From fef191df5e681c3c89166fdf3b0caff8e75399c8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 13:03:43 +0000 Subject: [PATCH 001/160] Add the RowFn scalar function framework Progress towards #9128. `RowFn` derives the whole of `ScalarFnVTable` from a row closure: name the element types, write the closure, and lifting supplies null propagation, constant folding, nullability, validity, and options serde. The input side is open. `InputElement::Elem` is a GAT, so an element can hand the closure borrowed variable-length data or drill through a wrapper to an extension array's storage. Covering a new type family is one impl. The output is always an `OutputSink`, allocated once per batch and handing the closure one row to write. `ElementSink` covers one owned `OutputElement` per row, and a custom sink carries runtime-shaped output such as a tensor whose width comes from its input dtype. Work that depends only on a batch-constant operand goes in `RowVisitor::visit_prepared_into`'s once-per-batch prepare step. Prepare must not be load-bearing for validation, because an empty batch decodes every operand as non-constant. Null-strategy selection is derived too. A nullable batch runs densely, by branch-and-skip, or by filtering, and the framework picks per batch. The one input an element controls is `InputElement::FILTERED_DECODE_COST`, set when decoding a column does expensive per-row work, so sparse batches keep the filter strategy's shrunken decode. Two things send a function to `ScalarFnVTable` instead, and no output sink covers either: a result that aliases an input, and a null result for a non-null row. The module docs on `scalar_fn` record the full choice between the two traits. Signed-off-by: Connor Tsui Co-authored-by: Claude --- vortex-array/src/scalar_fn/mod.rs | 65 ++ .../src/scalar_fn/row/element/bool.rs | 73 ++ .../src/scalar_fn/row/element/conformance.rs | 78 ++ vortex-array/src/scalar_fn/row/element/mod.rs | 167 +++++ .../src/scalar_fn/row/element/primitive.rs | 78 ++ .../src/scalar_fn/row/element/tuple.rs | 370 ++++++++++ vortex-array/src/scalar_fn/row/execute.rs | 195 +++++ vortex-array/src/scalar_fn/row/lift.rs | 688 ++++++++++++++++++ vortex-array/src/scalar_fn/row/mod.rs | 74 ++ vortex-array/src/scalar_fn/row/result.rs | 161 ++++ vortex-array/src/scalar_fn/row/row_fn.rs | 138 ++++ vortex-array/src/scalar_fn/row/sink.rs | 133 ++++ .../src/scalar_fn/row/tests/conformance.rs | 76 ++ .../scalar_fn/row/tests/constant_operands.rs | 187 +++++ .../scalar_fn/row/tests/decode_fallibility.rs | 98 +++ .../src/scalar_fn/row/tests/dispatched.rs | 77 ++ .../src/scalar_fn/row/tests/lifting.rs | 230 ++++++ vortex-array/src/scalar_fn/row/tests/mod.rs | 509 +++++++++++++ .../scalar_fn/row/tests/null_strategies.rs | 502 +++++++++++++ .../scalar_fn/row/tests/nullable_outputs.rs | 120 +++ .../src/scalar_fn/row/tests/prepared.rs | 144 ++++ vortex-array/src/scalar_fn/row/tests/sink.rs | 375 ++++++++++ vortex-array/src/scalar_fn/row/vtable.rs | 398 ++++++++++ vortex-array/src/scalar_fn/vtable.rs | 2 +- 24 files changed, 4937 insertions(+), 1 deletion(-) create mode 100644 vortex-array/src/scalar_fn/row/element/bool.rs create mode 100644 vortex-array/src/scalar_fn/row/element/conformance.rs create mode 100644 vortex-array/src/scalar_fn/row/element/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/element/primitive.rs create mode 100644 vortex-array/src/scalar_fn/row/element/tuple.rs create mode 100644 vortex-array/src/scalar_fn/row/execute.rs create mode 100644 vortex-array/src/scalar_fn/row/lift.rs create mode 100644 vortex-array/src/scalar_fn/row/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/result.rs create mode 100644 vortex-array/src/scalar_fn/row/row_fn.rs create mode 100644 vortex-array/src/scalar_fn/row/sink.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/conformance.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/constant_operands.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/decode_fallibility.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/dispatched.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/lifting.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/null_strategies.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/nullable_outputs.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/prepared.rs create mode 100644 vortex-array/src/scalar_fn/row/tests/sink.rs create mode 100644 vortex-array/src/scalar_fn/row/vtable.rs diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 003dbc2ca48..11bcefe0325 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -6,6 +6,68 @@ //! This module contains the [`ScalarFnVTable`] trait and all built-in scalar function //! implementations. Expressions ([`crate::expr::Expression`]) reference scalar functions //! at each node. +//! +//! # Choosing a trait +//! +//! Two traits reach this vtable, and [`RowFn`] derives the whole of [`ScalarFnVTable`] from a row +//! closure. Implement `RowFn` when the function fits it, and `ScalarFnVTable` when it does not. +//! +//! [`RowFn`] is for a kernel whose value at a row is determined by that row alone, and which has to +//! read every row anyway: the arithmetic operators over primitive columns, `vortex.tensor.l2_norm`, +//! `vortex.tensor.inner_product`, `vortex.tensor.cosine_similarity`, `vortex.geo.distance`, +//! `vortex.geo.contains`. Name the element types and write the row closure, and the rest is +//! derived, including which rows get visited. +//! +//! Its *input* side is open. [`InputElement::Elem`] is a GAT, so an element can hand the closure +//! borrowed variable-length data (a byte-string element yielding `&[u8]`) or drill through a wrapper +//! (`vortex-tensor`'s `TensorRow` yields a slice of an extension array's storage). Covering a new +//! type family, a list row included, is one impl. +//! +//! Its output is always an [`OutputSink`], allocated once per batch and handing the closure one row +//! to write. [`ElementSink`] is the standard sink for one owned [`OutputElement`] per row. A custom +//! sink carries runtime-shaped output, such as a tensor whose width comes from its input dtype, or a +//! future string transform appending every row into one shared byte buffer. +//! +//! When part of the kernel's work depends only on an operand that is constant for the batch (the +//! norm of a broadcast query vector, a prepared form of a constant geometry), do that work in +//! [`RowVisitor::visit_prepared_into`]'s once-per-batch prepare step. Pass `|_| ()` when there is +//! nothing to prepare. Prepare **must not** be load-bearing for validation: an empty batch decodes +//! every operand as non-constant, so a prepare that validated its constant would silently not run. +//! +//! Null handling is derived too, null-strategy selection included: a nullable batch runs densely +//! (compute every row, mask after), by branch-and-skip (decode full length, compute only the +//! conjoined-valid rows, mask after), or by filtering (shrink the inputs to the valid rows, +//! compute, scatter back), and the framework picks per batch. Function authors do nothing. The one +//! input to that choice an element controls is [`InputElement::FILTERED_DECODE_COST`]: set it when +//! decoding a column does expensive per-row work (parsing a geometry), so sparse batches keep the +//! filter strategy's shrunken decode. Costs from separate arguments are additive. +//! +//! Two things no output sink covers, and they are what actually send a function to +//! [`ScalarFnVTable`]: +//! +//! - **A result that aliases an input.** Sinks own their output bytes. Trimming strings is the +//! example, where the ideal kernel keeps the input's data buffer and writes new views over it, +//! copying no bytes, which only a columnar kernel can express. +//! - **A null result for a non-null row.** Sinks build an all-valid column, so +//! `vortex.list.sum` cannot be a row function: a valid empty list sums to null. +//! +//! [`ScalarFnVTable`] takes the whole column instead, and everything a row function gets derived is +//! then hand-written: null propagation, constant folding, nullability, validity, and options serde. +//! Besides the two cases above and the functions that are simply not strict (Kleene logic, or a +//! strictness that depends on the options), reach for it when a row loop *could* express the +//! function but would do avoidable work: +//! +//! - **The answer is already an array, or is one value for the whole column.** +//! `vortex.list.length` hands back a `ListViewArray`'s sizes child, and a single `ConstantArray` +//! for a `FixedSizeListArray`. A row loop would rebuild that one `u64` at a time, even given a +//! list-length element that reads the size out of the layout rather than the list. +//! - **A row is not the natural unit of work.** `vortex.not` is one `!` per 64-bit word, in place +//! when the bit buffer is unshared, against 64 loop iterations and 64 bit writes, and its +//! encoding-aware fallback pushes the inversion down instead of canonicalizing. +//! - **The row's value is cheaper to read than the row.** `vortex.byte_length` was tried as a row +//! function and measured 7.6x slower than its columnar implementation, because the length is a +//! field of the view and the row loop paid to resolve the bytes it never looked at. Being +//! row-determined is necessary but not sufficient. use vortex_session::registry::Id; @@ -35,6 +97,9 @@ pub use options::*; mod signature; pub use signature::*; +mod row; +pub use row::*; + pub mod fns; pub mod internal; pub mod session; diff --git a/vortex-array/src/scalar_fn/row/element/bool.rs b/vortex-array/src/scalar_fn/row/element/bool.rs new file mode 100644 index 00000000000..d4fc51c769c --- /dev/null +++ b/vortex-array/src/scalar_fn/row/element/bool.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::InputElement; +use crate::scalar_fn::OutputElement; +use crate::validity::Validity; + +impl InputElement for bool { + type Column = BitBuffer; + type Varying<'a> = &'a BitBuffer; + type Elem<'a> = bool; + + // Every bit of the buffer is readable, valid or not. + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + matches!(dtype, DType::Bool(_)), + "expected a Bool column, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_bit_buffer()) + } + + fn get(column: &Self::Column, index: usize) -> bool { + column.value(index) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> bool + where + Self: 'a, + { + column.value(index) + } +} + +impl OutputElement for bool { + fn element_dtype() -> DType { + DType::Bool(Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + // `From>` packs through the multiversioned SIMD path; `from_iter` would set one + // bit at a time, which measures 6.6-7.9x slower on the packing step alone. + BoolArray::new(BitBuffer::from(values), Validity::NonNullable).into_array() + } + + fn placeholder() -> Self { + false + } +} diff --git a/vortex-array/src/scalar_fn/row/element/conformance.rs b/vortex-array/src/scalar_fn/row/element/conformance.rs new file mode 100644 index 00000000000..45687ea9c1a --- /dev/null +++ b/vortex-array/src/scalar_fn/row/element/conformance.rs @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A shared conformance check every [`InputElement`] should be run through. + +use std::hint::black_box; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::InputElement; + +/// Assert that `E` honors its [`InputElement`] contract over `array`, and rejects `rejected_dtype`. +/// +/// The part worth checking mechanically is [`InputElement::DENSE_SAFE`]. An element claiming it will +/// be read at rows that are *null*, where an array guarantees nothing about the payload, and getting +/// the `const` wrong is either an out-of-bounds panic in production (the failure mode of +/// [#9090](https://github.com/vortex-data/vortex/issues/9090)) or an unnecessary valid-only +/// execution path. Nothing else verifies it, since the framework reads the `const` rather than +/// testing the claim. +/// +/// So `array` **must** contain at least one null row, and its payload behind those nulls **must** be +/// deliberately extreme rather than zeroed, or the check passes vacuously. Build that safely by +/// putting the extreme values in the array first and masking those rows afterwards, as the callers of +/// this function do. +/// +/// What this cannot check: [`DECODE_FALLIBLE`](InputElement::DECODE_FALLIBLE), which needs data that +/// is legal but malformed, and whether `validate` accepts everything it *should*, since only the +/// element knows its full dtype domain. Pass one representative rejection. +#[track_caller] +pub fn assert_element_conforms( + array: ArrayRef, + rejected_dtype: &DType, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let dtype = array.dtype().clone(); + E::validate(&dtype)?; + + assert!( + E::validate(rejected_dtype).is_err(), + "element accepted {rejected_dtype}, which it was expected to reject", + ); + + let len = array.len(); + let valid = array.validity()?.execute_mask(len, ctx)?; + assert!( + !valid.all_true(), + "conformance needs a null row to read behind, but every row of the {dtype} input is valid", + ); + + let column = E::decode(array, ctx)?; + let varying = E::varying(&column); + assert_eq!( + E::varying_len(&varying), + len, + "varying element view changed the decoded row count", + ); + + // The claim under test. Reading a null row may yield garbage, but it must not fault, so an + // element that secretly follows a per-row offset panics here instead of in production. + if E::DENSE_SAFE { + for index in 0..len { + black_box(E::get(&column, index)); + black_box(E::get_varying(&varying, index)); + } + } else { + for index in 0..len { + if valid.value(index) { + black_box(E::get(&column, index)); + black_box(E::get_varying(&varying, index)); + } + } + } + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/element/mod.rs b/vortex-array/src/scalar_fn/row/element/mod.rs new file mode 100644 index 00000000000..fc88690b6b0 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/element/mod.rs @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The element types a row function can read and produce. +//! +//! Both traits are open, and this module holds one file per type family, so covering a new one is a +//! sibling file and every row function gains it. The families are not confined to this crate: +//! `vortex-tensor`'s `TensorRow` drills through an extension wrapper into its storage. +//! +//! The two directions are deliberately asymmetric. [`InputElement::Elem`] is a GAT, so an input row +//! can borrow out of the decoded column, while an [`OutputElement`] is one owned value written into +//! an [`ElementSink`](crate::scalar_fn::ElementSink). Runtime-shaped output uses a custom +//! [`OutputSink`](crate::scalar_fn::OutputSink) instead. + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; + +mod bool; + +#[cfg(any(test, feature = "_test-harness"))] +mod conformance; +#[cfg(any(test, feature = "_test-harness"))] +pub use conformance::assert_element_conforms; + +mod primitive; + +mod tuple; +pub use tuple::ElementTuple; +pub(super) use tuple::batch_constant; + +/// An element type that can be read row-wise out of an input column. +pub trait InputElement: 'static { + /// The decoded column representation supporting `O(1)` row access. + type Column; + + /// The view of a varying decoded column read by the hot row loop. + /// + /// This may borrow a cheaper representation than [`Column`](Self::Column). Primitive elements, + /// for example, expose a slice so its pointer and length are loop invariants rather than + /// re-reading a [`Buffer`](vortex_buffer::Buffer) descriptor for every row. + type Varying<'a>; + + /// The borrowed element value handed to the row closure a [`RowFn`](crate::scalar_fn::RowFn) + /// visits with. + type Elem<'a>; + + /// Whether [`decode`](Self::decode) and [`get`](Self::get) tolerate rows that are null in the + /// input. + /// + /// Arrays only guarantee their contents for _valid_ rows, so this is `false` for any element + /// that follows an offset or pointer stored in the array: behind a null row that value is + /// arbitrary and may not address anything. Reading a whole value out of a flat buffer is `true`, + /// since the value is garbage but the read cannot fault. + /// + /// Dense execution requires this of every argument; otherwise the row layer executes only + /// valid rows. + const DENSE_SAFE: bool = false; + + /// Whether [`decode`](Self::decode) can fail on _legal_ input data. + /// + /// `false` for an element read straight out of a buffer: decoding can still fail for + /// infrastructural reasons (IO, allocation), but never because of the values. `true` for an + /// element that parses its bytes, since a malformed WKB geometry in a _valid_ row is a domain + /// error, which makes a function over that element + /// [fallible](crate::scalar_fn::ScalarFnVTable::is_fallible) however infallible its own row + /// computation is. + const DECODE_FALLIBLE: bool = true; + + /// A relative unit count for per-row decode work avoided by filtering this argument first. + /// + /// Use `1` for an element whose decode _parses_ every row (a geometry built from coordinate + /// storage): decoding only the survivors of a sparse validity mask is genuinely cheaper than + /// decoding everyone. Keep the default `0` for a bulk canonicalization (bytes, bools, + /// primitives), whose decode is a memcpy-shaped pass that filtering barely shrinks. Larger + /// values may express a proportionally more expensive decode. + /// + /// The lifting reads this when it picks a null strategy for a batch with a mixed + /// validity mask: filtering the inputs first only pays off when it shrinks a per-row decode, + /// so elements that leave this at zero always take the cheaper branch-and-skip strategy. + /// Getting it wrong is a performance bug, never a correctness bug. + const FILTERED_DECODE_COST: usize = 0; + + /// Validate that `dtype` is an acceptable input column dtype for this element type. + fn validate(dtype: &DType) -> VortexResult<()>; + + /// Decode `array` into its column representation. Called once per batch. + /// + /// This is where every per-batch cost belongs: resolving the dtype, downcasting the buffer, + /// checking the ptype, and anything else that does not vary by row. [`Column`](Self::Column) is + /// the type to widen if that means carrying more, since it is chosen by the element. + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Decode `array` _without_ assuming every row is valid, or `Ok(None)` when this element + /// cannot for this particular array. + /// + /// An element with [`DENSE_SAFE`](Self::DENSE_SAFE) set **should not** override this: its + /// ordinary decode already tolerates null payloads, so the default is already correct and an + /// override just restates it. Overriding is for an element that is *not* dense-safe but can + /// still write an arbitrary placeholder into null slots; the caller guarantees + /// [`get`](Self::get) is never called for such a row. It is what the branch-and-skip null + /// strategy decodes with. + /// + /// Return `Ok(None)` rather than an error when an array has no null-tolerant decode; the lifting + /// falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if Self::DENSE_SAFE { + Self::decode(array, ctx).map(Some) + } else { + Ok(None) + } + } + + /// Read the element at `index`, the one function called once per row. + /// + /// `O(1)` is necessary but **not sufficient**: this must not repeat work that is constant across + /// the batch, however cheap that work looks per call. An `O(1)` ptype check and buffer downcast + /// per row cost `l2_norm` 2x at width 2, invisible in the call because it read like a getter. Do + /// that work in [`decode`](Self::decode) and leave this an offset computation. + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>; + + /// Borrow the representation used when this argument varies within the batch. + /// + /// Called once before the hot loop. Constants do not use this view because the tuple adapter + /// keeps their one-row decoded representation separate. + fn varying(column: &Self::Column) -> Self::Varying<'_>; + + /// Number of rows addressable through a [`Varying`](Self::Varying) view. + fn varying_len(column: &Self::Varying<'_>) -> usize; + + /// Read one row from a [`Varying`](Self::Varying) view. + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a; +} + +/// An element type that a row computation can produce, buildable into an all-valid column. +/// +/// [`Clone`] is required so [`ElementSink`](crate::scalar_fn::ElementSink) can allocate through +/// `vec![placeholder; rows]`, which is what lets a zero placeholder reach the allocator's zeroed +/// path instead of costing a write pass over the output. +pub trait OutputElement: 'static + Sized + Clone { + /// The dtype of columns built from this element type. Must be non-nullable: nullability is + /// derived from the inputs by the lifting. + /// + /// Taking no arguments confines an element's dtype to a property of its Rust type, so an output + /// whose dtype depends on runtime data (a tensor, whose dtype carries its shape) cannot be an + /// element. Such an output uses an [`OutputSink`](crate::scalar_fn::OutputSink), whose + /// [`sink_dtype`](crate::scalar_fn::OutputSink::sink_dtype) does see the input dtypes. + fn element_dtype() -> DType; + + /// Build a column from one value per row. Called once per batch. + fn build(values: Vec) -> ArrayRef; + + /// An arbitrary value of this element, pre-filled into the output slots that the + /// branch-and-skip null strategy skips. + /// + /// The value is never observable: the lifting masks every slot holding it before the + /// result escapes. It only has to be cheap to construct and legal to + /// [`build`](Self::build) with. + fn placeholder() -> Self; +} diff --git a/vortex-array/src/scalar_fn/row/element/primitive.rs b/vortex-array/src/scalar_fn/row/element/primitive.rs new file mode 100644 index 00000000000..d54c922bf2d --- /dev/null +++ b/vortex-array/src/scalar_fn/row/element/primitive.rs @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::InputElement; +use crate::scalar_fn::OutputElement; +use crate::validity::Validity; + +impl InputElement for T { + type Column = Buffer; + type Varying<'a> = &'a [T]; + type Elem<'a> = T; + + // Every lane of the buffer holds a `T`, valid or not. + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + let expected = T::PTYPE; + let DType::Primitive(ptype, _) = dtype else { + vortex_bail!("expected a {expected} column, got {dtype}"); + }; + vortex_ensure_eq!( + *ptype, + expected, + "expected a {expected} column, got {dtype}" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_buffer::()) + } + + fn get(column: &Self::Column, index: usize) -> T { + column[index] + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column.as_slice() + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> T + where + Self: 'a, + { + column[index] + } +} + +impl OutputElement for T { + fn element_dtype() -> DType { + DType::Primitive(T::PTYPE, Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + PrimitiveArray::new(values, Validity::NonNullable).into_array() + } + + fn placeholder() -> Self { + T::default() + } +} diff --git a/vortex-array/src/scalar_fn/row/element/tuple.rs b/vortex-array/src/scalar_fn/row/element/tuple.rs new file mode 100644 index 00000000000..a3c31cfeb2e --- /dev/null +++ b/vortex-array/src/scalar_fn/row/element/tuple.rs @@ -0,0 +1,370 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Argument lists built from [`InputElement`]s, and the per-argument decode behind them. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Extension; +use crate::arrays::Masked; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::InputElement; + +mod private { + pub trait Sealed {} +} + +/// One decoded input column of an [`ElementTuple`]. +/// +/// A constant operand holds the same value in every row, so it is decoded once as a single row and +/// read at index 0 forever. That is what stops a constant argument costing one decode per row, which +/// matters whenever the decode is more than a buffer read: parsing a geometry from WKB, or +/// canonicalizing an extension row. +pub struct ArgColumn(ArgColumnKind); + +enum ArgColumnKind { + Varying(T::Column), + Constant(T::Column), +} + +impl ArgColumn { + /// Decode one input column, collapsing a constant operand to its single distinct row. + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // An empty input has no row 0 to slice, and its row loop runs zero times either way. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?))); + } + + Ok(Self(ArgColumnKind::Varying(T::decode(array, ctx)?))) + } + + /// Like [`decode`](Self::decode), but a varying column decodes null-tolerantly through + /// [`InputElement::decode_null_tolerant`]. `Ok(None)` means the element cannot, and the + /// caller falls back to the filter strategy. + /// + /// A constant operand still takes the ordinary decode: the lifting short-circuits null + /// constants before any strategy runs, so a constant reaching here is non-null. + fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Some(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?)))); + } + + Ok(T::decode_null_tolerant(array, ctx)? + .map(ArgColumnKind::Varying) + .map(Self)) + } + + /// Read the element at `index`, which for a constant operand is always its single row. + fn get(&self, index: usize) -> T::Elem<'_> { + match &self.0 { + ArgColumnKind::Varying(column) => T::get(column, index), + ArgColumnKind::Constant(column) => T::get(column, 0), + } + } + + /// The decoded full column, or `None` when this argument was collapsed to one constant row. + fn varying(&self) -> Option<&T::Column> { + match &self.0 { + ArgColumnKind::Varying(column) => Some(column), + ArgColumnKind::Constant(_) => None, + } + } + + /// Whether this argument addresses exactly `row_count` rows. + /// + /// A constant operand was collapsed to its one distinct row and is read at index 0 forever, so + /// it addresses any row count and is exempt. + fn addresses_rows(&self, row_count: usize) -> bool { + match &self.0 { + ArgColumnKind::Varying(column) => T::varying_len(&T::varying(column)) == row_count, + ArgColumnKind::Constant(_) => true, + } + } + + /// The single decoded element of a constant operand, or `None` for a real column. + /// + /// `Some` exactly when [`decode`](Self::decode) collapsed the operand to its one distinct row, + /// in which case the value returned is the element every row of the batch reads. + fn constant(&self) -> Option> { + match &self.0 { + ArgColumnKind::Varying(_) => None, + ArgColumnKind::Constant(column) => Some(T::get(column, 0)), + } + } +} + +/// The array whose every row holds one distinct value, when `array` is constant for the batch. +/// +/// Beyond the constant encoding itself this sees one level through two wrappers that spell "the +/// same value in every row" without being the constant encoding: +/// +/// - [`Masked`], how the compressor spells an all-same-with-nulls chunk: the child carries the +/// value, the wrapper carries only validity. Reading the child's value for a null row is sound +/// here because the lifting owns validity entirely; the row loop's output behind a null +/// row is masked away (dense) or never computed (filter), so which value the loop read there +/// cannot be observed. An all-null constant never reaches decode at all, since the lifting +/// short-circuits it to an all-null result first. +/// - [`Extension`] over constant storage, the shape an extension-typed builder produces before +/// `ExtensionConstantRule` normalizes it to a top-level constant. Every row wraps the same +/// storage value, so the whole array (sliced to one row, keeping its extension dtype) is the +/// constant. +pub(in crate::scalar_fn::row) fn batch_constant(array: &ArrayRef) -> Option { + if array.as_constant().is_some() { + return Some(array.clone()); + } + + if let Some(masked) = array.as_opt::() { + return Some(masked.child().clone()).filter(|child| child.as_constant().is_some()); + } + + array + .as_opt::() + .is_some_and(|ext| ext.storage_array().as_constant().is_some()) + .then(|| array.clone()) +} + +/// Tuples of [`InputElement`]s forming the typed argument list a [`RowFn`](crate::scalar_fn::RowFn) +/// visits with. Implemented for `()` and tuples of one through twelve elements. This trait is +/// framework-only; add a new decode primitive by implementing [`InputElement`], then use it inside +/// one of those tuples. +/// +/// The arities past the widest function in tree are deliberate. This trait is **sealed**, so a +/// downstream crate cannot add the one it needs, and an unused arity costs only its own macro +/// expansion: no monomorphization happens until something instantiates it. +pub trait ElementTuple: 'static + private::Sealed { + /// The decoded column representations. + type Columns; + + /// Direct references to decoded columns when every argument varies within the batch. + type VaryingColumns<'a>; + + /// The borrowed row of element values. + type Elems<'a>; + + /// The batch-constant element values: [`Elems`](Self::Elems) with every argument wrapped in + /// `Option`. + /// + /// `Some` marks an argument whose operand is constant for the batch and carries the element + /// every row reads; `None` marks one that varies by row. This is what + /// [`visit_prepared_into`](crate::scalar_fn::RowVisitor::visit_prepared_into) hands to its prepare + /// closure, so a kernel can hoist work that depends only on a constant argument out of the + /// row loop. + type ConstElems<'a>; + + /// The number of arguments. + const ARITY: usize; + + /// Whether every argument is [`InputElement::DENSE_SAFE`]. + const DENSE_SAFE: bool; + + /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. + const DECODE_FALLIBLE: bool; + + /// The additive cost of per-row decode work avoided by filtering the arguments first. + const FILTERED_DECODE_COST: usize; + + /// Validate the input dtypes, including that `dtypes` has exactly `ARITY` entries. + /// + /// The expression layer checks the count against [`Arity`](crate::scalar_fn::Arity) before it + /// builds a call, but this is also the entry point of the public + /// [`return_dtype`](crate::scalar_fn::ScalarFnVTable::return_dtype), so the count is enforced + /// here rather than assumed. + fn validate(dtypes: &[DType]) -> VortexResult<()>; + + /// Decode every input column once. Called once per batch. + fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Decode every input column once, tolerating null rows, or `Ok(None)` when some argument + /// cannot. Called once per batch by the branch-and-skip null strategy. + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; + + /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; + + /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// + /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple + /// gives the optimizer ordinary contiguous column access without a per-row constant check. + fn varying(columns: &Self::Columns) -> Option>; + + /// Whether every varying column contains exactly `row_count` rows. + fn varying_len_matches(columns: &Self::VaryingColumns<'_>, row_count: usize) -> bool; + + /// Whether every argument that varies within the batch contains exactly `row_count` rows. + /// + /// The same guarantee as [`varying_len_matches`](Self::varying_len_matches), for the mixed case + /// [`varying`](Self::varying) declines: a batch-constant argument is exempt because it was + /// collapsed to one row, while every argument beside it still has to address the whole batch. + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; + + /// Read one row from columns already known to vary within the batch. + fn get_varying<'a>(columns: &Self::VaryingColumns<'a>, index: usize) -> Self::Elems<'a>; + + /// Read the batch-constant elements out of the decoded columns. Called once per batch. + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; +} + +impl private::Sealed for () {} + +impl ElementTuple for () { + type Columns = (); + type VaryingColumns<'a> = (); + type Elems<'a> = (); + type ConstElems<'a> = (); + + const ARITY: usize = 0; + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + const FILTERED_DECODE_COST: usize = 0; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + 0, + "expected 0 argument dtypes, got {}", + dtypes.len(), + ); + Ok(()) + } + + fn decode(_args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(()) + } + + fn decode_null_tolerant( + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(())) + } + + fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} + + fn varying(_columns: &Self::Columns) -> Option> { + Some(()) + } + + fn varying_len_matches(_columns: &Self::VaryingColumns<'_>, _row_count: usize) -> bool { + true + } + + fn decoded_lens_match(_columns: &Self::Columns, _row_count: usize) -> bool { + true + } + + fn get_varying<'a>(_columns: &Self::VaryingColumns<'a>, _index: usize) -> Self::Elems<'a> {} + + fn constants(_columns: &Self::Columns) -> Self::ConstElems<'_> {} +} + +macro_rules! element_tuple { + ($arity:literal; $($t:ident : $idx:tt),+) => { + impl<$($t: InputElement),+> private::Sealed for ($($t,)+) {} + + impl<$($t: InputElement),+> ElementTuple for ($($t,)+) { + type Columns = ($(ArgColumn<$t>,)+); + type VaryingColumns<'a> = ($($t::Varying<'a>,)+); + type Elems<'a> = ($($t::Elem<'a>,)+); + type ConstElems<'a> = ($(Option<$t::Elem<'a>>,)+); + + const ARITY: usize = $arity; + const DENSE_SAFE: bool = $($t::DENSE_SAFE &&)+ true; + const DECODE_FALLIBLE: bool = $($t::DECODE_FALLIBLE ||)+ false; + const FILTERED_DECODE_COST: usize = $($t::FILTERED_DECODE_COST +)+ 0; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + $arity, + "expected {} argument dtypes, got {}", + $arity, + dtypes.len(), + ); + + $($t::validate(&dtypes[$idx])?;)+ + Ok(()) + } + + fn decode( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(($(ArgColumn::<$t>::decode(args.get($idx)?, ctx)?,)+)) + } + + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(( + $(match ArgColumn::<$t>::decode_null_tolerant(args.get($idx)?, ctx)? { + Some(column) => column, + None => return Ok(None), + },)+ + ))) + } + + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_> { + ($(columns.$idx.get(index),)+) + } + + fn varying(columns: &Self::Columns) -> Option> { + Some(($($t::varying(columns.$idx.varying()?),)+)) + } + + fn varying_len_matches( + columns: &Self::VaryingColumns<'_>, + row_count: usize, + ) -> bool { + $($t::varying_len(&columns.$idx) == row_count &&)+ true + } + + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { + $(columns.$idx.addresses_rows(row_count) &&)+ true + } + + fn get_varying<'a>( + columns: &Self::VaryingColumns<'a>, + index: usize, + ) -> Self::Elems<'a> { + ($($t::get_varying(&columns.$idx, index),)+) + } + + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { + ($(columns.$idx.constant(),)+) + } + } + }; +} + +element_tuple!(1; A:0); +element_tuple!(2; A:0, B:1); +element_tuple!(3; A:0, B:1, C:2); +element_tuple!(4; A:0, B:1, C:2, D:3); +element_tuple!(5; A:0, B:1, C:2, D:3, E:4); +element_tuple!(6; A:0, B:1, C:2, D:3, E:4, F:5); +element_tuple!(7; A:0, B:1, C:2, D:3, E:4, F:5, G:6); +element_tuple!(8; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7); +element_tuple!(9; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8); +element_tuple!(10; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9); +element_tuple!(11; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10); +element_tuple!(12; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11); diff --git a/vortex-array/src/scalar_fn/row/execute.rs b/vortex-array/src/scalar_fn/row/execute.rs new file mode 100644 index 00000000000..345bed991fe --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute.rs @@ -0,0 +1,195 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The pieces every row function is built from, whatever its `dispatch` chooses. +//! +//! These back the blanket impls in [`row_fn`](super::row_fn) and are deliberately not public: +//! [`RowFn`](crate::scalar_fn::RowFn) is the abstraction, these are its internals. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::DeferredError; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::SinkResult; + +/// The value path out of a row executor, keeping a deferred row error distinct from structural +/// execution errors so nullable lifting retries only the former. +pub(super) enum RowExecution { + /// A successfully built output column. + Output(ArrayRef), + + /// A batch-wide row error that nullable lifting may retry over only the valid rows. + DeferredError(VortexError), +} + +impl RowExecution { + /// Return the output or surface its deferred row error. + pub(super) fn into_result(self) -> VortexResult { + match self { + Self::Output(output) => Ok(output), + Self::DeferredError(error) => Err(error), + } + } +} + +/// Validate the input dtypes of a sink-writing row function and return the dtype its sink builds. +/// +/// The output dtype may be a function of the inputs. A sink can also own a batch-wide builder, such +/// as the shared byte and view buffers of a future string transform. +pub(super) fn validate_row_sink( + args: &[DType], +) -> VortexResult { + A::validate(args)?; + let dtype = S::sink_dtype(args)?; + vortex_ensure!( + !dtype.is_nullable(), + "row output sinks must declare a non-nullable dtype, got {dtype}", + ); + Ok(dtype) +} + +/// Decode every input column once, allocate the sink once, then write one row at a time. +/// +/// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and the loop keeps the +/// unconditional shape that lets it vectorize. Monomorphic in `A`, `S` and `R`, so `apply` and +/// [`OutputSink::row`] both inline. +pub(super) fn execute_row_sink_prepared( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, +) -> VortexResult { + let row_count = args.row_count(); + let mut sink = S::with_capacity(row_count, sink_dtype)?; + let columns = A::decode(args, ctx)?; + let state = prepare(A::constants(&columns)); + let mut accumulated = R::Accumulated::default(); + + { + let mut rows = sink.rows(); + vortex_ensure!( + S::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + if let Some(varying) = A::varying(&columns) { + vortex_ensure!( + A::varying_len_matches(&varying, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + for index in 0..row_count { + apply( + &state, + A::get_varying(&varying, index), + S::row(&mut rows, index), + ) + .accumulate(&mut accumulated)?; + } + } else { + vortex_ensure!( + A::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + for index in 0..row_count { + apply(&state, A::get(&columns, index), S::row(&mut rows, index)) + .accumulate(&mut accumulated)?; + } + } + } + + finish_sink(sink, DeferredError::new(R::occurred(accumulated))) +} + +/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +pub(super) fn execute_row_sink_branch( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + valid: &Mask, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, +) -> VortexResult> { + if !S::SUPPORTS_SKIPPED_ROWS { + return Ok(None); + } + + let Some(columns) = A::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + let state = prepare(A::constants(&columns)); + let row_count = args.row_count(); + let mut sink = S::with_capacity(row_count, sink_dtype)?; + let mut accumulated = R::Accumulated::default(); + + let AllOr::Some(valid) = valid.bit_buffer() else { + vortex_bail!("execute_row_sink_branch requires a mixed mask"); + }; + + { + let mut rows = sink.rows(); + vortex_ensure!( + S::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + let varying = A::varying(&columns); + let lens_match = match &varying { + Some(varying) => A::varying_len_matches(varying, row_count), + None => A::decoded_lens_match(&columns, row_count), + }; + vortex_ensure!( + lens_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut error = None; + valid.for_each_set_index(|index| { + if error.is_some() { + return; + } + + let result = match &varying { + Some(varying) => apply( + &state, + A::get_varying(varying, index), + S::row(&mut rows, index), + ), + None => apply(&state, A::get(&columns, index), S::row(&mut rows, index)), + }; + if let Err(err) = result.accumulate(&mut accumulated) { + error = Some(err); + } + }); + + if let Some(error) = error { + return Err(error); + } + } + + finish_sink(sink, DeferredError::new(R::occurred(accumulated))).map(Some) +} + +/// Finish a sink while preserving whether its error came from the deferred row accumulator. +fn finish_sink( + sink: S, + deferred_error: DeferredError, +) -> VortexResult { + match sink.finish(deferred_error) { + Ok(output) => Ok(RowExecution::Output(output)), + Err(error) if deferred_error.occurred() => Ok(RowExecution::DeferredError(error)), + Err(error) => Err(error), + } +} diff --git a/vortex-array/src/scalar_fn/row/lift.rs b/vortex-array/src/scalar_fn/row/lift.rs new file mode 100644 index 00000000000..a842ef5deef --- /dev/null +++ b/vortex-array/src/scalar_fn/row/lift.rs @@ -0,0 +1,688 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lifting a kernel over non-null values into a full [`ScalarFnVTable::execute`]. +//! +//! A [`RowFn`] hands the framework a kernel that only ever computes rows valid in every argument. +//! Everything between that kernel and [`ScalarFnVTable::execute`] lives here: null propagation, +//! constant folding, nullability widening, output dtype reconciliation, and the per-batch choice +//! between dense execution and the two mechanisms that execute only valid rows. +//! +//! This is machinery, not an interface. It takes the kernel as a pair of closures rather than a +//! trait because the one trait that ever occupied the slot (a public `StrictScalarFnVTable`, with +//! [`RowFn`] blanket-implementing it) never found a second implementor, and the indirection cost +//! more than it explained. Extract a trait if and when a non-row user appears. +//! +//! [`RowFn`]: crate::scalar_fn::RowFn +//! [`ScalarFnVTable::execute`]: crate::scalar_fn::ScalarFnVTable::execute + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::SinkResult; +use crate::scalar_fn::row::element::batch_constant; +use crate::scalar_fn::row::execute::RowExecution; +use crate::validity::Validity; + +struct BorrowedExecutionArgs<'a> { + inputs: &'a [ArrayRef], + row_count: usize, +} + +impl<'a> BorrowedExecutionArgs<'a> { + fn new(inputs: &'a [ArrayRef], row_count: usize) -> Self { + Self { inputs, row_count } + } +} + +impl ExecutionArgs for BorrowedExecutionArgs<'_> { + fn get(&self, index: usize) -> VortexResult { + self.inputs.get(index).cloned().ok_or_else(|| { + vortex_error::vortex_err!( + "Input index {} out of bounds (num_inputs={})", + index, + self.inputs.len() + ) + }) + } + + fn num_inputs(&self) -> usize { + self.inputs.len() + } + + fn row_count(&self) -> usize { + self.row_count + } +} + +/// The arguments handed to one kernel invocation. +/// +/// `arrays` may be filtered or sliced, while `dtypes` and `sink_dtype` always describe the original +/// planned batch. Keeping them together prevents an execution path from accidentally pairing an +/// input view with unrelated planning metadata. +#[derive(Clone, Copy)] +pub(super) struct KernelArgs<'a> { + /// The executor-facing view, including the row count for this invocation. + pub(super) execution: &'a dyn ExecutionArgs, + + /// The same inputs as concrete arrays for encoding-aware rewrites. + pub(super) arrays: &'a [ArrayRef], + + /// The original input dtypes used to select the row implementation. + pub(super) dtypes: &'a [DType], + + /// The non-nullable dtype allocated by the selected output sink. + pub(super) sink_dtype: &'a DType, +} + +/// The execution policy and output dtype selected by a planning visit. +pub(super) struct BatchPlan { + /// The non-nullable dtype built by the selected sink. + pub(super) sink_dtype: DType, + + /// How this concrete dispatch executes nullable rows. + pub(super) policy: RowPolicy, +} + +/// The nullable execution policy derived from one concrete dispatch. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum RowPolicy { + /// Evaluate all rows and mask the result. + Dense, + + /// Evaluate all rows, retrying only valid rows if a deferred error is raised. + DenseWithRetry, + + /// Execute only valid rows, choosing branch-and-skip or filtering from the mask and decode + /// cost. + ValidOnly { filtered_decode_cost: usize }, +} + +impl RowPolicy { + /// The policy one concrete dispatch executes nullable rows under. + /// + /// Note what is deliberately **not** read here: [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. Hoisting + /// it into the plan so that a non-skipping sink never enters the branch path looks like a free + /// win, and #9130 records it as one, but it is not: the branch path probes + /// [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) against the _original_ arrays + /// before it ever consults the sink, and that is the only probe that sees them still encoded. + /// Skipping the path early would leave such a function with only the filtered probe, whose + /// canonical arrays match no encoding fast path. For a function whose reduction is defined to + /// answer differently from its row loop, that is a wrong answer rather than a slow one. + /// + /// [`OutputSink::SUPPORTS_SKIPPED_ROWS`]: crate::scalar_fn::OutputSink::SUPPORTS_SKIPPED_ROWS + pub(super) const fn for_dispatch() -> Self { + if A::DENSE_SAFE && !A::DECODE_FALLIBLE && !R::FALLIBLE { + if R::DEFERRED { + Self::DenseWithRetry + } else { + Self::Dense + } + } else { + Self::ValidOnly { + filtered_decode_cost: A::FILTERED_DECODE_COST, + } + } + } +} + +/// How far [`Batch::resolve_validity`] got before a mixed-mask strategy became necessary. +enum ResolvedMask { + /// The batch was answered without one: every row valid, or every row null. + Decided(ArrayRef), + + /// A mask with both set and unset bits, which a strategy must now execute. + Mixed(Mask), +} + +/// One batch of inputs, with everything the lifting reads off them before the kernel runs. +pub(super) struct Batch<'a> { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The arguments as the execution layer handed them over. Every path but the filter strategy + /// gives the kernel these untouched, so it sees the original encodings. + args: &'a dyn ExecutionArgs, + + /// The input columns, collected once: constant folding inspects them and the filter strategy + /// filters them. + inputs: SmallVec<[ArrayRef; 4]>, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: SmallVec<[DType; 4]>, + + /// The conjoined input validity, so a row of the output is valid iff it is valid in every + /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. + validity: Validity, + + /// The dtype the function declares for these inputs, which the kernel's output is reconciled + /// against. Already widened to nullable if any input is nullable. + result_dtype: DType, + + /// The non-nullable dtype the dispatched sink builds, computed once while planning. + sink_dtype: DType, + + /// How the concrete dispatch executes nullable rows. + policy: RowPolicy, +} + +impl<'a> Batch<'a> { + /// Collect `args` and read the lifting's facts off them, `return_dtype` being the function's + /// declared return dtype for the input dtypes it is handed. + /// + /// **Not** for a nullary function: with no inputs there is no validity to propagate and no + /// per-row work to fold, and the all-constant check below would vacuously pass. + pub(super) fn new( + id: ScalarFnId, + args: &'a dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|i| args.get(i)) + .collect::>()?; + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let nullability = plan.sink_dtype.nullability() + | Nullability::from(arg_dtypes.iter().any(DType::is_nullable)); + let result_dtype = plan.sink_dtype.with_nullability(nullability); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + args, + inputs, + arg_dtypes, + validity, + result_dtype, + sink_dtype: plan.sink_dtype, + policy: plan.policy, + }) + } + + /// Run `kernel` over this batch, adding everything the kernel does not do: the null-constant + /// short circuit, the all-constant fold, and the null handling. + /// + /// `kernel` computes the whole column from the arguments it is handed. Those are this batch's + /// arguments untouched, except under the filter strategy, where they are filtered copies, and + /// in the all-constant fold, where they are one row each. What it may assume: + /// + /// - No input is a null constant, and the inputs are not all constant. + /// - Under valid-only execution, every row of every input is valid. + /// - Under dense execution, rows behind nulls hold arbitrary values, and their results are + /// discarded. + /// + /// Either way the kernel can ignore input validity, and its output **must** equal + /// `return_dtype` up to nullability. A kernel that returns nulls of its own keeps them, unioned + /// with the ones the lifting applies, which requires its declared dtype to be nullable. + /// + /// `branch` computes only the rows set in the conjoined mask, over the _unfiltered_ arguments, + /// writing an arbitrary placeholder everywhere else; `Ok(None)` means it cannot for these + /// inputs, which sends the batch to the filter strategy. It is only ever called with a mixed + /// mask, and it **must not** run its row computation (nor any per-row fallible decode) on an + /// unset row, since those rows hold arbitrary values and a fallible kernel would spuriously + /// fail on them. + pub(super) fn execute( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + branch: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: any null-constant input forces an all-null result without evaluating the + // kernel. + if self + .inputs + .iter() + .any(|input| input.as_constant().is_some_and(|scalar| scalar.is_null())) + { + return Ok(self.all_null()); + } + + // All inputs constant, and their conjoined validity proves every row non-null. This sees + // through extension and masked wrappers just like argument decoding does. + if self.args.row_count() > 0 + && self.validity.definitely_no_nulls() + && self + .inputs + .iter() + .all(|input| batch_constant(input).is_some()) + { + return self.broadcast_one_row(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, false, ctx), + RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), + RowPolicy::ValidOnly { + filtered_decode_cost, + } => self.execute_filtered(kernel, branch, filtered_decode_cost, ctx), + } + } + + /// Evaluate a single row of all-constant inputs and broadcast its value. + /// + /// Reconciling the row's dtype before reading the scalar keeps this path on the same + /// kernel/declaration agreement check as the dense and filter paths, rather than letting `cast` + /// paper over a disagreement. + fn broadcast_one_row( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one_row: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.slice(0..1)) + .collect::>()?; + + let args = BorrowedExecutionArgs::new(&one_row, 1); + let result = kernel(self.kernel_args(&args, &one_row), ctx)?.into_result()?; + let scalar = self.with_return_dtype(result, 1)?.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.args.row_count()).into_array()) + } + + /// Run the kernel over every row, including the rows behind nulls, then mask its result. + /// + /// The arguments reach the kernel untouched, so the inputs keep their original encoding, and + /// the conjoined validity is handed to `mask` as an array rather than materialized into a + /// [`Mask`] first. + fn execute_dense( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + retry_deferred_error: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Every row is null, so the kernel has nothing to contribute. + if matches!(self.validity, Validity::AllInvalid) { + return Ok(self.all_null()); + } + + let values = match kernel(self.kernel_args(self.args, &self.inputs), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self + .validity + .clone() + .execute_mask(self.args.row_count(), ctx)?; + + // The same shortcut pair as `resolve_validity`, with different outcomes: every + // row valid means some valid row genuinely failed, and no row valid means every + // failure was behind a null. An empty mask is both all-true and all-false, but + // cannot reach this arm: a zero-row loop accumulates no evidence, so a zero-row + // batch never reports a deferred error. + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + // Filtering unconditionally, rather than consulting `branch_beats_filter`. Not + // because branch-and-skip is unavailable in principle: `ERRORS_ARE_DEFERRED` and + // `SUPPORTS_SKIPPED_ROWS` are independent, and a sink may legally set both. It is + // that `execute_dense` is not handed the `branch` closure at all, so filtering is + // the only strategy reachable from here. This is the cold path, taken only after a + // batch has already reported an error, so the choice has not been worth plumbing + // for. + return self.filter_and_scatter(kernel, &valid, ctx); + } + RowExecution::DeferredError(error) => return Err(error), + }; + + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.with_return_dtype(values, self.args.row_count()) + } + Validity::Array(valid) => { + self.with_return_dtype(values.mask(valid)?, self.args.row_count()) + } + // Handled by the guard above, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } + + /// Materialize the conjoined validity and resolve everything that does not need a mixed-mask + /// strategy, so that the production selector and the forced-strategy test seam cannot drift + /// apart on the shortcuts they share. The deferred-error retry in + /// [`execute_dense`](Self::execute_dense) repeats the same materialize-then-shortcut shape + /// with different outcomes — all-true is an error there, all-false is all-null — so it stays + /// open-coded, with its own note on why the ordering is safe. + fn resolve_validity( + &self, + kernel: &impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self + .validity + .clone() + .execute_mask(self.args.row_count(), ctx)?; + + // Check all-true before all-false: an empty mask is both, and must not be treated as + // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). + if valid.all_true() { + return self + .with_return_dtype( + kernel(self.kernel_args(self.args, &self.inputs), ctx)?.into_result()?, + self.args.row_count(), + ) + .map(ResolvedMask::Decided); + } + + if valid.all_false() { + return Ok(ResolvedMask::Decided(self.all_null())); + } + + Ok(ResolvedMask::Mixed(valid)) + } + + /// Materialize the conjoined validity once, take the all-true and all-false shortcuts, and + /// pick a strategy per batch for a mixed mask. + /// + /// Two strategies can execute a mixed mask, and neither is visible to the kernel: + /// + /// - **Branch-and-skip** ([`execute_branched`](Self::execute_branched)): hand the _unfiltered_ + /// arguments plus the mask to `branch`, which computes only the valid rows, then mask the + /// full-length result exactly as the dense path does. This skips the filter and the scatter + /// entirely, at the price of decoding full-length columns. + /// - **Filter** ([`filter_and_scatter`](Self::filter_and_scatter)): filter every input down to + /// the conjoined-valid rows, run the kernel over those, and scatter its results back into a + /// null-padded output. Always available, never encoding-preserving. + /// + /// Branch-and-skip is preferred whenever [`branch_beats_filter`] says so, and the filter + /// strategy is also the fallback for a kernel with no branch execution. + fn execute_filtered( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + branch: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + filtered_decode_cost: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedMask::Decided(result) => return Ok(result), + ResolvedMask::Mixed(valid) => valid, + }; + + if branch_beats_filter(filtered_decode_cost, &valid) + && let Some(result) = self.execute_branched(branch, &valid, ctx)? + { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Try the branch-and-skip strategy for a mixed mask: the kernel computes only the rows set in + /// `valid` over the unfiltered inputs, and the full-length result is masked exactly as the + /// dense path masks. `Ok(None)` means the kernel has no branch execution for these inputs, and + /// the caller falls back to [`filter_and_scatter`](Self::filter_and_scatter). + fn execute_branched( + &self, + branch: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(values) = branch(self.kernel_args(self.args, &self.inputs), valid, ctx)? else { + return Ok(None); + }; + let values = values.into_result()?; + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + self.with_return_dtype(values.mask(mask)?, valid.len()) + .map(Some) + } + + /// The filter strategy for a mixed mask: filter every input down to the rows set in `valid`, + /// run the kernel over those, and scatter its results back into a null-padded output. + fn filter_and_scatter( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(valid.clone())) + .collect::>()?; + + let args = BorrowedExecutionArgs::new(&filtered, valid.true_count()); + let values = kernel(self.kernel_args(&args, &filtered), ctx)?.into_result()?; + + self.with_return_dtype(self.scatter_valid(values, valid)?, valid.len()) + } + + /// An all-null result of the function's declared return dtype. + fn all_null(&self) -> ArrayRef { + ConstantArray::new( + Scalar::null(self.result_dtype.clone()), + self.args.row_count(), + ) + .into_array() + } + + /// Pair an input view with this batch's planning metadata. + fn kernel_args<'b>( + &'b self, + execution: &'b dyn ExecutionArgs, + arrays: &'b [ArrayRef], + ) -> KernelArgs<'b> { + KernelArgs { + execution, + arrays, + dtypes: &self.arg_dtypes, + sink_dtype: &self.sink_dtype, + } + } + + /// Reconcile the kernel's output dtype with the function's declared return dtype. + /// + /// The kernel may ignore nullability, so a nullability difference is cast away. Any other + /// difference means the declared dtype and the kernel disagree, which is a bug worth naming + /// rather than silently casting away. + fn with_return_dtype(&self, values: ArrayRef, expected_len: usize) -> VortexResult { + reconcile_return(self.id, &self.result_dtype, expected_len, values) + } + + /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set + /// bits, producing an array of length `valid.len()` that is null at every unset position. + fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult { + vortex_ensure_eq!( + values.len(), + valid.true_count(), + "the {} kernel produced {} rows for {} filtered rows", + self.id, + values.len(), + valid.true_count(), + ); + + let AllOr::Some(slices) = valid.slices() else { + // The caller handles the all-true and all-false masks. + vortex_bail!("scatter_valid requires a mixed mask"); + }; + + // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index + // 0, and any in-bounds index would do since they are masked out below (values is non-empty + // here). + let mut indices = vec![0u64; valid.len()]; + let mut rank = 0u64; + for &(start, end) in slices { + for index in &mut indices[start..end] { + *index = rank; + rank += 1; + } + } + let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array(); + + let scattered = values.take(indices)?; + + // A kernel that produced nulls of its own (only `reduce_encoded` may) cannot be wrapped, + // since a `Masked` child must be all valid. Those nulls have to be unioned with the + // lifting's, which is what the general masking pass does. + if scattered.dtype().is_nullable() { + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + return scattered.mask(mask); + } + + // Attaching the mask as validity rather than masking again: the gathered values are + // already all valid, so recording which rows survive is the whole job and a `Masked` + // wrapper says exactly that. Worth 1.13-1.53x here, growing with null density + // (`null_strategy_bytes`, 65536 rows, divan fastest and median of 100 samples, best of two + // runs, Apple M4 Max). The same substitution on the dense path measured no difference, so + // it is deliberately confined to the scatter. + Ok(MaskedArray::try_new( + scattered, + Validity::from_mask(valid.clone(), Nullability::Nullable), + )? + .into_array()) + } +} + +/// Validate the row count and reconcile nullability against a row function's declared dtype. +pub(super) fn reconcile_return( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, +) -> VortexResult { + vortex_ensure_eq!( + values.len(), + expected_len, + "the {id} kernel produced {} rows for {expected_len} input rows", + values.len(), + ); + vortex_ensure!( + values.dtype().eq_ignore_nullability(result_dtype), + "the {id} kernel produced {} but the function declares {result_dtype}", + values.dtype(), + ); + + if values.dtype() == result_dtype { + Ok(values) + } else { + values.cast(result_dtype.clone()) + } +} + +/// The minimum surviving-row fraction (`true_count / len` of the conjoined mask) at which +/// branch-and-skip is still chosen for one filtered decode unit. +/// +/// From the branch-and-skip measurements (65536 rows, divan fastest of 100 samples, two runs on a +/// shared 4-vCPU VM). A kernel with a _bulk_ decode never lost under branch: `byte_length` over +/// a byte-string element ran 1.8-5.9x faster than filter at every null density from 1% to 90%, so +/// such kernels skip this check entirely. A kernel with a _per-row_ decode (geo `contains`, which +/// arrow-exports and parses one geometry per row) pays that decode over the full column under +/// branch but only over the survivors under filter, so filter wins once validity is sparse: +/// +/// - polygons CONTAINS constant point: branch won 1.07-1.18x at 1-50% nulls; filter won 1.38x at +/// 90% nulls (10% of rows surviving). +/// - polygons CONTAINS points, independent nulls on both: branch won up to ~10% null density +/// (~81% surviving); filter won 1.2x at ~56% surviving, 1.9x at ~25%, 11.3x at ~1%. +/// +/// A single nullable operand still favored branch at 50% surviving, while two independent nullable +/// operands favored filtering at 81% surviving. Keep those cases distinct instead of collapsing +/// every per-row decode into one boolean. There is not yet enough evidence to distinguish two from +/// three or more decode units, so they share the conservative multi-decode threshold. +pub(super) const ONE_DECODE_BRANCH_MIN_SURVIVING_FRACTION: f64 = 0.50; +pub(super) const MULTI_DECODE_BRANCH_MIN_SURVIVING_FRACTION: f64 = 0.85; + +/// Whether the branch-and-skip strategy should be preferred over filtering for the mixed mask +/// `valid`. A zero cost always branches; otherwise the survivor threshold grows when filtering +/// avoids more than one unit of per-row decode work. +pub(super) fn branch_beats_filter(filtered_decode_cost: usize, valid: &Mask) -> bool { + if filtered_decode_cost == 0 { + return true; + } + + let minimum = if filtered_decode_cost == 1 { + ONE_DECODE_BRANCH_MIN_SURVIVING_FRACTION + } else { + MULTI_DECODE_BRANCH_MIN_SURVIVING_FRACTION + }; + valid.true_count() as f64 >= valid.len() as f64 * minimum +} + +/// Which null strategy a forced execution takes for a mixed validity mask. +/// +/// A test and benchmark seam: pinning a strategy is how the two are compared and how their +/// agreement is asserted. Production execution selects per batch inside the lifting and never +/// names one. See [`execute_row_fn_with_strategy`](super::execute_row_fn_with_strategy). +#[cfg(any(test, feature = "_test-harness"))] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum NullStrategy { + /// Filter the inputs down to the conjoined-valid rows, run the kernel, and scatter back. + Filter, + + /// Decode the unfiltered inputs null-tolerantly, compute only the conjoined-valid rows, and + /// mask the full-length result. + BranchAndSkip, +} + +#[cfg(any(test, feature = "_test-harness"))] +impl Batch<'_> { + /// Execute this batch with a forced null strategy, bypassing the per-batch selection. + /// + /// A test and benchmark seam only. It mirrors [`execute_filtered`](Self::execute_filtered) + /// (conjoined validity, the all-true and all-false shortcuts, output dtype reconciliation) but + /// takes the strategy from the caller instead of the selection rule, and it skips the + /// null-constant and all-constant folds, so do not pass such inputs. `Ok(None)` means + /// [`NullStrategy::BranchAndSkip`] was forced on a kernel with no branch execution, which the + /// caller reports rather than silently falling back. + pub(super) fn execute_with_strategy( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + branch: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + strategy: NullStrategy, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedMask::Decided(result) => return Ok(Some(result)), + ResolvedMask::Mixed(valid) => valid, + }; + + match strategy { + NullStrategy::Filter => self.filter_and_scatter(kernel, &valid, ctx).map(Some), + NullStrategy::BranchAndSkip => self.execute_branched(branch, &valid, ctx), + } + } +} diff --git a/vortex-array/src/scalar_fn/row/mod.rs b/vortex-array/src/scalar_fn/row/mod.rs new file mode 100644 index 00000000000..f62d8bb8d49 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Defining scalar functions one row at a time. +//! +//! This is the derived way to write a scalar function, and the right default for a kernel that has +//! to read every row anyway. See [choosing a trait](crate::scalar_fn#choosing-a-trait) for when to +//! drop to [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) instead. +//! +//! [`RowFn`] names its arguments and provides a [`dispatch`](RowFn::dispatch) that picks the +//! concrete element and sink types for a batch. Everything else (dtype checks, output dtype, null +//! handling, constants, and validity) is derived from that dispatch. +//! +//! When the element types are fixed, `dispatch` is a single visit at those types. When one function +//! ID has to cover several (`l2_norm` accepts `f16`, `f32` and `f64` columns), `dispatch` matches on +//! the input dtypes and visits at the chosen width. Kernel fallibility is declared separately +//! because callers need it before dispatch. +//! +//! [`RowFn`] does not say how a row is _stored_, which is the element's job: `vortex-tensor` adds a +//! `TensorRow` [`InputElement`] and writes ordinary kernels over it. +//! +//! Output always goes through [`RowVisitor::visit_prepared_into`]. [`ElementSink`] covers one owned +//! [`OutputElement`] per row; custom [`OutputSink`] implementations cover runtime-shaped rows. The +//! prepare closure sees every batch-constant input and returns shared state for the row loop. Pass +//! `|_| ()` when there is nothing to prepare. +//! +//! A kernel that can safely write a provisional value uses [`DeferredError`] instead of returning +//! a per-row result. The executor vector-reduces those bits and hands one batch-wide error to the +//! sink. With nullable fixed-width inputs it runs densely and retries only valid rows on the cold +//! error path. +//! +//! Null handling is derived and executed by the [lifting](lift), never by the row closure, which +//! only ever computes rows valid in every argument. A batch with a mixed validity mask executes by +//! one of two strategies, selected per batch: _branch-and-skip_ (decode the unfiltered columns +//! null-tolerantly via [`InputElement::decode_null_tolerant`], compute only the valid rows a word +//! of the mask at a time, mask the result) whenever it can, and _filter_ (shrink every input to +//! the surviving rows, compute, scatter back) when an argument has no null-tolerant decode for its +//! array or when a per-row decode makes filtering cheaper at sparse validity. Authors do nothing; +//! an element whose decode does expensive per-row work reports that work through +//! [`InputElement::FILTERED_DECODE_COST`]. The costs of all arguments are added together so the +//! batch selector can distinguish one expensive decode from several. A sink opts into +//! branch-and-skip with [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. + +mod element; +pub use element::ElementTuple; +pub use element::InputElement; +pub use element::OutputElement; +#[cfg(any(test, feature = "_test-harness"))] +pub use element::assert_element_conforms; + +mod result; +pub use result::DeferredError; +pub use result::SinkResult; + +mod sink; +pub use sink::ElementSink; +pub use sink::OutputSink; + +mod execute; + +mod lift; +#[cfg(any(test, feature = "_test-harness"))] +pub use lift::NullStrategy; + +mod row_fn; +pub use row_fn::RowFn; +pub use row_fn::RowVisitor; + +mod vtable; +#[cfg(any(test, feature = "_test-harness"))] +pub use vtable::execute_row_fn_with_strategy; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/row/result.rs b/vortex-array/src/scalar_fn/row/result.rs new file mode 100644 index 00000000000..86b90db41f5 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/result.rs @@ -0,0 +1,161 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What a sink-writing row closure may return. + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +mod private { + pub trait Sealed {} +} + +/// A value-dependent failure bit reduced across the whole row loop and handed to the output sink. +/// +/// Unlike [`VortexResult`], this never exits the loop. It is for kernels such as checked addition +/// that can safely write a provisional value for every row and report any failure once at the end. +/// +/// **The reduction is one byte wide on purpose.** It is OR-reduced once per row alongside the +/// kernel's own arithmetic, so a wider accumulator caps how many rows a vector of the reduction +/// covers, whatever the element width. Carrying the bit in an `i64` instead cost the primitive +/// `Mul` kernel 3.1x at `i8`, 1.9x at `i16` and 1.2x at `i32`, and nothing at `i64` where the two +/// widths already agree (`binary_ops`, 65536 rows, divan fastest of 100 samples, best of two runs, +/// Apple M4 Max). +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct DeferredError(bool); + +impl DeferredError { + /// Record whether this row encountered an error. + pub const fn new(failed: bool) -> Self { + Self(failed) + } + + /// Whether any row accumulated into this value failed. + pub const fn occurred(self) -> bool { + self.0 + } +} + +impl BitOrAssign for DeferredError { + fn bitor_assign(&mut self, rhs: Self) { + self.0 |= rhs.0; + } +} + +/// What a row computation that _writes_ into an [`OutputSink`](crate::scalar_fn::OutputSink) may +/// produce: nothing, an early [`VortexResult`] error, or non-branching failure evidence. +/// +/// The value is already in the sink by the time the closure returns, so the only thing left to +/// report is failure. +/// +/// [`Accumulated`](Self::Accumulated) is the word the executor OR-reduces in a **local**, which is +/// what keeps the reduction in a register and the row loop vectorizable. It exists so that evidence +/// can be wider than one bit when narrowing it per row would cost more than carrying it: unsigned +/// multiplication hands back the discarded high half of its product, because comparing that half +/// against zero per row is what LLVM folds into `llvm.umul.with.overflow`, which has no vector form. +/// **The word must be no wider than the element**, or the reduction, rather than the arithmetic, +/// bounds how many rows a vector covers. +/// +/// The sink never sees this. It is handed a plain [`DeferredError`] once, after the loop. +/// +/// This trait is framework-only. Row functions choose one of the supplied return forms; custom +/// output representation belongs in [`OutputSink`](crate::scalar_fn::OutputSink). +pub trait SinkResult: 'static + private::Sealed { + /// The word this result reduces into, kept in a loop-local by the executor. + type Accumulated: 'static + Copy + Default; + + /// Whether this return type can carry an error. + const FALLIBLE: bool; + + /// Whether this result carries non-branching failure evidence for the sink. + const DEFERRED: bool; + + /// Merge this row's outcome into the batch-wide reduction. + fn accumulate(self, accumulated: &mut Self::Accumulated) -> VortexResult<()>; + + /// Whether the finished reduction means some row failed. + fn occurred(accumulated: Self::Accumulated) -> bool; +} + +impl private::Sealed for () {} + +impl SinkResult for () { + type Accumulated = (); + + const FALLIBLE: bool = false; + const DEFERRED: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + Ok(()) + } + + fn occurred(_accumulated: ()) -> bool { + false + } +} + +impl private::Sealed for VortexResult<()> {} + +impl SinkResult for VortexResult<()> { + type Accumulated = (); + + const FALLIBLE: bool = true; + const DEFERRED: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + self + } + + fn occurred(_accumulated: ()) -> bool { + false + } +} + +/// The evidence widths a row closure may reduce. `bool` is the ordinary answer; the unsigned +/// integers exist for a kernel whose per-row comparison would cost it its vectorization. +macro_rules! impl_sink_result_word { + ($($word:ty),+ $(,)?) => { + $( + impl private::Sealed for $word {} + + impl SinkResult for $word { + type Accumulated = $word; + + const FALLIBLE: bool = false; + const DEFERRED: bool = true; + + fn accumulate(self, accumulated: &mut $word) -> VortexResult<()> { + *accumulated |= self; + Ok(()) + } + + fn occurred(accumulated: $word) -> bool { + accumulated != <$word>::default() + } + } + )+ + }; +} + +impl_sink_result_word!(bool, u8, u16, u32, u64); + +#[cfg(test)] +mod tests { + use super::DeferredError; + + #[test] + fn one_failing_row_is_enough() { + let mut error = DeferredError::default(); + assert!(!error.occurred()); + + error |= DeferredError::new(false); + assert!(!error.occurred()); + + error |= DeferredError::new(true); + assert!(error.occurred()); + + error |= DeferredError::new(false); + assert!(error.occurred()); + } +} diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs new file mode 100644 index 00000000000..f2a63636b1d --- /dev/null +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -0,0 +1,138 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions computed one row at a time. + +use std::fmt::Debug; +use std::fmt::Display; +use std::hash::Hash; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::SinkResult; + +/// A scalar function computed one row at a time. +/// +/// An implementor declares its argument names, then [`dispatch`](Self::dispatch) picks the concrete +/// element and sink types for a batch. The planning visit reads dense safety, decode fallibility, +/// and decode cost from that concrete choice; no representative element types are needed. +/// +/// A function whose kernel is columnar rather than row-at-a-time (negating a whole bit buffer, a +/// zero-copy unwrap) is not a `RowFn`, and implements +/// [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) directly. +pub trait RowFn: 'static + Sized + Clone + Send + Sync { + /// Options for this function, if any. Use [`EmptyOptions`](crate::scalar_fn::EmptyOptions) + /// for none. + type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash; + + /// The arguments in display order. Its length is the function's exact arity. + const ARG_NAMES: &'static [&'static str]; + + /// Whether any legal dispatch can fail while decoding or computing a row. + /// + /// The framework verifies that every fallible dispatched element or result implies this value. + /// A conservative `true` is allowed when only some dtype choices are fallible. + const FALLIBLE: bool = false; + + /// Returns the ID of the scalar function. + fn id(&self) -> ScalarFnId; + + /// Serialize this function's options, or return `None` when the function is not serializable. + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + _ = options; + Ok(None) + } + + /// Restore options written by [`serialize`](Self::serialize). + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + vortex_bail!("Expression {} is not deserializable", self.id()) + } + + /// Choose element types for these input dtypes and visit the framework with them. + /// + /// This is where a per-batch width match lives (`match_each_float_ptype!` and friends panic + /// outside their width class, so check the class first), and where cross-argument dtype + /// constraints belong, since per-argument validation runs inside the visit. Plan time and run + /// time both come through here, so the choice **must** be a pure function of `options` and + /// `args`. + fn dispatch( + &self, + options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult; + + /// An encoding-aware rewrite, tried on the input arrays before the row loop. + /// + /// `Some` skips the row loop entirely, which makes this the escape hatch for a function that is + /// row-shaped in general but has a bulk answer for some encodings: reading stored values back out + /// of a wrapper encoding, or handing back a child array whole. The result may be lazy and + /// nullable, but its nulls **must** be a subset of the rows the lifting will mask, and it + /// **must** have one row per row of `args`, which on the filter strategy is the _filtered_ count + /// rather than the original one. Size the result from `args`, which are filtered to match, and + /// never from a length captured elsewhere. + /// + /// Whether the arrays still carry their original encoding depends on the execution path. + /// Dense execution always passes them through untouched. Valid-only execution does too when + /// no row is null; for a mixed mask, branch-and-skip also passes them through untouched (full + /// length, with the result masked afterwards), while filtering hands over filtered copies, + /// which are canonical and so match no encoding fast path. + /// + /// A non-nullable operand therefore reaches an encoding fast path under either. Note also that + /// filtering a constant yields a constant, so a fast path keyed on + /// [`as_constant`](ArrayRef::as_constant) still fires even for a filtered batch. + fn reduce_encoded( + &self, + options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + _ = (options, args, ctx); + Ok(None) + } +} + +/// One use of a [`RowFn`] at concrete element types. +/// +/// The framework hands a visitor to [`RowFn::dispatch`], which calls one of the visit methods with +/// the element types it chose: at plan time the visit validates dtypes, at run time it executes the row +/// loop. Only the framework implements this trait, and a function only ever _calls_ a visit. +/// +/// The function names one output sink and one preparation step. Passing `|_| ()` is the no-prepare +/// case. +pub trait RowVisitor: private::Sealed { + /// What this visit produces. + type Out; + + /// Visit at argument tuple `A`, preparing shared state once and writing every output row into + /// sink `S`. + /// + /// `prepare` receives [`A::ConstElems`](ElementTuple::ConstElems): the element value of every + /// argument whose operand is constant for the batch, and `None` for each one that varies by + /// row. Whatever it returns is handed to every `apply` call by shared reference. + /// + /// `A` **must** have the arity declared by [`RowFn::ARG_NAMES`]. A fallible element or result + /// also requires [`RowFn::FALLIBLE`] to be `true`; the reverse is not required. A deferred result + /// must be paired with a sink whose [`OutputSink::ERRORS_ARE_DEFERRED`] is `true`. + fn visit_prepared_into( + self, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, + ) -> VortexResult; +} + +pub(super) mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/row/sink.rs b/vortex-array/src/scalar_fn/row/sink.rs new file mode 100644 index 00000000000..50d07d1b7ff --- /dev/null +++ b/vortex-array/src/scalar_fn/row/sink.rs @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The column builders a row function can write its output into. + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::DeferredError; +use crate::scalar_fn::OutputElement; + +/// A column allocated once per batch that a row closure writes into, one row at a time. +/// +/// Every [`RowFn`](crate::scalar_fn::RowFn) writes through one. [`ElementSink`] covers an ordinary +/// owned value per row. A custom sink covers output whose width is runtime data or whose rows append +/// into one batch-wide builder. +/// +/// Two properties of the contract are worth stating, since both are load-bearing: +/// +/// - **The row loop, not the closure, holds the sink.** [`row`](Self::row) is called by the framework +/// and its result passed in, so a writing closure stays [`Fn`] and captures nothing mutable. +/// Relaxing the row closure to `FnMut` instead was measured at 8 to 11%, because a captured `&mut` +/// inhibits vectorization of the loop. +/// - **[`sink_dtype`](Self::sink_dtype) sees the input dtypes**, unlike +/// [`OutputElement::element_dtype`](crate::scalar_fn::OutputElement::element_dtype), which takes +/// none. That is the whole reason a runtime-shaped output fits here: the width comes out of the +/// arguments. +/// +/// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; +/// branch-and-skip may omit null rows when [`SUPPORTS_SKIPPED_ROWS`](Self::SUPPORTS_SKIPPED_ROWS) +/// is `true`. +pub trait OutputSink: 'static + Sized { + /// Whether this sink accepts [`DeferredError`] from its row closure instead of requiring a + /// per-row [`VortexResult`]. + /// + /// The executor OR-reduces the row error words and passes the result to + /// [`finish`](Self::finish). When the arguments are safe to read behind nulls, this lets the + /// lifting optimistically run a dense loop. If `finish` reports the deferred error for a + /// nullable batch, the lifting retries over only the valid rows: success means the error came + /// exclusively from null rows, while another deferred error is real. + /// + /// A supporting sink must return an error from `finish` when its `error` argument occurred. + const ERRORS_ARE_DEFERRED: bool = false; + + /// Whether this sink can finish a full-length output when some rows were never visited. + /// + /// A supporting sink must leave a legal arbitrary value at every skipped row. The lifting masks + /// those rows before the result escapes, so that value is never observable. + const SUPPORTS_SKIPPED_ROWS: bool = false; + + /// A loop-local view of all output rows. + /// + /// Borrowed once before execution so the sink's buffer descriptor and shape become loop + /// invariants rather than being re-read through `&mut Self` for every row. + type Rows<'a> + where + Self: 'a; + + /// The place a row closure writes one row through, borrowed from the sink. + type Row<'a> + where + Self: 'a; + + /// The dtype of the column this sink builds, given the function's input dtypes. + /// + /// Must be non-nullable: nullability is derived from the inputs by the lifting, which + /// widens the result and masks the null rows itself. + fn sink_dtype(args: &[DType]) -> VortexResult; + + /// Allocate a sink for `rows` rows of `dtype`, which is this sink's own + /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; + + /// Borrow all output rows for the hot loop. + fn rows(&mut self) -> Self::Rows<'_>; + + /// Whether every index in `0..row_count` is addressable through [`row`](Self::row). + /// + /// Called once before the hot loop. Besides validating the sink contract, this gives the + /// optimizer the output bounds it needs to remove the bounds check hidden in each row accessor. + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool; + + /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; + + /// Finish into the built column, whose dtype **must** be this sink's + /// [`sink_dtype`](Self::sink_dtype). Called once per batch with the OR of every row's deferred + /// error bit. + fn finish(self, error: DeferredError) -> VortexResult; +} + +/// The standard output sink for one owned [`OutputElement`] per row. +pub struct ElementSink { + values: Vec, +} + +impl OutputSink for ElementSink { + const SUPPORTS_SKIPPED_ROWS: bool = true; + + type Rows<'a> = &'a mut [T]; + type Row<'a> = &'a mut T; + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(T::element_dtype()) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + // `vec![placeholder; rows]` rather than `resize_with(rows, placeholder)`: the former hands + // a zeroable placeholder (every primitive, `false`) straight to `alloc_zeroed`, while the + // latter always writes one element at a time. Only branch-and-skip ever reads a + // placeholder back, so on the dense and filter paths that write is pure waste. + Ok(Self { + values: vec![T::placeholder(); rows], + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.values + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + fn finish(self, _error: DeferredError) -> VortexResult { + Ok(T::build(self.values)) + } +} diff --git a/vortex-array/src/scalar_fn/row/tests/conformance.rs b/vortex-array/src/scalar_fn/row/tests/conformance.rs new file mode 100644 index 00000000000..8df438a734c --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/conformance.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Conformance tests for every [`InputElement`](crate::scalar_fn::InputElement) in this crate. + +use std::sync::Arc; + +use vortex_buffer::BitBuffer; +use vortex_buffer::ByteBuffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::assert_element_conforms; +use crate::scalar_fn::row::tests::TestBytes; +use crate::validity::Validity; + +/// A `Utf8` column whose single null row carries a view naming a buffer that does not exist, at +/// an offset far past the end of the data. Reading its _bytes_ densely panics; reading its +/// _length_ does not, which is exactly the distinction `DENSE_SAFE` encodes. +fn hostile_views() -> VortexResult { + let views = buffer![ + BinaryView::make_view(b"a longer string here", 0, 0), + BinaryView::new_ref(64, *b"junk", 9, 4096), + ]; + Ok(VarBinViewArray::try_new( + views, + Arc::from([ByteBuffer::copy_from(b"a longer string here")]), + DType::Utf8(Nullability::Nullable), + Validity::from_iter([true, false]), + )? + .into_array()) +} + +#[test] +fn primitive_element_conforms() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + // The extremes sit at the rows that are then marked null. + let array = PrimitiveArray::new( + buffer![i32::MAX, 1, i32::MIN, 2], + Validity::from_iter([false, true, false, true]), + ) + .into_array(); + + assert_element_conforms::(array, &DType::Utf8(Nullability::NonNullable), &mut ctx) +} + +#[test] +fn bool_element_conforms() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = BoolArray::new( + BitBuffer::from(vec![true, true, false, true]), + Validity::from_iter([false, true, true, false]), + ) + .into_array(); + + assert_element_conforms::(array, &DType::Utf8(Nullability::NonNullable), &mut ctx) +} + +#[test] +fn test_bytes_element_conforms() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + assert_element_conforms::( + hostile_views()?, + &DType::Bool(Nullability::NonNullable), + &mut ctx, + ) +} diff --git a/vortex-array/src/scalar_fn/row/tests/constant_operands.rs b/vortex-array/src/scalar_fn/row/tests/constant_operands.rs new file mode 100644 index 00000000000..5e86b925587 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/constant_operands.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests that constant operands are decoded once and broadcast across the batch. + +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use vortex_buffer::Buffer; + +use super::*; + +/// Total rows handed to [`CountedI64::decode`] across one execution. Sound as a global because +/// each test binary runs one test per process. +static DECODED_ROWS: AtomicUsize = AtomicUsize::new(0); + +/// Stands in for an element whose decode is expensive per row, recording how wide a column each +/// decode was actually given. +struct CountedI64; + +impl InputElement for CountedI64 { + type Column = Buffer; + type Varying<'a> = ::Varying<'a>; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + DECODED_ROWS.fetch_add(array.len(), Ordering::Relaxed); + ::decode(array, ctx) + } + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + ::varying(column) + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + ::varying_len(column) + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_varying(column, index) + } +} + +#[derive(Clone)] +struct AddCounted; + +impl RowFn for AddCounted { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.add_counted"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(CountedI64, CountedI64), ElementSink, _, _>( + |_| (), + |&(), (lhs, rhs), output| *output = lhs + rhs, + ) + } +} + +/// An element whose decode drops the last row, standing in for a buggy element implementation. +struct ShortDecodeI64; + +impl InputElement for ShortDecodeI64 { + type Column = Buffer; + type Varying<'a> = ::Varying<'a>; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let column = ::decode(array, ctx)?; + Ok(column.slice(0..column.len().saturating_sub(1))) + } + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + ::varying(column) + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + ::varying_len(column) + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_varying(column, index) + } +} + +/// Pairs a short-decoding argument with an ordinary one, so a batch-constant second operand takes +/// the mixed constant-and-varying read path rather than the all-varying one. +#[derive(Clone)] +struct AddShort; + +impl RowFn for AddShort { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.add_short"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(ShortDecodeI64, i64), ElementSink, _, _>( + |_| (), + |&(), (lhs, rhs), output| *output = lhs + rhs, + ) + } +} + +/// A constant operand makes the whole tuple decline the all-varying read path, so the row loop +/// indexes each [`ArgColumn`](crate::scalar_fn::ArgColumn) directly. The decoded length still has to +/// be checked there, or a short column reaches an out-of-bounds row read. +#[test] +fn a_short_decode_beside_a_constant_operand_is_rejected() { + let mut ctx = array_session().create_execution_ctx(); + let column = PrimitiveArray::from_iter(0..64i64).into_array(); + let constant = ConstantArray::new(Scalar::from(10i64), 64).into_array(); + + let error = apply(AddShort, [column, constant], &mut ctx).unwrap_err(); + + assert!( + error + .to_string() + .contains("does not address exactly 64 rows"), + "{error}" + ); +} + +#[test] +fn a_constant_operand_is_decoded_once() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let column = PrimitiveArray::from_iter(0..64i64).into_array(); + let constant = ConstantArray::new(Scalar::from(10i64), 64).into_array(); + + let result = apply(AddCounted, [column, constant], &mut ctx)?; + + // 64 rows for the real column, plus exactly one for the constant. + assert_eq!(DECODED_ROWS.load(Ordering::Relaxed), 65); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter((0..64i64).map(|value| value + 10)), + &mut ctx + ); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/tests/decode_fallibility.rs b/vortex-array/src/scalar_fn/row/tests/decode_fallibility.rs new file mode 100644 index 00000000000..e3ade44f680 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/decode_fallibility.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for fallible input decoding and its effect on execution strategy. + +use super::*; + +/// Stands in for an element that _parses_ its bytes, like a WKB geometry: malformed bytes in a +/// valid row are a domain error, so decoding can fail on otherwise legal input. +struct ParsedBytes; + +impl InputElement for ParsedBytes { + type Column = VarBinViewArray; + type Varying<'a> = &'a VarBinViewArray; + type Elem<'a> = usize; + + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = true; + + fn validate(_dtype: &DType) -> VortexResult<()> { + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + array.execute::(ctx) + } + + fn get(column: &Self::Column, index: usize) -> usize { + column.views()[index].len() as usize + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> usize + where + Self: 'a, + { + Self::get(column, index) + } +} + +/// Its row computation is total; only the decode can fail. +#[derive(Clone)] +struct TotalKernelOverParsedInput; + +impl RowFn for TotalKernelOverParsedInput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.total_over_parsed"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(ParsedBytes,), ElementSink, _, _>( + |_| (), + |&(), (len,), output| *output = len as u64, + ) + } +} + +/// The row closure is infallible, so reading only the kernel declaration would report +/// `false` and let dict pushdown speculatively evaluate the parse over unreferenced values. +#[test] +fn a_fallible_decode_makes_the_function_fallible() { + assert!(ScalarFnVTable::is_fallible( + &TotalKernelOverParsedInput, + &EmptyOptions + )); +} + +/// And it must not run densely: rows behind nulls would be parsed too. +#[test] +fn a_fallible_decode_forces_filtering() { + assert_eq!( + policy( + &TotalKernelOverParsedInput, + &[DType::Binary(Nullability::Nullable)] + ), + RowPolicy::ValidOnly { + filtered_decode_cost: 0 + } + ); +} diff --git a/vortex-array/src/scalar_fn/row/tests/dispatched.rs b/vortex-array/src/scalar_fn/row/tests/dispatched.rs new file mode 100644 index 00000000000..ca377bea4b1 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/dispatched.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for row functions that choose their element types per batch. + +use vortex_error::vortex_ensure; + +use super::*; +use crate::match_each_integer_ptype; + +#[derive(Clone)] +struct Max; + +impl RowFn for Max { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.int_max"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let DType::Primitive(ptype, _) = args[0] else { + vortex_bail!("int_max requires primitive inputs, got {}", args[0]); + }; + vortex_ensure!( + ptype.is_int(), + "int_max requires integer inputs, got {ptype}" + ); + + match_each_integer_ptype!(ptype, |T| { + visitor.visit_prepared_into::<(T, T), ElementSink, _, _>( + |_| (), + |&(), (a, b), output| *output = a.max(b), + ) + }) + } +} + +#[rstest] +#[case::i16(buffer![1i16, 9, 3].into_array(), buffer![4i16, 2, 3].into_array(), buffer![4i16, 9, 3].into_array())] +#[case::i64(buffer![1i64, 9, 3].into_array(), buffer![4i64, 2, 3].into_array(), buffer![4i64, 9, 3].into_array())] +#[case::u8(buffer![1u8, 9, 3].into_array(), buffer![4u8, 2, 3].into_array(), buffer![4u8, 9, 3].into_array())] +fn dispatches_at_each_integer_width( + #[case] lhs: ArrayRef, + #[case] rhs: ArrayRef, + #[case] expected: ArrayRef, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let result = apply(Max, [lhs, rhs], &mut ctx)?; + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} + +#[test] +fn rejects_a_float_width() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = buffer![1.0f64].into_array(); + let rhs = buffer![2.0f64].into_array(); + + let error = apply(Max, [lhs, rhs], &mut ctx) + .expect_err("a float width must be rejected at construction"); + + assert!( + error.to_string().contains("integer inputs"), + "unexpected error: {error}" + ); +} diff --git a/vortex-array/src/scalar_fn/row/tests/lifting.rs b/vortex-array/src/scalar_fn/row/tests/lifting.rs new file mode 100644 index 00000000000..f7d2fba4d04 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/lifting.rs @@ -0,0 +1,230 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for null propagation, constant folding, nullability widening, and options serde. + +use super::*; +use crate::dtype::Nullability; +use crate::dtype::PType; + +/// An `i32` element that is [dense-safe] iff `DENSE`, and otherwise the plain `i32` element in +/// every respect. Dense-safety is what decides the null-handling path, so a pair of these is +/// how one kernel gets run under both. +/// +/// [dense-safe]: InputElement::DENSE_SAFE +struct MaybeDenseI32; + +impl InputElement for MaybeDenseI32 { + type Column = ::Column; + type Varying<'a> = ::Varying<'a>; + type Elem<'a> = i32; + + const DENSE_SAFE: bool = DENSE; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + ::decode(array, ctx) + } + + fn get(column: &Self::Column, index: usize) -> i32 { + ::get(column, index) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + ::varying(column) + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + ::varying_len(column) + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> i32 + where + Self: 'a, + { + ::get_varying(column, index) + } +} + +/// Wrapping addition over two [`MaybeDenseI32`] columns. +#[derive(Clone)] +struct Add; + +impl RowFn for Add { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + if DENSE { + static ID: CachedId = CachedId::new("vortex.test.add.dense"); + *ID + } else { + static ID: CachedId = CachedId::new("vortex.test.add.filter"); + *ID + } + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::< + (MaybeDenseI32, MaybeDenseI32), + ElementSink, + _, + _, + >( + |_| (), + |&(), (lhs, rhs), output| *output = lhs.wrapping_add(rhs), + ) + } +} + +/// Adds `lhs` to `rhs` under both null-handling paths and asserts each result equals +/// `expected`, which is what every case below does. +/// +/// Forcing a _strategy_ within the filter contract is a separate axis, covered in +/// [`null_strategies`](super::null_strategies). +fn assert_add(lhs: ArrayRef, rhs: ArrayRef, expected: ArrayRef) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let dense = apply(Add::, [lhs.clone(), rhs.clone()], &mut ctx)?; + let filtered = apply(Add::, [lhs, rhs], &mut ctx)?; + + let args = [ + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Primitive(PType::I32, Nullability::NonNullable), + ]; + assert_eq!(policy(&Add::, &args), RowPolicy::Dense); + assert_eq!( + policy(&Add::, &args), + RowPolicy::ValidOnly { + filtered_decode_cost: 0 + } + ); + assert_arrays_eq!(dense, expected, &mut ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); + Ok(()) +} + +#[test] +fn no_nulls() -> VortexResult<()> { + assert_add( + PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), + PrimitiveArray::from_iter([10i32, 20, 30]).into_array(), + PrimitiveArray::from_iter([11i32, 22, 33]).into_array(), + ) +} + +#[test] +fn nulls_propagate() -> VortexResult<()> { + assert_add( + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(), + PrimitiveArray::from_option_iter([Some(10i32), Some(20), None, None]).into_array(), + PrimitiveArray::from_option_iter([Some(11i32), None, None, None]).into_array(), + ) +} + +/// Strictness: a null constant makes the whole output null without the kernel running at all. +#[test] +fn null_constant_short_circuits() -> VortexResult<()> { + let null = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); + + assert_add( + PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), + ConstantArray::new(null, 3).into_array(), + PrimitiveArray::from_option_iter([Option::::None, None, None]).into_array(), + ) +} + +/// All-constant inputs evaluate one row and broadcast it. +#[test] +fn all_constants_broadcast() -> VortexResult<()> { + assert_add( + ConstantArray::new(Scalar::from(2i32), 4).into_array(), + ConstantArray::new(Scalar::from(40i32), 4).into_array(), + PrimitiveArray::from_iter([42i32, 42, 42, 42]).into_array(), + ) +} + +#[test] +fn mixed_constant_and_column() -> VortexResult<()> { + assert_add( + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), + ConstantArray::new(Scalar::from(10i32), 3).into_array(), + PrimitiveArray::from_option_iter([Some(11i32), None, Some(13)]).into_array(), + ) +} + +/// An empty batch is neither all-valid nor all-null, and a zero-length non-nullable execution +/// keeps its non-nullable dtype. +#[test] +fn empty_input_keeps_dtype() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let empty = || PrimitiveArray::from_iter(Vec::::new()).into_array(); + + let result = apply(Add::, [empty(), empty()], &mut ctx)?; + + assert_eq!(result.len(), 0); + assert!(!result.dtype().is_nullable()); + Ok(()) +} + +/// The output element dtype is non-nullable, and the lifting widens it iff an input is +/// nullable, which is what makes strictness's dtype contract hold by construction. +#[test] +fn return_dtype_unions_nullability() -> VortexResult<()> { + let non_nullable = DType::Primitive(PType::I32, Nullability::NonNullable); + let nullable = non_nullable.as_nullable(); + + assert_eq!( + ScalarFnVTable::return_dtype( + &Add::, + &EmptyOptions, + &[non_nullable.clone(), non_nullable.clone()] + )?, + non_nullable + ); + assert_eq!( + ScalarFnVTable::return_dtype( + &Add::, + &EmptyOptions, + &[non_nullable, nullable.clone()] + )?, + nullable + ); + Ok(()) +} + +#[test] +fn a_row_fn_is_strict() { + assert!(ScalarFnVTable::is_strict(&Add::, &EmptyOptions)); +} + +/// Output sinks build an all-valid column, so the output validity is exactly the child +/// conjunction and the planner never has to execute the function to learn which rows are null. +#[test] +fn validity_is_the_child_conjunction() -> VortexResult<()> { + let expr = Add::.new_expr(EmptyOptions, [root(), root()]); + + assert!(ScalarFnVTable::validity(&Add::, &EmptyOptions, &expr)?.is_some()); + Ok(()) +} + +/// A row function is not serializable until the function opts into a wire representation. +#[test] +fn options_are_not_serializable_by_default() -> VortexResult<()> { + assert_eq!( + ScalarFnVTable::serialize(&Add::, &EmptyOptions)?, + None + ); + assert!(ScalarFnVTable::deserialize(&Add::, &[], &array_session()).is_err()); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/tests/mod.rs b/vortex-array/src/scalar_fn/row/tests/mod.rs new file mode 100644 index 00000000000..c8f2ab77aee --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/mod.rs @@ -0,0 +1,509 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! End-to-end tests for row function execution. + +use rstest::rstest; +use vortex_buffer::ByteBuffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_session::registry::CachedId; + +use super::lift::RowPolicy; +use super::vtable::row_policy; +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::scalar_fn::ScalarFnFactoryExt; +use crate::arrays::varbinview::BinaryView; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::root; +use crate::scalar::Scalar; +use crate::scalar_fn::*; + +mod conformance; +mod constant_operands; +mod decode_fallibility; +mod dispatched; +mod lifting; +mod null_strategies; +mod nullable_outputs; +mod prepared; +mod sink; + +/// Builds `scalar_fn` over `args` and executes it end to end, which is what every test below does. +fn apply>( + scalar_fn: F, + args: impl IntoIterator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = args.into_iter().collect::>(); + let rows = args.first().map_or(0, |arg| arg.len()); + + Ok(scalar_fn + .try_new_array(rows, EmptyOptions, args)? + .execute::(ctx)? + .into_array()) +} + +/// A binary row function over fixed primitive types: `hypot(x, y)`. +#[derive(Clone)] +struct Hypot; + +impl RowFn for Hypot { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["x", "y"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.hypot"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(f64, f64), ElementSink, _, _>( + |_| (), + |&(), (x, y), output| *output = x.hypot(y), + ) + } +} + +/// A byte-string element that resolves each row's view into a data buffer, which is only +/// meaningful for a valid row. +/// +/// This is the crate's only non-dense-safe element, so it exercises valid-only execution, +/// branch-and-skip, and the agreement between the two strategies. It lives here rather than beside +/// the framework because no production row function reads bytes yet. +struct TestBytes; + +/// The canonical views array plus its resolved data buffers. +struct TestBytesColumn { + array: VarBinViewArray, + buffers: Vec, +} + +/// Resolve one view, which is either inlined or an offset into `buffers`. +fn read_view<'a>(view: &'a BinaryView, buffers: &'a [ByteBuffer]) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let view = view.as_view(); + &buffers[view.buffer_index as usize].as_slice()[view.as_range()] + } +} + +impl InputElement for TestBytes { + type Column = TestBytesColumn; + // The views slice, not the array: `VarBinViewArray::views` resolves a host buffer and its + // `vortex_expect` is a side effect the optimizer cannot hoist out of the row loop. + type Varying<'a> = (&'a [BinaryView], &'a [ByteBuffer]); + type Elem<'a> = &'a [u8]; + + const DENSE_SAFE: bool = false; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + matches!(dtype, DType::Utf8(_) | DType::Binary(_)), + "expected a Utf8 or Binary column, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let array = array.execute::(ctx)?; + let buffers = (0..array.data_buffers().len()) + .map(|idx| array.buffer(idx).clone()) + .collect(); + Ok(TestBytesColumn { array, buffers }) + } + + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Self::decode(array, ctx).map(Some) + } + + fn get(column: &Self::Column, index: usize) -> &[u8] { + read_view(&column.array.views()[index], &column.buffers) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + (column.array.views(), &column.buffers) + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.0.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [u8] + where + Self: 'a, + { + read_view(&column.0[index], column.1) + } +} + +impl OutputElement for String { + fn element_dtype() -> DType { + DType::Utf8(Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + VarBinViewArray::from_iter_str(values).into_array() + } + + fn placeholder() -> Self { + String::new() + } +} + +/// A unary row function over strings: uppercased text, exercising [`TestBytes`] input and +/// [`String`] output. +#[derive(Clone)] +struct Shout; + +impl RowFn for Shout { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.shout"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(TestBytes,), ElementSink, _, _>( + |_| (), + |&(), (text,), output| { + *output = String::from_utf8_lossy(text).to_uppercase(); + }, + ) + } +} + +/// A fallible row function: integer division, undefined at a zero divisor. +#[derive(Clone)] +struct CheckedDiv; + +impl RowFn for CheckedDiv { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.checked_div"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64, i64), ElementSink, _, _>( + |_| (), + |&(), (lhs, rhs), output| { + if rhs == 0 { + vortex_bail!("division by zero"); + } + *output = lhs / rhs; + Ok(()) + }, + ) + } +} + +#[test] +fn hypot_columns() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let x = buffer![3.0f64, 5.0].into_array(); + let y = buffer![4.0f64, 12.0].into_array(); + + let result = apply(Hypot, [x, y], &mut ctx)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([5.0f64, 13.0]), &mut ctx); + Ok(()) +} + +#[test] +fn hypot_propagates_nulls_and_constants() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let x = PrimitiveArray::from_option_iter([Some(3.0f64), None, Some(8.0)]).into_array(); + let y = ConstantArray::new(Scalar::from(4.0f64), 3).into_array(); + + let result = apply(Hypot, [x, y], &mut ctx)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(5.0f64), None, Some((80.0f64).sqrt())]), + &mut ctx + ); + Ok(()) +} + +#[test] +fn shout_strings() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = + VarBinViewArray::from_iter_nullable_str([Some("hello"), None, Some("Vortex")]).into_array(); + + let result = apply(Shout, [input], &mut ctx)?; + + let expected = + VarBinViewArray::from_iter_nullable_str([Some("HELLO"), None, Some("VORTEX")]).into_array(); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} + +#[test] +fn display_names_the_function_id() { + let expr = Hypot.new_expr(EmptyOptions, [root(), root()]); + assert_eq!(expr.to_string(), "vortex.test.hypot($, $)"); +} + +#[derive(Clone)] +struct WrongLength; + +impl RowFn for WrongLength { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.wrong_length"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64,), ElementSink, _, _>( + |_| (), + |&(), (value,), output| *output = value, + ) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(PrimitiveArray::from_iter([0i64]).into_array())) + } +} + +#[test] +fn kernel_result_length_is_validated() { + let mut ctx = array_session().create_execution_ctx(); + let input = buffer![1i64, 2, 3].into_array(); + + let error = apply(WrongLength, [input], &mut ctx).unwrap_err(); + + assert!( + error + .to_string() + .contains("produced 1 rows for 3 input rows"), + "{error}" + ); +} + +#[derive(Clone)] +struct FortyTwo; + +impl RowFn for FortyTwo { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.forty_two"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(), ElementSink, _, _>( + |()| (), + |&(), (), output| *output = 42, + ) + } +} + +#[test] +fn nullary_row_fn_executes_requested_rows() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let result = FortyTwo + .try_new_array(3, EmptyOptions, [])? + .execute::(&mut ctx)?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([42i64; 3]), &mut ctx); + Ok(()) +} + +#[derive(Clone)] +struct SumFour; + +impl RowFn for SumFour { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["a", "b", "c", "d"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.sum_four"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64, i64, i64, i64), ElementSink, _, _>( + |_| (), + |&(), (a, b, c, d), output| *output = a + b + c + d, + ) + } +} + +#[test] +fn four_argument_row_fn_executes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let result = apply( + SumFour, + [ + buffer![1i64, 2].into_array(), + buffer![10i64, 20].into_array(), + buffer![100i64, 200].into_array(), + buffer![1000i64, 2000].into_array(), + ], + &mut ctx, + )?; + + assert_arrays_eq!(result, PrimitiveArray::from_iter([1111i64, 2222]), &mut ctx); + Ok(()) +} + +#[test] +fn tuples_are_supported_through_arity_twelve() { + type TwelveI64s = (i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64); + + assert_eq!(<() as ElementTuple>::ARITY, 0); + assert_eq!(::ARITY, 12); +} + +#[test] +fn kernel_flag_decides_fallibility() { + assert!(!ScalarFnVTable::is_fallible(&Hypot, &EmptyOptions)); + assert!(ScalarFnVTable::is_fallible(&CheckedDiv, &EmptyOptions)); +} + +#[test] +fn fallible_apply_propagates_its_error() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = buffer![10i64, 10].into_array(); + let rhs = buffer![2i64, 0].into_array(); + + let error = apply(CheckedDiv, [lhs, rhs], &mut ctx) + .expect_err("a zero divisor must fail the execution"); + + assert!( + error.to_string().contains("division by zero"), + "unexpected error: {error}" + ); + Ok(()) +} + +/// The divisor's null slot holds a zero, which a dense pass would divide by. Filtering keeps the +/// fallible kernel away from it. +#[test] +fn fallible_apply_never_sees_rows_behind_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lhs = buffer![10i64, 10].into_array(); + let rhs = PrimitiveArray::from_option_iter([Some(2i64), None]).into_array(); + + let result = apply(CheckedDiv, [lhs, rhs], &mut ctx)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(5i64), None]), + &mut ctx + ); + Ok(()) +} + +/// The internal nullable execution policy selected by a concrete dispatch. +fn policy>(row_fn: &F, args: &[DType]) -> RowPolicy { + row_policy(row_fn, &EmptyOptions, args).expect("test dispatch must produce a policy") +} + +/// The function never declares a policy: the dispatched arguments and result decide it. +#[test] +fn null_handling_follows_from_args_and_fallibility() { + // Primitive arguments, infallible: nothing behind a null row can fault. + assert_eq!( + policy( + &Hypot, + &[ + DType::Primitive(PType::F64, Nullability::NonNullable), + DType::Primitive(PType::F64, Nullability::NonNullable), + ] + ), + RowPolicy::Dense + ); + // `TestBytes` resolves a view into a data buffer, which is only meaningful for valid rows. + assert_eq!( + policy(&Shout, &[DType::Utf8(Nullability::Nullable)]), + RowPolicy::ValidOnly { + filtered_decode_cost: 0 + } + ); + // Fallible: a garbage row could raise an error of its own. + assert_eq!( + policy( + &CheckedDiv, + &[ + DType::Primitive(PType::I64, Nullability::NonNullable), + DType::Primitive(PType::I64, Nullability::NonNullable), + ] + ), + RowPolicy::ValidOnly { + filtered_decode_cost: 0 + } + ); +} diff --git a/vortex-array/src/scalar_fn/row/tests/null_strategies.rs b/vortex-array/src/scalar_fn/row/tests/null_strategies.rs new file mode 100644 index 00000000000..d8a7fe18816 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/null_strategies.rs @@ -0,0 +1,502 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Equivalence and selection tests for filtering and branch-and-skip null execution. + +use std::sync::Arc; + +use vortex_buffer::ByteBuffer; + +use super::*; +use crate::arrays::varbinview::BinaryView; +use crate::dtype::Nullability; +use crate::validity::Validity; + +/// Executes `scalar_fn` over `args` with `strategy` forced, canonicalized like [`apply`]. +fn apply_forced>( + scalar_fn: &F, + args: &[ArrayRef], + strategy: NullStrategy, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let rows = args.first().map_or(0, |arg| arg.len()); + + Ok( + execute_row_fn_with_strategy(scalar_fn, &EmptyOptions, args.to_vec(), rows, strategy, ctx)? + .execute::(ctx)? + .into_array(), + ) +} + +/// Runs `scalar_fn` under forced filter, forced branch-and-skip, and the automatic per-batch +/// selection, and asserts all three produce identical arrays. +fn assert_strategies_agree>( + scalar_fn: F, + args: Vec, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let filtered = apply_forced(&scalar_fn, &args, NullStrategy::Filter, &mut ctx)?; + let branched = apply_forced(&scalar_fn, &args, NullStrategy::BranchAndSkip, &mut ctx)?; + let auto = apply(scalar_fn, args, &mut ctx)?; + + assert_arrays_eq!(branched, filtered, &mut ctx); + assert_arrays_eq!(auto, filtered, &mut ctx); + Ok(()) +} + +/// A `Utf8` column whose null rows carry views naming a buffer that does not exist, at +/// offsets far out of bounds. Resolving such a row's bytes panics, so strategy agreement +/// proves the branch loop never calls `get` behind a null. +fn hostile_nullable_strings() -> VortexResult { + let views = buffer![ + BinaryView::make_view(b"a longer string here", 0, 0), + BinaryView::new_ref(64, *b"junk", 9, 4096), + BinaryView::make_view(b"another non-inlined string", 1, 0), + BinaryView::new_ref(64, *b"junk", 7, 1 << 20), + ]; + + Ok(VarBinViewArray::try_new( + views, + Arc::from([ + ByteBuffer::copy_from(b"a longer string here"), + ByteBuffer::copy_from(b"another non-inlined string"), + ]), + DType::Utf8(Nullability::Nullable), + Validity::from_iter([true, false, true, false]), + )? + .into_array()) +} + +/// `Bytes` is not dense-safe, so `Shout` uses valid-only execution; both strategies must produce +/// the same array without resolving the hostile views behind the nulls. +#[test] +fn branch_matches_filter_for_bytes() -> VortexResult<()> { + assert_strategies_agree(Shout, vec![hostile_nullable_strings()?]) +} + +/// A fallible kernel with a poison value (zero divisor) behind every null: the branch loop +/// must skip those rows rather than spuriously failing on them. +#[test] +fn branch_never_applies_a_fallible_kernel_behind_nulls() -> VortexResult<()> { + let lhs = buffer![10i64, 10, 12, 9].into_array(); + let rhs = PrimitiveArray::new( + buffer![2i64, 0, 3, 0], + Validity::from_iter([true, false, true, false]), + ) + .into_array(); + + assert_strategies_agree(CheckedDiv, vec![lhs, rhs]) +} + +/// Nulls in both operands: the branch loop must honor the _conjoined_ mask, not either +/// input's own validity. +#[test] +fn branch_conjoins_validities() -> VortexResult<()> { + let lhs = + PrimitiveArray::from_option_iter([Some(10i64), None, Some(12), Some(9), None]).into_array(); + let rhs = PrimitiveArray::new( + buffer![2i64, 0, 3, 0, 0], + Validity::from_iter([true, true, true, false, false]), + ) + .into_array(); + + assert_strategies_agree(CheckedDiv, vec![lhs, rhs]) +} + +/// A constant operand under the branch strategy still hoists through the stride-0 decode. +#[test] +fn branch_handles_constant_operands() -> VortexResult<()> { + let lhs = PrimitiveArray::from_option_iter([Some(10i64), None, Some(12)]).into_array(); + let rhs = ConstantArray::new(Scalar::from(2i64), 3).into_array(); + + assert_strategies_agree(CheckedDiv, vec![lhs, rhs]) +} + +/// An error from a _valid_ row still propagates under the branch strategy. +#[test] +fn branch_propagates_real_errors() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = buffer![10i64, 10, 12].into_array(); + let rhs = PrimitiveArray::new( + buffer![2i64, 3, 0], + Validity::from_iter([true, false, true]), + ) + .into_array(); + + let error = apply_forced( + &CheckedDiv, + &[lhs, rhs], + NullStrategy::BranchAndSkip, + &mut ctx, + ) + .expect_err("a zero divisor in a valid row must fail"); + + assert!( + error.to_string().contains("division by zero"), + "unexpected error: {error}" + ); +} + +/// The automatic per-batch selection, observed through elements that record which decode ran +/// on how many rows: the branch strategy decodes null-tolerantly at full length, the filter +/// strategy decodes ordinarily over the survivors. +mod selection { + use std::cell::Cell; + use std::cell::RefCell; + + use vortex_buffer::Buffer; + use vortex_error::vortex_err; + use vortex_mask::Mask; + + use super::*; + use crate::scalar_fn::row::lift::branch_beats_filter; + + thread_local! { + /// What the last varying-column decode did: `(null_tolerant, rows)`. Thread-local so + /// concurrent tests in one process cannot race it; execution runs on the calling + /// thread. + static LAST_DECODE: Cell> = const { Cell::new(None) }; + } + + /// An i64 element that records its decodes and reports `COST` units of filtered decode work. + /// It is not dense-safe, so strategy selection actually happens. + struct TrackedI64; + + impl InputElement for TrackedI64 { + type Column = Buffer; + type Varying<'a> = ::Varying<'a>; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = false; + const DECODE_FALLIBLE: bool = false; + const FILTERED_DECODE_COST: usize = COST; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + LAST_DECODE.set(Some((false, array.len()))); + ::decode(array, ctx) + } + + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + LAST_DECODE.set(Some((true, array.len()))); + ::decode(array, ctx).map(Some) + } + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + ::varying(column) + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + ::varying_len(column) + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_varying(column, index) + } + } + + /// Negation over one tracked column. + #[derive(Clone)] + struct TrackedNegate; + + impl RowFn for TrackedNegate { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + if COST == 0 { + static ID: CachedId = CachedId::new("vortex.test.tracked_negate.bulk"); + *ID + } else { + static ID: CachedId = CachedId::new("vortex.test.tracked_negate.per_row"); + *ID + } + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(TrackedI64,), ElementSink, _, _>( + |_| (), + |&(), (value,), output| *output = -value, + ) + } + } + + /// A 32-row nullable i64 column whose first `valid_count` rows are valid. + fn column_with_survivors(valid_count: usize) -> ArrayRef { + PrimitiveArray::from_option_iter( + (0..32u16).map(|i| (usize::from(i) < valid_count).then_some(i64::from(i))), + ) + .into_array() + } + + /// Executes the tracked function through the full pipeline and returns what the decode + /// recorded: whether it was null-tolerant, and how many rows it saw. + fn run(valid_count: usize) -> VortexResult<(bool, usize)> { + let mut ctx = array_session().create_execution_ctx(); + LAST_DECODE.set(None); + + apply( + TrackedNegate::, + [column_with_survivors(valid_count)], + &mut ctx, + )?; + + LAST_DECODE + .get() + .ok_or_else(|| vortex_err!("no decode ran")) + } + + /// A bulk-decoded element takes branch-and-skip on a mixed mask however sparse the + /// survivors: the decode is null-tolerant and full length. + #[test] + fn bulk_decode_branches_at_any_density() -> VortexResult<()> { + assert_eq!(run::<0>(31)?, (true, 32)); + assert_eq!(run::<0>(4)?, (true, 32)); + Ok(()) + } + + /// One per-row decode still branches when half the rows survive, matching the measured + /// single-nullable-input crossover. + #[test] + fn per_row_decode_filters_when_sparse() -> VortexResult<()> { + // 30/32 surviving: branch, full-length null-tolerant decode. + assert_eq!(run::<1>(30)?, (true, 32)); + // 16/32 = 50% surviving sits exactly on the threshold: still branch. + assert_eq!(run::<1>(16)?, (true, 32)); + // Below 50% surviving: filter, ordinary decode over the survivors. + assert_eq!(run::<1>(15)?, (false, 15)); + Ok(()) + } + + /// An all-true mask short-circuits to the plain kernel and an all-false mask to an + /// all-null constant, before any strategy is selected. + #[test] + fn degenerate_masks_bypass_the_selection() -> VortexResult<()> { + assert_eq!(run::<1>(32)?, (false, 32)); + + let mut ctx = array_session().create_execution_ctx(); + LAST_DECODE.set(None); + apply(TrackedNegate::<1>, [column_with_survivors(0)], &mut ctx)?; + assert_eq!(LAST_DECODE.get(), None); + Ok(()) + } + + /// An i64 element that omits `decode_null_tolerant`: the conservative default refuses, so + /// the batch must fall back to the filter strategy even though the selection preferred + /// branch. + struct RefusesNullTolerant; + + impl InputElement for RefusesNullTolerant { + type Column = Buffer; + type Varying<'a> = ::Varying<'a>; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = false; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + LAST_DECODE.set(Some((false, array.len()))); + ::decode(array, ctx) + } + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + ::varying(column) + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + ::varying_len(column) + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_varying(column, index) + } + } + + #[derive(Clone)] + struct RefusingNegate; + + impl RowFn for RefusingNegate { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.refusing_negate"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(RefusesNullTolerant,), ElementSink, _, _>( + |_| (), + |&(), (value,), output| *output = -value, + ) + } + } + + /// The fallback is silent and correct: the ordinary decode runs over the survivors and + /// the result matches the expected negation. + #[test] + fn missing_null_tolerant_decode_falls_back_to_filter() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + LAST_DECODE.set(None); + + let result = apply( + RefusingNegate, + [PrimitiveArray::from_option_iter([Some(3i64), None, Some(5)]).into_array()], + &mut ctx, + )?; + + assert_eq!(LAST_DECODE.get(), Some((false, 2))); + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(-3i64), None, Some(-5)]), + &mut ctx + ); + Ok(()) + } + + thread_local! { + /// Every `row_count` `reduce_encoded` was handed, in call order. + static REDUCE_ROW_COUNTS: RefCell> = const { RefCell::new(Vec::new()) }; + } + + /// [`RefusingNegate`] with an encoding-aware rewrite that declines, recording the row count it + /// was offered. + #[derive(Clone)] + struct ProbingNegate; + + impl RowFn for ProbingNegate { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.probing_negate"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(RefusesNullTolerant,), ElementSink, _, _>( + |_| (), + |&(), (value,), output| *output = -value, + ) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + REDUCE_ROW_COUNTS.with_borrow_mut(|counts| counts.push(args[0].len())); + Ok(None) + } + } + + /// A `reduce_encoded` rewrite must be sized from the arrays it was handed, which under the + /// filter strategy hold the surviving rows rather than the whole batch. This is the only + /// mixed-mask path where the two differ, so nothing else would catch a rewrite sized from a + /// length captured elsewhere. + /// + /// This also pins the double probe as deliberate. The first call sees the original arrays at + /// full length, and is the only one that does; the second sees filtered, canonical copies. An + /// "optimization" that skipped the first because the batch will end up filtering would take an + /// encoding-aware rewrite away from every function whose sink cannot skip rows. + #[test] + fn reduce_encoded_is_probed_before_and_after_filtering() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + REDUCE_ROW_COUNTS.with_borrow_mut(Vec::clear); + + let result = apply( + ProbingNegate, + [PrimitiveArray::from_option_iter([Some(3i64), None, Some(5)]).into_array()], + &mut ctx, + )?; + + assert_eq!( + REDUCE_ROW_COUNTS.with_borrow(|counts| counts.clone()), + vec![3, 2], + "expected an unfiltered probe at the batch length, then a filtered one at the \ + surviving count", + ); + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(-3i64), None, Some(-5)]), + &mut ctx + ); + Ok(()) + } + + /// The rule itself, at and around the threshold, without going through an execution. + #[rstest] + #[case::bulk_dense_mask(0, 99, 100, true)] + #[case::bulk_sparse_mask(0, 1, 100, true)] + #[case::one_decode_dense_mask(1, 99, 100, true)] + #[case::one_decode_at_threshold(1, 50, 100, true)] + #[case::one_decode_below_threshold(1, 49, 100, false)] + #[case::two_decodes_at_old_boolean_choice(2, 81, 100, false)] + #[case::two_decodes_dense_mask(2, 90, 100, true)] + fn selects_branch_per_the_measured_rule( + #[case] filtered_decode_cost: usize, + #[case] true_count: usize, + #[case] len: usize, + #[case] expect_branch: bool, + ) { + let valid = Mask::from_indices(len, 0..true_count); + assert_eq!( + branch_beats_filter(filtered_decode_cost, &valid), + expect_branch, + ); + } + + #[test] + fn planning_adds_decode_cost_across_arguments() { + assert_eq!( + RowPolicy::for_dispatch::<(TrackedI64<1>, TrackedI64<1>), ()>(), + RowPolicy::ValidOnly { + filtered_decode_cost: 2 + } + ); + } +} diff --git a/vortex-array/src/scalar_fn/row/tests/nullable_outputs.rs b/vortex-array/src/scalar_fn/row/tests/nullable_outputs.rs new file mode 100644 index 00000000000..50ea77ef015 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/nullable_outputs.rs @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests that row output dtypes cannot introduce their own nulls. + +use super::*; +use crate::dtype::Nullability; +use crate::dtype::PType; + +#[derive(Clone)] +struct NullableI64(i64); + +impl OutputElement for NullableI64 { + fn element_dtype() -> DType { + DType::Primitive(PType::I64, Nullability::Nullable) + } + + fn build(values: Vec) -> ArrayRef { + PrimitiveArray::from_option_iter(values.into_iter().map(|value| Some(value.0))).into_array() + } + + fn placeholder() -> Self { + Self(0) + } +} + +struct NullableSink(usize); + +impl OutputSink for NullableSink { + type Rows<'a> = usize; + type Row<'a> = (); + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::Primitive(PType::I64, Nullability::Nullable)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(rows)) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0 + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + *rows == row_count + } + + fn row<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} + + fn finish(self, _error: DeferredError) -> VortexResult { + Ok(PrimitiveArray::from_option_iter(Vec::>::new()).into_array()) + } +} + +#[derive(Clone)] +struct NullableElementFn; + +impl RowFn for NullableElementFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.nullable_element"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64,), ElementSink, _, _>( + |_| (), + |&(), (value,), output| *output = NullableI64(value), + ) + } +} + +#[derive(Clone)] +struct NullableSinkFn; + +impl RowFn for NullableSinkFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.nullable_sink"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64,), NullableSink, _, _>(|_| (), |&(), _, ()| {}) + } +} + +#[test] +fn nullable_element_dtype_is_rejected() { + let input = DType::Primitive(PType::I64, Nullability::NonNullable); + let error = + ScalarFnVTable::return_dtype(&NullableElementFn, &EmptyOptions, &[input]).unwrap_err(); + + assert!(error.to_string().contains("non-nullable dtype"), "{error}"); +} + +#[test] +fn nullable_sink_dtype_is_rejected() { + let input = DType::Primitive(PType::I64, Nullability::NonNullable); + let error = ScalarFnVTable::return_dtype(&NullableSinkFn, &EmptyOptions, &[input]).unwrap_err(); + + assert!(error.to_string().contains("non-nullable dtype"), "{error}"); +} diff --git a/vortex-array/src/scalar_fn/row/tests/prepared.rs b/vortex-array/src/scalar_fn/row/tests/prepared.rs new file mode 100644 index 00000000000..64ed82c8d2c --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/prepared.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for preparing batch-constant state once before the row loop. + +use std::cell::Cell; + +use super::*; +use crate::validity::Validity; + +thread_local! { + /// Which operands the last `prepare` saw as constant, as a bitmask (bit 0 for `x`, bit 1 + /// for `y`). Thread-local rather than a process global so concurrent tests in one process + /// (plain `cargo test`) cannot race it; execution runs on the calling thread. + static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; +} + +/// `sqrt(x^2 + y^2)` through [`RowVisitor::visit_prepared_into`]: the square of any constant +/// operand is hoisted out of the row loop, and recorded in [`SEEN_CONSTANTS`]. +#[derive(Clone)] +struct PreparedHypot; + +impl RowFn for PreparedHypot { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["x", "y"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.prepared_hypot"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(f64, f64), ElementSink, _, _>( + |(x, y)| { + SEEN_CONSTANTS.set(u8::from(x.is_some()) | (u8::from(y.is_some()) << 1)); + (x.map(|x| x * x), y.map(|y| y * y)) + }, + |&(x_sq, y_sq), (x, y), output| { + *output = (x_sq.unwrap_or(x * x) + y_sq.unwrap_or(y * y)).sqrt(); + }, + ) + } +} + +/// A constant operand reaches `prepare` as `Some`, and the result is identical to the same +/// value expanded into a full column, which reaches `prepare` as `None`. +#[test] +fn a_constant_operand_matches_its_expanded_column() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let x = buffer![3.0f64, 5.0, 8.0].into_array(); + + let constant = ConstantArray::new(Scalar::from(4.0f64), 3).into_array(); + let from_constant = apply(PreparedHypot, [x.clone(), constant], &mut ctx)?; + assert_eq!(SEEN_CONSTANTS.get(), 0b10); + + let expanded = buffer![4.0f64, 4.0, 4.0].into_array(); + let from_expanded = apply(PreparedHypot, [x, expanded], &mut ctx)?; + assert_eq!(SEEN_CONSTANTS.get(), 0b00); + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// A masked constant (the same value in every row, some rows null, how the compressor spells +/// an all-same-with-nulls chunk) is a batch constant too: the wrapper carries only validity, +/// which the lifting owns, so `prepare` sees the child's value and the null rows stay +/// null in the result. +#[test] +fn a_masked_constant_operand_is_seen_as_constant() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let x = buffer![3.0f64, 5.0, 8.0].into_array(); + + let masked_constant = MaskedArray::try_new( + ConstantArray::new(Scalar::from(4.0f64), 3).into_array(), + Validity::from_iter([true, false, true]), + )? + .into_array(); + let result = apply(PreparedHypot, [x, masked_constant], &mut ctx)?; + + assert_eq!(SEEN_CONSTANTS.get(), 0b10); + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(5.0f64), None, Some((80.0f64).sqrt())]), + &mut ctx + ); + Ok(()) +} + +/// With no constant operand every `ConstElems` slot is `None` and the loop computes exactly +/// what unit preparation would. +#[test] +fn all_varying_operands_prepare_nothing() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let x = buffer![3.0f64, 5.0].into_array(); + let y = buffer![4.0f64, 12.0].into_array(); + + let result = apply(PreparedHypot, [x, y], &mut ctx)?; + + assert_eq!(SEEN_CONSTANTS.get(), 0b00); + assert_arrays_eq!(result, PrimitiveArray::from_iter([5.0f64, 13.0]), &mut ctx); + Ok(()) +} + +/// Two constant operands are folded to a single-row execution by the lifting, and that +/// row still goes through `prepare`, seeing both constants. +#[test] +fn all_constant_operands_fold_and_still_prepare() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let x = ConstantArray::new(Scalar::from(3.0f64), 4).into_array(); + let y = ConstantArray::new(Scalar::from(4.0f64), 4).into_array(); + + let result = apply(PreparedHypot, [x, y], &mut ctx)?; + + assert_eq!(SEEN_CONSTANTS.get(), 0b11); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([5.0f64, 5.0, 5.0, 5.0]), + &mut ctx + ); + Ok(()) +} + +/// Null rows pass through the prepared path exactly as through unit preparation. +#[test] +fn nulls_propagate_through_the_prepared_path() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let x = PrimitiveArray::from_option_iter([Some(3.0f64), None, Some(8.0)]).into_array(); + let y = ConstantArray::new(Scalar::from(4.0f64), 3).into_array(); + + let result = apply(PreparedHypot, [x, y], &mut ctx)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(5.0f64), None, Some((80.0f64).sqrt())]), + &mut ctx + ); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/tests/sink.rs b/vortex-array/src/scalar_fn/row/tests/sink.rs new file mode 100644 index 00000000000..95b5e2df05f --- /dev/null +++ b/vortex-array/src/scalar_fn/row/tests/sink.rs @@ -0,0 +1,375 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for row functions that write into a batch-wide output sink. + +use std::sync::Arc; + +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; + +use super::*; +use crate::arrays::FixedSizeListArray; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::validity::Validity; + +/// Builds a `FixedSizeList` column, presenting each row as the `&mut [T]` slice to fill. +/// +/// Its element dtype comes from the input rather than from `T` alone, so it exercises +/// [`OutputSink::sink_dtype`] actually reading `args`. +struct SpreadSink { + dtype: DType, + rows: usize, + elements: BufferMut, +} + +impl OutputSink for SpreadSink { + type Rows<'a> = (&'a mut [T], usize); + type Row<'a> = &'a mut [T]; + + fn sink_dtype(args: &[DType]) -> VortexResult { + let element = args + .first() + .ok_or_else(|| vortex_err!("a spread sink takes its element dtype from its input"))?; + ::validate(element)?; + Ok(DType::FixedSizeList( + Arc::new(element.as_nonnullable()), + u32::try_from(W).vortex_expect("test width fits in u32"), + Nullability::NonNullable, + )) + } + + fn with_capacity(rows: usize, dtype: &DType) -> VortexResult { + Ok(Self { + dtype: dtype.clone(), + rows, + elements: BufferMut::zeroed(rows * W), + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + (self.elements.as_mut_slice(), self.rows) + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.1 == row_count && row_count.checked_mul(W) == Some(rows.0.len()) + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> &'a mut [T] { + &mut rows.0[index * W..][..W] + } + + fn finish(self, _error: DeferredError) -> VortexResult { + vortex_ensure_eq!( + self.dtype, + Self::sink_dtype(&[DType::Primitive(T::PTYPE, self.dtype.nullability())])?, + "the sink must build the dtype it named", + ); + Ok(FixedSizeListArray::try_new( + PrimitiveArray::new(self.elements.freeze(), Validity::NonNullable).into_array(), + u32::try_from(W).vortex_expect("test width fits in u32"), + Validity::NonNullable, + self.rows, + )? + .into_array()) + } +} + +/// A sink that reports a data-dependent error only after every row has been written. +struct NonNegativeSink { + values: BufferMut, +} + +struct NonNegativeRow<'a> { + value: &'a mut i64, +} + +impl NonNegativeRow<'_> { + fn write(self, value: i64) -> bool { + *self.value = value; + value < 0 + } +} + +impl OutputSink for NonNegativeSink { + const ERRORS_ARE_DEFERRED: bool = true; + + type Rows<'a> = &'a mut [i64]; + type Row<'a> = NonNegativeRow<'a>; + + fn sink_dtype(args: &[DType]) -> VortexResult { + let dtype = args + .first() + .ok_or_else(|| vortex_err!("a non-negative sink requires one input"))?; + ::validate(dtype)?; + Ok(DType::Primitive(PType::I64, Nullability::NonNullable)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self { + values: BufferMut::zeroed(rows), + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.values.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + NonNegativeRow { + value: &mut rows[index], + } + } + + fn finish(self, error: DeferredError) -> VortexResult { + if error.occurred() { + vortex_bail!("negative output"); + } + Ok(PrimitiveArray::new(self.values.freeze(), Validity::NonNullable).into_array()) + } +} + +/// Broadcasts each input value across a fixed-size list row: `spread(x) == [x, x, x]`. +#[derive(Clone)] +struct Spread; + +impl RowFn for Spread { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.spread"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64,), SpreadSink, _, _>( + |_| (), + |&(), (x,), out| out.fill(x), + ) + } +} + +/// The same, but refusing negative inputs, so its row closure returns `VortexResult<()>`. +#[derive(Clone)] +struct SpreadNonNegative; + +impl RowFn for SpreadNonNegative { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.spread_non_negative"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64,), SpreadSink, _, _>( + |_| (), + |&(), (x,), out| { + if x < 0 { + vortex_bail!("negative input {x}"); + } + out.fill(x); + Ok(()) + }, + ) + } +} + +/// Writes infallibly and lets its output sink report the error after the loop. +#[derive(Clone)] +struct DeferredNonNegative; + +impl RowFn for DeferredNonNegative { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.deferred_non_negative"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64,), NonNegativeSink, _, _>( + |_| (), + |&(), (value,), output| output.write(value), + ) + } +} + +/// `SpreadSink`'s three-element rows, built from `values`. +fn spread_rows(values: impl IntoIterator>) -> VortexResult { + let values = values.into_iter().collect::>(); + let rows = values.len(); + let flat = values + .iter() + .flat_map(|value| [value.unwrap_or(0); 3]) + .collect::>(); + let validity = if values.iter().all(Option::is_some) { + Validity::NonNullable + } else { + Validity::from_iter(values.iter().map(Option::is_some)) + }; + + Ok(FixedSizeListArray::try_new( + PrimitiveArray::new(flat, Validity::NonNullable).into_array(), + 3, + validity, + rows, + )? + .into_array()) +} + +/// The output dtype is the sink's, with its width, and every row holds the written slice. +#[test] +fn writes_one_row_at_a_time() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = buffer![7i64, -2, 0].into_array(); + + let result = apply(Spread, [input], &mut ctx)?; + + assert_arrays_eq!(result, spread_rows([Some(7), Some(-2), Some(0)])?, &mut ctx); + Ok(()) +} + +/// A null input row is written densely and masked away afterwards, exactly as on the value path. +#[test] +fn nulls_are_masked_after_the_sink() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = PrimitiveArray::from_option_iter([Some(7i64), None, Some(4)]).into_array(); + + let result = apply(Spread, [input], &mut ctx)?; + + assert!(result.dtype().is_nullable()); + assert_arrays_eq!(result, spread_rows([Some(7), None, Some(4)])?, &mut ctx); + Ok(()) +} + +/// A sink whose closure cannot fail is dense, while one whose closure can fail is valid-only. +#[test] +fn null_handling_follows_from_declared_fallibility() { + let args = [DType::Primitive(PType::I64, Nullability::NonNullable)]; + assert_eq!(policy(&Spread, &args), RowPolicy::Dense); + assert!(!ScalarFnVTable::is_fallible(&Spread, &EmptyOptions)); + + assert_eq!( + policy(&SpreadNonNegative, &args), + RowPolicy::ValidOnly { + filtered_decode_cost: 0 + } + ); + assert!(ScalarFnVTable::is_fallible( + &SpreadNonNegative, + &EmptyOptions + )); +} + +/// An error from a writing closure aborts the batch rather than being written into the sink. +#[test] +fn a_failing_row_propagates() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = buffer![1i64, -5, 3].into_array(); + + let error = apply(SpreadNonNegative, [input], &mut ctx).unwrap_err(); + + assert!(error.to_string().contains("negative input -5"), "{error}"); + Ok(()) +} + +/// A sink may accumulate a failure while its row closure remains infallible. +#[test] +fn a_sink_can_defer_its_error_until_finish() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = buffer![1i64, -5, 3].into_array(); + + let error = apply(DeferredNonNegative, [input], &mut ctx).unwrap_err(); + + assert!(error.to_string().contains("negative output"), "{error}"); + Ok(()) +} + +/// A deferred failure behind a null triggers a valid-row retry and is then discarded. +#[test] +fn a_deferred_error_behind_a_null_is_ignored() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = PrimitiveArray::new( + buffer![1i64, -5, 3], + Validity::from_iter([true, false, true]), + ) + .into_array(); + + let result = apply(DeferredNonNegative, [input], &mut ctx)?; + + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]), + &mut ctx + ); + Ok(()) +} + +/// Being fallible, `SpreadNonNegative` is filtered, so its closure never sees the value behind a +/// null row. A negative payload there must therefore not raise. +#[test] +fn a_failing_row_is_never_reached_behind_a_null() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = PrimitiveArray::new( + buffer![1i64, -5, 3], + Validity::from_iter([true, false, true]), + ) + .into_array(); + + let result = apply(SpreadNonNegative, [input], &mut ctx)?; + + assert_arrays_eq!(result, spread_rows([Some(1), None, Some(3)])?, &mut ctx); + Ok(()) +} + +/// The sink names its output dtype from the input, so a wrong input dtype is rejected at plan +/// time rather than producing a mis-typed column. +#[test] +fn the_sink_dtype_validates_its_input() { + let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ScalarFnVTable::return_dtype(&Spread, &EmptyOptions, &[dtype]).is_err()); +} + +/// The width the sink declares is the width it builds, over the element dtype it read off the +/// input. +#[test] +fn the_return_dtype_is_the_sinks() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); + assert_eq!( + ScalarFnVTable::return_dtype(&Spread, &EmptyOptions, std::slice::from_ref(&dtype))?, + DType::FixedSizeList(Arc::new(dtype), 3, Nullability::NonNullable), + ); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs new file mode 100644 index 00000000000..cff4831aa8f --- /dev/null +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -0,0 +1,398 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Blanket scalar-function implementation and execution visitors for row functions. + +use std::marker::PhantomData; + +use vortex_error::VortexResult; +#[cfg(any(test, feature = "_test-harness"))] +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use super::row_fn::RowFn; +use super::row_fn::RowVisitor; +use super::row_fn::private; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::expr::Expression; +use crate::expr::union_child_validities; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::ExecutionArgs; +#[cfg(any(test, feature = "_test-harness"))] +use crate::scalar_fn::NullStrategy; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::SinkResult; +#[cfg(any(test, feature = "_test-harness"))] +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::row::execute::RowExecution; +use crate::scalar_fn::row::execute::execute_row_sink_branch; +use crate::scalar_fn::row::execute::execute_row_sink_prepared; +use crate::scalar_fn::row::execute::validate_row_sink; +use crate::scalar_fn::row::lift::Batch; +use crate::scalar_fn::row::lift::BatchPlan; +use crate::scalar_fn::row::lift::KernelArgs; +use crate::scalar_fn::row::lift::RowPolicy; +use crate::scalar_fn::row::lift::reconcile_return; + +/// Compile-time check that a dispatched `(A, S, R)` agrees with `F`'s public metadata. Evaluated by +/// monomorphizing +/// [`visit_prepared_into`](RowVisitor::visit_prepared_into), so even a dispatch arm that never runs +/// is checked. +const fn assert_dispatch_agrees() { + assert!( + A::ARITY == F::ARG_NAMES.len(), + "dispatch visited a tuple whose arity differs from RowFn::ARG_NAMES", + ); + // Dictionary pushdown treats an infallible function as safe to evaluate over values no code + // references, so every dispatch must fit the function-wide declaration. + assert!( + !A::DECODE_FALLIBLE || F::FALLIBLE, + "dispatch decoded fallibly without declaring RowFn::FALLIBLE", + ); + assert!( + !R::FALLIBLE || F::FALLIBLE, + "dispatch returned an error without declaring RowFn::FALLIBLE", + ); + assert!( + !R::DEFERRED || F::FALLIBLE, + "dispatch deferred an error without declaring RowFn::FALLIBLE", + ); + assert!( + S::ERRORS_ARE_DEFERRED == R::DEFERRED, + "a deferred-error sink and row closure must be used together", + ); +} + +/// The plan-time visit: validate the dtypes and derive execution from the concrete sink and row +/// closure selected by dispatch. +struct PlanRows<'a, F> { + args: &'a [DType], + + /// The visited function, carried only so the dispatch check can name its contract. + row_fn: PhantomData, +} + +impl private::Sealed for PlanRows<'_, F> {} + +impl RowVisitor for PlanRows<'_, F> { + type Out = BatchPlan; + + fn visit_prepared_into( + self, + _prepare: impl FnOnce(A::ConstElems<'_>) -> P, + _apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, + ) -> VortexResult { + const { assert_dispatch_agrees::() }; + + Ok(BatchPlan { + sink_dtype: validate_row_sink::(self.args)?, + policy: RowPolicy::for_dispatch::(), + }) + } +} + +/// The run-time visit: decode every column once and run the row loop. +struct ExecuteRows<'a, 'b, F> { + args: &'a dyn ExecutionArgs, + + /// The sink dtype computed by the planning visit. + sink_dtype: &'a DType, + + ctx: &'b mut ExecutionCtx, + + /// The visited function, carried only so the dispatch check can name its contract. + row_fn: PhantomData, +} + +impl private::Sealed for ExecuteRows<'_, '_, F> {} + +impl RowVisitor for ExecuteRows<'_, '_, F> { + type Out = RowExecution; + + fn visit_prepared_into( + self, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, + ) -> VortexResult { + const { assert_dispatch_agrees::() }; + execute_row_sink_prepared::( + self.args, + self.sink_dtype, + self.ctx, + prepare, + apply, + ) + } +} + +/// The run-time visit for the branch-and-skip null strategy: compute only the conjoined-valid +/// rows over unfiltered columns. +/// +/// `Ok(None)` means the visit cannot take that strategy because the sink cannot skip rows or an +/// argument has no null-tolerant decode, and the lifting falls back to the filter strategy. +struct ExecuteRowsBranch<'a, 'b, F> { + args: &'a dyn ExecutionArgs, + + /// The sink dtype computed by the planning visit. + sink_dtype: &'a DType, + + /// The conjoined validity, materialized by the lifting and guaranteed mixed. + valid: &'a Mask, + + ctx: &'b mut ExecutionCtx, + + /// The visited function, carried only so the dispatch check can name its contract. + row_fn: PhantomData, +} + +impl private::Sealed for ExecuteRowsBranch<'_, '_, F> {} + +impl RowVisitor for ExecuteRowsBranch<'_, '_, F> { + type Out = Option; + + fn visit_prepared_into( + self, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, + ) -> VortexResult> { + const { assert_dispatch_agrees::() }; + execute_row_sink_branch::( + self.args, + self.sink_dtype, + self.valid, + self.ctx, + prepare, + apply, + ) + } +} + +/// The kernel the lifting runs: the encoding-aware rewrite if it answers, otherwise the row loop +/// over whichever arguments the lifting hands over. +fn execute_rows( + row_fn: &F, + options: &F::Options, + args: KernelArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if let Some(reduced) = row_fn.reduce_encoded(options, args.arrays, ctx)? { + return Ok(RowExecution::Output(reduced)); + } + + row_fn.dispatch( + options, + args.dtypes, + ExecuteRows:: { + args: args.execution, + sink_dtype: args.sink_dtype, + ctx, + row_fn: PhantomData, + }, + ) +} + +/// The branch-and-skip kernel: compute only the rows set in `valid`, over the unfiltered `args`. +/// +/// `Ok(None)` sends the batch to the filter strategy instead. +fn execute_rows_branch( + row_fn: &F, + options: &F::Options, + args: KernelArgs<'_>, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + // The encoding-aware rewrite runs before the row loop exactly as in [`execute_rows`]. Here it + // sees the original (unfiltered) encodings, and its full-length result is masked by the caller + // like any other branch result. + if let Some(reduced) = row_fn.reduce_encoded(options, args.arrays, ctx)? { + return Ok(Some(RowExecution::Output(reduced))); + } + + row_fn.dispatch( + options, + args.dtypes, + ExecuteRowsBranch:: { + args: args.execution, + sink_dtype: args.sink_dtype, + valid, + ctx, + row_fn: PhantomData, + }, + ) +} + +/// The batch facts for `row_fn` over `args`, derived from its dispatched elements and sink. +fn lift_batch<'a, F: RowFn>( + row_fn: &F, + options: &F::Options, + args: &'a dyn ExecutionArgs, +) -> VortexResult> { + Batch::new(RowFn::id(row_fn), args, |arg_dtypes| { + let plan = row_fn.dispatch( + options, + arg_dtypes, + PlanRows:: { + args: arg_dtypes, + row_fn: PhantomData, + }, + )?; + Ok(plan) + }) +} + +/// The nullable execution policy selected by one concrete dispatch. +#[cfg(test)] +pub(super) fn row_policy( + row_fn: &F, + options: &F::Options, + args: &[DType], +) -> VortexResult { + row_fn + .dispatch( + options, + args, + PlanRows:: { + args, + row_fn: PhantomData, + }, + ) + .map(|plan| plan.policy) +} + +/// Every [`RowFn`] is a [`ScalarFnVTable`], the row loop lifted by `Batch`. +/// +/// This impl is why a [`RowFn`] cannot also implement [`ScalarFnVTable`] itself: coherence forbids +/// the second impl. Nothing in tree needs to, since everything a row function can vary lives on +/// [`RowFn`]; mirror another [`ScalarFnVTable`] method onto it when something actually does. +impl ScalarFnVTable for F { + type Options = F::Options; + + fn id(&self) -> ScalarFnId { + RowFn::id(self) + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + RowFn::serialize(self, options) + } + + fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { + RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(F::ARG_NAMES.len()) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + ChildName::from(F::ARG_NAMES[child_idx]) + } + + /// The visited output element's dtype, widened to nullable iff any input is nullable, which is + /// what makes the strictness dtype contract hold by construction. + fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { + let plan = self.dispatch( + options, + args, + PlanRows:: { + args, + row_fn: PhantomData, + }, + )?; + + let nullability = + plan.sink_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); + Ok(plan.sink_dtype.with_nullability(nullability)) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Nullary functions have no input values that could be null, so there is nothing to lift. + if args.num_inputs() == 0 { + let result_dtype = ScalarFnVTable::return_dtype(self, options, &[])?; + let values = execute_rows( + self, + options, + KernelArgs { + execution: args, + arrays: &[], + dtypes: &[], + sink_dtype: &result_dtype, + }, + ctx, + )? + .into_result()?; + return reconcile_return(RowFn::id(self), &result_dtype, args.row_count(), values); + } + + lift_batch(self, options, args)?.execute( + |args, ctx| execute_rows(self, options, args, ctx), + |args, valid, ctx| execute_rows_branch(self, options, args, valid, ctx), + ctx, + ) + } + + /// Output sinks build an all-valid column, so a row kernel cannot turn a wholly non-null row into + /// a null and the output validity is exactly the conjunction of the inputs'. Letting a sink + /// produce nulls would invalidate this. + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + /// A row kernel maps a null input row to a null output row, and computes non-null outputs from + /// non-null inputs alone, which is exactly strictness. The lifting is what makes it true. + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + F::FALLIBLE + } +} + +/// Execute `row_fn` over `inputs` with a forced null strategy, bypassing the per-batch selection. +/// +/// A test and benchmark seam only, and the only way to name a strategy from outside: it is how the +/// two are compared and how their agreement is asserted. It skips the null-constant and +/// all-constant folds, so do not pass such inputs. Forcing [`NullStrategy::BranchAndSkip`] on a +/// dispatch with no branch execution is an error rather than a silent fallback to filtering. +#[cfg(any(test, feature = "_test-harness"))] +pub fn execute_row_fn_with_strategy( + row_fn: &F, + options: &F::Options, + inputs: Vec, + row_count: usize, + strategy: NullStrategy, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(inputs, row_count); + + lift_batch(row_fn, options, &args)? + .execute_with_strategy( + |args, ctx| execute_rows(row_fn, options, args, ctx), + |args, valid, ctx| execute_rows_branch(row_fn, options, args, valid, ctx), + strategy, + ctx, + )? + .ok_or_else(|| { + vortex_err!( + "{} has no branch-and-skip execution for these inputs", + RowFn::id(row_fn), + ) + }) +} diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 5d3561ff039..e5b074e9cf3 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -361,7 +361,7 @@ impl ExecutionArgs for VecExecutionArgs { } } -#[derive(Clone, Debug, PartialEq, Eq, Hash)] +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] pub struct EmptyOptions; impl Display for EmptyOptions { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { From ae099e8909d387b881ed932689467955818349d6 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 13:03:43 +0000 Subject: [PATCH 002/160] Benchmark the row executor and the null strategies Progress towards #9128. `row_fn_executor` measures the derived execution loop against a hand-written columnar kernel of the same arithmetic, so the cost of going through `RowFn` is separated from the cost of the operation. `strict_validity` measures the three null strategies against each other across validity densities, which is the evidence behind the per-batch selection rule. `like` gains a pair. `like_per_row_repeated_patterns` carries one infix pattern on every row so the compile cache always hits, and `like_per_row_distinct_patterns` varies that pattern so it never hits. Both compile the same shape and match the same way, so the difference between them is compilation alone, which is what a kernel that cannot cache across rows pays. The existing `like_per_row_patterns` keeps its input, so its measurements stay comparable with develop. Signed-off-by: Connor Tsui Co-authored-by: Claude --- vortex-array/Cargo.toml | 8 + vortex-array/benches/like.rs | 44 ++- vortex-array/benches/row_fn_executor.rs | 419 ++++++++++++++++++++++++ vortex-array/benches/strict_validity.rs | 217 ++++++++++++ 4 files changed, 683 insertions(+), 5 deletions(-) create mode 100644 vortex-array/benches/row_fn_executor.rs create mode 100644 vortex-array/benches/strict_validity.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index d00b811a387..2da7f59309f 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -129,6 +129,10 @@ harness = false name = "binary_ops" harness = false +[[bench]] +name = "row_fn_executor" +harness = false + [[bench]] name = "kleene_bool" harness = false @@ -282,3 +286,7 @@ harness = false [[bench]] name = "slice_dict_primitive" harness = false + +[[bench]] +name = "strict_validity" +harness = false diff --git a/vortex-array/benches/like.rs b/vortex-array/benches/like.rs index 68219724717..e83fae69b28 100644 --- a/vortex-array/benches/like.rs +++ b/vortex-array/benches/like.rs @@ -87,13 +87,9 @@ fn like_regex(bencher: Bencher) { bench_like(bencher, "h_llo%w%d", LikeOptions::default()); } -#[divan::bench] -fn like_per_row_patterns(bencher: Bencher) { +fn bench_per_row_patterns(bencher: Bencher, patterns: ArrayRef) { let session = vortex_array::array_session(); let array = strings(); - // A non-constant pattern child takes the per-row path; repeated patterns hit the - // compile cache. - let patterns = VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| "hello%")).into_array(); bencher .with_inputs(|| { ( @@ -109,6 +105,44 @@ fn like_per_row_patterns(bencher: Bencher) { .bench_values(|(array, mut ctx)| array.execute::(&mut ctx).unwrap()); } +#[divan::bench] +fn like_per_row_patterns(bencher: Bencher) { + // A non-constant pattern child takes the per-row path; repeated patterns hit the + // compile cache. + let patterns = VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| "hello%")).into_array(); + bench_per_row_patterns(bencher, patterns); +} + +/// The per-row path with the compile cache hit on every row, carrying the infix pattern that +/// [`like_per_row_distinct_patterns`] varies. Both compile the same shape and match the same way, +/// so the only difference between them is how often a pattern is compiled. +#[divan::bench] +fn like_per_row_repeated_patterns(bencher: Bencher) { + let patterns = VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| "%aaa%")).into_array(); + bench_per_row_patterns(bencher, patterns); +} + +/// The per-row path with the compile cache defeated: every row carries a distinct pattern of the +/// same shape, so each row pays one [`LikePattern`] compilation. +/// +/// Paired with [`like_per_row_repeated_patterns`] this isolates the cost of compiling a pattern from +/// the cost of matching against it, which is what any kernel that cannot cache across rows pays. +#[divan::bench] +fn like_per_row_distinct_patterns(bencher: Bencher) { + let patterns = VarBinViewArray::from_iter_str( + (0..ARRAY_SIZE).map(|i| format!("%{}%", distinct_trigram(i))), + ) + .into_array(); + bench_per_row_patterns(bencher, patterns); +} + +/// A distinct three-letter lowercase infix per row, so `ARRAY_SIZE` rows never repeat a pattern +/// while every pattern keeps the same shape and compiles the same way. +fn distinct_trigram(i: usize) -> String { + let letter = |shift: usize| char::from(b'a' + u8::try_from((i >> shift) % 26).unwrap()); + [letter(0), letter(5), letter(10)].iter().collect() +} + #[divan::bench] fn ilike_contains(bencher: Bencher) { bench_like( diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs new file mode 100644 index 00000000000..c66865ef26d --- /dev/null +++ b/vortex-array/benches/row_fn_executor.rs @@ -0,0 +1,419 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares cheap primitive row functions with the specialized columnar implementation. + +#![expect(clippy::unwrap_used)] + +use std::mem::MaybeUninit; +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::DeferredError; +use vortex_array::scalar_fn::ElementSink; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::OutputSink; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const ROWS: usize = 65_536; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +#[derive(Clone)] +struct RowWrappingAdd; + +impl RowFn for RowWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_wrapping_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64, i64), ElementSink, _, _>( + |_| (), + |&(), (lhs, rhs), output| *output = lhs.wrapping_add(rhs), + ) + } +} + +#[derive(Clone)] +struct RowCheckedAdd; + +struct CheckedAddSink { + values: BufferMut, + row_count: usize, +} + +struct CheckedAddRows<'a> { + values: &'a mut [MaybeUninit], +} + +struct CheckedAddRow<'a> { + value: &'a mut MaybeUninit, +} + +impl CheckedAddRow<'_> { + fn write(self, lhs: i64, rhs: i64) -> bool { + let value = lhs.wrapping_add(rhs); + let error = (lhs ^ value) & (rhs ^ value); + self.value.write(value); + error < 0 + } +} + +impl OutputSink for CheckedAddSink { + const ERRORS_ARE_DEFERRED: bool = true; + + type Rows<'a> = CheckedAddRows<'a>; + type Row<'a> = CheckedAddRow<'a>; + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self { + values: BufferMut::with_capacity(rows), + row_count: rows, + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + CheckedAddRows { + values: &mut self.values.spare_capacity_mut()[..self.row_count], + } + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.values.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + CheckedAddRow { + value: &mut rows.values[index], + } + } + + fn finish(mut self, error: DeferredError) -> VortexResult { + if error.occurred() { + return Err(vortex_err!("integer overflow in row checked add")); + } + + // SAFETY: dense execution writes every slot before `finish` is called. This sink does not + // support branch-and-skip, and filtered execution allocates exactly one slot per valid row. + unsafe { self.values.set_len(self.row_count) }; + Ok(PrimitiveArray::new(self.values.freeze(), Validity::NonNullable).into_array()) + } +} + +impl RowFn for RowCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_checked_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64, i64), CheckedAddSink, _, _>( + |_| (), + |&(), (lhs, rhs), output| output.write(lhs, rhs), + ) + } +} + +struct I64Sink(BufferMut); + +impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + fn finish(self, _error: DeferredError) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +#[derive(Clone)] +struct RowSinkWrappingAdd; + +impl RowFn for RowSinkWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_wrapping_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i64, i64), I64Sink, _, _>( + |_| (), + |&(), (lhs, rhs), out| { + *out = lhs.wrapping_add(rhs); + }, + ) + } +} + +fn inputs() -> (ArrayRef, ArrayRef) { + let lhs = (0..ROWS) + .map(|index| index as i64) + .collect::>() + .into_array(); + let rhs = (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>() + .into_array(); + (lhs, rhs) +} + +fn constant_inputs() -> (ArrayRef, ArrayRef) { + let (lhs, _) = inputs(); + let rhs = ConstantArray::new(Scalar::from(7i64), ROWS).into_array(); + (lhs, rhs) +} + +fn nullable_inputs() -> (ArrayRef, ArrayRef) { + let lhs = PrimitiveArray::new( + (0..ROWS).map(|index| index as i64).collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(5))), + ) + .into_array(); + let rhs = PrimitiveArray::new( + (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(7))), + ) + .into_array(); + (lhs, rhs) +} + +fn bench_row_fn(bencher: Bencher, row_fn: F) +where + F: RowFn, +{ + bencher + .with_inputs(inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + row_fn + .clone() + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn specialized_checked_add(bencher: Bencher) { + bencher + .with_inputs(inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + lhs.binary(rhs, Operator::Add) + .unwrap() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd); +} + +#[divan::bench] +fn row_sink_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd); +} + +#[divan::bench] +fn handrolled_sink_wrapping_add(bencher: Bencher) { + bencher + .with_inputs(inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = lhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let rhs = rhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let mut output = BufferMut::zeroed(ROWS); + for ((out, lhs), rhs) in output + .as_mut_slice() + .iter_mut() + .zip(lhs.as_slice()) + .zip(rhs.as_slice()) + { + *out = lhs.wrapping_add(*rhs); + } + PrimitiveArray::new(output.freeze(), Validity::NonNullable).into_array() + }); +} + +#[divan::bench] +fn row_checked_add(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd); +} + +#[divan::bench] +fn specialized_checked_add_constant(bencher: Bencher) { + bencher + .with_inputs(constant_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + lhs.binary(rhs, Operator::Add) + .unwrap() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_wrapping_add_constant(bencher: Bencher) { + bencher + .with_inputs(constant_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + RowWrappingAdd + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_checked_add_constant(bencher: Bencher) { + bencher + .with_inputs(constant_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + RowCheckedAdd + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn specialized_checked_add_nullable(bencher: Bencher) { + bencher + .with_inputs(nullable_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + lhs.binary(rhs, Operator::Add) + .unwrap() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_checked_add_nullable(bencher: Bencher) { + bencher + .with_inputs(nullable_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + RowCheckedAdd + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_wrapping_add_nullable(bencher: Bencher) { + bencher + .with_inputs(nullable_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + RowWrappingAdd + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs new file mode 100644 index 00000000000..c4e7ca77d8e --- /dev/null +++ b/vortex-array/benches/strict_validity.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks how the row lifting applies validity to a dense kernel. +//! +//! Both arms run the same kernel over the same nullable input and differ only in how the output's +//! nulls are restored: +//! +//! - `lazy` is what the lifting behind [`RowFn`] does: conjoin the input validities (itself lazy) +//! and hand the resulting boolean array straight to `mask`. +//! - `eager` is the strategy it replaced: materialize the conjunction into a `Mask`, copy it into a +//! `BitBuffer`, wrap that in a `BoolArray`, and mask with it. +//! +//! `chain` composes three of the same function, which is where a per-call materialization compounds. + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ElementSink; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +const SIZES: &[usize] = &[65536, 1 << 20]; + +/// The shared kernel: double every lane, ignoring validity. +fn doubled(input: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let input = input.clone().execute::(ctx)?; + let values: Buffer = input + .as_slice::() + .iter() + .map(|value| value.wrapping_mul(2)) + .collect(); + Ok(PrimitiveArray::new(values, Validity::NonNullable).into_array()) +} + +/// Dense row function: the lifting decides how validity is applied. +/// +/// Its [`reduce_encoded`](RowFn::reduce_encoded) answers with [`doubled`] before the row loop is +/// reached, so this arm runs exactly the kernel [`EagerDouble`] runs and the two differ only in +/// validity. The row closure below is what makes it a [`RowFn`] at all, and never executes. +#[derive(Clone)] +struct LazyDouble; + +impl RowFn for LazyDouble { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.lazy_double"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(i32,), ElementSink, _, _>( + |_| (), + |&(), (value,), output| *output = value.wrapping_mul(2), + ) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + doubled(&args[0], ctx).map(Some) + } +} + +/// The same function, applying validity the way the adapter used to: materialize a mask first. +#[derive(Clone)] +struct EagerDouble; + +impl ScalarFnVTable for EagerDouble { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.eager_double"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { + ChildName::from("input") + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + Ok(DType::Primitive(PType::I32, args[0].nullability())) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let valid = input.validity()?.execute_mask(args.row_count(), ctx)?; + let values = doubled(&input, ctx)?; + + if valid.all_true() { + return values.cast(DType::Primitive(PType::I32, input.dtype().nullability())); + } + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + values.mask(mask) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// A nullable i32 column with array-backed validity (~10% nulls). +fn nullable_input(len: usize) -> ArrayRef { + PrimitiveArray::new( + (0..len as i32).collect::>(), + Validity::from_iter((0..len).map(|i| !i.is_multiple_of(10))), + ) + .into_array() +} + +fn bench_depth(bencher: Bencher, scalar_fn: V, len: usize, depth: usize) +where + V: ScalarFnVTable + Clone, +{ + bencher + .with_inputs(|| nullable_input(len)) + .bench_local_values(|input| { + let mut ctx = SESSION.create_execution_ctx(); + let mut array = input; + for _ in 0..depth { + array = scalar_fn + .clone() + .try_new_array(len, EmptyOptions, [array]) + .unwrap() + .into_array(); + } + array.execute::(&mut ctx).unwrap() + }); +} + +#[divan::bench(args = SIZES)] +fn lazy(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn eager(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn lazy_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 3); +} + +#[divan::bench(args = SIZES)] +fn eager_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 3); +} From b324f3e266774eb8ae79b44a464c2240bc87bebf Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 13:03:43 +0000 Subject: [PATCH 003/160] Port the numeric binary kernels to RowFn Progress towards #9128. `vortex.numeric` becomes a `RowFn` over a primitive pair, taking `NumericOperator` as its options, which deletes the hand-written null propagation, constant folding, and validity logic in `numeric/primitive.rs`. Checked arithmetic reports overflow as evidence the row closure returns, rather than as a comparison the caller re-derives. Decimal keeps its own columnar implementation. `PrimitiveOperand` was defined in `numeric/primitive.rs` and shared out of `numeric/mod.rs`. The port drops its numeric caller, so it moves into `compare/primitive.rs`, its only remaining user. `NumericOperator` gains `Hash`, which a `RowFn`'s options require. `map_checked_into` in `vortex-compute` loses its last caller with the port and is deleted. Also adds a `list_length` test pinning that a non-nullable fixed-size list keeps a constant result rather than materializing one `u64` per row, which is the reason `vortex.list.length` stays on `ScalarFnVTable`. Signed-off-by: Connor Tsui Co-authored-by: Claude --- .../typed_view/primitive/numeric_operator.rs | 2 +- .../scalar_fn/fns/binary/compare/primitive.rs | 69 +++- .../scalar_fn/fns/binary/numeric/checked.rs | 93 +---- .../src/scalar_fn/fns/binary/numeric/mod.rs | 8 +- .../scalar_fn/fns/binary/numeric/primitive.rs | 389 ++++-------------- .../src/scalar_fn/fns/binary/numeric/row.rs | 269 ++++++++++++ .../src/scalar_fn/fns/binary/numeric/tests.rs | 8 +- vortex-array/src/scalar_fn/fns/list_length.rs | 19 + vortex-compute/src/lane_kernels/map_into.rs | 73 ---- 9 files changed, 461 insertions(+), 469 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/numeric/row.rs diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index 6b389a0554b..d7bcc8d0b4c 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_err; use crate::scalar_fn::fns::operators::Operator; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. pub enum NumericOperator { /// Binary element-wise addition of two arrays or of two scalars. diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 3bfb11a266e..58fb22ed73e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -4,6 +4,7 @@ //! Native comparison of primitive arrays via bit-packing lane kernels. use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -11,18 +12,20 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::BoolArray; +use crate::arrays::Constant; use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_native_ptype; use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::PrimitiveOperand; use crate::scalar_fn::fns::binary::compare::collect_bits; use crate::scalar_fn::fns::binary::compare::collect_zip_bits; use crate::scalar_fn::fns::binary::compare::compare_validity; use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; /// Compare two primitive arrays of the same [`PType`]. /// @@ -128,3 +131,67 @@ fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), } } + +/// A primitive binary-operator operand: a materialized buffer, a non-null constant, or an +/// all-null constant. +/// +/// Splitting the constant out of the buffer is what lets the lane kernels above hoist it into a +/// register instead of reading it back per lane. +enum PrimitiveOperand { + /// A decoded column, one value per row. + Array { + values: Buffer, + validity: Validity, + }, + + /// The same non-null value in every row. + Constant { + value: T, + len: usize, + validity: Validity, + }, + + /// A null in every row, carrying only the row count. + Null(usize), +} + +impl PrimitiveOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + return Ok( + match constant.scalar().as_primitive().try_typed_value::()? { + Some(value) => Self::Constant { + value, + len: array.len(), + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }, + None => Self::Null(array.len()), + }, + ); + } + + let array = array.clone().execute::(ctx)?; + let validity = array.validity()?; + let values = array.into_buffer::(); + Ok(Self::Array { values, validity }) + } + + fn len(&self) -> usize { + match self { + Self::Array { values, .. } => values.len(), + Self::Constant { len, .. } | Self::Null(len) => *len, + } + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } => validity.clone(), + Self::Constant { validity, .. } => validity.clone(), + Self::Null(_) => Validity::AllInvalid, + } + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index afb709abe1c..47c7d1351b9 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Checked-lane execution for numeric kernels, driven by the shared +//! Checked-lane execution for the decimal kernels, driven by the shared //! `vortex-compute` lane kernels. - -use std::ops::BitOrAssign; +//! +//! The primitive widths do not come through here: they are computed one row at a time by +//! [`row`](super::row), which writes a value for every row and reduces failure as one bit rather +//! than scanning lanes. use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -13,34 +15,22 @@ use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_mask::AllOr; use vortex_mask::Mask; -/// Evidence that a lane failed, anything other than [`Default`] meaning failure. -/// -/// `bool` is the ordinary choice; [`map_checked_into`] explains why the wider members exist and -/// asserts the width bound that membership here does **not** imply. -/// -/// [`map_checked_into`]: IndexedSourceExt::map_checked_into -pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - -impl Failure for bool {} -impl Failure for u8 {} -impl Failure for u16 {} -impl Failure for u32 {} -impl Failure for u64 {} - /// Apply the fallible `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it returns `None` on a _valid_ lane. +/// `Err(first_failing_valid_lane)` only when it returns `None` on a valid lane. /// /// `apply` also runs on invalid lanes, whose failures are masked out and whose values are /// unspecified, so it must be total: no panics or side effects on any stored lane value. /// -/// This drives the one-pass early-exit kernels, which abort at the end of the enclosing 64-lane -/// chunk. Use it when per-lane failure handling is cheap relative to the operation, as in integer -/// division. Prefer [`checked_apply_lanes`] when the failure check itself vectorizes. +/// This drives the one-pass early-exit kernels: failures abort at the end of the enclosing +/// 64-lane chunk. It suits an operation whose per-lane failure handling is cheap relative to the +/// operation itself, which is what the decimal kernels and their per-lane casts are. /// -/// `#[inline]`: this must inline into the caller that builds the closure, so that a captured -/// constant operand flattens into a register rather than living behind a pointer the loop reloads -/// on every lane, which blocks vectorization. -#[inline] +/// `#[inline(always)]`: this wrapper and its kernel calls must inline into the caller that +/// constructs the closure, so the closure environment (e.g. a captured constant operand) +/// flattens into registers. Left to its own devices under `codegen-units > 1`, the compiler +/// keeps the environment behind a pointer, and reloading a captured constant on every lane +/// blocks vectorization of the whole loop. +#[inline(always)] pub(super) fn checked_lanes( source: S, valid_rows: &Mask, @@ -61,7 +51,6 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, @@ -73,55 +62,3 @@ where Ok(values.freeze()) } - -/// Apply the split value/failure `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it flags a _valid_ lane. -/// -/// The hot pass writes every value unconditionally and OR-reduces one piece of evidence, leaving -/// the loop free of per-lane selects and per-chunk exit branches. Only if a lane flagged does a -/// cold second pass re-run `apply` through the early-exit kernels, which drop null-lane failures -/// and attribute the first valid one. -/// -/// Like [`checked_lanes`], `apply` runs on invalid lanes and must be total. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -pub(super) fn checked_apply_lanes( - source: S, - valid_rows: &Mask, - mut apply: Apply, -) -> Result, usize> -where - S: IndexedSource + Copy, - T: Copy + Default, - Fail: Failure, - Apply: FnMut(S::Item) -> (T, Fail), -{ - let len = source.len(); - debug_assert_eq!(len, valid_rows.len()); - - let valid_bits = match valid_rows.bit_buffer() { - AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), - AllOr::Some(valid_bits) => Some(valid_bits), - }; - - let mut values = BufferMut::::with_capacity(len); - - let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &mut apply) != Fail::default() { - let mut checked = |item: S::Item| { - let (value, failure) = apply(item); - (failure == Fail::default()).then_some(value) - }; - match valid_bits { - None => source.try_map_into(out, &mut checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?, - } - } - - // SAFETY: the kernels initialize every lane in `out`. - unsafe { values.set_len(len) }; - - Ok(values.freeze()) -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 6622e08f82b..21db5dc8c8e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,17 +4,21 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! +//! The primitive widths are computed by a [`RowFn`](crate::scalar_fn::RowFn), which owns null +//! handling, constants and validity for them; see [`row`]. Decimal keeps its own columnar +//! implementation in [`decimal`]. +//! //! [`Binary`]: super::Binary mod checked; mod decimal; mod primitive; +mod row; #[cfg(test)] mod tests; use decimal::execute_numeric_decimal; -pub(crate) use primitive::PrimitiveOperand; -use primitive::execute_numeric_primitive; +use row::execute_numeric_primitive; use vortex_error::VortexResult; use vortex_error::vortex_ensure; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 357547f25b8..6de544eaaa2 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -1,65 +1,56 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::checked::Failure; -use super::checked::checked_apply_lanes; -use super::checked::checked_lanes; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; +//! The checked arithmetic one row of a primitive column is computed with. +//! +//! Each operator is a type implementing [`CheckedPrimitiveOp`] at every native width, and each +//! width implements [`CheckedArithmetic`] with the value and failure evidence written separately. +//! Keeping them apart is what lets [`row`](super::row) write a value for every row and reduce the +//! evidence without a branch, so the loop vectorizes. + use crate::dtype::NativePType; -use crate::dtype::PType; use crate::dtype::half::f16; -use crate::match_each_native_ptype; -use crate::scalar::NumericOperator; -use crate::scalar::Scalar; -use crate::validity::Validity; +use crate::scalar_fn::SinkResult; -struct CheckedAdd; +/// Checked addition, failing on integer overflow. +pub(super) struct CheckedAdd; -struct CheckedSub; +/// Checked subtraction, failing on integer overflow. +pub(super) struct CheckedSub; -struct CheckedMul; +/// Checked multiplication, failing on integer overflow. +pub(super) struct CheckedMul; -struct CheckedDiv; +/// Checked division, failing on integer division by zero and on `MIN / -1`. +pub(super) struct CheckedDiv; -trait CheckedPrimitiveOp: Sized { - /// Message for the error raised when this operation fails on a valid lane. +/// Evidence that some row failed, in a form that OR-reduces across the batch. +/// +/// A plain `bool` is the obvious choice and the right one for most operations. Unsigned +/// multiplication is the exception: deriving a `bool` from the widened product costs a comparison, +/// and LLVM rewrites that comparison plus the product into `llvm.umul.with.overflow`, which has no +/// vector form and scalarizes the whole loop. Carrying the discarded high half instead means the row +/// never compares, so the multiply stays a widening vector multiply and the reduction stays a +/// vector OR. **The width must not exceed the element's**, or the reduction becomes the loop's +/// bottleneck instead of the arithmetic. +pub(super) trait Failure: SinkResult + Copy + Default {} + +impl + Copy + Default> Failure for T {} + +/// One arithmetic operator at one width, as a value and its failure evidence. +/// +/// The pair rather than an `Option` is what a row can write unconditionally: the value is stored +/// whatever the evidence says, and a failing row is either masked away as null or turned into a +/// batch error before anything reads it. +pub(super) trait CheckedPrimitiveOp: 'static + Sized { + /// The error reported for a batch in which some valid row failed. const ERROR: &'static str; - /// Whether to check inside the value loop and exit early, rather than reducing evidence over - /// the whole batch and re-running only once some lane has flagged. - const CHECKED_VALUE_LOOP: bool = false; - - /// How this operation reports a failing lane. See [`Failure`]. + /// How this operation reports a failing row. See [`Failure`]. type Failure: Failure; - /// Compute a lane's value and its failure evidence together. - /// - /// Overflowing-style rather than `Option`, so the vectorizable kernels can write every - /// lane's value unconditionally and reduce the evidence separately. [`Self::checked`] is the - /// shape for the scalar and early-exit paths. + /// The result of this operation, paired with evidence of whether the row failed. fn apply(lhs: T, rhs: T) -> (T, Self::Failure); - - /// [`Self::apply`] folded into the `Option` that the early-exit kernels take. - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - let (value, failed) = Self::apply(lhs, rhs); - - (failed == Self::Failure::default()).then_some(value) - } } impl CheckedPrimitiveOp for CheckedAdd { @@ -97,12 +88,6 @@ impl CheckedPrimitiveOp for CheckedMul { impl CheckedPrimitiveOp for CheckedDiv { const ERROR: &'static str = "integer division by zero or overflow in checked div"; - // Integer division still lowers to scalar divides, so the split - // value/error-scan loop used to auto-vectorize add/sub/mul only adds a - // second full scan. Use the one-pass early-exit checked kernel for integer - // division, matching Arrow/Velox. Float division has no checked errors and - // stays on the split/vectorizable default path. - const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; type Failure = bool; @@ -116,207 +101,21 @@ impl CheckedPrimitiveOp for CheckedDiv { }; (value, failed) } - - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - lhs.div_checked(rhs) - } -} - -/// Execute a numeric operation between two primitive-typed arrays. -pub(super) fn execute_numeric_primitive( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - - match_each_native_ptype!(ptype, |T| { - match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), - } - }) -} - -fn execute_checked_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType, - Op: CheckedPrimitiveOp, - Scalar: From, - Scalar: From>, -{ - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - let len = lhs.len(); - debug_assert_eq!(len, rhs.len()); - - let validity = lhs.validity().and(rhs.validity())?; - let valid_rows = validity.execute_mask(len, ctx)?; - - let values = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_op_lanes::<_, T, Op>( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - &valid_rows, - |(lhs, rhs)| (lhs, rhs), - ), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Capture the constant by value so it stays hoisted out of the lane loop. - let rhs = *rhs; - checked_op_lanes::<_, T, Op>(lhs.as_slice(), &valid_rows, move |lhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let lhs = *lhs; - checked_op_lanes::<_, T, Op>(rhs.as_slice(), &valid_rows, move |rhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let value = Op::checked(*lhs, *rhs) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - return Ok(constant_result_array(value, len, &result_dtype)); - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => Ok(Buffer::zeroed(len)), - } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - - primitive_result_array::(values, validity, &result_dtype) } -/// Run `Op` over the lanes of `source` in the loop shape it declares: the one-pass early-exit -/// kernel when `Op::CHECKED_VALUE_LOOP` is set, and the split value/failure kernel otherwise. +/// The per-width arithmetic behind [`CheckedPrimitiveOp`], with each operation split into the value +/// it produces and whether producing it failed. /// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -fn checked_op_lanes( - source: S, - valid_rows: &Mask, - mut to_operands: impl FnMut(S::Item) -> (T, T), -) -> Result, usize> -where - S: IndexedSource + Copy, - T: NativePType, - Op: CheckedPrimitiveOp, -{ - if Op::CHECKED_VALUE_LOOP { - checked_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::checked(lhs, rhs) - }) - } else { - checked_apply_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::apply(lhs, rhs) - }) - } -} - -fn primitive_result_array( - values: Buffer, - validity: Validity, - dtype: &DType, -) -> VortexResult { - let array = PrimitiveArray::new(values, validity).into_array(); - if array.dtype() == dtype { - return Ok(array); - } - array.cast(dtype.clone()) -} - -fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef -where - T: NativePType, - Scalar: From + From>, -{ - if dtype.is_nullable() { - ConstantArray::new(Some(value), len).into_array() - } else { - ConstantArray::new(value, len).into_array() - } -} - -/// A primitive binary-operator operand: a materialized buffer, a non-null constant, or an -/// all-null constant. -pub(crate) enum PrimitiveOperand { - Array { - values: Buffer, - validity: Validity, - }, - Constant { - value: T, - len: usize, - validity: Validity, - }, - Null(usize), -} - -impl PrimitiveOperand { - pub(crate) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - if let Some(constant) = array.as_opt::() { - return Ok( - match constant.scalar().as_primitive().try_typed_value::()? { - Some(value) => Self::Constant { - value, - len: array.len(), - validity: if constant.scalar().dtype().is_nullable() { - Validity::AllValid - } else { - Validity::NonNullable - }, - }, - None => Self::Null(array.len()), - }, - ); - } - - let array = array.clone().execute::(ctx)?; - let validity = array.validity()?; - let values = array.into_buffer::(); - Ok(Self::Array { values, validity }) - } - - pub(crate) fn len(&self) -> usize { - match self { - Self::Array { values, .. } => values.len(), - Self::Constant { len, .. } | Self::Null(len) => *len, - } - } - - pub(crate) fn validity(&self) -> Validity { - match self { - Self::Array { validity, .. } => validity.clone(), - Self::Constant { validity, .. } => validity.clone(), - Self::Null(_) => Validity::AllInvalid, - } - } -} - -trait CheckedArithmetic: NativePType { - const DIV_CHECKS_IN_VALUE_LOOP: bool; - - /// How multiplication reports a failing lane, which is width-dependent in a way the other - /// operations are not. `impl_checked_unsigned!` and `impl_checked_signed!` say why each family - /// reports what it reports. +/// Every `_value` method **must** be total: it is called for rows behind nulls, whose operands are +/// arbitrary, so it may not panic or trap. Integer division is the one that needs care, and +/// [`CheckedDiv`] supplies the default instead of dividing when the divisor is rejected. +pub(super) trait CheckedArithmetic: NativePType { + /// How multiplication reports a failing row. + /// + /// `Self` for the unsigned widths that have a widening multiply, so the row can hand back the + /// discarded high half rather than comparing. `bool` everywhere else: the narrow signed widths + /// already vectorize through a two-sided range check, floats never overflow, and the 64-bit + /// widths use a full-width evidence word. type MulFailure: Failure; fn add_value(self, rhs: Self) -> Self; @@ -327,16 +126,10 @@ trait CheckedArithmetic: NativePType { fn mul_failure(self, rhs: Self) -> Self::MulFailure; fn div_value(self, rhs: Self) -> Self; fn div_error(self, rhs: Self) -> bool; - fn div_checked(self, rhs: Self) -> Option; } /// The integer arithmetic every width shares, given the two things that actually differ between -/// them: how multiplication reports a failing lane, and how add/sub/div detect one. -/// -/// Only `mul_failure` genuinely varies per width, so it is the macro's parameter and everything -/// else is written once. The `$mul_failure_ty` and body are supplied by the caller because the -/// choice between a widened high half and a `bool` is exactly the vectorization decision described -/// on [`Failure`]. +/// them: how multiplication reports a failing row, and how add/sub/div detect one. macro_rules! impl_checked_integer { ( $ty:ty, @@ -347,8 +140,6 @@ macro_rules! impl_checked_integer { = |$mf_lhs:ident, $mf_rhs:ident| $mul_failure:expr, ) => { impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = true; - type MulFailure = $mul_failure_ty; #[inline(always)] @@ -395,19 +186,12 @@ macro_rules! impl_checked_integer { let ($div_lhs, $div_rhs) = (self, rhs); $div_error } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - self.checked_div(rhs) - } } }; } -/// The unsigned widths. `add`, `sub` and `div` are written once here, and `widening_mul` covers -/// every width including the 64-bit one, which widens into `u128`: the discarded high half of the -/// widened product is the failure evidence, and costs none of the comparison LLVM folds into -/// `umul.with.overflow`. +/// The unsigned widths. The discarded high half of the widened product is the failure evidence, +/// and costs none of the comparison LLVM folds into `umul.with.overflow`. macro_rules! impl_checked_unsigned { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_integer!( @@ -420,12 +204,8 @@ macro_rules! impl_checked_unsigned { }; } -/// The signed widths, on the same principle. `widening_mul` is the shorthand for the narrow widths, -/// whose two-sided range check over a wider product is not the shape LLVM folds into an overflow -/// intrinsic, so they vectorize while reporting a plain `bool`. The 64-bit width cannot: deriving a -/// `bool` there costs the comparison that scalarizes the loop, so `high_half_mul` hands back the -/// discarded high half as a word. Both arms derive every shift and bound from `$ty`, so -/// instantiating one at a new width cannot silently keep another width's constants. +/// The signed widths. The narrow widths report a two-sided range check as `bool`; the 64-bit width +/// reports the discarded high half as a word so deriving the evidence does not scalarize the loop. macro_rules! impl_checked_signed { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_signed!($ty, mul_failure: bool = |lhs, rhs| { @@ -433,9 +213,6 @@ macro_rules! impl_checked_signed { product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide }); }; - // Zero exactly when the product fits: a signed multiply overflows iff the high half of the true - // product differs from the sign extension of the half that was kept, so the two XOR to zero on - // the lanes that fit. `tests::test_multiply_overflow_boundaries` pins the boundaries. ($ty:ty, high_half_mul: $wide:ty => $failure:ty) => { impl_checked_signed!($ty, mul_failure: #[expect( clippy::cast_possible_truncation, @@ -451,7 +228,7 @@ macro_rules! impl_checked_signed { ( $ty:ty, mul_failure: $(#[$mul_failure_attr:meta])* $mul_failure_ty:ty - = |$l:ident, $r:ident| $mul_failure:expr + = |$lhs:ident, $rhs:ident| $mul_failure:expr ) => { impl_checked_integer!( $ty, @@ -464,7 +241,7 @@ macro_rules! impl_checked_signed { ((lhs ^ rhs) & (lhs ^ value)) < 0 }, div_error: |lhs, rhs| rhs == 0 || (lhs == <$ty>::MIN && rhs == -1), - mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$l, $r| $mul_failure, + mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$lhs, $rhs| $mul_failure, ); }; } @@ -473,8 +250,6 @@ macro_rules! impl_checked_float { ($($ty:ty),+ $(,)?) => { $( impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = false; - type MulFailure = bool; #[inline(always)] @@ -516,11 +291,6 @@ macro_rules! impl_checked_float { fn div_error(self, _rhs: Self) -> bool { false } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - Some(self / rhs) - } } )+ }; @@ -539,34 +309,34 @@ impl_checked_float!(f16, f32, f64); #[cfg(test)] mod tests { use super::CheckedArithmetic; + use crate::scalar_fn::SinkResult; /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change /// pivots, and a spread of magnitudes that straddles the 64-bit split. const PROBES: &[i64] = &[ - 0, // - 1, // - -1, // - 2, // - -2, // - 3, // - i64::MIN, // - i64::MIN + 1, // - i64::MAX, // - i64::MAX - 1, // - 1 << 31, // - 1 << 32, // - 1 << 62, // - -(1 << 62), // - 0x7FFF_FFFF, // - -0x8000_0000, // + 0, + 1, + -1, + 2, + -2, + 3, + i64::MIN, + i64::MIN + 1, + i64::MAX, + i64::MAX - 1, + 1 << 31, + 1 << 32, + 1 << 62, + -(1 << 62), + 0x7FFF_FFFF, + -0x8000_0000, ]; - /// Every `mul_failure` impl is either a bit trick or a two-sided range check, so hold each - /// against the obvious reference: `reference` is the width's own `checked_mul`, whose `None` - /// _is_ the definition of overflow. + /// Every `mul_failure` implementation is either a bit trick or a two-sided range check, so + /// hold each against `checked_mul`, whose `None` is the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { - let failed = lhs.mul_failure(rhs) != ::default(); + let failed = ::occurred(lhs.mul_failure(rhs)); assert_eq!(failed, reference.is_none(), "{lhs:?} * {rhs:?}"); } @@ -578,14 +348,13 @@ mod tests { assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); let (lhs, rhs) = (lhs as u64, rhs as u64); - assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); } } } - /// The 8-bit widths are cheap enough to check exhaustively, which pins the shift of the - /// unsigned formula and the range check of the signed one against every product that exists. + /// The 8-bit widths are cheap enough to check exhaustively, pinning the unsigned shift and the + /// signed range check against every product that exists. #[test] fn mul_failure_agrees_with_checked_mul_exhaustively_at_8_bits() { for lhs in u8::MIN..=u8::MAX { diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs new file mode 100644 index 00000000000..aab36bf196b --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -0,0 +1,269 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The primitive arithmetic operators as a [`RowFn`]. +//! +//! [`Binary`] keeps its ID, its options serialization, and its strictness, fallibility and validity +//! contracts, and delegates only the _execution_ of `Add`, `Sub`, `Mul` and `Div` over primitive +//! columns to [`NumericBinary`]. Delegation rather than conversion is what makes the port possible +//! at all: `Binary` also covers Kleene `And`/`Or`, which are not strict, and the six comparisons, +//! which are infallible, so no single [`RowFn`] can stand in for the whole function. +//! +//! [`NumericBinary`] is not registered and appears in no serialized expression. It is reached only +//! through the [`ScalarFnVTable::execute`] that the blanket [`RowFn`] implementation provides, so +//! it needs no rewrite rule, no ID in the registry, and no wire format of its own. +//! +//! Everything the previous hand-written implementation did around the arithmetic itself now comes +//! from the lifting: input decoding, the constant operand collapse, the all-constant fold, the +//! null-constant short circuit, output allocation, nullability widening, and masking. What is left +//! here is the per-type checked operation and the sink that carries its overflow bit. +//! +//! [`Binary`]: crate::scalar_fn::fns::binary::Binary + +use std::marker::PhantomData; +use std::mem::MaybeUninit; + +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::primitive::CheckedAdd; +use super::primitive::CheckedDiv; +use super::primitive::CheckedMul; +use super::primitive::CheckedPrimitiveOp; +use super::primitive::CheckedSub; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::NumericOperator; +use crate::scalar_fn::DeferredError; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::RowVisitor; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::validity::Validity; + +/// Execute a numeric operation between two primitive-typed arrays. +/// +/// The caller has already established that both operands are primitive, of the same type, and of +/// the same length. +pub(super) fn execute_numeric_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + ScalarFnVTable::execute(&NumericBinary, &op, &args, ctx) +} + +/// The four arithmetic operators of [`Binary`] over primitive columns, as one row function per +/// operator and width. +/// +/// [`Binary`]: crate::scalar_fn::fns::binary::Binary +#[derive(Clone)] +struct NumericBinary; + +impl RowFn for NumericBinary { + type Options = NumericOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + /// Only the integer widths can overflow, and only integer division can divide by zero, but + /// fallibility is declared without input dtypes. The float widths are therefore covered by the + /// same `true`, which costs them nothing: a deferred error keeps the batch on the dense path. + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.numeric_binary"); + *ID + } + + fn dispatch( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = operand_ptype(args)?; + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_checked::(visitor), + } + }) + } +} + +/// The width both operands are read at. +/// +/// Only the left operand is inspected. `(T, T)` validates each argument against the chosen width, +/// so a right operand of a different type is rejected by the visit rather than here. +fn operand_ptype(args: &[DType]) -> VortexResult { + let lhs = args + .first() + .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?; + + PType::try_from(lhs) +} + +/// Visit at two `T` columns, applying `Op` per row into the sink that defers its overflow bit. +/// +/// The const block enforces, at monomorphization time, the width rule stated on +/// [`Failure`](super::primitive::Failure): evidence wider than the element would make the +/// OR-reduction rather than the arithmetic decide how many rows fit in a vector. +fn visit_checked(visitor: V) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + V: RowVisitor, +{ + const { + assert!( + size_of::() <= size_of::(), + "failure evidence must be no wider than the value, or it bounds the vector width" + ) + }; + + visitor.visit_prepared_into::<(T, T), CheckedSink, _, _>( + |_| (), + |&(), (lhs, rhs), output| output.write(lhs, rhs), + ) +} + +/// The output column of one checked arithmetic batch, reporting failure once after the row loop. +/// +/// Deferring the failure is what keeps a fallible kernel on the dense path: every row writes a +/// value unconditionally and OR-reduces its failure evidence, so the loop holds no branch and no +/// `Result` discriminant. The lifting retries a nullable batch over only its valid rows if that +/// reduction is non-zero, which is what makes an overflow behind a null row invisible. +/// +/// The reduction lives in the sink rather than in the row closure's return type so that its width +/// is [`Op::Failure`](CheckedPrimitiveOp::Failure), the operator's choice, rather than one bit. That +/// is what lets unsigned multiplication report its discarded high half instead of a comparison, and +/// so stay vectorized. +/// +/// **The storage is deliberately uninitialized, not zeroed.** Substituting `BufferMut::zeroed` to +/// make the sink safe was measured at **1.65 to 1.71x** the cost of allocate-and-fill, stable across +/// two runs and every batch size from 8 KiB to 2 MiB, because `alloc_zeroed` does not avoid the +/// write: below glibc's mmap threshold `calloc` recycles a dirty chunk and memsets it, and above it +/// the first touch of each fresh page faults instead. The row loop overwrites every slot regardless, +/// so that pass is pure duplicate work on the hottest kernel in the system. This is the case the +/// repository's "avoid `unsafe` unless it is necessary" rule leaves room for: the safe spelling +/// exists, and it costs a second pass over the output. +/// +/// Rows are written into uninitialized storage, so this sink cannot finish a batch whose rows were +/// not all visited, and leaves [`OutputSink::SUPPORTS_SKIPPED_ROWS`] at `false`. Nothing is lost: +/// `SUPPORTS_SKIPPED_ROWS` is what makes branch-and-skip unavailable, which is the guard that keeps +/// the uninitialized slots sound. Note this is _not_ implied by the dispatch policy alone: a +/// deferred result still reaches the executor's valid-only policy whenever its arguments are not +/// dense-safe, so the `false` here is load-bearing rather than a restatement. +struct CheckedSink> { + /// The result values, initialized one row at a time up to `row_count`. + values: BufferMut, + + /// The batch length, which is the capacity `values` was allocated with. + row_count: usize, + + /// The operation applied to every row, which names the error reported by + /// [`finish`](OutputSink::finish). + op: PhantomData, +} + +/// The uninitialized output slots of a [`CheckedSink`], borrowed once for the row loop. +struct CheckedRows<'a, T: NativePType, Op: CheckedPrimitiveOp> { + values: &'a mut [MaybeUninit], + op: PhantomData, +} + +/// One output slot of a [`CheckedSink`]. +struct CheckedRow<'a, T: NativePType, Op: CheckedPrimitiveOp> { + value: &'a mut MaybeUninit, + op: PhantomData, +} + +impl> CheckedRow<'_, T, Op> { + /// Apply `Op` to one row, writing its value and handing back its failure evidence. + /// + /// The value is written whether or not the operation failed, since a failing row is either + /// masked away as null or turned into a batch error before it can be read. The evidence is + /// returned rather than reduced here so the executor can keep the reduction in a register, and + /// it is `Op`'s own width so the row never has to compare. + fn write(self, lhs: T, rhs: T) -> Op::Failure { + let (value, failure) = Op::apply(lhs, rhs); + self.value.write(value); + + failure + } +} + +impl> OutputSink for CheckedSink { + const ERRORS_ARE_DEFERRED: bool = true; + + type Rows<'a> + = CheckedRows<'a, T, Op> + where + Self: 'a; + type Row<'a> + = CheckedRow<'a, T, Op> + where + Self: 'a; + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::Primitive(T::PTYPE, Nullability::NonNullable)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self { + values: BufferMut::with_capacity(rows), + row_count: rows, + op: PhantomData, + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + let row_count = self.row_count; + CheckedRows { + values: &mut self.values.spare_capacity_mut()[..row_count], + op: PhantomData, + } + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.values.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + CheckedRow { + value: &mut rows.values[index], + op: PhantomData, + } + } + + fn finish(mut self, error: DeferredError) -> VortexResult { + if error.occurred() { + return Err(vortex_err!(InvalidArgument: "{}", Op::ERROR)); + } + + // SAFETY: the sink reports `SUPPORTS_SKIPPED_ROWS = false`, so every path that reaches + // `finish` without an error has written all `row_count` slots: dense execution visits + // `0..row_count`, and the valid-row retry runs densely over a sink allocated for exactly + // the filtered rows. + unsafe { self.values.set_len(self.row_count) }; + + Ok(PrimitiveArray::new(self.values.freeze(), Validity::NonNullable).into_array()) + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 7ee98ae717e..08d3e2daca9 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -297,13 +297,13 @@ fn test_multiply_overflow_on_null_lane_ignored( Ok(()) } -/// The hot pass OR-reduces evidence across whole 64-lane chunks before anything looks at it, so an -/// overflow in a late chunk must still be caught, and must still be suppressed when its lane is -/// null. Every other test here fits in a single chunk and cannot show either. +/// The hot pass OR-reduces evidence across the whole row loop before anything looks at it, so an +/// overflow late in the batch must still be caught, and must still be suppressed when its lane is +/// null. #[rstest] #[case::reported(true)] #[case::suppressed_by_null(false)] -fn test_multiply_overflow_survives_chunk_reduction( +fn test_multiply_overflow_survives_batch_reduction( #[case] lane_is_valid: bool, ) -> VortexResult<()> { const LEN: u32 = 1000; diff --git a/vortex-array/src/scalar_fn/fns/list_length.rs b/vortex-array/src/scalar_fn/fns/list_length.rs index 415a65416db..a971fe71141 100644 --- a/vortex-array/src/scalar_fn/fns/list_length.rs +++ b/vortex-array/src/scalar_fn/fns/list_length.rs @@ -350,6 +350,25 @@ mod tests { Ok(()) } + /// A non-nullable fixed-size list has one length for the whole column, so the result stays a + /// constant rather than materializing one `u64` per row. + #[test] + fn test_fixed_size_list_length_stays_constant() -> VortexResult<()> { + let fsl = create_fixed_size_list(Validity::NonNullable); + let mut ctx = array_session().create_execution_ctx(); + + let result = fsl + .apply(&list_length(root()))? + .execute::(&mut ctx)?; + + assert_eq!( + result.as_constant(), + Some(Scalar::primitive(2u64, Nullability::NonNullable)), + "expected a constant length column" + ); + Ok(()) + } + #[test] fn test_fixed_size_list_length_nullable() -> VortexResult<()> { let fsl = create_fixed_size_list(Validity::Array( diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs index c1e1107b1b9..258913e9fc7 100644 --- a/vortex-compute/src/lane_kernels/map_into.rs +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -5,7 +5,6 @@ //! caller-provided `&mut [MaybeUninit]`. use std::mem::MaybeUninit; -use std::ops::BitOrAssign; use vortex_buffer::BitBuffer; @@ -218,58 +217,6 @@ pub trait IndexedSourceExt: IndexedSource + Sized { } } - /// Split value/failure map with **no validity awareness at all**: write every lane's value - /// unconditionally and OR-reduce its failure evidence into the return. - /// - /// The fastest checked shape, running at the speed of the unchecked [`map_into`] in exchange - /// for reporting only _that_ some lane failed and never exiting early. Re-run the now known - /// cold input through [`try_map_into`] or [`try_map_masked_into`] to attribute the failure or - /// to drop the null-lane ones. The evidence reduces inside the kernel because a captured `&mut` - /// becomes a loop-carried memory dependence that blocks vectorization. - /// - /// Anything other than [`Default`] means failure, and `bool` is the ordinary `Fail`. Wider - /// words exist for operations where deriving a `bool` costs the vectorization it guards. - /// **`Fail` must be no wider than `R`**, asserted below, or the reduction rather than the - /// operation decides how many lanes fit in a vector. - /// - /// [`map_into`]: IndexedSourceExt::map_into - /// [`try_map_into`]: IndexedSourceExt::try_map_into - /// [`try_map_masked_into`]: IndexedSourceExt::try_map_masked_into - /// - /// # Panics - /// - /// Panics if `out.len() != self.len()`. - #[inline] - fn map_checked_into(self, out: &mut [MaybeUninit], mut apply: Apply) -> Fail - where - Fail: Copy + Default + BitOrAssign, - Apply: FnMut(Self::Item) -> (R, Fail), - { - const { - assert!( - size_of::() <= size_of::(), - "failure evidence must be no wider than the value, or it bounds the vector width" - ) - }; - - let values = self; - let len = values.len(); - assert_eq!(out.len(), len, "out must have the same length as values"); - - let mut failed = Fail::default(); - for idx in 0..len { - // SAFETY: idx < len by the loop bound, and out.len() == len. - let val = unsafe { values.get_unchecked(idx) }; - - let (result, failure) = apply(val); - failed |= failure; - - // SAFETY: idx < len == out.len(). - unsafe { out.get_unchecked_mut(idx).write(result) }; - } - failed - } - /// Fallible map with **no validity awareness at all** — every `None` returned /// by the closure is treated as a failure, even at null lanes. /// @@ -599,26 +546,6 @@ mod tests { assert!(res.is_ok(), "null lane should bypass the range check"); } - #[test] - fn map_checked_into_writes_all_lanes_and_reduces_flag() { - let mut values: Vec = (0..130).collect(); - let mut out = vec![MaybeUninit::::uninit(); 130]; - let failed = values - .as_slice() - .map_checked_into(&mut out, |v| (v as u32, v > u32::MAX as u64)); - assert!(!failed); - assert_eq!(write_t(out), (0..130u32).collect::>()); - - values[77] = (u32::MAX as u64) + 1; - let mut out = vec![MaybeUninit::::uninit(); 130]; - let failed = values - .as_slice() - .map_checked_into(&mut out, |v| (v as u32, v > u32::MAX as u64)); - assert!(failed); - // Failing lanes still write their (wrapped) value. - assert_eq!(write_t(out)[76], 76); - } - #[test] fn map_bits_into_packs_full_and_remainder_words() { let values: Vec = (0..130).collect(); From aebe3caf772948c4e3290dbcb1680e1d03150830 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 13:03:43 +0000 Subject: [PATCH 004/160] Port the tensor scalar functions to RowFn Progress towards #9128. `vortex.tensor.l2_norm`, `vortex.tensor.inner_product`, and `vortex.tensor.cosine_similarity` become row functions over a `TensorRow` element, which yields a slice of the extension array's storage straight to the closure. Lifting supplies the null propagation, constant folding, nullability, validity, and options serde that the three kernels each wrote by hand. Cosine similarity hoists the norm of a broadcast query vector into `visit_prepared_into`'s prepare step. The prepared and per-row arms must agree bit for bit, which only holds while both accumulate in the same order, so `l2_norm_row` moves into `utils.rs` and both call it. `BinaryTensorOpMetadata` and `build_tensor_array` move there too, shared by the two binary operators and by the normalized encoding. Constant folding through `try_build_constant_normalized` is now derived from the row closure, so the export is gone. The tests move out of the three kernel modules and into `scalar_fns/tests/`. Signed-off-by: Connor Tsui Co-authored-by: Claude --- vortex-tensor/benches/cosine_similarity.rs | 8 +- vortex-tensor/benches/inner_product.rs | 8 +- vortex-tensor/benches/l2_norm.rs | 6 +- .../src/encodings/normalized/execute.rs | 22 +- vortex-tensor/src/encodings/normalized/mod.rs | 1 - .../src/scalar_fns/cosine_similarity.rs | 918 +++++------------- vortex-tensor/src/scalar_fns/inner_product.rs | 490 ++-------- vortex-tensor/src/scalar_fns/l2_norm.rs | 380 +------- vortex-tensor/src/scalar_fns/mod.rs | 4 + vortex-tensor/src/scalar_fns/row.rs | 131 +++ .../src/scalar_fns/tests/cosine_similarity.rs | 586 +++++++++++ .../src/scalar_fns/tests/inner_product.rs | 253 +++++ vortex-tensor/src/scalar_fns/tests/l2_norm.rs | 295 ++++++ vortex-tensor/src/scalar_fns/tests/mod.rs | 9 + vortex-tensor/src/scalar_fns/tests/row.rs | 113 +++ vortex-tensor/src/utils.rs | 224 +++-- vortex-tensor/src/vector_search.rs | 4 +- 17 files changed, 1882 insertions(+), 1570 deletions(-) create mode 100644 vortex-tensor/src/scalar_fns/row.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/inner_product.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/l2_norm.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/mod.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/row.rs diff --git a/vortex-tensor/benches/cosine_similarity.rs b/vortex-tensor/benches/cosine_similarity.rs index 6cc5eb867ef..fef94a0aa91 100644 --- a/vortex-tensor/benches/cosine_similarity.rs +++ b/vortex-tensor/benches/cosine_similarity.rs @@ -22,10 +22,12 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity; @@ -85,9 +87,9 @@ fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { bencher .with_inputs(|| { ( - CosineSimilarity::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + CosineSimilarity + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) diff --git a/vortex-tensor/benches/inner_product.rs b/vortex-tensor/benches/inner_product.rs index 796e9b648d6..c0918f87ba2 100644 --- a/vortex-tensor/benches/inner_product.rs +++ b/vortex-tensor/benches/inner_product.rs @@ -20,6 +20,8 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::inner_product::InnerProduct; @@ -62,9 +64,9 @@ fn bench_inner_product(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { .counter(ItemsCount::new(lhs.len())) .with_inputs(|| { ( - InnerProduct::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + InnerProduct + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) diff --git a/vortex-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index 6f597084113..d96e4877af9 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -20,6 +20,8 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::l2_norm::L2Norm; @@ -60,7 +62,9 @@ fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { .counter(ItemsCount::new(input.len())) .with_inputs(|| { ( - L2Norm::try_new_array(input.clone()).unwrap().into_array(), + L2Norm + .try_new_array(input.len(), EmptyOptions, [input.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) diff --git a/vortex-tensor/src/encodings/normalized/execute.rs b/vortex-tensor/src/encodings/normalized/execute.rs index 637c8c07117..b96a03113b7 100644 --- a/vortex-tensor/src/encodings/normalized/execute.rs +++ b/vortex-tensor/src/encodings/normalized/execute.rs @@ -14,7 +14,6 @@ use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; @@ -24,6 +23,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::matcher::AnyTensor; +use crate::utils::build_tensor_array; use crate::utils::extract_flat_elements; use crate::utils::unit_norm_tolerance; @@ -115,26 +115,6 @@ fn denormalize_constant_norms( Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) } -/// Rebuilds a tensor-like extension array from flat primitive elements. -fn build_tensor_array( - dtype: DType, - tensor_flat_size: usize, - row_count: usize, - validity: Validity, - elements: Buffer, -) -> VortexResult { - let list_size = - u32::try_from(tensor_flat_size).vortex_expect("tensor flat size must fit into `u32`"); - - // SAFETY: Tensor elements are always non-nullable, so the validity carries no length. - let elements = unsafe { PrimitiveArray::new_unchecked(elements, Validity::NonNullable) }; - - let storage = - FixedSizeListArray::try_new(elements.into_array(), list_size, validity, row_count)?; - - Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) -} - /// Returns the flattened element count of each row of a tensor-like extension dtype. fn tensor_flat_size(dtype: &DType) -> usize { dtype diff --git a/vortex-tensor/src/encodings/normalized/mod.rs b/vortex-tensor/src/encodings/normalized/mod.rs index 545236bba7d..2c9a72d2988 100644 --- a/vortex-tensor/src/encodings/normalized/mod.rs +++ b/vortex-tensor/src/encodings/normalized/mod.rs @@ -31,7 +31,6 @@ pub use array::NormalizedSlots; mod compress; pub use compress::NormalizedScheme; pub use compress::normalize; -pub(crate) use compress::try_build_constant_normalized; mod execute; diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index ca8fcf0efd4..e786d49b47c 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -1,48 +1,49 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Cosine similarity expression for tensor-like types. +//! Cosine similarity between two tensor columns. +use num_traits::Float; use num_traits::Zero; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; +use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ElementSink; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::encodings::normalized::try_build_constant_normalized; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_norm::L2Norm; +use crate::scalar_fns::row::TensorRow; +#[cfg(test)] +use crate::scalar_fns::row::probe; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; +use crate::utils::l2_norm_row; /// Cosine similarity between two columns. /// /// Computes `dot(a, b) / (||a|| * ||b||)` over the flat backing buffer of each tensor or vector. /// The shape and permutation do not affect the result because cosine similarity only depends on the -/// element values, not their logical arrangement. +/// element values, not their logical arrangement. A zero norm on either side yields `0.0`. /// /// Both inputs must be tensor-like extension arrays ([`FixedShapeTensor`] or [`Vector`]) with the /// same dtype and a float element type. The output is a float column of the same float type. @@ -55,143 +56,79 @@ use crate::utils::validate_binary_tensor_float_inputs; /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct CosineSimilarity; -impl CosineSimilarity { - /// Creates a new [`TypedScalarFnInstance`] wrapping the cosine similarity operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(CosineSimilarity, EmptyOptions) - } - - /// Constructs a [`ScalarFnArray`] that lazily computes the cosine similarity between `lhs` and - /// `rhs`. - /// - /// # Errors - /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). - pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(CosineSimilarity::new().erased(), vec![lhs, rhs]) - } -} - -impl ScalarFnVTable for CosineSimilarity { +impl RowFn for CosineSimilarity { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.cosine_similarity"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("CosineSimilarity must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_prepared_into::<(TensorRow, TensorRow), ElementSink, _, _>( + |(lhs, rhs)| { + #[cfg(test)] + probe::record(lhs.is_some(), rhs.is_some()); + ConstNorms { + lhs: lhs.map(l2_norm_row), + rhs: rhs.map(l2_norm_row), + } + }, + |norms, (lhs, rhs), output| { + *output = cosine_similarity_row_prepared(norms, lhs, rhs); + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands make the *stored* norms and normalized children + /// authoritative: `cos(D(x, s), D(y, t)) = dot(x, y)` and `cos(D(x, s), y) = dot(x, y) / + /// ||y||`, in both cases forced to `0.0` on rows where any authoritative norm is `0.0` (even + /// for lossy children whose decoded coordinates are nonzero). + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, + args: &[ArrayRef], ctx: &mut ExecutionCtx, - ) -> VortexResult { - let mut lhs_ref = args.get(0)?; - let mut rhs_ref = args.get(1)?; - let len = args.row_count(); - - // If either side is a constant tensor-like extension array, eagerly normalize the single - // stored row and re-encode it as an `Normalized` whose children are both `ConstantArray`s. - // The `Normalized` fast path below then picks it up. - if let Some(normalized_array) = try_build_constant_normalized(&lhs_ref, len, ctx)? { - lhs_ref = normalized_array.into_array(); - } - if let Some(normalized_array) = try_build_constant_normalized(&rhs_ref, len, ctx)? { - rhs_ref = normalized_array.into_array(); - } + ) -> VortexResult> { + let lhs = args[0].clone(); + let rhs = args[1].clone(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { + match NormalizedOrientation::classify(&lhs, &rhs) { NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); + cosine_both_normalized(lhs, rhs, ctx).map(Some) } NormalizedOrientation::One { normalized_array, plain, - } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); - } - NormalizedOrientation::Neither => {} + } => cosine_one_normalized(normalized_array, plain, ctx).map(Some), + NormalizedOrientation::Neither => Ok(None), } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm::try_new_array(lhs_ref.clone())?; - let norm_rhs_arr = L2Norm::try_new_array(rhs_ref.clone())?; - let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; - - // Execute to get the inner product and norms of the arrays. We only fully decompress - // because we need to perform special logic (guard against 0) during division. - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - let norm_l: PrimitiveArray = norm_lhs_arr.into_array().execute(ctx)?; - let norm_r: PrimitiveArray = norm_rhs_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norm_l.as_slice::(); - let norms_r = norm_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - let denom = norms_l[i] * norms_r[i]; - - if denom == T::zero() { - T::zero() - } else { - dots[i] / denom - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false } } @@ -221,578 +158,177 @@ impl ScalarFnArrayVTable for CosineSimilarity { } } -impl CosineSimilarity { - /// Both sides are [`Normalized`]-encoded: treat the normalized children as authoritative, so - /// `cosine_similarity = dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - // `Normalized` makes the normalized children authoritative, so their dot product is the - // cosine similarity even for lossy storage wrappers, except that a zero stored norm still - // represents a zero vector. - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norms_l.as_slice::(); - let norms_r = norms_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if norms_l[i] == T::zero() || norms_r[i] == T::zero() { - T::zero() - } else { - dots[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: treat the normalized child as authoritative, so - /// `cosine_similarity = dot(n, b) / ||b||`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, normalized_norms) = extract_normalized_children(normalized_ref); - - let dot_arr = InnerProduct::try_new_array(normalized, plain_ref.clone())?; - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - - let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - - let norm_arr = L2Norm::try_new_array(plain_ref.clone())?; - let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let normalized_norms = normalized_norms.as_slice::(); - let plain_norms = plain_norm.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if normalized_norms[i] == T::zero() || plain_norms[i] == T::zero() { - T::zero() - } else { - dots[i] / plain_norms[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } +/// Per-batch state for the cosine row kernel: the L2 norm of each operand that is constant for +/// the batch. +/// +/// A broadcast query vector holds the same elements in every row, so its norm is the same in +/// every row too. Computing it in the prepare step hoists an `O(width)` pass and a `sqrt` per row +/// out of the row loop. `None` marks an operand that varies by row, whose norm the row closure +/// computes exactly as it did before the hoist. +struct ConstNorms { + /// The norm of the lhs when it is batch-constant. + lhs: Option, + + /// The norm of the rhs when it is batch-constant. + rhs: Option, } -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::cosine_similarity::CosineSimilarity; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::constant_tensor_array; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. - fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[test] - fn unit_vectors_1d() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 0.0, 1.0, 0.0, // Tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 1.0, 0.0, 0.0, // Tensor 2 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - /// Single-row cosine similarity for various vector pairs. - #[rstest] - // Antiparallel -> -1.0. - #[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] - // dot=24, both magnitudes=5 -> 24/25 = 0.96. - #[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] - // Zero vector -> guarded to 0.0. - #[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); - Ok(()) - } - - /// Self-similarity across various tensor shapes should always produce 1.0. - #[rstest] - // 2x3 matrix, flattened to 6 elements. - #[case::matrix_2d( - &[2, 3], - &[ - 1.0, 0.0, 0.0, // row 0 - 0.0, 0.0, 0.0, // row 1 - ], - )] - // 2x2x2 tensor, 8 elements. - #[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] - fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { - let lhs = tensor_array(shape, elements)?; - let rhs = tensor_array(shape, elements)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn scalar_0d() -> VortexResult<()> { - // 0-dimensional tensor: each "tensor" is a single scalar value. - let lhs = tensor_array(&[], &[5.0, 3.0])?; - let rhs = tensor_array(&[], &[5.0, -3.0])?; - - // Same sign -> 1.0, opposite sign -> -1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); - Ok(()) - } - - #[test] - fn many_rows() -> VortexResult<()> { - // 5 tensors of shape [4] compared against themselves -> all 1.0. - let lhs = tensor_array( - &[4], - &[ - 1.0, 2.0, 3.0, 4.0, // tensor 0 - 0.0, 1.0, 0.0, 0.0, // tensor 1 - 5.0, 0.0, 5.0, 0.0, // tensor 2 - 1.0, 1.0, 1.0, 1.0, // tensor 3 - 0.0, 0.0, 0.0, 7.0, // tensor 4 - ], - )?; - let rhs = lhs.clone(); - - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0, 1.0, 1.0, 1.0, 1.0], - ); - Ok(()) - } - - #[test] - fn constant_query_tensor() -> VortexResult<()> { - // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. - let data = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 0.0, 1.0, 0.0, // tensor 1 - 0.0, 0.0, 1.0, // tensor 2 - 1.0, 0.0, 0.0, // tensor 3 - ], - )?; - let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn vector_unit_vectors() -> VortexResult<()> { - let lhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 1.0, 0.0, 0.0, // vector 1 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn vector_constant_query() -> VortexResult<()> { - let data = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - 0.0, 0.0, 1.0, // vector 2 - 1.0, 0.0, 0.0, // vector 3 - ], - )?; - let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. - let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; - let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; - let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: self-similarity = 1.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_self_similarity() -> VortexResult<()> { - // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. - // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Self-similarity should always be 1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); - Ok(()) - } - - #[test] - fn both_normalized_orthogonal() -> VortexResult<()> { - // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. - // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn both_normalized_zero_norm() -> VortexResult<()> { - // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS is plain [3.0, 4.0]. - // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[3.0, 4.0])?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 0.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); - Ok(()) - } - - #[test] - fn both_normalized_null_norms() -> VortexResult<()> { - // Row 0: valid, row 1: null (via nullable norms on rhs). - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_r = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); - let rhs = Normalized::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine - // similarity for that row must be `0.0` even though the dot product of the normalized - // children is nonzero. - let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; - let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally violates the unit-norm invariant by pairing a nonzero normalized row - // with a stored norm of `0.0`, mimicking lossy storage. - // SAFETY: The children are structurally valid. - let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l) }.into_array(); - - let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; - let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Same as above for the rhs operand. - // SAFETY: The children are structurally valid. - let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r) }.into_array(); - - // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both - // `0.0`, so cosine similarity must be `0.0`. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. The plain side is a normal nonzero - // tensor with positive norm. cosine similarity must still be `0.0` because the - // authoritative stored norm on the normalized_array side is `0.0`. - let normalized = tensor_array(&[2], &[0.6, 0.8])?; - let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally pairs a nonzero normalized row with a stored norm of `0.0`, mimicking - // lossy storage where the stored norm is authoritative. - // SAFETY: The children are structurally valid. - let normalized_array = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); - - let plain = tensor_array(&[2], &[1.0, 0.0])?; - - // Normalized encoding on the lhs: `One { normalized_array: lhs, plain: rhs }`. - assert_close( - &eval_cosine_similarity(normalized_array.clone(), plain.clone())?, - &[0.0], - ); - - // Normalized encoding on the rhs: `One { normalized_array: rhs, plain: lhs }`. The same - // zero-norm guard must fire regardless of operand order. - assert_close(&eval_cosine_similarity(plain, normalized_array)?, &[0.0]); - Ok(()) - } - - #[test] - fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { - // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. - // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. - let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 - 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 - 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 - 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { - // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn both_constant_tensors() -> VortexResult<()> { - // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). - let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; - let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; - let expected = 1.0 / 2.0_f64.sqrt(); - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[expected, expected, expected], - ); - Ok(()) - } - - #[test] - fn constant_zero_norm_query() -> VortexResult<()> { - // A zero-norm constant query must produce `0.0` for every row via the zero-norm guard in - // `execute_one_normalized` and `execute_both_normalized`. - let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 2.0, 3.0, // - 4.0, 5.0, 6.0, // - 7.0, 8.0, 9.0, // - ], - )?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); - Ok(()) - } - - #[test] - fn constant_self_similarity_nonunit() -> VortexResult<()> { - // A non-unit constant query compared to itself must produce `1.0`. This exercises the - // helper's division: after normalization, both sides must be exactly unit so the - // Normalized fast path's inner product yields 1. - let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); - Ok(()) - } - - #[test] - fn vector_constant_matches_plain() -> VortexResult<()> { - // Exercise the `Vector` extension variant through the new pre-pass. - let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[rstest] - #[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] - #[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = CosineSimilarity::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("CosineSimilarity serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) +/// Computes the cosine similarity of one row, taking any hoisted norm from `norms` and computing +/// the rest exactly as [`cosine_similarity_row`] does. +/// +/// Each arm accumulates the same values in the same order as [`cosine_similarity_row`], and the +/// denominator keeps its lhs-times-rhs order, so the result is bit-identical whether a norm was +/// hoisted or not. The match costs one predictable branch per row: the arm is the same for the +/// whole batch. +fn cosine_similarity_row_prepared( + norms: &ConstNorms, + a: &[T], + b: &[T], +) -> T { + match (norms.lhs, norms.rhs) { + (None, None) => cosine_similarity_row(a, b), + (Some(norm_a), None) => { + let mut dot = T::zero(); + let mut norm_sq_b = T::zero(); + for (&x, &y) in a.iter().zip(b.iter()) { + dot = dot + x * y; + norm_sq_b = norm_sq_b + y * y; + } + cosine_from_parts(dot, norm_a * norm_sq_b.sqrt()) + } + (None, Some(norm_b)) => { + let mut dot = T::zero(); + let mut norm_sq_a = T::zero(); + for (&x, &y) in a.iter().zip(b.iter()) { + dot = dot + x * y; + norm_sq_a = norm_sq_a + x * x; + } + cosine_from_parts(dot, norm_sq_a.sqrt() * norm_b) + } + (Some(norm_a), Some(norm_b)) => { + let mut dot = T::zero(); + for (&x, &y) in a.iter().zip(b.iter()) { + dot = dot + x * y; + } + cosine_from_parts(dot, norm_a * norm_b) + } } +} - fn cosine_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") - } +/// Computes the cosine similarity of two equal-length float slices. +/// +/// Returns `dot(a, b) / (||a|| * ||b||)`, or `0.0` when either norm is zero. +fn cosine_similarity_row(a: &[T], b: &[T]) -> T { + let mut dot = T::zero(); + let mut norm_sq_a = T::zero(); + let mut norm_sq_b = T::zero(); + for (&x, &y) in a.iter().zip(b.iter()) { + dot = dot + x * y; + norm_sq_a = norm_sq_a + x * x; + norm_sq_b = norm_sq_b + y * y; + } + + cosine_from_parts(dot, norm_sq_a.sqrt() * norm_sq_b.sqrt()) +} - fn cosine_vector_rhs() -> ArrayRef { - vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +/// The shared tail of every cosine arm: `dot / denom`, guarded to `0.0` when the denominator is +/// zero. +fn cosine_from_parts(dot: T, denom: T) -> T { + if denom == T::zero() { + T::zero() + } else { + dot / denom } +} - fn cosine_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") - } +/// Both sides are [`Normalized`]-encoded: the normalized children are authoritative, so their dot +/// product is the cosine similarity, except that a row with a zero *stored* norm is a zero vector. +/// +/// Unlike [`InnerProduct::reduce_encoded`], which composes lazy `Mul` arrays over the norm columns, +/// this executes and materializes. The zero-norm guard is a conditional per row rather than an +/// arithmetic factor, so there is no lazy array that expresses it; the norm columns are one value +/// per row rather than one per coordinate, so materializing them is cheap next to the decode this +/// avoids. +/// +/// [`InnerProduct::reduce_encoded`]: InnerProduct::reduce_encoded +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_both_normalized( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized_l, normalized_r])? + .execute(ctx)?; + let norms_l: PrimitiveArray = norms_l.execute(ctx)?; + let norms_r: PrimitiveArray = norms_r.execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let norms_l = norms_l.as_slice::(); + let norms_r = norms_r.as_slice::(); + // Zipped rather than indexed by `0..len`: one bounds check per iterator instead of three + // per row. A length disagreement between the children shortens the result, which the + // lifting reports against the batch row count rather than panicking mid-loop. + let buffer: Buffer = dots + .iter() + .zip(norms_l) + .zip(norms_r) + .map(|((&dot, &norm_l), &norm_r)| { + if norm_l.is_zero() || norm_r.is_zero() { + T::zero() + } else { + dot + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) +} - fn cosine_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") - } +/// One side is [`Normalized`]-encoded: `cos = dot(normalized, plain) / ||plain||`, forced to `0.0` +/// on rows where the stored norm or the plain norm is `0.0`. +/// +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_one_normalized( + normalized_array: &ArrayRef, + plain: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = normalized_array.len(); + let (normalized, normalized_norms) = extract_normalized_children(normalized_array); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized, plain.clone()])? + .execute(ctx)?; + let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; + let plain_norm: PrimitiveArray = L2Norm + .try_new_array(len, EmptyOptions, [plain.clone()])? + .execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let normalized_norms = normalized_norms.as_slice::(); + let plain_norms = plain_norm.as_slice::(); + // Zipped for the same reason as [`cosine_both_normalized`]. + let buffer: Buffer = dots + .iter() + .zip(normalized_norms) + .zip(plain_norms) + .map(|((&dot, &stored_norm), &plain_norm)| { + if stored_norm.is_zero() || plain_norm.is_zero() { + T::zero() + } else { + dot / plain_norm + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) } diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 53ae82eb4a2..3d8255f3599 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -6,40 +6,30 @@ use num_traits::Float; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ElementSink; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::matcher::AnyTensor; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; -use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; /// Inner product (dot product) between two columns. /// @@ -52,131 +42,82 @@ use crate::utils::validate_binary_tensor_float_inputs; /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct InnerProduct; -impl InnerProduct { - /// Creates a new [`TypedScalarFnInstance`] wrapping the inner product operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(InnerProduct, EmptyOptions) - } - - /// Constructs a [`ScalarFnArray`] that lazily computes the inner product between `lhs` and - /// `rhs`. - /// - /// # Errors - /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). - pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(InnerProduct::new().erased(), vec![lhs, rhs]) - } -} - -impl ScalarFnVTable for InnerProduct { +impl RowFn for InnerProduct { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.inner_product"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("InnerProduct must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - // TODO(connor): relax the float-only gate once integer tensors are supported. - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_prepared_into::<(TensorRow, TensorRow), ElementSink, _, _>( + |_| (), + |&(), (lhs, rhs), output| *output = inner_product_row(lhs, rhs), + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands factor through their stored norms: with `D(x, s)` denoting + /// `x * s` rowwise, `dot(D(x, s), D(y, t)) = s * t * dot(x, y)` and + /// `dot(D(x, s), y) = s * dot(x, y)`. The rewrite is expressed with lazy [`Operator::Mul`] + /// arrays over the (much smaller) norm columns, so no denormalized coordinates are decoded. + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let lhs_ref = args.get(0)?; - let rhs_ref = args.get(1)?; - let len = args.row_count(); + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let len = args[0].len(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { + Ok(match NormalizedOrientation::classify(&args[0], &args[1]) { NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized_l, normalized_r])?; + Some( + dot.binary(norms_l, Operator::Mul)? + .binary(norms_r, Operator::Mul)?, + ) } NormalizedOrientation::One { normalized_array, plain, } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); + let (normalized, norms) = extract_normalized_children(normalized_array); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized, plain.clone()])?; + Some(dot.binary(norms, Operator::Mul)?) } - NormalizedOrientation::Neither => {} - } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Canonicalize so we can perform the math directly. - let lhs: ExtensionArray = lhs_ref.execute(ctx)?; - let rhs: ExtensionArray = rhs_ref.execute(ctx)?; - - // We validated that both inputs have the same type. - let ext = lhs.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let dimensions = tensor_match.list_size() as usize; - - // Extract the storage array from each extension input. We pass the storage (FSL) rather - // than the extension array to avoid canonicalizing the extension wrapper. - let lhs_storage = lhs.storage_array(); - let rhs_storage = rhs.storage_array(); - - let lhs_flat = extract_flat_elements(lhs_storage, dimensions, ctx)?; - let rhs_flat = extract_flat_elements(rhs_storage, dimensions, ctx)?; - - match_each_float_ptype!(lhs_flat.ptype(), |T| { - let buffer: Buffer = (0..len) - .map(|i| inner_product_row(lhs_flat.row::(i), rhs_flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) + NormalizedOrientation::Neither => None, }) } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false - } } impl ScalarFnArrayVTable for InnerProduct { @@ -205,72 +146,6 @@ impl ScalarFnArrayVTable for InnerProduct { } } -impl InnerProduct { - /// Both sides are [`Normalized`]-encoded: `inner_product = s_l * s_r * dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let nl = norms_l.as_slice::(); - let nr = norms_r.as_slice::(); - let buffer: Buffer = (0..len).map(|i| nl[i] * nr[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: `inner_product = s * dot(n, other)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, norms) = extract_normalized_children(normalized_ref); - let normalized_norms: PrimitiveArray = norms.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized, plain_ref.clone())? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let ns = normalized_norms.as_slice::(); - let buffer: Buffer = (0..len).map(|i| ns[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } -} - /// Computes the inner product (dot product) of two equal-length float slices. /// /// Returns `sum(a_i * b_i)`. @@ -280,254 +155,3 @@ fn inner_product_row(a: &[T], b: &[T]) -> T { .map(|(&x, &y)| x * y) .fold(T::zero(), |acc, v| acc + v) } - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::inner_product::InnerProduct; - use crate::tests::SESSION; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates inner product between two tensor arrays and returns the result as `Vec`. - fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - /// Single-row inner product for various vector pairs. - #[rstest] - // Orthogonal: [1, 0] . [0, 1] = 0. - #[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] - // Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. - #[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] - // Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. - #[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] - // Scaled: [2, 0] . [3, 0] = 6. - #[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_inner_product(lhs, rhs)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 3.0, 4.0, 0.0, // tensor 1 - 1.0, 1.0, 1.0, // tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 0.0, 1.0, 0.0, // tensor 0: dot = 0 - 3.0, 4.0, 0.0, // tensor 1: dot = 25 - 2.0, 2.0, 2.0, // tensor 2: dot = 6 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); - Ok(()) - } - - #[test] - fn vector_inner_product() -> VortexResult<()> { - let lhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0 - 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0: dot = 25 - 0.0, 1.0, // vector 1: dot = 0 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. - let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; - let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; - let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert!(prim.is_valid(2, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[23.0]); - assert_close(&[prim.as_slice::()[2]], &[127.0]); - Ok(()) - } - - #[test] - fn rejects_non_extension_dtype() { - let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); - let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - } - - #[test] - fn rejects_mismatched_dtypes() -> VortexResult<()> { - let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; - let rhs = vector_array(2, &[3.0_f64, 4.0])?; - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn both_normalized() -> VortexResult<()> { - // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). - // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). - // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; - - // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. - assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); - Ok(()) - } - - #[test] - fn both_normalized_multiple_rows() -> VortexResult<()> { - // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. - // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS: plain [1.0, 2.0]. - // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[1.0, 2.0])?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS: plain [1.0, 2.0]. - // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 2.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn both_normalized_null_norms() -> VortexResult<()> { - // Row 0: valid, row 1: null (via nullable norms on lhs). - let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_l = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - let lhs = Normalized::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[25.0]); - Ok(()) - } - - #[rstest] - #[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] - #[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = InnerProduct::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(InnerProduct); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("InnerProduct serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn inner_product_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } - - fn inner_product_vector_rhs() -> ArrayRef { - vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") - } - - fn inner_product_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") - } - - fn inner_product_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") - } -} diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index b7e9060ed3f..433a6527636 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -3,40 +3,23 @@ //! L2 norm expression for tensor-like types. -use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::Constant; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; -use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; -use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ElementSink; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -44,9 +27,10 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::Normalized; -use crate::matcher::AnyTensor; -use crate::utils::extract_flat_elements; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::extract_normalized_children; +use crate::utils::l2_norm_row; use crate::utils::validate_tensor_float_input; /// L2 norm (Euclidean norm) of a tensor or vector column. @@ -62,139 +46,62 @@ use crate::utils::validate_tensor_float_input; /// of the storage contract, not a separate lossy-compute mode. /// /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct L2Norm; -impl L2Norm { - /// Creates a new [`TypedScalarFnInstance`] wrapping the L2 norm operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(L2Norm, EmptyOptions) - } - - /// Constructs a [`ScalarFnArray`] that lazily computes the L2 norm over `child`. - /// - /// # Errors - /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). - pub fn try_new_array(child: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(L2Norm::new().erased(), vec![child]) - } -} - -impl ScalarFnVTable for L2Norm { +impl RowFn for L2Norm { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["input"]; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.l2_norm"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(1) - } - - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("input"), - _ => unreachable!("L2Norm must have exactly one child"), - } + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let input_dtype = &arg_dtypes[0]; - let tensor_match = validate_tensor_float_input(input_dtype)?; - let ptype = tensor_match.element_ptype(); - - let nullability = Nullability::from(input_dtype.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn execute( + fn dispatch( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let input_ref = args.get(0)?; - let row_count = args.row_count(); - - let ext = input_ref.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let tensor_flat_size = tensor_match.list_size() as usize; - let element_ptype = tensor_match.element_ptype(); - - let norm_dtype = DType::Primitive(element_ptype, ext.nullability()); - - // L2Norm over a `Normalized`-encoded column is defined to read back the authoritative stored - // norms. Callers of lossy encodings opt into that storage semantics instead of forcing a - // decode-and-recompute path here. - if input_ref.is::() { - let (_, norms) = extract_normalized_children(&input_ref); - vortex_ensure_eq!(norms.dtype(), &norm_dtype); - return Ok(norms); - } - - // Optimize for the constant array case. - if let Some(array) = input_ref.as_opt::() { - let scalar = array.scalar().as_extension().to_storage_scalar(); - - let Some(elements) = scalar.as_list().elements() else { - return Ok(ConstantArray::new(Scalar::null(norm_dtype), row_count).into_array()); - }; - - let norm_scalar = match_each_float_ptype!(element_ptype, |T| { - let values: Vec = elements - .iter() - .map(|s| { - s.as_primitive() - .as_::() - .vortex_expect("element was somehow not the correct float") - }) - .collect(); - let norm = l2_norm_row::(&values); - - Scalar::try_new(norm_dtype, Some(norm.into())) - })?; - - let norms = ConstantArray::new(norm_scalar, row_count).into_array(); - return Ok(norms); - } - - let input: ExtensionArray = input_ref.execute(ctx)?; - let validity = input.as_ref().validity()?; - - let storage = input.storage_array(); - let flat = extract_flat_elements(storage, tensor_flat_size, ctx)?; - - match_each_float_ptype!(flat.ptype(), |T| { - let buffer: Buffer = (0..row_count) - .map(|i| l2_norm_row(flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_prepared_into::<(TensorRow,), ElementSink, _, _>( + |_| (), + |&(), (row,), output| *output = l2_norm_row(row), + ) }) } - fn validity( + /// `L2Norm` over a [`Normalized`]-encoded column is defined to read back the authoritative + /// stored norms. Callers of lossy encodings opt into that storage semantics instead of forcing + /// a decode-and-recompute path here. + fn reduce_encoded( &self, _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if the input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let input = &args[0]; + if !input.is::() { + return Ok(None); + } + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let (_, norms) = extract_normalized_children(input); + vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); + Ok(Some(norms)) } } @@ -240,206 +147,3 @@ impl ScalarFnArrayVTable for L2Norm { }) } } - -/// Computes the L2 norm (Euclidean norm) of a float slice. -/// -/// Returns `sqrt(sum(v_i^2))`. A zero-length or all-zero input produces `0.0`. -fn l2_norm_row(v: &[T]) -> T { - let mut sum_sq = T::zero(); - for &x in v { - sum_sq = sum_sq + x * x; - } - sum_sq.sqrt() -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::EmptyMetadata; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Constant; - use vortex_array::arrays::ConstantArray; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::dtype::extension::ExtDType; - use vortex_array::scalar::Scalar; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::scalar_fns::l2_norm::L2Norm; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::literal_vector_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. - fn eval_l2_norm(input: ArrayRef) -> VortexResult> { - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[rstest] - #[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] - #[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] - #[case::single_element(&[1], &[7.0], &[7.0])] - #[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] - fn known_norms( - #[case] shape: &[usize], - #[case] elements: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let arr = tensor_array(shape, elements)?; - assert_close(&eval_l2_norm(arr)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let arr = tensor_array( - &[3], - &[ - 3.0, 4.0, 0.0, // norm = 5.0 - 0.0, 0.0, 0.0, // norm = 0.0 - 1.0, 1.0, 1.0, // norm = sqrt(3) - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); - Ok(()) - } - - #[test] - fn vector_multiple_rows() -> VortexResult<()> { - let arr = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // norm = 1.0 - 3.0, 4.0, 0.0, // norm = 5.0 - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 is masked as null. - let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; - let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: norm = 5.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is a non-null tensor should short-circuit to a - /// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so - /// execution stops at the [`Constant`] encoding instead of canonicalizing into a - /// [`PrimitiveArray`]. - #[test] - fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { - let input = literal_vector_array(&[3.0f64, 4.0], 4); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let output = result.execute_until::(&mut ctx)?; - - let constant = output - .as_opt::() - .expect("L2Norm over a constant input must produce a constant output"); - assert_eq!(constant.len(), 4); - let norm = constant - .scalar() - .as_primitive() - .as_::() - .expect("norm scalar must be a non-null primitive"); - assert_close(&[norm], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of - /// the correct primitive dtype and length. - #[test] - fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { - let storage_dtype = DType::FixedSizeList( - DType::Primitive(PType::F64, Nullability::NonNullable).into(), - 2, - Nullability::Nullable, - ); - let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); - let null_scalar = Scalar::null(DType::Extension(ext_dtype)); - let input = ConstantArray::new(null_scalar, 3).into_array(); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let output = result.execute_until::(&mut ctx)?; - - let constant = output - .as_opt::() - .expect("null constant input must produce a constant output"); - assert_eq!(constant.len(), 3); - assert!(constant.scalar().is_null()); - assert_eq!( - constant.dtype(), - &DType::Primitive(PType::F64, Nullability::Nullable) - ); - Ok(()) - } - - #[rstest] - #[case::fixed_shape_tensor(l2_norm_tensor_child())] - #[case::vector(l2_norm_vector_child())] - fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { - let original = L2Norm::try_new_array(child.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(L2Norm); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("L2Norm serialize must produce metadata"); - - let children = vec![child]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn l2_norm_tensor_child() -> ArrayRef { - tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") - } - - fn l2_norm_vector_child() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } -} diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 68f10ca6b01..706392d3b25 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -6,3 +6,7 @@ pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; +pub mod row; + +#[cfg(test)] +mod tests; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs new file mode 100644 index 00000000000..3c02a1d2615 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the tensor scalar functions add to the row-function machinery: an element type that reads a +//! tensor row and the width rule they share. + +use std::marker::PhantomData; + +use num_traits::Float; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::scalar_fn::InputElement; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use crate::utils::extract_flat_elements; +use crate::utils::validate_tensor_float_input; +use crate::utils::validate_tensor_float_inputs; + +/// The width rule the tensor scalar functions share: every argument is the same float tensor dtype, +/// and the width is its element ptype. +pub(crate) fn tensor_element_ptype(args: &[DType]) -> VortexResult { + Ok(validate_tensor_float_inputs(args)?.element_ptype()) +} + +/// Marker for tensor-valued input elements: accepts any tensor-like extension column whose +/// elements are `T`, and presents each row as its flat elements, `&[T]`. +pub struct TensorRow(PhantomData); + +/// The decoded form of a [`TensorRow`] column: one flat typed buffer plus the stride to read it at. +/// +/// Typed at decode time rather than per row. `FlatElements::row` re-derives its typed slice on every +/// call, which costs a ptype check and a buffer downcast per row; a row loop reads every row, so it +/// pays that once here instead. +pub struct TensorRows { + /// Every row's elements, back to back. + elements: Buffer, + + /// Number of logical tensor rows, stored so zero-width tensors retain their length. + rows: usize, + + /// Elements per row, the length of each row slice. + list_size: usize, + + /// `list_size` for a full column and `0` for constant-backed storage, so `index * stride` pins a + /// constant to its single materialized row without a branch in the loop. + stride: usize, +} + +impl InputElement for TensorRow { + type Column = TensorRows; + type Varying<'a> = &'a TensorRows; + type Elem<'a> = &'a [T]; + + // Tensor storage is a fully materialized non-nullable primitive buffer, so the elements behind + // a null row are arbitrary values rather than an unresolvable reference. + const DENSE_SAFE: bool = true; + // Tensor storage is a primitive buffer; reading it cannot fail on account of its values. + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + let tensor_match = validate_tensor_float_input(dtype)?; + let expected = T::PTYPE; + vortex_ensure_eq!( + tensor_match.element_ptype(), + expected, + "expected a tensor of {expected} elements, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let rows = array.len(); + let list_size = validate_tensor_float_input(array.dtype())?.list_size() as usize; + let ext: ExtensionArray = array.execute(ctx)?; + let flat = extract_flat_elements(ext.storage_array(), list_size, ctx)?; + + Ok(TensorRows { + rows, + list_size: flat.list_size(), + stride: flat.row_stride(), + elements: flat.into_buffer::(), + }) + } + + fn get(column: &Self::Column, index: usize) -> &[T] { + let start = index * column.stride; + &column.elements.as_slice()[start..start + column.list_size] + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.rows + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + Self::get(column, index) + } +} + +/// Test-only probe recording which operands the last `prepare` step saw as batch-constant, so a +/// test can assert its inputs took the stride-0 decode path rather than merely producing the right +/// values through the varying path. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + thread_local! { + /// Bitmask of the constant operands the last `prepare` saw (bit 0 for the lhs, bit 1 for + /// the rhs). Thread-local rather than a process global so concurrent tests in one process + /// (plain `cargo test`) cannot race it; execution runs on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Record which operands `prepare` saw as constant. + pub(crate) fn record(lhs_constant: bool, rhs_constant: bool) { + SEEN_CONSTANTS.set(u8::from(lhs_constant) | (u8::from(rhs_constant) << 1)); + } +} diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs new file mode 100644 index 00000000000..60e75792109 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -0,0 +1,586 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::assert_arrays_eq; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::cosine_similarity::CosineSimilarity; +use crate::scalar_fns::row::probe; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::constant_tensor_array; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; + +/// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. +fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +/// Like [`eval_cosine_similarity`], but returns the executed array for exact array comparisons. +fn eval_cosine_similarity_array( + lhs: ArrayRef, + rhs: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + Ok(result + .into_array() + .execute::(ctx)? + .into_array()) +} + +#[test] +fn unit_vectors_1d() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 0.0, 1.0, 0.0, // Tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 1.0, 0.0, 0.0, // Tensor 2 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +/// Single-row cosine similarity for various vector pairs. +#[rstest] +// Antiparallel -> -1.0. +#[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] +// dot=24, both magnitudes=5 -> 24/25 = 0.96. +#[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] +// Zero vector -> guarded to 0.0. +#[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); + Ok(()) +} + +/// Self-similarity across various tensor shapes should always produce 1.0. +#[rstest] +// 2x3 matrix, flattened to 6 elements. +#[case::matrix_2d( + &[2, 3], + &[ + 1.0, 0.0, 0.0, // row 0 + 0.0, 0.0, 0.0, // row 1 + ], +)] +// 2x2x2 tensor, 8 elements. +#[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] +fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { + let lhs = tensor_array(shape, elements)?; + let rhs = tensor_array(shape, elements)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn scalar_0d() -> VortexResult<()> { + // 0-dimensional tensor: each "tensor" is a single scalar value. + let lhs = tensor_array(&[], &[5.0, 3.0])?; + let rhs = tensor_array(&[], &[5.0, -3.0])?; + + // Same sign -> 1.0, opposite sign -> -1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); + Ok(()) +} + +#[test] +fn many_rows() -> VortexResult<()> { + // 5 tensors of shape [4] compared against themselves -> all 1.0. + let lhs = tensor_array( + &[4], + &[ + 1.0, 2.0, 3.0, 4.0, // tensor 0 + 0.0, 1.0, 0.0, 0.0, // tensor 1 + 5.0, 0.0, 5.0, 0.0, // tensor 2 + 1.0, 1.0, 1.0, 1.0, // tensor 3 + 0.0, 0.0, 0.0, 7.0, // tensor 4 + ], + )?; + let rhs = lhs.clone(); + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0, 1.0, 1.0, 1.0, 1.0], + ); + Ok(()) +} + +#[test] +fn constant_query_tensor() -> VortexResult<()> { + // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. + let data = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 0.0, 1.0, 0.0, // tensor 1 + 0.0, 0.0, 1.0, // tensor 2 + 1.0, 0.0, 0.0, // tensor 3 + ], + )?; + let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn vector_unit_vectors() -> VortexResult<()> { + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 1.0, 0.0, 0.0, // vector 1 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn vector_constant_query() -> VortexResult<()> { + let data = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + 0.0, 0.0, 1.0, // vector 2 + 1.0, 0.0, 0.0, // vector 3 + ], + )?; + let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. + let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; + let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; + let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: self-similarity = 1.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_self_similarity() -> VortexResult<()> { + // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. + // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Self-similarity should always be 1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); + Ok(()) +} + +#[test] +fn both_normalized_orthogonal() -> VortexResult<()> { + // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. + // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn both_normalized_zero_norm() -> VortexResult<()> { + // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS is plain [3.0, 4.0]. + // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[3.0, 4.0])?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 0.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); + Ok(()) +} + +#[test] +fn both_normalized_null_norms() -> VortexResult<()> { + // Row 0: valid, row 1: null (via nullable norms on rhs). + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_r = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let rhs = Normalized::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine + // similarity for that row must be `0.0` even though the dot product of the normalized + // children is nonzero. + let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; + let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally violates the unit-norm invariant by + // pairing a nonzero normalized row with a stored norm of `0.0`, mimicking lossy storage. + let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l) }.into_array(); + + let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; + let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: Same as above for the rhs operand. + let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r) }.into_array(); + + // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both + // `0.0`, so cosine similarity must be `0.0`. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. The plain side is a normal nonzero + // tensor with positive norm. cosine similarity must still be `0.0` because the + // authoritative stored norm on the denorm side is `0.0`. + let normalized = tensor_array(&[2], &[0.6, 0.8])?; + let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally pairs a nonzero normalized row with a + // stored norm of `0.0`, mimicking lossy storage where the stored norm is authoritative. + let denorm = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); + + let plain = tensor_array(&[2], &[1.0, 0.0])?; + + // Denorm on the lhs: `One { denorm: lhs, plain: rhs }`. + assert_close( + &eval_cosine_similarity(denorm.clone(), plain.clone())?, + &[0.0], + ); + + // Denorm on the rhs: `One { denorm: rhs, plain: lhs }`. The same zero-norm guard must + // fire regardless of operand order. + assert_close(&eval_cosine_similarity(plain, denorm)?, &[0.0]); + Ok(()) +} + +#[test] +fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { + // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. + // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. + let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 + 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 + 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 + 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 + ], + )?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { + // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn both_constant_tensors() -> VortexResult<()> { + // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). + let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; + let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[expected, expected, expected], + ); + Ok(()) +} + +#[test] +fn constant_zero_norm_query() -> VortexResult<()> { + // A zero-norm constant query must produce `0.0` for every row via the zero-norm guard in + // `cosine_one_normalized` and `execute_both_normalized`. + let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 2.0, 3.0, // + 4.0, 5.0, 6.0, // + 7.0, 8.0, 9.0, // + ], + )?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +#[test] +fn constant_self_similarity_nonunit() -> VortexResult<()> { + // A non-unit constant query compared to itself must produce `1.0`. This exercises the + // helper's division: after normalization, both sides must be exactly unit so the + // Normalized fast path's inner product yields 1. + let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); + Ok(()) +} + +/// An extension array over constant storage (what [`Vector::constant_array`] builds) is a batch +/// constant like any other: the row layer sees through the wrapper, so `prepare` hoists its norm +/// exactly as it does for the literal shape. This used to be intercepted by a hand-written +/// `reduce_encoded` rewrite into `Normalized`, deleted in favor of the framework path. +#[test] +fn vector_constant_matches_plain() -> VortexResult<()> { + let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + assert_eq!( + probe::SEEN_CONSTANTS.get(), + 0b01, + "the extension-over-constant lhs must reach prepare as a batch constant", + ); + Ok(()) +} + +/// The literal-constant shape (a [`ConstantArray`] over a [`Vector`] extension scalar, what a +/// `lit(query)` expression produces) reaches the row loop, unlike an extension-wrapped constant, +/// which `reduce_encoded` rewrites into `Normalized`. There the prepared kernel hoists the query's +/// norm once per batch, and the result must be exactly the result of expanding the same query +/// into a full column, which hoists nothing. +/// +/// [`ConstantArray`]: vortex_array::arrays::ConstantArray +#[test] +fn literal_constant_rhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(lhs.clone(), literal_vector_array(&query, 4), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(lhs, vector_array(3, &query.repeat(4))?, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// The mirror of [`literal_constant_rhs_matches_expanded_column`], exercising the hoisted-lhs arm +/// of the prepared kernel. +#[test] +fn literal_constant_lhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(literal_vector_array(&query, 4), rhs.clone(), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(vector_array(3, &query.repeat(4))?, rhs, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// A zero-norm literal constant query must be guarded to `0.0` on every row by the prepared row +/// kernel, exactly as the unprepared kernel guards it. +#[test] +fn literal_constant_zero_norm_query_yields_zero() -> VortexResult<()> { + let lhs = vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = literal_vector_array(&[0.0f64, 0.0, 0.0], 2); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0]); + Ok(()) +} + +/// Two literal constants are folded to a single-row execution by the row lifting, and that row +/// still runs the prepared kernel with both norms hoisted. +#[test] +fn both_literal_constants() -> VortexResult<()> { + let lhs = literal_vector_array(&[1.0f64, 0.0, 0.0], 3); + let rhs = literal_vector_array(&[1.0f64, 1.0, 0.0], 3); + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[expected; 3]); + Ok(()) +} + +#[rstest] +#[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] +#[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("CosineSimilarity serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn cosine_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_vector_rhs() -> ArrayRef { + vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn cosine_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/inner_product.rs b/vortex-tensor/src/scalar_fns/tests/inner_product.rs new file mode 100644 index 00000000000..af7fbb7bc1a --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/inner_product.rs @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::inner_product::InnerProduct; +use crate::tests::SESSION; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; + +/// Evaluates inner product between two tensor arrays and returns the result as `Vec`. +fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +/// Single-row inner product for various vector pairs. +#[rstest] +// Orthogonal: [1, 0] . [0, 1] = 0. +#[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] +// Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. +#[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] +// Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. +#[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] +// Scaled: [2, 0] . [3, 0] = 6. +#[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_inner_product(lhs, rhs)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 3.0, 4.0, 0.0, // tensor 1 + 1.0, 1.0, 1.0, // tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 0.0, 1.0, 0.0, // tensor 0: dot = 0 + 3.0, 4.0, 0.0, // tensor 1: dot = 25 + 2.0, 2.0, 2.0, // tensor 2: dot = 6 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); + Ok(()) +} + +#[test] +fn vector_inner_product() -> VortexResult<()> { + let lhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0 + 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0: dot = 25 + 0.0, 1.0, // vector 1: dot = 0 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. + let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; + let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert!(prim.is_valid(2, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[23.0]); + assert_close(&[prim.as_slice::()[2]], &[127.0]); + Ok(()) +} + +#[test] +fn rejects_non_extension_dtype() { + let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); + let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); +} + +#[test] +fn rejects_mismatched_dtypes() -> VortexResult<()> { + let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; + let rhs = vector_array(2, &[3.0_f64, 4.0])?; + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn both_normalized() -> VortexResult<()> { + // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). + // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). + // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; + + // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. + assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); + Ok(()) +} + +#[test] +fn both_normalized_multiple_rows() -> VortexResult<()> { + // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. + // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS: plain [1.0, 2.0]. + // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[1.0, 2.0])?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS: plain [1.0, 2.0]. + // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 2.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn both_normalized_null_norms() -> VortexResult<()> { + // Row 0: valid, row 1: null (via nullable norms on lhs). + let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_l = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let lhs = Normalized::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[25.0]); + Ok(()) +} + +#[rstest] +#[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] +#[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(InnerProduct); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("InnerProduct serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn inner_product_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} + +fn inner_product_vector_rhs() -> ArrayRef { + vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") +} + +fn inner_product_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn inner_product_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/l2_norm.rs b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs new file mode 100644 index 00000000000..a9fda0326d8 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs @@ -0,0 +1,295 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::l2_norm::L2Norm; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; + +/// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. +fn eval_l2_norm(input: ArrayRef) -> VortexResult> { + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[rstest] +#[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] +#[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] +#[case::single_element(&[1], &[7.0], &[7.0])] +#[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] +fn known_norms( + #[case] shape: &[usize], + #[case] elements: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let arr = tensor_array(shape, elements)?; + assert_close(&eval_l2_norm(arr)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let arr = tensor_array( + &[3], + &[ + 3.0, 4.0, 0.0, // norm = 5.0 + 0.0, 0.0, 0.0, // norm = 0.0 + 1.0, 1.0, 1.0, // norm = sqrt(3) + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); + Ok(()) +} + +#[test] +fn vector_multiple_rows() -> VortexResult<()> { + let arr = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // norm = 1.0 + 3.0, 4.0, 0.0, // norm = 5.0 + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 is masked as null. + let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: norm = 5.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is a non-null tensor should short-circuit to a +/// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so +/// execution stops at the [`Constant`] encoding instead of canonicalizing into a +/// [`PrimitiveArray`]. +#[test] +fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { + let input = literal_vector_array(&[3.0f64, 4.0], 4); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over a constant input must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// An extension array over constant storage is folded just like a top-level constant instead of +/// recomputing the same norm once per row. +#[test] +fn extension_backed_constant_yields_constant_output() -> VortexResult<()> { + let input = Vector::constant_array(&[3.0f64, 4.0], 4)?; + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over constant-backed extension storage must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of +/// the correct primitive dtype and length. +#[test] +fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { + let storage_dtype = DType::FixedSizeList( + DType::Primitive(PType::F64, Nullability::NonNullable).into(), + 2, + Nullability::Nullable, + ); + let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); + let null_scalar = Scalar::null(DType::Extension(ext_dtype)); + let input = ConstantArray::new(null_scalar, 3).into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("null constant input must produce a constant output"); + assert_eq!(constant.len(), 3); + assert!(constant.scalar().is_null()); + assert_eq!( + constant.dtype(), + &DType::Primitive(PType::F64, Nullability::Nullable) + ); + Ok(()) +} + +/// An `f32` column must dispatch at `f32` and produce an `f32` result, which is the property that +/// makes width polymorphism load-bearing rather than decorative. +#[rstest] +#[case::f32(&[3.0f32, 4.0], PType::F32)] +#[case::f64(&[3.0f64, 4.0], PType::F64)] +fn dispatches_at_input_width( + #[case] elements: &[T], + #[case] expected: PType, +) -> VortexResult<()> { + let arr = tensor_array(&[2], elements)?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L2Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .execute(&mut ctx)?; + assert_eq!(prim.ptype(), expected); + Ok(()) +} + +/// `L2Norm(Normalized(normalized, norms))` reads back the authoritative stored norms rather than +/// recomputing over decoded coordinates. The normalized child here is deliberately *not* +/// unit-norm, mimicking lossy storage, so readthrough and recompute disagree: row 0 decodes to +/// `[6, 8]` (norm `10`) and row 1 to `[6, 0]` (norm `6`), while the stored norms are `5` and `2`. +#[test] +fn normalized_readthrough_returns_stored_norms() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 2.0]).into_array(); + // SAFETY: A focused test of the lossy storage contract: the stored norms are authoritative + // even though this normalized child violates the unit-norm invariant. + let denorm = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); + + assert_close(&eval_l2_norm(denorm)?, &[5.0, 2.0]); + Ok(()) +} + +/// The readthrough must survive a partially-null column. +/// +/// This pins the dense policy the row contract derives. Filtering could hand `reduce_encoded` a +/// filtered input, which is no longer an `ExactScalarFn`, silently falling back to +/// decode-and-recompute. For a lossy child that changes the answer: row 0 below would come back as +/// `10` (recomputed from `[6, 8]`) instead of the authoritative stored `5`. +#[test] +fn normalized_readthrough_survives_null_rows() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + // SAFETY: Intentionally lossy, as in `normalized_readthrough_returns_stored_norms`, so that + // a recompute fallback is observable. + let denorm = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// The readthrough must still propagate nulls carried by the `norms` child. +#[test] +fn normalized_readthrough_propagates_null_norms() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let denorm = Normalized::try_new(normalized, norms, &mut ctx)?.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +#[rstest] +#[case::fixed_shape_tensor(l2_norm_tensor_child())] +#[case::vector(l2_norm_vector_child())] +fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { + let original = L2Norm.try_new_array(child.len(), EmptyOptions, [child.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(L2Norm); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("L2Norm serialize must produce metadata"); + + let children = vec![child]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn l2_norm_tensor_child() -> ArrayRef { + tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") +} + +fn l2_norm_vector_child() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/mod.rs b/vortex-tensor/src/scalar_fns/tests/mod.rs new file mode 100644 index 00000000000..bb3726e9329 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/mod.rs @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for the tensor scalar functions. + +mod cosine_similarity; +mod inner_product; +mod l2_norm; +mod row; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs new file mode 100644 index 00000000000..f08f614cfef --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::Float; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::match_each_float_ptype; +use vortex_array::scalar_fn::ElementSink; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::assert_element_conforms; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::tensor_array; + +/// The marginal cost of a new tensor scalar function is this entire definition. Everything else +/// (null propagation, constants, validity, f16/f32/f64 dispatch, dtype checks, and constructors) is +/// derived. +#[derive(Clone, Debug, Default)] +struct L1Norm; + +impl RowFn for L1Norm { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.l1_norm"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_prepared_into::<(TensorRow,), ElementSink, _, _>( + |_| (), + |&(), (row,), output| *output = l1_norm_row(row), + ) + }) + } +} + +fn l1_norm_row(row: &[T]) -> T { + row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) +} + +#[test] +fn derived_fn_executes_with_nulls() -> VortexResult<()> { + let arr = tensor_array(&[2], &[3.0, -4.0, 1.0, 1.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L1Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[7.0]); + Ok(()) +} + +/// A kernel written once serves every float width. +#[test] +fn derived_fn_dispatches_at_input_width() -> VortexResult<()> { + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + + let f32_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f32, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f32_result.ptype(), PType::F32); + assert_eq!(f32_result.as_slice::(), &[7.0f32]); + + let f64_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f64, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f64_result.ptype(), PType::F64); + Ok(()) +} + +/// Runs the out-of-crate [`TensorRow`] element through `vortex-array`'s shared element conformance +/// check, with `NaN` and infinities sitting behind the null row so a wrong `DENSE_SAFE` would be +/// read rather than skipped. +#[test] +fn tensor_row_element_conforms() -> VortexResult<()> { + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + let arr = tensor_array(&[2], &[3.0, -4.0, f64::NAN, f64::INFINITY])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + assert_element_conforms::>( + arr, + &DType::Primitive(PType::F64, Nullability::NonNullable), + &mut ctx, + ) +} diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 488694bd47f..460dde82ea7 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -1,13 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Shared helpers for the tensor scalar functions. + use half::f16; +use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn; @@ -20,6 +24,8 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::dtype::proto::dtype as pb; use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -58,6 +64,20 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { SAFETY_FACTOR as f64 * machine_epsilon * dimensions_root } +/// The L2 norm of one row: `sqrt(sum(v_i^2))`. A zero-length or all-zero row gives `0.0`. +/// +/// Shared by `l2_norm` and by cosine similarity's hoisted constant norm. The accumulation order is +/// part of the contract rather than an implementation detail: cosine's prepared and per-row arms +/// must agree bit for bit, which only holds while both sum in this order. Keeping one copy is what +/// stops the two drifting apart. +pub(crate) fn l2_norm_row(v: &[T]) -> T { + let mut sum_sq = T::zero(); + for &x in v { + sum_sq = sum_sq + x * x; + } + sum_sq.sqrt() +} + /// Extracts the `(normalized, norms)` children of a [`Normalized`]-encoded array. /// /// # Panics @@ -97,17 +117,78 @@ pub fn validate_tensor_float_input(input_dtype: &DType) -> VortexResult( - lhs: &'a DType, - rhs: &DType, -) -> VortexResult> { - vortex_ensure!( - lhs.eq_ignore_nullability(rhs), - "binary tensor expression expects inputs to have the same dtype, got {lhs} and {rhs}" - ); - validate_tensor_float_input(lhs) +pub fn validate_tensor_float_inputs(args: &[DType]) -> VortexResult> { + let (first, rest) = args + .split_first() + .ok_or_else(|| vortex_err!("tensor expression expects at least one input"))?; + for arg in rest { + vortex_ensure!( + first.eq_ignore_nullability(arg), + "tensor expression expects inputs to have the same dtype, got {first} and {arg}" + ); + } + validate_tensor_float_input(first) +} + +/// Metadata for a serialized binary tensor-op array (shared by [`InnerProduct`] and +/// [`CosineSimilarity`]). Both operands share the same extension dtype up to nullability +/// (enforced by their `return_dtype` checks), but their individual nullabilities are lost in the +/// parent's unioned output, so both are persisted. +/// +/// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity +/// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct +#[derive(Clone, prost::Message)] +pub(crate) struct BinaryTensorOpMetadata { + #[prost(message, optional, tag = "1")] + pub(crate) lhs_dtype: Option, + #[prost(message, optional, tag = "2")] + pub(crate) rhs_dtype: Option, +} + +impl BinaryTensorOpMetadata { + /// Encodes the two children of `view` into a [`BinaryTensorOpMetadata`] byte blob. + pub(crate) fn encode_from_view( + view: &ScalarFnArrayView, + ) -> VortexResult> { + let scalar_fn_array = view.as_::(); + let lhs_dtype = Some(scalar_fn_array.child_at(0).dtype().try_into()?); + let rhs_dtype = Some(scalar_fn_array.child_at(1).dtype().try_into()?); + Ok(Self { + lhs_dtype, + rhs_dtype, + } + .encode_to_vec()) + } + + /// Decodes `metadata` and fetches both children from `children` using the decoded dtypes, + /// validating that `lhs` and `rhs` are compatible tensor operands. + pub(crate) fn decode_children( + metadata: &[u8], + len: usize, + children: &dyn vortex_array::serde::ArrayChildren, + session: &VortexSession, + ) -> VortexResult> { + let metadata = Self::decode(metadata) + .map_err(|e| vortex_err!("Failed to decode BinaryTensorOpMetadata: {e}"))?; + let lhs_pb = metadata + .lhs_dtype + .as_ref() + .ok_or_else(|| vortex_err!("metadata missing lhs_dtype"))?; + let rhs_pb = metadata + .rhs_dtype + .as_ref() + .ok_or_else(|| vortex_err!("metadata missing rhs_dtype"))?; + + let lhs_dtype = DType::from_proto(lhs_pb, session)?; + let rhs_dtype = DType::from_proto(rhs_pb, session)?; + validate_tensor_float_inputs(&[lhs_dtype.clone(), rhs_dtype.clone()])?; + + let lhs = children.get(0, &lhs_dtype, len)?; + let rhs = children.get(1, &rhs_dtype, len)?; + Ok(vec![lhs, rhs]) + } } /// The flat primitive elements of a tensor storage array, with typed row access. @@ -132,12 +213,58 @@ impl FlatElements { /// /// When the source was a constant-backed storage, all indices resolve to the single stored /// row. + /// + /// This re-derives the typed slice on every call, which costs a ptype check and a buffer + /// downcast per row. A caller reading every row in a loop should take [`into_buffer`](Self::into_buffer) + /// instead and pay that once. #[must_use] pub fn row(&self, i: usize) -> &[T] { let row_idx = if self.is_constant { 0 } else { i }; let slice = self.elems.as_slice::(); &slice[row_idx * self.list_size..][..self.list_size] } + + /// Elements per row. + #[must_use] + pub fn list_size(&self) -> usize { + self.list_size + } + + /// The row stride: `list_size` for a full column, and `0` for constant-backed storage, whose + /// single materialized row every index reads. + #[must_use] + pub fn row_stride(&self) -> usize { + if self.is_constant { 0 } else { self.list_size } + } + + /// The elements as a typed buffer, checking the ptype once instead of once per row. + pub fn into_buffer(self) -> Buffer { + self.elems.into_buffer::() + } +} + +/// Rebuilds a tensor-like extension array from flat primitive elements. +/// +/// # Errors +/// +/// Returns an error if `elements` does not hold exactly `tensor_flat_size * row_count` values. +pub(crate) fn build_tensor_array( + dtype: DType, + tensor_flat_size: usize, + row_count: usize, + validity: Validity, + elements: Buffer, +) -> VortexResult { + let list_size = + u32::try_from(tensor_flat_size).vortex_expect("tensor flat size must fit into `u32`"); + + // SAFETY: Tensor elements are always non-nullable, so the validity carries no length. + let elements = unsafe { PrimitiveArray::new_unchecked(elements, Validity::NonNullable) }; + + let storage = + FixedSizeListArray::try_new(elements.into_array(), list_size, validity, row_count)?; + + Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) } /// Extracts the flat primitive elements from a tensor storage array (FixedSizeList). @@ -161,10 +288,10 @@ pub fn extract_flat_elements( let fsl: FixedSizeListArray = source.execute(ctx)?; let elems: PrimitiveArray = fsl.elements().clone().execute(ctx)?; + let dtype = elems.dtype(); vortex_ensure!( !elems.nullability().is_nullable(), - "tensor storage elements must be non-nullable, got {}", - elems.dtype(), + "tensor storage elements must be non-nullable, got {dtype}", ); Ok(FlatElements { elems, @@ -216,73 +343,14 @@ pub fn extract_constant_flat_row( let single = ConstantArray::new(constant.scalar().clone(), 1).into_array(); let fsl: FixedSizeListArray = single.execute(ctx)?; let elems: PrimitiveArray = fsl.elements().clone().execute(ctx)?; + let dtype = elems.dtype(); vortex_ensure!( !elems.nullability().is_nullable(), - "tensor storage elements must be non-nullable, got {}", - elems.dtype(), + "tensor storage elements must be non-nullable, got {dtype}", ); Ok(FlatRow { elems }) } -/// Metadata for a serialized binary tensor-op array (shared by [`InnerProduct`] and -/// [`CosineSimilarity`]). Both operands share the same extension dtype up to nullability -/// (enforced by their `return_dtype` checks), but their individual nullabilities are lost in the -/// parent's unioned output, so both are persisted. -/// -/// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity -/// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct -#[derive(Clone, prost::Message)] -pub(crate) struct BinaryTensorOpMetadata { - #[prost(message, optional, tag = "1")] - pub(crate) lhs_dtype: Option, - #[prost(message, optional, tag = "2")] - pub(crate) rhs_dtype: Option, -} - -impl BinaryTensorOpMetadata { - /// Encodes the two children of `view` into a [`BinaryTensorOpMetadata`] byte blob. - pub(crate) fn encode_from_view( - view: &ScalarFnArrayView, - ) -> VortexResult> { - let scalar_fn_array = view.as_::(); - let lhs_dtype = Some(scalar_fn_array.child_at(0).dtype().try_into()?); - let rhs_dtype = Some(scalar_fn_array.child_at(1).dtype().try_into()?); - Ok(Self { - lhs_dtype, - rhs_dtype, - } - .encode_to_vec()) - } - - /// Decodes `metadata` and fetches both children from `children` using the decoded dtypes, - /// validating that `lhs` and `rhs` are compatible tensor operands. - pub(crate) fn decode_children( - metadata: &[u8], - len: usize, - children: &dyn vortex_array::serde::ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = Self::decode(metadata) - .map_err(|e| vortex_err!("Failed to decode BinaryTensorOpMetadata: {e}"))?; - let lhs_pb = metadata - .lhs_dtype - .as_ref() - .ok_or_else(|| vortex_err!("metadata missing lhs_dtype"))?; - let rhs_pb = metadata - .rhs_dtype - .as_ref() - .ok_or_else(|| vortex_err!("metadata missing rhs_dtype"))?; - - let lhs_dtype = DType::from_proto(lhs_pb, session)?; - let rhs_dtype = DType::from_proto(rhs_pb, session)?; - validate_binary_tensor_float_inputs(&lhs_dtype, &rhs_dtype)?; - - let lhs = children.get(0, &lhs_dtype, len)?; - let rhs = children.get(1, &rhs_dtype, len)?; - Ok(vec![lhs, rhs]) - } -} - #[cfg(test)] pub mod test_helpers { use vortex_array::ArrayRef; @@ -358,9 +426,9 @@ pub mod test_helpers { } /// Builds a [`ConstantArray`] whose scalar is itself a [`Vector`] extension scalar, broadcast - /// to `len` rows. This is the shape produced by an `lit(vector_scalar)` literal expression — - /// the constant lives at the extension level rather than inside the FSL storage, in contrast - /// to [`Vector::constant_array`]. + /// to `len` rows. This is the shape produced by an `lit(vector_scalar)` literal expression, where + /// the constant lives at the extension level rather than inside the FSL storage, in contrast to + /// [`Vector::constant_array`]. pub fn literal_vector_array>( elements: &[T], len: usize, @@ -401,10 +469,10 @@ pub mod test_helpers { if a.is_nan() && e.is_nan() { continue; } + let diff = (a - e).abs(); assert!( (a - e).abs() < 1e-10, - "element {i}: got {a}, expected {e} (diff = {})", - (a - e).abs() + "element {i}: got {a}, expected {e} (diff = {diff})" ); } } diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index ad3b96d1bff..492bc837b89 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -35,11 +35,13 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; @@ -79,7 +81,7 @@ pub fn build_similarity_search_tree>( let num_rows = data.len(); let query_vec = Vector::constant_array(query, num_rows)?; - let cosine = CosineSimilarity::try_new_array(data, query_vec)?.into_array(); + let cosine = CosineSimilarity.try_new_array(num_rows, EmptyOptions, [data, query_vec])?; let threshold_scalar = Scalar::primitive(threshold, Nullability::NonNullable); let threshold_array = ConstantArray::new(threshold_scalar, num_rows).into_array(); From 6c13e8516a0440f17ed74ccaeb87d68bbb38d13b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 13:03:43 +0000 Subject: [PATCH 005/160] Port the geo predicates to RowFn Progress towards #9128. `vortex.geo.distance`, `vortex.geo.contains`, and `vortex.geo.intersects` become row functions, which deletes `scalar_fn/execute.rs` and the shared columnar execution it held. Decoding a geometry does expensive per-row work, so the geo elements set `InputElement::FILTERED_DECODE_COST` and sparse batches keep the filter strategy's shrunken decode. Branch-and-skip needs a decode that tolerates null rows without parsing them, so `geometries_null_tolerant` writes a placeholder geometry into null slots for `Point` and `Polygon`. It returns `Ok(None)` for any other geometry type, and the caller falls back to the filter strategy, which never decodes a null row. `contains` prepares a constant geometry once per batch rather than re-preparing it per row. The row layer sees through extension-over-constant, so the hand-written rewrite that used to uncover the constant is gone. `geo` is pinned to `=0.31.0` in the workspace manifest. `contains_route` transcribes geo's `impl_contains_from_relate!` dispatch table, so any bump that moves a row silently changes containment verdicts while the tests stay green wherever relate and the direct algorithm agree. A caret requirement would let `cargo update` take 0.31.x with no diff to review. `null_strategies` benchmarks the three strategies against `GeoContains` across validity densities, which is where the crossover between filtering and branch-and-skip was measured. Signed-off-by: Connor Tsui Co-authored-by: Claude --- Cargo.toml | 8 +- vortex-spatial/Cargo.toml | 8 +- vortex-spatial/benches/null_strategies.rs | 199 ++++++ vortex-spatial/src/extension/mod.rs | 41 ++ vortex-spatial/src/extension/point.rs | 18 + vortex-spatial/src/extension/polygon.rs | 18 + vortex-spatial/src/scalar_fn/contains.rs | 649 +++++++++++++++--- vortex-spatial/src/scalar_fn/distance.rs | 108 +-- vortex-spatial/src/scalar_fn/execute.rs | 19 +- .../src/scalar_fn/execute/binary.rs | 334 --------- .../src/scalar_fn/execute/geo_types.rs | 144 ---- vortex-spatial/src/scalar_fn/execute/unary.rs | 2 - vortex-spatial/src/scalar_fn/intersects.rs | 240 +++++-- vortex-spatial/src/scalar_fn/mod.rs | 1 + vortex-spatial/src/scalar_fn/row.rs | 170 +++++ vortex-spatial/src/test_harness.rs | 2 +- 16 files changed, 1224 insertions(+), 737 deletions(-) create mode 100644 vortex-spatial/benches/null_strategies.rs delete mode 100644 vortex-spatial/src/scalar_fn/execute/binary.rs delete mode 100644 vortex-spatial/src/scalar_fn/execute/geo_types.rs create mode 100644 vortex-spatial/src/scalar_fn/row.rs diff --git a/Cargo.toml b/Cargo.toml index 36aa5b2ac9e..5328905cc95 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -160,7 +160,13 @@ flatbuffers = "25.2.10" fsst-rs = "0.6.0" futures = { version = "0.3.31", default-features = false } fuzzy-matcher = "0.3" -geo = "0.31.0" +# `vortex-spatial`'s `contains_route` transcribes geo's `impl_contains_from_relate!` dispatch table, so +# any bump that moves a row silently changes containment verdicts — the tests stay green wherever +# relate and the direct algorithm agree. Pinned exactly so that taking any new geo, patch releases +# included, is a deliberate edit of this line that re-verifies the table; a caret requirement would +# let `cargo update` (or automated lockfile maintenance) take 0.31.x with no diff to review. See +# `vortex-spatial/src/scalar_fn/contains.rs`. +geo = "=0.31.0" geo-traits = "0.3.0" geo-types = "0.7.19" geoarrow = "0.8.0" diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 8b790da3c05..139f28181d6 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -47,11 +47,15 @@ name = "envelope" harness = false [[bench]] -name = "predicate_bbox" +name = "binary_predicates" harness = false [[bench]] -name = "binary_predicates" +name = "null_strategies" +harness = false + +[[bench]] +name = "predicate_bbox" harness = false [[bench]] diff --git a/vortex-spatial/benches/null_strategies.rs b/vortex-spatial/benches/null_strategies.rs new file mode 100644 index 00000000000..1ef453d645b --- /dev/null +++ b/vortex-spatial/benches/null_strategies.rs @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Null-strategy comparison for the geo `contains` kernel, whose per-row geometry decode is what +//! the selection threshold exists for. +//! +//! Arms: `filter` and `branch` force one strategy through the test-harness seam +//! ([`execute_row_fn_with_strategy`]); `auto` executes the full pipeline and lets the per-batch +//! selection choose, which should track the faster forced arm on both sides of the crossover +//! (branch at dense validity, filter at sparse). +//! +//! Workloads: a column of small polygons CONTAINS a constant point, and polygon column CONTAINS +//! point column with independent nulls on both, each at null densities 0/1/5/10/25/50/90 percent +//! over 65536 rows, nulls placed by a seeded splitmix hash. +//! +//! Run with `cargo bench -p vortex-spatial --bench null_strategies`. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::MaskedArray; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::NullStrategy; +use vortex_array::scalar_fn::execute_row_fn_with_strategy; +use vortex_array::validity::Validity; +use vortex_spatial::scalar_fn::contains::SpatialContains; +use vortex_spatial::test_harness::geo_session; +use vortex_spatial::test_harness::point_column; +use vortex_spatial::test_harness::polygon_column; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(geo_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +const ROWS: usize = 65536; + +/// Null densities in percent. +const DENSITIES: &[usize] = &[0, 1, 5, 10, 25, 50, 90]; + +/// Deterministic pseudo-random value in `[0, 1)` (same generator as `binary_predicates`). +fn unit(i: usize) -> f64 { + ((i.wrapping_mul(2654435761) >> 8) % 10_000) as f64 / 10_000.0 +} + +/// splitmix64, for seeded random null placement. +fn splitmix64(mut x: u64) -> u64 { + x = x.wrapping_add(0x9E3779B97F4A7C15); + x = (x ^ (x >> 30)).wrapping_mul(0xBF58476D1CE4E5B9); + x = (x ^ (x >> 27)).wrapping_mul(0x94D049BB133111EB); + x ^ (x >> 31) +} + +/// A small square (side 2) centered at `(cx, cy)`. +fn square(cx: f64, cy: f64) -> Vec> { + vec![vec![ + (cx - 1.0, cy - 1.0), + (cx + 1.0, cy - 1.0), + (cx + 1.0, cy + 1.0), + (cx - 1.0, cy + 1.0), + (cx - 1.0, cy - 1.0), + ]] +} + +/// [`ROWS`] small squares spread over roughly `[-150, 150)^2`; a handful contain the origin, and +/// each row's verdict is a direct point-in-polygon test. +fn squares() -> ArrayRef { + let rows = (0..ROWS) + .map(|i| square(300.0 * unit(i) - 150.0, 300.0 * unit(i + 1) - 150.0)) + .collect(); + polygon_column(rows).unwrap() +} + +/// [`ROWS`] points over the same region. +fn points() -> ArrayRef { + let xs = (0..ROWS).map(|i| 300.0 * unit(i + 7) - 150.0).collect(); + let ys = (0..ROWS).map(|i| 300.0 * unit(i + 8) - 150.0).collect(); + point_column(xs, ys).unwrap() +} + +/// The constant point operand, at the origin so some squares contain it. +fn constant_point(ctx: &mut ExecutionCtx) -> ArrayRef { + let scalar = point_column(vec![0.0], vec![0.0]) + .unwrap() + .execute_scalar(0, ctx) + .unwrap(); + ConstantArray::new(scalar, ROWS).into_array() +} + +/// Wrap `array` with seeded random nulls at `density` percent. Zero density stays unwrapped, as a +/// non-nullable column would. +fn with_nulls(array: ArrayRef, seed: u64, density: usize) -> ArrayRef { + if density == 0 { + return array; + } + + let valid = (0..ROWS).map(|i| (splitmix64(seed ^ i as u64) % 100) >= density as u64); + MaskedArray::try_new(array, Validity::from_iter(valid)) + .unwrap() + .into_array() +} + +/// One arm over the operand pair: `Some` forces a strategy through the harness seam, `None` runs +/// the full pipeline with the per-batch selection. +fn bench_contains(bencher: Bencher, a: ArrayRef, b: ArrayRef, strategy: Option) { + let mut ctx = SESSION.create_execution_ctx(); + + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| match strategy { + None => SpatialContains::try_new_array(a.clone(), b.clone()) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap(), + Some(strategy) => execute_row_fn_with_strategy( + &SpatialContains, + &EmptyOptions, + vec![a.clone(), b.clone()], + ROWS, + strategy, + &mut ctx, + ) + .unwrap() + .execute::(&mut ctx) + .unwrap(), + }); +} + +/// Column of polygons CONTAINS constant point, nulls on the polygon column. +mod polygons_x_constant_point { + use super::*; + + fn operands(density: usize) -> (ArrayRef, ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + (with_nulls(squares(), 1, density), constant_point(&mut ctx)) + } + + #[divan::bench(args = DENSITIES)] + fn filter(bencher: Bencher, density: usize) { + let (a, b) = operands(density); + bench_contains(bencher, a, b, Some(NullStrategy::Filter)); + } + + #[divan::bench(args = DENSITIES)] + fn branch(bencher: Bencher, density: usize) { + let (a, b) = operands(density); + bench_contains(bencher, a, b, Some(NullStrategy::BranchAndSkip)); + } + + #[divan::bench(args = DENSITIES)] + fn auto(bencher: Bencher, density: usize) { + let (a, b) = operands(density); + bench_contains(bencher, a, b, None); + } +} + +/// Column of polygons CONTAINS column of points, independent nulls on both, so the conjoined +/// valid fraction is roughly `(1 - d)^2`. +mod polygons_x_points { + use super::*; + + fn operands(density: usize) -> (ArrayRef, ArrayRef) { + ( + with_nulls(squares(), 1, density), + with_nulls(points(), 2, density), + ) + } + + #[divan::bench(args = DENSITIES)] + fn filter(bencher: Bencher, density: usize) { + let (a, b) = operands(density); + bench_contains(bencher, a, b, Some(NullStrategy::Filter)); + } + + #[divan::bench(args = DENSITIES)] + fn branch(bencher: Bencher, density: usize) { + let (a, b) = operands(density); + bench_contains(bencher, a, b, Some(NullStrategy::BranchAndSkip)); + } + + #[divan::bench(args = DENSITIES)] + fn auto(bencher: Bencher, density: usize) { + let (a, b) = operands(density); + bench_contains(bencher, a, b, None); + } +} diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index d1e2c37ebf4..e89ec15b500 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -178,6 +178,47 @@ pub(crate) fn geometries( } } +/// The geometry a null row decodes to under [`geometries_null_tolerant`]. Arbitrary: the caller +/// guarantees null rows are never read. +pub(crate) fn placeholder_geometry() -> Geometry { + Geometry::Point(geo_types::Point::new(0.0, 0.0)) +} + +/// Decode a native geometry column that may contain null rows, writing [`placeholder_geometry`] +/// into their slots. The caller guarantees null rows are never read. +/// +/// `Ok(None)` means this geometry type has no null-tolerant decode yet (`Point` and `Polygon` are +/// covered), and the caller falls back to the filter strategy, which never decodes a null row. A +/// column with definitely no nulls delegates to the ordinary [`geometries`] for any type. +pub(crate) fn geometries_null_tolerant( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>>> { + if array.validity()?.definitely_no_nulls() { + return geometries(array, ctx).map(Some); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + + if ext.is::() { + point_geometries_null_tolerant(&storage, ctx).map(Some) + } else if ext.is::() { + polygon_geometries_null_tolerant(&storage, ctx).map(Some) + } else { + Ok(None) + } +} + /// Decode a constant operand scalar to one geometry, a constant of any /// supported geometry type is decoded exactly like a column. pub(crate) fn single_geometry( diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index e8f3ad3c169..8189fcba7bf 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -52,6 +52,7 @@ use super::coordinate::coordinate_from_struct; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A single location: `geoarrow.point`, stored as `Struct` of non-nullable `f64`. @@ -149,6 +150,23 @@ pub(crate) fn point_geometries( .collect() } +/// Like [`point_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn point_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + point_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + impl ArrowExportVTable for Point { fn arrow_ext_id(&self) -> Id { *ARROW_POINT diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index dcfa8514ff3..362dfe311e9 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -52,6 +52,7 @@ use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A polygon: `geoarrow.polygon`, stored as `List>>` (rings of vertices). @@ -131,6 +132,23 @@ pub(crate) fn polygon_geometries( .collect() } +/// Like [`polygon_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn polygon_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + polygon_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + /// Build a geoarrow `PolygonArray` from a `Polygon`'s `List>` storage. fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let polygon_type = polygon_type( diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 599c0eee2be..54316678f5f 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -3,44 +3,30 @@ //! `ST_Contains`: OGC containment test between two native geometries. +use std::cell::OnceCell; + +use geo::BoundingRect; use geo::Contains; +use geo::PreparedGeometry; +use geo::Relate; +use geo_types::Geometry; +use geo_types::Rect; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ElementSink; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Contains`. -fn validate_contains_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: contains requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: contains operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; +#[cfg(test)] +use crate::scalar_fn::row::probe; /// OGC `ST_Contains` between two native geometry operands, each a column or a constant /// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone @@ -59,83 +45,297 @@ impl SpatialContains { } } -impl ScalarFnVTable for SpatialContains { +impl RowFn for SpatialContains { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.contains"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { Ok(EmptyOptions) } - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) + /// Containment is not symmetric, so `a` is always the container and `b` the contained. + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), ElementSink, _, _>( + |(a, b)| { + #[cfg(test)] + probe::record(a.is_some(), b.is_some()); + ConstOperands { + a: a.map(PreparedOperand::new), + b: b.map(PreparedOperand::new), + } + }, + |operands, (a, b), output| *output = contains_row_prepared(operands, a, b), + ) } +} + +/// Per-batch state for the contains row kernel: the prepared form of whichever operand is +/// constant for the batch. `None` marks an operand that varies by row. +struct ConstOperands { + /// Operand `a` (the container) when it is batch-constant. + a: Option, + + /// Operand `b` (the contained) when it is batch-constant. + b: Option, +} + +/// One batch-constant operand: the geometry cloned out of its decoded column (the state must not +/// borrow from the columns), plus its [`PreparedGeometry`], built on the first row whose pairing +/// routes through relate. +/// +/// The build is lazy because preparation (self-noding the topology graph plus an R*-tree over the +/// edges) costs `O(edges log edges)` and pays off only on relate-routed pairings; a batch of +/// point rows against a constant polygon never touches it, and preparing a large constant eagerly +/// would charge such a batch for nothing. +struct PreparedOperand { + /// The constant's decoded geometry, owned so [`prepared`](Self::prepared) can be `'static`. + geometry: Geometry, + + /// The constant's bounding rectangle, folded once for conservative row rejection. + bbox: Option>, + + /// The lazily built prepared form of [`geometry`](Self::geometry). + prepared: OnceCell, f64>>, +} - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("contains has exactly two children"), +impl PreparedOperand { + fn new(geometry: &Geometry) -> Self { + Self { + geometry: geometry.clone(), + bbox: geometry.bounding_rect(), + prepared: OnceCell::new(), } } - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_contains_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Bool(nullability)) + /// The prepared geometry, built on first use. + fn get(&self) -> &PreparedGeometry<'static, Geometry, f64> { + self.prepared + .get_or_init(|| PreparedGeometry::from(self.geometry.clone())) } +} - fn execute( - &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Containment is not symmetric: `a` is always the container and `b` the contained. A - // container's rect must cover the contained's rect (`Rect::contains` is the closed - // test), so a contained rect poking outside proves the row false. - execute_binary_geo_types( - &a, - &b, - |a, b| a.contains(b), - Some(|ra, rb| (!ra.contains(rb)).then_some(false)), - ctx, - ) - } +/// How geo's `a.contains(b)` computes its verdict for a pairing. +enum ContainsRoute { + /// `a.relate(b).is_contains()`. + ForwardRelate, - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) + /// `b.relate(a).is_within()`, how geo phrases relate for `MultiPolygon` containers. + ReversedRelate, + + /// A direct algorithm (coordinate position, point arithmetic); nothing to prepare. + Direct, +} + +/// The route geo 0.31's `Contains` dispatch takes for `a.contains(b)`. +/// +/// The prepared substitution in [`contains_row_prepared`] **must** run relate exactly where geo +/// runs relate, with the same argument order, because geo's direct algorithms are not everywhere +/// bit-identical to a relate matrix query (they resolve degenerate and boundary cases with +/// different arithmetic). The relate rows below transcribe geo's `impl_contains_from_relate!` +/// lists per container type; everything else, notably every `Point`/`MultiPoint` contained side +/// and every `Point` container, is direct. +/// +/// **This table is coupled to the geo version.** It transcribes a dispatch that geo is free to +/// reshuffle in any release, and a wrong row is a silently wrong verdict rather than a build error. +/// The workspace therefore pins `geo = "=0.31.0"`: taking any new geo, patch releases included, is +/// a deliberate edit of that line, and the edit must re-verify this table against +/// `impl_contains_from_relate!`. +/// +/// `constant_operands_agree_with_columns` is the mechanical check, and it is **not** complete: it +/// compares the prepared route against plain `a.contains(b)` only for the container types it has +/// cases for. `routes_agree_with_geo_for_every_container` covers the rest, one representative +/// pairing per container variant, and is the one to extend when geo grows a geometry type. Both +/// stay green wherever relate and the direct algorithm agree, so neither replaces the pin. +fn contains_route(a: &Geometry, b: &Geometry) -> ContainsRoute { + use Geometry as G; + + match (a, b) { + // Line contains [Polygon, MultiLineString, MultiPolygon, GeometryCollection, Rect, + // Triangle]. + ( + G::Line(_), + G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // LineString contains [Polygon, MultiPoint, MultiLineString, MultiPolygon, + // GeometryCollection, Rect, Triangle]. + | ( + G::LineString(_), + G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // MultiLineString contains everything except Point. + | ( + G::MultiLineString(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // MultiPoint contains [Line, LineString, Polygon, MultiLineString, MultiPolygon, + // GeometryCollection, Rect, Triangle]. + | ( + G::MultiPoint(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // Polygon contains everything except Point and MultiPoint. + | ( + G::Polygon(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // Rect contains [Line, LineString, MultiPoint, MultiLineString, MultiPolygon, + // GeometryCollection, Triangle]; Rect contains Rect and Polygon are direct. + | ( + G::Rect(_), + G::Line(_) + | G::LineString(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Triangle(_), + ) + // Triangle and GeometryCollection contain everything except Point. + | ( + G::Triangle(_) | G::GeometryCollection(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) => ContainsRoute::ForwardRelate, + + // MultiPolygon contains everything except Point and MultiPoint, phrased reversed. + ( + G::MultiPolygon(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) => ContainsRoute::ReversedRelate, + + _ => ContainsRoute::Direct, } +} - fn is_strict(&self, _: &Self::Options) -> bool { - true +/// Computes one row of contains, substituting a prepared graph for a constant operand on the +/// pairings geo itself answers through relate. +/// +/// [`PreparedGeometry`] carries the operand's self-noded topology graph and edge R*-tree, so a +/// relate against it skips rebuilding both and reads its bounding rect from cache; geo asserts +/// the cached graph equal to a freshly built one (its `swap_arg_index` test), which is what makes +/// the substitution result-preserving. Before dispatch, a disjoint constant-side bounding rect +/// conservatively rejects the row, matching the columnar implementation's #9076 optimization. +/// All other rows delegate to the same direct or relate route as `a.contains(b)`. +fn contains_row_prepared(operands: &ConstOperands, a: &Geometry, b: &Geometry) -> bool { + let rejected = match (&operands.a, &operands.b) { + (None, None) => false, + (Some(const_a), Some(const_b)) => const_a + .bbox + .zip(const_b.bbox) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + (Some(const_a), None) => const_a + .bbox + .zip(b.bounding_rect()) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + (None, Some(const_b)) => a + .bounding_rect() + .zip(const_b.bbox) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + }; + + if rejected { + return false; } - fn is_fallible(&self, _: &Self::Options) -> bool { - false + match contains_route(a, b) { + ContainsRoute::Direct => a.contains(b), + ContainsRoute::ForwardRelate => match (&operands.a, &operands.b) { + (Some(const_a), Some(const_b)) => const_a.get().relate(const_b.get()).is_contains(), + (Some(const_a), None) => const_a.get().relate(b).is_contains(), + (None, Some(const_b)) => a.relate(const_b.get()).is_contains(), + (None, None) => a.contains(b), + }, + ContainsRoute::ReversedRelate => match (&operands.a, &operands.b) { + (Some(const_a), Some(const_b)) => const_b.get().relate(const_a.get()).is_within(), + (Some(const_a), None) => b.relate(const_a.get()).is_within(), + (None, Some(const_b)) => const_b.get().relate(a).is_within(), + (None, None) => a.contains(b), + }, } } #[cfg(test)] mod tests { + use geo::Contains; + use geo_types::Coord; use geo_types::Geometry; + use geo_types::GeometryCollection; + use geo_types::Line; use geo_types::LineString; + use geo_types::MultiLineString; + use geo_types::MultiPoint; + use geo_types::MultiPolygon; use geo_types::Point; use geo_types::Polygon; + use geo_types::Rect; + use geo_types::Triangle; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -144,23 +344,31 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::MaskedArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::NullStrategy; use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::scalar_fn::execute_row_fn_with_strategy; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_err; use wkb::writer::WriteOptions; + use super::ConstOperands; + use super::PreparedOperand; use super::SpatialContains; + use super::contains_row_prepared; + use crate::scalar_fn::row::probe::assert_prepared_agrees_with_columns; use crate::test_harness::linestring_column; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; + use crate::test_harness::polygon_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { @@ -244,6 +452,20 @@ mod tests { assert_contains(container, points, [true, false, false]) } + /// Constant container vs a linestring column: a row whose bounding rect pokes outside the + /// container's is not contained, while one wholly inside is. Carried over from the columnar + /// bounding-rect rejection in #9076, since it constrains the verdict rather than the mechanism. + #[test] + fn constant_container_vs_row_rect_poking_outside() -> VortexResult<()> { + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let lines = linestring_column(vec![ + vec![(1.0, 1.0), (3.0, 3.0)], + vec![(1.0, 1.0), (9.0, 1.0)], + vec![(5.0, 5.0), (9.0, 9.0)], + ])?; + assert_contains(container, lines, [true, false, false]) + } + /// Polygon column vs constant point: only the polygon around the point contains it. #[test] fn polygon_column_vs_constant_point() -> VortexResult<()> { @@ -264,20 +486,6 @@ mod tests { assert_contains(away, point, [false; 2]) } - /// Constant container vs a linestring column: a row whose bounding rect pokes outside the - /// container's rect is proven false by the rect pre-check alone; a fully inside row still - /// needs (and passes) the exact test. - #[test] - fn constant_container_vs_row_rect_poking_outside() -> VortexResult<()> { - let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; - let lines = linestring_column(vec![ - vec![(1.0, 1.0), (3.0, 3.0)], - vec![(1.0, 1.0), (9.0, 1.0)], - vec![(5.0, 5.0), (9.0, 9.0)], - ])?; - assert_contains(container, lines, [true, false, false]) - } - /// Column vs column pairs rows: each polygon row is tested against the point row at the /// same position. #[test] @@ -408,6 +616,117 @@ mod tests { Ok(()) } + /// A nullable polygon column: unit squares at `centers`, the rows where `nulls` is true + /// masked out, spelled as `Masked` over non-nullable storage. + fn nullable_squares(centers: &[(f64, f64)], nulls: &[bool]) -> VortexResult { + let squares = centers + .iter() + .map(|&(x, y)| { + vec![vec![ + (x - 1.0, y - 1.0), + (x + 1.0, y - 1.0), + (x + 1.0, y + 1.0), + (x - 1.0, y + 1.0), + (x - 1.0, y - 1.0), + ]] + }) + .collect(); + let polygons = polygon_column(squares)?; + + Ok( + MaskedArray::try_new(polygons, Validity::from_iter(nulls.iter().map(|n| !n)))? + .into_array(), + ) + } + + /// Executes `SpatialContains(a, b)` with a forced null strategy, canonicalized. + fn contains_forced( + a: &ArrayRef, + b: &ArrayRef, + strategy: NullStrategy, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(execute_row_fn_with_strategy( + &SpatialContains, + &EmptyOptions, + vec![a.clone(), b.clone()], + a.len(), + strategy, + ctx, + )? + .execute::(ctx)? + .into_array()) + } + + /// The branch-and-skip and filter strategies, plus the automatic per-batch selection, must + /// return identical arrays for nullable geometry operands: `Masked` polygons against nullable + /// points, with independent nulls conjoined. + #[test] + fn branch_matches_filter_for_nullable_geometries() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let centers = [(0.0, 0.0), (5.0, 5.0), (0.5, -0.2), (9.0, 9.0), (0.0, 1.0)]; + let nulls = [false, true, false, false, true]; + let polygons = nullable_squares(¢ers, &nulls)?; + let points = nullable_point_column(vec![ + Some((0.0, 0.0)), + Some((5.0, 5.0)), + None, + Some((0.0, 0.0)), + Some((0.0, 1.0)), + ])?; + + let filtered = contains_forced(&polygons, &points, NullStrategy::Filter, &mut ctx)?; + let branched = contains_forced(&polygons, &points, NullStrategy::BranchAndSkip, &mut ctx)?; + let auto = SpatialContains::try_new_array(polygons, points)? + .into_array() + .execute::(&mut ctx)? + .into_array(); + + assert_arrays_eq!(branched, filtered, &mut ctx); + assert_arrays_eq!(auto, filtered, &mut ctx); + Ok(()) + } + + /// Geometry types without a null-tolerant decode refuse the branch strategy: forcing it is an + /// error, and the automatic selection (which prefers branch at this density) silently falls + /// back to filtering with the correct result. + #[test] + fn unsupported_geometry_falls_back_to_filter() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + // Four rows with one null: 75% surviving, so the selection prefers branch. + let lines = MaskedArray::try_new( + linestring_column(vec![ + vec![(0.0, 0.0), (4.0, 4.0)], + vec![(0.0, 0.0), (1.0, 1.0)], + vec![(2.0, 2.0), (3.0, 3.0)], + vec![(0.0, 4.0), (4.0, 0.0)], + ])?, + Validity::from_iter([true, false, true, true]), + )? + .into_array(); + let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 4)?; + + let error = contains_forced(&lines, &point, NullStrategy::BranchAndSkip, &mut ctx) + .expect_err("a linestring column with nulls has no branch decode"); + assert!( + error.to_string().contains("branch-and-skip"), + "unexpected error: {error}" + ); + + let filtered = contains_forced(&lines, &point, NullStrategy::Filter, &mut ctx)?; + let auto = SpatialContains::try_new_array(lines, point)? + .into_array() + .execute::(&mut ctx)? + .into_array(); + + assert_arrays_eq!(auto, filtered, &mut ctx); + Ok(()) + } + /// A non-geometry operand dtype is rejected up front, before execution. #[test] fn non_geometry_operand_is_rejected() -> VortexResult<()> { @@ -417,4 +736,166 @@ mod tests { assert!(result.is_err()); Ok(()) } + + // The prepared-vs-expanded agreement grid: every constant arrangement of a pairing must + // return exactly what the fully expanded columns return. + + /// A point geometry. + fn point(x: f64, y: f64) -> Geometry { + Geometry::Point(Point::new(x, y)) + } + + /// A linestring geometry through `coords`. + fn line(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::LineString(LineString::from(coords)) + } + + /// A multipoint geometry over `coords`. + fn multipoint(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::MultiPoint(MultiPoint::from(coords)) + } + + /// A two-point line segment geometry, the `Line` container variant. + fn line_geometry(start: (f64, f64), end: (f64, f64)) -> Geometry { + Geometry::Line(Line::new( + Coord { + x: start.0, + y: start.1, + }, + Coord { x: end.0, y: end.1 }, + )) + } + + /// A multilinestring geometry over one linestring per entry of `parts`. + fn multilinestring(parts: Vec>) -> Geometry { + Geometry::MultiLineString(MultiLineString::new( + parts.into_iter().map(LineString::from).collect(), + )) + } + + /// A geometry collection wrapping `parts`. + fn collection(parts: Vec) -> Geometry { + Geometry::GeometryCollection(GeometryCollection::from(parts)) + } + + /// An axis-aligned rectangle geometry, the `Rect` container variant. + fn rect_geometry(x0: f64, y0: f64, x1: f64, y1: f64) -> Geometry { + Geometry::Rect(Rect::new(Coord { x: x0, y: y0 }, Coord { x: x1, y: y1 })) + } + + /// A triangle geometry large enough to contain the small test polygons. + fn triangle_geometry() -> Geometry { + Geometry::Triangle(Triangle::new( + Coord { x: 0.0, y: 0.0 }, + Coord { x: 8.0, y: 0.0 }, + Coord { x: 0.0, y: 8.0 }, + )) + } + + /// A two-part multipolygon: `4x4` squares at the origin and at `(10, 10)`. + fn two_part_multipolygon() -> Geometry { + Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 4.0, 4.0), + rect_polygon(10.0, 10.0, 14.0, 14.0), + ])) + } + + /// Every container variant `contains_route` distinguishes, checked against plain + /// `a.contains(b)` in all four constant arrangements. + /// + /// Every case is a containment geo answers `true`, which the test asserts: a pairing that is + /// false regardless of route (a lower-dimensional container, say) also agrees regardless of + /// route, and pins nothing. A true case fails when the prepared substitution diverges from + /// geo — a table row whose relate phrasing disagrees with geo's dispatch on this input, or a + /// bounding-rect prescreen that wrongly rejects a contained row. It is **not** a version + /// tripwire: a geo release that reshuffles its dispatch stays green wherever relate and the + /// direct algorithm agree, which is why the workspace pins `geo` exactly. + /// + /// This is the table's own regression, and the one to extend when geo grows a geometry type: + /// `constant_operands_agree_with_columns` below goes through real arrays and so is the better + /// end-to-end check, but it only covers the container types it has cases for, and WKB decoding + /// limits which types those can be. The MultiPoint and Line containers route relate only for + /// contained types a MultiPoint or Line can rarely contain, so their true cases lean on + /// `GeometryCollection` membership and collinear `MultiLineString` parts respectively. + #[rstest] + #[case::point(point(1.0, 1.0), point(1.0, 1.0))] + #[case::line(line_geometry((0.0, 0.0), (4.0, 4.0)), point(2.0, 2.0))] + #[case::line_x_multilinestring(line_geometry((0.0, 0.0), (4.0, 4.0)), multilinestring(vec![vec![(1.0, 1.0), (2.0, 2.0)]]))] + #[case::linestring(line(vec![(0.0, 0.0), (4.0, 4.0)]), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::multipoint(multipoint(vec![(0.0, 0.0), (2.0, 2.0), (4.0, 4.0)]), collection(vec![point(2.0, 2.0)]))] + #[case::multilinestring(multilinestring(vec![vec![(0.0, 0.0), (4.0, 4.0)]]), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipolygon(two_part_multipolygon(), rect_polygon(1.0, 1.0, 3.0, 3.0).into())] + #[case::geometrycollection(collection(vec![rect_polygon(0.0, 0.0, 8.0, 8.0).into()]), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::rect(rect_geometry(0.0, 0.0, 8.0, 8.0), line(vec![(2.0, 2.0), (4.0, 4.0)]))] + #[case::triangle(triangle_geometry(), rect_polygon(1.0, 1.0, 2.0, 2.0).into())] + fn routes_agree_with_geo_for_every_container(#[case] a: Geometry, #[case] b: Geometry) { + let expected = a.contains(&b); + assert!( + expected, + "route cases must be containments geo answers true, or every route agrees vacuously", + ); + + let arrangements = [ + (None, None), + (Some(PreparedOperand::new(&a)), None), + (None, Some(PreparedOperand::new(&b))), + ( + Some(PreparedOperand::new(&a)), + Some(PreparedOperand::new(&b)), + ), + ]; + + for (index, (const_a, const_b)) in arrangements.into_iter().enumerate() { + let operands = ConstOperands { + a: const_a, + b: const_b, + }; + assert_eq!( + contains_row_prepared(&operands, &a, &b), + expected, + "arrangement {index} disagrees with geo's own contains", + ); + } + } + + /// Constant arrangements agree with expanded columns across the routes the prepared kernel + /// distinguishes: forward relate (polygon, linestring and multipoint containers), reversed + /// relate (multipolygon containers), and the direct pairings (a point on either side, + /// multipoint over multipoint, polygon over multipoint), including boundary contact, + /// crossing, disjoint and empty cases. + #[rstest] + #[case::polygon_nested_polygon(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::polygon_touching_from_inside(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(0.0, 2.0, 2.0, 4.0).into())] + #[case::polygon_overlapping_polygon(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(2.0, 2.0, 6.0, 6.0).into())] + #[case::polygon_disjoint_polygon(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::polygon_x_point_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(2.0, 2.0))] + #[case::polygon_x_point_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(0.0, 2.0))] + #[case::polygon_x_point_outside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(20.0, 20.0))] + #[case::polygon_x_nan_point(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(f64::NAN, 2.0))] + #[case::polygon_x_linestring_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon_x_linestring_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(0.0, 1.0), (0.0, 3.0)]))] + #[case::polygon_x_linestring_crossing(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(-2.0, 2.0), (2.0, 2.0)]))] + #[case::polygon_x_empty_linestring(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![]))] + #[case::polygon_x_multipoint_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon_x_multipoint_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), multipoint(vec![(0.0, 1.0), (0.0, 3.0)]))] + #[case::linestring_x_multipoint_on_line(line(vec![(0.0, 0.0), (4.0, 4.0)]), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipoint_x_multipoint_subset(multipoint(vec![(0.0, 0.0), (2.0, 2.0), (4.0, 4.0)]), multipoint(vec![(2.0, 2.0)]))] + #[case::multipoint_x_linestring_between_points(multipoint(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipolygon_x_polygon_in_one_part(two_part_multipolygon(), rect_polygon(1.0, 1.0, 3.0, 3.0).into())] + #[case::multipolygon_x_polygon_straddling(two_part_multipolygon(), rect_polygon(3.0, 3.0, 11.0, 11.0).into())] + #[case::multipolygon_x_polygon_disjoint(two_part_multipolygon(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::multipolygon_x_point_inside(two_part_multipolygon(), point(11.0, 11.0))] + #[case::point_x_point_equal(point(1.0, 1.0), point(1.0, 1.0))] + #[case::point_x_polygon(point(2.0, 2.0), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + fn constant_operands_agree_with_columns( + #[case] a: Geometry, + #[case] b: Geometry, + ) -> VortexResult<()> { + assert_prepared_agrees_with_columns( + SpatialContains::try_new_array, + geometry_constant(&a, 3)?, + geometry_constant(&b, 3)?, + ) + } } diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index dfd3d09ed23..e8de808bd32 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -6,43 +6,19 @@ use geo::Distance; use geo::Euclidean; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ElementSink; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Distance`. -fn validate_distance_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: distance requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: distance operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -60,66 +36,45 @@ impl SpatialDistance { } } -impl ScalarFnVTable for SpatialDistance { +impl RowFn for SpatialDistance { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.distance"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { - Ok(EmptyOptions) - } - - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("distance has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_distance_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Primitive(PType::F64, nullability)) - } - - fn execute( + fn deserialize( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Distance is a value, not a verdict: no bounding-rect test can decide it. - execute_binary_geo_types(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx) + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn validity( + /// Deliberately uses unit preparation: a batch-constant operand offers nothing + /// sound to hoist. geo computes linestring and polygon distances through its private + /// `nearest_neighbour_distance`, which builds the R*-trees for *both* sides inside each call, + /// and the point pairings are single expressions; reusing a tree across rows would mean + /// reimplementing geo's internals. A batch where both operands are constant already folds to + /// a single-row execution before the row loop. + fn dispatch( &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } - - fn is_strict(&self, _: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _: &Self::Options) -> bool { - false + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), ElementSink, _, _>( + |_| (), + |&(), (a, b), output| *output = Euclidean.distance(a, b), + ) } } @@ -196,8 +151,9 @@ mod tests { Ok(()) } - /// Distance passes no bounding-rect rejection: a point far outside a constant polygon's - /// bounding rect still gets its true distance, alongside an inside point at distance zero. + /// Distance is a value rather than a verdict, so no bounding-rect rejection may fire for it: a + /// point far outside a constant polygon's rect still gets its true distance. Carried over from + /// #9076, which added the rejection to the predicates but deliberately not to this function. #[test] fn distance_to_constant_polygon_is_exact() -> VortexResult<()> { let session = vortex_array::array_session(); diff --git a/vortex-spatial/src/scalar_fn/execute.rs b/vortex-spatial/src/scalar_fn/execute.rs index ca5b4018249..836577ec26e 100644 --- a/vortex-spatial/src/scalar_fn/execute.rs +++ b/vortex-spatial/src/scalar_fn/execute.rs @@ -1,24 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Shared execution for native geometry scalar functions. -//! -//! [`dispatch_unary`] and the binary dispatcher handle constant/column operands and strict null -//! propagation without prescribing how a kernel represents geometries or builds its output. -//! Native columnar kernels such as `ST_Envelope` use the unary dispatcher directly. -//! -//! [`execute_binary_geo_types`] adapts row-oriented algorithms from the `geo` ecosystem. It decodes -//! valid inputs into `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`], such -//! as an `f64` or boolean array. +//! Shared unary execution for native geometry scalar functions. -mod binary; -mod geo_types; mod unary; -pub(crate) use binary::execute_binary_geo_types; pub(crate) use unary::dispatch_unary; use vortex_array::ArrayRef; -use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_mask::Mask; @@ -31,9 +19,6 @@ pub(crate) enum Operand { } /// Shared batch state presented to a null-propagating geometry kernel with `N` operands. -/// -/// Binary kernels use the default materialized [`Mask`]. Unary columnar kernels can instead -/// retain a lazy [`vortex_array::validity::Validity`] until they need row-wise access. pub(crate) struct Execution { /// Constant/column shape of each operand. pub(crate) operands: [Operand; N], @@ -41,6 +26,4 @@ pub(crate) struct Execution { pub(crate) valid: V, /// Number of output rows. pub(crate) len: usize, - /// Output nullability from the scalar function's return dtype. - pub(crate) nullability: Nullability, } diff --git a/vortex-spatial/src/scalar_fn/execute/binary.rs b/vortex-spatial/src/scalar_fn/execute/binary.rs deleted file mode 100644 index f2c03bd1beb..00000000000 --- a/vortex-spatial/src/scalar_fn/execute/binary.rs +++ /dev/null @@ -1,334 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Binary pairwise dispatch, plus an adapter for row-oriented `geo_types` kernels. - -use geo::BoundingRect; -use geo_types::Geometry; -use geo_types::Rect; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::Constant; -use vortex_array::arrays::ConstantArray; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::scalar::Scalar; -use vortex_error::VortexResult; -use vortex_mask::Mask; - -use super::Execution; -use super::Operand; -use super::geo_types::GeoTypesOutput; -use super::geo_types::eval_column; -use super::geo_types::eval_column_pair; -use crate::extension::single_geometry; - -/// Dispatch a binary strict geometry kernel over constants and columns. -/// -/// A null constant or an empty combined validity mask short-circuits to an all-null constant -/// output. Otherwise, `kernel` receives both operand shapes and the mask of rows where both are -/// valid. Two columns are always paired by row index. The kernel remains responsible for physical -/// input interpretation and Vortex output construction. -pub(crate) fn dispatch_binary( - left: &ArrayRef, - right: &ArrayRef, - output_dtype: DType, - kernel: K, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - K: FnOnce(Execution<2>, &mut ExecutionCtx) -> VortexResult, -{ - let len = left.len(); - for operand in [left, right] { - if operand - .as_opt::() - .is_some_and(|constant| constant.scalar().is_null()) - { - return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); - } - } - - let (left, right, valid) = match (left.as_opt::(), right.as_opt::()) { - (Some(left), Some(right)) => ( - Operand::Constant(left.scalar().clone()), - Operand::Constant(right.scalar().clone()), - Mask::new_true(len), - ), - (Some(left), None) => ( - Operand::Constant(left.scalar().clone()), - Operand::Column(right.clone()), - right.validity()?.execute_mask(len, ctx)?, - ), - (None, Some(right)) => ( - Operand::Column(left.clone()), - Operand::Constant(right.scalar().clone()), - left.validity()?.execute_mask(len, ctx)?, - ), - (None, None) => { - let left_valid = left.validity()?.execute_mask(len, ctx)?; - let right_valid = right.validity()?.execute_mask(len, ctx)?; - ( - Operand::Column(left.clone()), - Operand::Column(right.clone()), - &left_valid & &right_valid, - ) - } - }; - - if len != 0 && valid.all_false() { - return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); - } - kernel( - Execution { - operands: [left, right], - valid, - len, - nullability: output_dtype.nullability(), - }, - ctx, - ) -} - -/// A bounding-rectangle pre-check for [`execute_binary_geo_types`]'s one-constant paths. -/// -/// Called per row with rectangles in operand order, it returns `Some(result)` when they prove the -/// result and `None` when the exact kernel must run. -pub(crate) type BboxPrecheck = fn(&Rect, &Rect) -> Option; - -/// Run a binary row-oriented kernel whose inputs are decoded to `geo_types::Geometry`. -/// -/// The `geo_types` name describes the values passed to `compute`, not the output. `T` is converted -/// into a Vortex array before this function returns. Nulls propagate from either operand. With -/// exactly one constant operand, `bbox_precheck` may prove a result from the fixed constant -/// bounding rectangle and the current row's rectangle before the exact kernel runs. -pub(crate) fn execute_binary_geo_types( - left: &ArrayRef, - right: &ArrayRef, - compute: F, - bbox_precheck: Option>, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T + Copy, -{ - let nullability = Nullability::from(left.dtype().is_nullable() || right.dtype().is_nullable()); - dispatch_binary( - left, - right, - T::dtype(nullability), - |execution, ctx| match execution.operands { - [Operand::Constant(left), Operand::Constant(right)] => { - let left = single_geometry(&left, ctx)?; - let right = single_geometry(&right, ctx)?; - Ok(ConstantArray::new( - compute(&left, &right).into_scalar(execution.nullability), - execution.len, - ) - .into_array()) - } - [Operand::Constant(left), Operand::Column(right)] => { - let left = single_geometry(&left, ctx)?; - let prescreen = bbox_precheck.zip(left.bounding_rect()); - eval_column( - &right, - &execution.valid, - |right| { - prescreen - .and_then(|(precheck, fixed)| precheck(&fixed, &right.bounding_rect()?)) - .unwrap_or_else(|| compute(&left, right)) - }, - execution.nullability, - ctx, - ) - } - [Operand::Column(left), Operand::Constant(right)] => { - let right = single_geometry(&right, ctx)?; - let prescreen = bbox_precheck.zip(right.bounding_rect()); - eval_column( - &left, - &execution.valid, - |left| { - prescreen - .and_then(|(precheck, fixed)| precheck(&left.bounding_rect()?, &fixed)) - .unwrap_or_else(|| compute(left, &right)) - }, - execution.nullability, - ctx, - ) - } - [Operand::Column(left), Operand::Column(right)] => eval_column_pair( - &left, - &right, - &execution.valid, - compute, - execution.nullability, - ctx, - ), - }, - ctx, - ) -} - -#[cfg(test)] -mod tests { - use std::cell::Cell; - - use geo::Contains; - use geo::Intersects; - use geo_types::Geometry; - use vortex_array::ArrayRef; - use vortex_array::ExecutionCtx; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::ConstantArray; - use vortex_array::assert_arrays_eq; - use vortex_array::validity::Validity; - use vortex_buffer::BitBuffer; - use vortex_error::VortexResult; - - use super::BboxPrecheck; - use super::execute_binary_geo_types; - use crate::test_harness::linestring_column; - use crate::test_harness::nullable_point_column; - use crate::test_harness::point_column; - use crate::test_harness::polygon_column; - - const DISJOINT_PRECHECK: BboxPrecheck = - |left, right| (!left.intersects(right)).then_some(false); - - fn triangle_constant(len: usize, ctx: &mut ExecutionCtx) -> VortexResult { - let ring = vec![(0.0, 0.0), (10.0, 0.0), (0.0, 10.0), (0.0, 0.0)]; - let scalar = polygon_column(vec![vec![ring]])?.execute_scalar(0, ctx)?; - Ok(ConstantArray::new(scalar, len).into_array()) - } - - fn counting_intersects( - counter: &Cell, - ) -> impl Fn(&Geometry, &Geometry) -> bool + Copy { - move |left, right| { - counter.set(counter.get() + 1); - left.intersects(right) - } - } - - #[test] - fn bbox_precheck_skips_exact_test() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(3, &mut ctx)?; - let probes = point_column(vec![50.0, 8.0, 2.0], vec![50.0, 8.0, 2.0])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - #[test] - fn bbox_precheck_leaves_nulls_alone() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(3, &mut ctx)?; - let probes = nullable_point_column(vec![Some((50.0, 50.0)), None, Some((2.0, 2.0))])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - let expected = BoolArray::new( - BitBuffer::from_iter([false, false, true]), - Validity::from_iter([true, false, true]), - ) - .into_array(); - - assert_arrays_eq!(result, expected, &mut ctx); - assert_eq!(exact_runs.get(), 1); - Ok(()) - } - - #[test] - fn bbox_precheck_sees_rects_in_operand_order() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let triangle = triangle_constant(2, &mut ctx)?; - let exact_runs = Cell::new(0); - let counted = |left: &Geometry, right: &Geometry| { - exact_runs.set(exact_runs.get() + 1); - left.contains(right) - }; - - let result = execute_binary_geo_types( - &probes, - &triangle, - counted, - Some(|left, right| (!left.contains(right)).then_some(false)), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 0); - Ok(()) - } - - #[test] - fn empty_constant_falls_through_to_exact() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let scalar = linestring_column(vec![vec![]])?.execute_scalar(0, &mut ctx)?; - let empty = ConstantArray::new(scalar, 2).into_array(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &empty, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - #[test] - fn bbox_precheck_matches_exact_results() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(6, &mut ctx)?; - let probes = nullable_point_column(vec![ - Some((50.0, 50.0)), - Some((8.0, 8.0)), - Some((2.0, 2.0)), - None, - Some((0.0, 0.0)), - Some((10.0, 0.0)), - ])?; - let exact = |left: &Geometry, right: &Geometry| left.intersects(right); - - let with_precheck = - execute_binary_geo_types(&triangle, &probes, exact, Some(DISJOINT_PRECHECK), &mut ctx)?; - let exact_only = execute_binary_geo_types(&triangle, &probes, exact, None, &mut ctx)?; - - assert_arrays_eq!(with_precheck, exact_only, &mut ctx); - Ok(()) - } -} diff --git a/vortex-spatial/src/scalar_fn/execute/geo_types.rs b/vortex-spatial/src/scalar_fn/execute/geo_types.rs deleted file mode 100644 index 038aca46502..00000000000 --- a/vortex-spatial/src/scalar_fn/execute/geo_types.rs +++ /dev/null @@ -1,144 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Shared input decoding and Vortex output construction for `geo_types` kernels. -//! -//! `geo_types` is the row representation consumed by the kernel. These helpers always construct -//! and return Vortex arrays; they do not expose `geo_types` values as scalar-function outputs. - -use geo_types::Geometry; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::BoolArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::scalar::Scalar; -use vortex_array::validity::Validity; -use vortex_buffer::BitBuffer; -use vortex_error::VortexResult; -use vortex_mask::AllOr; -use vortex_mask::Mask; - -use crate::extension::geometries; - -/// A primitive result produced after kernel inputs are decoded to `geo_types`. -pub(crate) trait GeoTypesOutput: Copy { - /// The Vortex dtype used to represent this output. - fn dtype(nullability: Nullability) -> DType; - - /// Convert one computed value into a Vortex scalar for constant output. - fn into_scalar(self, nullability: Nullability) -> Scalar; - - /// Scatter values computed for valid rows into a full-length output array. - fn build_array( - len: usize, - valid: &Mask, - values: Vec, - nullability: Nullability, - ) -> ArrayRef; -} - -impl GeoTypesOutput for f64 { - fn dtype(nullability: Nullability) -> DType { - DType::Primitive(PType::F64, nullability) - } - - fn into_scalar(self, nullability: Nullability) -> Scalar { - Scalar::primitive(self, nullability) - } - - fn build_array( - len: usize, - valid: &Mask, - values: Vec, - nullability: Nullability, - ) -> ArrayRef { - let validity = Validity::from_mask(valid.clone(), nullability); - match valid.indices() { - AllOr::All => PrimitiveArray::new(values, validity).into_array(), - AllOr::None => PrimitiveArray::new(vec![0.0f64; len], validity).into_array(), - AllOr::Some(rows) => { - let mut data = vec![0.0f64; len]; - for (&row, value) in rows.iter().zip(values) { - data[row] = value; - } - PrimitiveArray::new(data, validity).into_array() - } - } - } -} - -impl GeoTypesOutput for bool { - fn dtype(nullability: Nullability) -> DType { - DType::Bool(nullability) - } - - fn into_scalar(self, nullability: Nullability) -> Scalar { - Scalar::bool(self, nullability) - } - - fn build_array( - len: usize, - valid: &Mask, - values: Vec, - nullability: Nullability, - ) -> ArrayRef { - let validity = Validity::from_mask(valid.clone(), nullability); - match valid.indices() { - AllOr::All => BoolArray::new(BitBuffer::from_iter(values), validity).into_array(), - AllOr::None => BoolArray::new(BitBuffer::new_unset(len), validity).into_array(), - AllOr::Some(rows) => { - let mut data = vec![false; len]; - for (&row, value) in rows.iter().zip(values) { - data[row] = value; - } - BoolArray::new(BitBuffer::from_iter(data), validity).into_array() - } - } - } -} - -/// Evaluate a decoded kernel over each valid row of one geometry column. -pub(super) fn eval_column( - column: &ArrayRef, - valid: &Mask, - compute: F, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry) -> T, -{ - let len = column.len(); - let decoded = geometries(&column.filter(valid.clone())?, ctx)?; - let values = decoded.iter().map(compute).collect(); - Ok(T::build_array(len, valid, values, nullability)) -} - -/// Evaluate a decoded kernel over rows where both geometry columns are valid. -pub(super) fn eval_column_pair( - left: &ArrayRef, - right: &ArrayRef, - valid: &Mask, - compute: F, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T, -{ - let len = left.len(); - let left = geometries(&left.filter(valid.clone())?, ctx)?; - let right = geometries(&right.filter(valid.clone())?, ctx)?; - let values = left - .iter() - .zip(&right) - .map(|(left, right)| compute(left, right)) - .collect(); - Ok(T::build_array(len, valid, values, nullability)) -} diff --git a/vortex-spatial/src/scalar_fn/execute/unary.rs b/vortex-spatial/src/scalar_fn/execute/unary.rs index bdbbd0b33ac..478c62eef50 100644 --- a/vortex-spatial/src/scalar_fn/execute/unary.rs +++ b/vortex-spatial/src/scalar_fn/execute/unary.rs @@ -40,7 +40,6 @@ where operands: [Operand::Constant(constant.scalar().clone())], valid: Validity::AllValid, len, - nullability: output_dtype.nullability(), }, ctx, ); @@ -55,7 +54,6 @@ where operands: [Operand::Column(array.clone())], valid, len, - nullability: output_dtype.nullability(), }, ctx, ) diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index bdabd2b9967..3694303e954 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -3,44 +3,26 @@ //! `ST_Intersects`: OGC intersection test between two native geometries. +use geo::BoundingRect; use geo::Intersects; +use geo_types::Geometry; +use geo_types::Rect; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ElementSink; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Intersects`. -fn validate_intersects_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: intersects requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: intersects operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; +#[cfg(test)] +use crate::scalar_fn::row::probe; /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry /// operands, each a column or a constant literal. @@ -58,74 +40,97 @@ impl SpatialIntersects { } } -impl ScalarFnVTable for SpatialIntersects { +impl RowFn for SpatialIntersects { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.intersects"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { Ok(EmptyOptions) } - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("intersects has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_intersects_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Bool(nullability)) - } - - fn execute( + fn dispatch( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Disjoint bounding rects prove the geometries disjoint; rect contact (closed test) - // falls through to the exact test. - execute_binary_geo_types( - &a, - &b, - |x, y| x.intersects(y), - Some(|ra, rb| (!ra.intersects(rb)).then_some(false)), - ctx, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), ElementSink, _, _>( + |(a, b)| { + #[cfg(test)] + probe::record(a.is_some(), b.is_some()); + ConstBboxes::new(a, b) + }, + |bboxes, (a, b), output| *output = intersects_row_prepared(bboxes, a, b), ) } +} - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } +/// Per-batch state for the intersects row kernel: the bounding rect of each operand that is +/// constant for the batch. +/// +/// geo opens many intersects pairings with `has_disjoint_bboxes`, an early-out that folds +/// [`bounding_rect`] over both operands. For a batch-constant operand that fold recomputes the +/// same rect every row, so it is hoisted here and [`intersects_row_prepared`] replays the +/// comparison with the hoisted value. `None` marks an operand that varies by row or has no +/// bounding rect (an empty geometry); both mean no early-out, exactly as `has_disjoint_bboxes` +/// treats a missing rect. +/// +/// [`bounding_rect`]: BoundingRect::bounding_rect +struct ConstBboxes { + /// The bounding rect of operand `a` when it is batch-constant. + a: Option>, + + /// The bounding rect of operand `b` when it is batch-constant. + b: Option>, +} - fn is_strict(&self, _: &Self::Options) -> bool { - true +impl ConstBboxes { + fn new(a: Option<&Geometry>, b: Option<&Geometry>) -> Self { + Self { + a: a.and_then(BoundingRect::bounding_rect), + b: b.and_then(BoundingRect::bounding_rect), + } } +} - fn is_fallible(&self, _: &Self::Options) -> bool { - false - } +/// Computes one row of intersects, spending any bounding rect hoisted into `bboxes`. +/// +/// Disjoint bounding rectangles conservatively prove that the geometries do not intersect. The +/// fall-through delegates to the unchanged `a.intersects(b)`, which refolds both rects internally, +/// so a batch where every row overlaps pays one extra `bounding_rect` fold over the row operand; +/// the win concentrates where most rows are disjoint, the usual spatial-filter shape. +fn intersects_row_prepared(bboxes: &ConstBboxes, a: &Geometry, b: &Geometry) -> bool { + let disjoint = match (bboxes.a, bboxes.b) { + (None, None) => false, + (Some(bbox_a), Some(bbox_b)) => !bbox_a.intersects(&bbox_b), + (Some(bbox_a), None) => b + .bounding_rect() + .is_some_and(|bbox_b| !bbox_a.intersects(&bbox_b)), + (None, Some(bbox_b)) => a + .bounding_rect() + .is_some_and(|bbox_a| !bbox_a.intersects(&bbox_b)), + }; + + if disjoint { + return false; + } + + a.intersects(b) } #[cfg(test)] @@ -133,7 +138,9 @@ mod tests { use geo_types::Coord; use geo_types::Geometry; use geo_types::LineString; + use geo_types::MultiPoint; use geo_types::MultiPolygon; + use geo_types::Point; use geo_types::Polygon; use rstest::rstest; use vortex_array::ArrayRef; @@ -157,8 +164,10 @@ mod tests { use wkb::writer::WriteOptions; use super::SpatialIntersects; + use crate::scalar_fn::row::probe::assert_prepared_agrees_with_columns; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; + use crate::test_harness::rect_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { @@ -439,4 +448,85 @@ mod tests { assert!(result.is_err()); Ok(()) } + + // The prepared-vs-expanded agreement grid: every constant arrangement of a pairing must + // return exactly what the fully expanded columns return. + + /// A point geometry. + fn point(x: f64, y: f64) -> Geometry { + Geometry::Point(Point::new(x, y)) + } + + /// A linestring geometry through `coords`. + fn line(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::LineString(LineString::from(coords)) + } + + /// A multipoint geometry over `coords`. + fn multipoint(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::MultiPoint(MultiPoint::from(coords)) + } + + /// Constant arrangements agree with expanded columns across the pairing classes the prepared + /// kernel treats differently: bbox-prechecked pairs (polygon x polygon, linestring x + /// anything, multipolygon blankets), direct pairs (points), the excluded `MultiPoint` route, + /// and an empty geometry whose bounding rect does not exist. + #[rstest] + #[case::polygons_overlapping(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(2.0, 2.0, 6.0, 6.0).into())] + #[case::polygons_touching_edge(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(4.0, 0.0, 8.0, 4.0).into())] + #[case::polygons_touching_corner(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(4.0, 4.0, 8.0, 8.0).into())] + #[case::polygons_disjoint(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::polygons_nested(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::polygon_x_point_inside(donut(), point(2.0, 2.0))] + #[case::polygon_x_point_on_boundary(donut(), point(0.0, 5.0))] + #[case::polygon_x_point_in_hole(donut(), point(5.0, 5.0))] + #[case::point_outside_x_polygon(point(20.0, 20.0), donut())] + #[case::nan_point_x_polygon(point(f64::NAN, 2.0), donut())] + #[case::polygon_x_nan_point(donut(), point(f64::NAN, 2.0))] + #[case::points_equal(point(1.0, 1.0), point(1.0, 1.0))] + #[case::points_distinct(point(1.0, 1.0), point(2.0, 1.0))] + #[case::linestring_crossing_polygon(line(vec![(-2.0, -2.0), (2.0, 2.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::linestring_disjoint_polygon(line(vec![(-2.0, -2.0), (-6.0, -6.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::linestring_in_polygon_hole(line(vec![(4.5, 4.5), (5.5, 5.5)]), donut())] + #[case::linestrings_crossing(line(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(0.0, 4.0), (4.0, 0.0)]))] + #[case::linestrings_disjoint(line(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(10.0, 10.0), (14.0, 14.0)]))] + #[case::empty_linestring_x_polygon(line(vec![]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipoint_straddling_polygon(multipoint(vec![(2.0, 2.0), (20.0, 20.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipoint_outside_polygon(multipoint(vec![(20.0, 20.0), (30.0, 30.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipolygon_disjoint_polygon( + Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 2.0, 2.0), + rect_polygon(10.0, 10.0, 12.0, 12.0), + ])), + rect_polygon(20.0, 20.0, 24.0, 24.0).into() + )] + fn constant_operands_agree_with_columns( + #[case] a: Geometry, + #[case] b: Geometry, + ) -> VortexResult<()> { + assert_prepared_agrees_with_columns( + SpatialIntersects::try_new_array, + geometry_constant(&a, 3)?, + geometry_constant(&b, 3)?, + ) + } + + /// `Rect` has no WKB form, so its constant comes from a one-row rect column; its conservative + /// bbox early-out and exact fall-through must agree with the expanded form like the rest. + #[test] + fn rect_operand_agrees_with_columns() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let rect_scalar = rect_column(vec![(0.0, 0.0, 4.0, 4.0)])?.execute_scalar(0, &mut ctx)?; + let rect_constant = ConstantArray::new(rect_scalar, 3).into_array(); + let polygon_constant = + geometry_constant(&Geometry::Polygon(rect_polygon(2.0, 2.0, 6.0, 6.0)), 3)?; + + assert_prepared_agrees_with_columns( + SpatialIntersects::try_new_array, + rect_constant, + polygon_constant, + ) + } } diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index e6770be4fff..bcdb15e51e6 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -8,3 +8,4 @@ pub mod distance; pub mod envelope; mod execute; pub mod intersects; +pub(crate) mod row; diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs new file mode 100644 index 00000000000..ce7c34072ff --- /dev/null +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -0,0 +1,170 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the geo scalar functions add to the row-function machinery: an element type that decodes a +//! native geometry column into `geo_types` geometries. + +use geo_types::Geometry; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::InputElement; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::extension::geometries; +use crate::extension::geometries_null_tolerant; +use crate::extension::is_native_geometry; + +/// Marker for native geometry input elements: accepts any native geometry column and presents each +/// row as a decoded `geo_types` geometry. +/// +/// The two operands of a binary geo function need not share a geometry type, since distance, +/// containment and intersection across types are all meaningful, so this validates only that the +/// column is *some* native geometry. +pub struct GeometryRow; + +impl InputElement for GeometryRow { + type Column = Vec>; + type Varying<'a> = &'a [Geometry]; + type Elem<'a> = &'a Geometry; + + // A geometry row is decoded from its coordinate storage, which behind a null row holds arbitrary + // coordinates that need not describe a well-formed geometry. + const DENSE_SAFE: bool = false; + // Decoding builds a geometry from stored coordinates, and a malformed one in a *valid* row is a + // domain error rather than an infrastructural failure. + const DECODE_FALLIBLE: bool = true; + // Decoding arrow-exports the column and parses one geometry per row, so filtering the column + // first shrinks the decode itself, not just the row loop. + const FILTERED_DECODE_COST: usize = 1; + + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + is_native_geometry(dtype), + "spatial: operand {dtype} is not a native geometry type" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + geometries(&array, ctx) + } + + fn get(column: &Self::Column, index: usize) -> &Geometry { + &column[index] + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column.as_slice() + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a Geometry + where + Self: 'a, + { + &column[index] + } + + /// Null rows decode to a placeholder geometry that the branch-and-skip row loop never reads. + /// `Point` and `Polygon` columns are covered; other geometry types return `Ok(None)` and the + /// batch falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + geometries_null_tolerant(&array, ctx) + } +} + +/// Test-only support for the prepared geo row kernels: a probe recording which operands a +/// `prepare` step saw as batch-constant, and the shared prepared-vs-expanded agreement check +/// built on it. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::ScalarFnArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + + thread_local! { + /// Which operands the last `prepare` saw as constant, as a bitmask (bit 0 for `a`, bit 1 + /// for `b`). Thread-local rather than a process global so concurrent tests in one process + /// (plain `cargo test`) cannot race it; execution runs on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Record which operands `prepare` saw as constant. + pub(crate) fn record(a_constant: bool, b_constant: bool) { + SEEN_CONSTANTS.set(u8::from(a_constant) | (u8::from(b_constant) << 1)); + } + + /// Execute `build(a, b)` and assert that `prepare` saw exactly `expect_seen` as its constant + /// operands, so the test knows which decode path the inputs took. + fn run_probed( + build: &impl Fn(ArrayRef, ArrayRef) -> VortexResult, + a: ArrayRef, + b: ArrayRef, + expect_seen: u8, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + SEEN_CONSTANTS.set(u8::MAX); + let result = build(a, b)? + .into_array() + .execute::(ctx)? + .into_array(); + + assert_eq!( + SEEN_CONSTANTS.get(), + expect_seen, + "prepare saw the wrong constant operands", + ); + Ok(result) + } + + /// Assert that every constant-operand arrangement of `build(a, b)` returns exactly what the + /// fully expanded columns return, and that each arrangement's constness really reached + /// `prepare` (so the constants exercised the stride-0 path rather than a decoded column). + /// + /// Arrangements: `a` constant, `b` constant, and both constant with `a` masked. A plain + /// constant pair folds to a single-row execution before the row loop, so masking one side is + /// what drives the both-hoisted arm across rows; that run is compared against the same mask + /// over the expanded column. + pub(crate) fn assert_prepared_agrees_with_columns( + build: impl Fn(ArrayRef, ArrayRef) -> VortexResult, + const_a: ArrayRef, + const_b: ArrayRef, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let col_a = const_a.clone().execute::(&mut ctx)?.into_array(); + let col_b = const_b.clone().execute::(&mut ctx)?.into_array(); + + let baseline = run_probed(&build, col_a.clone(), col_b.clone(), 0b00, &mut ctx)?; + let a_hoisted = run_probed(&build, const_a.clone(), col_b.clone(), 0b01, &mut ctx)?; + let b_hoisted = run_probed(&build, col_a.clone(), const_b.clone(), 0b10, &mut ctx)?; + assert_arrays_eq!(a_hoisted, baseline, &mut ctx); + assert_arrays_eq!(b_hoisted, baseline, &mut ctx); + + let validity = Validity::from_iter((0..col_a.len()).map(|row| row != 1)); + let masked_const_a = MaskedArray::try_new(const_a, validity.clone())?.into_array(); + let masked_col_a = MaskedArray::try_new(col_a, validity)?.into_array(); + let both_hoisted = run_probed(&build, masked_const_a, const_b, 0b11, &mut ctx)?; + let masked_baseline = run_probed(&build, masked_col_a, col_b, 0b00, &mut ctx)?; + assert_arrays_eq!(both_hoisted, masked_baseline, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-spatial/src/test_harness.rs b/vortex-spatial/src/test_harness.rs index 7b471bdf2c4..d8175d14f53 100644 --- a/vortex-spatial/src/test_harness.rs +++ b/vortex-spatial/src/test_harness.rs @@ -251,7 +251,7 @@ pub fn nullable_rect_column(boxes: Vec>) -> VortexR Ok(ExtensionArray::try_new(ext.erased(), storage)?.into_array()) } -/// Decode a [`Coordinate`] from an extension-typed point scalar (unwrapped to its coordinate +/// Decode a `Coordinate` from an extension-typed point scalar (unwrapped to its coordinate /// storage) or a bare coordinate `Struct` scalar — used to read back a single point in assertions. pub fn coordinate_from_scalar(scalar: &Scalar) -> VortexResult { match scalar.as_extension_opt() { From b541d4bdcb1ab684316766e2b961e39d694546bf Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 13:03:43 +0000 Subject: [PATCH 006/160] Record the row scalar function research and handoff notes Progress towards #9128. `STRICT_SCALAR_FN_RESEARCH.md` holds the null-strategy measurements behind the per-batch selection rule, the three verdicts and the crossover, and the record of why `vortex.byte_length`, `vortex.not`, and `vortex.list.sum` stay on `ScalarFnVTable`. `SCALAR_FN_HANDOFF.md` records the current state of the work, including which API proposals from the review were backed out and why. `NUMERIC_ROWFN_PLAN.md` is the plan and measured outcome for the numeric port. `docs/strictness-and-validity-pushdown.typ` writes up strictness and validity pushdown, which is the property the whole derivation rests on. These are working notes rather than published documentation, and they are separated here so they are easy to drop before this ships. Signed-off-by: Connor Tsui Co-authored-by: Claude --- NUMERIC_ROWFN_PLAN.md | 256 +++ SCALAR_FN_HANDOFF.md | 443 +++++ STRICT_SCALAR_FN_RESEARCH.md | 1790 +++++++++++++++++++++ docs/strictness-and-validity-pushdown.typ | 243 +++ 4 files changed, 2732 insertions(+) create mode 100644 NUMERIC_ROWFN_PLAN.md create mode 100644 SCALAR_FN_HANDOFF.md create mode 100644 STRICT_SCALAR_FN_RESEARCH.md create mode 100644 docs/strictness-and-validity-pushdown.typ diff --git a/NUMERIC_ROWFN_PLAN.md b/NUMERIC_ROWFN_PLAN.md new file mode 100644 index 00000000000..2f69b708b01 --- /dev/null +++ b/NUMERIC_ROWFN_PLAN.md @@ -0,0 +1,256 @@ + + + +# Plan: fit the numeric binary operators onto `RowFn` + +Working note, branch-only, like `SCALAR_FN_HANDOFF.md`. Written so this survives a conversation +compaction: everything needed to start is here, and nothing below depends on chat history. + +## Where things stand + +Branch `ct/row-fn`, at `4becc863ae` after the final API +simplification. Issues #9128, #9129, and #9130 match the implementation. The public-path benchmark +baseline from #9136 is now in the repository. + +This document preserves the original spike plan and the measurements that answered it. The current +API has no witnesses, persistence is function-owned, executor-only helper traits are sealed, and +filtered decode cost is additive per input. Read the outcome and final API/codegen sections before +following an earlier step literally. + +`byte_length` is no longer a row function, and `Bytes`/`BytesLen` are deleted. It measured 7.6-7.7x +slower than develop and is the case #9128 already excludes. + +## Goal of this spike + +Prove, or disprove, that the four arithmetic operators can move onto `RowFn` without changing the +`RowFn` API, without a second scalar function ID, and without touching serialization. Doing this +first is deliberate: it is the change most likely to force an API change, and discovering that after +tensor and geo are ported would mean reworking them. + +## The design + +`Binary` keeps everything and delegates only execution: + +```rust +Operator::Add => ScalarFnVTable::execute(&NumericBinary, &NumericOperator::Add, args, ctx), +``` + +`NumericBinary` is a `RowFn` with `Options = NumericOperator` and `FALLIBLE = true`. It is **not** +registered as a public scalar function, so it needs no ID in the registry and appears in no +serialized expression. It is reached through the `ScalarFnVTable::execute` that the blanket impl +already provides. + +Why this works, and each of these was verified against the code rather than assumed: + +- **Nothing is lost.** `BooleanKernel` and `CompareKernel` exist with per-encoding pushdown; there is + no `NumericKernel`. Unlike `not`, a numeric port gives up no encoding fast path. +- **The seam is already numeric-only.** All four arithmetic arms of `Binary::execute` funnel into + `execute_numeric(lhs, rhs, NumericOperator, ctx)`, and `NumericOperator` is already its own enum in + `crate::scalar`, so it is a ready-made `RowFn::Options`. +- **Fallibility is uniform.** `Binary::is_fallible` is false for the six comparisons plus `And`/`Or` + and true for exactly the four arithmetic operators, so `FALLIBLE = true` on a numeric-only `RowFn` + is exactly right. The options-independence of `RowFn::is_fallible` only bites when one function + spans both families. +- **Strictness stays where it belongs.** `Binary::is_strict` is `!matches!(op, And | Or)` because + Kleene `false AND null` is a valid `false`. `Binary` keeps owning that; `NumericBinary` never sees + a boolean operator. +- **Decimal fits.** `OutputSink::sink_dtype(args)` sees the input dtypes, which is what + `numeric_op_result_decimal_dtype(decimal_dtype, op)` needs. + +## Steps + +1. **Primitive path only, `Add` only.** A `NumericBinary` `RowFn` over `(T, T)` for one integer + width, with a deferred-error sink that writes the wrapping sum and ORs an overflow bit. Delegate + only `Operator::Add` from `Binary::execute` and leave the other three on `execute_numeric`. + Success is: the existing `binary/numeric/tests.rs` suite passes unchanged. +2. **Widen to every primitive ptype**, through `match_each_native_ptype!` in `dispatch`. Confirm the + compile-time witness check tolerates it, as it does for tensor widths. +3. **Add `Sub`, `Mul`, `Div`.** `Div` is the awkward one: see the risk below. +4. **Decide decimal.** Either a decimal input element plus a sink that carries the result precision + and scale, or leave `DType::Decimal` on `execute_numeric` and delegate only the primitive path. + Leaving it is a legitimate outcome for the spike and possibly for the first PR. +5. **Delete the replaced code** only once benchmarks agree, not before. + +## Risks, in the order they are likely to bite + +- **`Div` already has a per-type strategy.** `primitive.rs` carries `CHECKED_VALUE_LOOP` and + `DIV_CHECKS_IN_VALUE_LOOP`, set per type, so division checking is not uniform. A single row closure + may not express it, and `Div` may have to stay behind. +- **The existing implementation is tuned, not naive.** `checked.rs` has `checked_lanes` and + `checked_apply_lanes` taking a `valid_rows: &Mask` and returning `Result, usize>` with the + failing index. The port is replacing real engineering, so parity is not a given. This is the reason + the CodSpeed gate on the `binary_ops` names from #9136 matters. +- **Two declarations of the result dtype must agree.** `Binary::return_dtype` is what the expression + layer uses, while `reconcile_return` checks the kernel output against `NumericBinary`'s + sink-derived dtype. Cover every operator and dtype pair with a test that asserts they match. +- **Error messages are part of the contract.** `primitive.rs` defines `ERROR` per operator, such as + `"integer overflow in checked add"`, and `numeric/tests.rs` asserts on failures. The deferred-error + sink reports once from `finish`, so the message must be preserved and the error must still be + raised for the same inputs. +- **Overflow behind a null row must stay invisible.** `numeric/tests.rs` has + `test_decimal_overflow_on_null_lane_ignored`. The lifting's deferred-error retry over valid rows is + exactly this behavior, so the test should pass, but it is the first thing to check. + +## Verification + +```bash +cargo nextest run -p vortex-array +cargo clippy --all-targets --all-features -p vortex-array +cargo +nightly fmt --all +cargo test --doc -p vortex-array +``` + +The numeric suite specifically: + +```bash +cargo nextest run -p vortex-array scalar_fn::fns::binary +``` + +Performance gate is CodSpeed on the stable `binary_ops` names from #9136. Locally, use +`cargo bench -p vortex-array --bench binary_ops` with two runs, fastest and median, machine stated. + +## What this spike is not + +Not a PR. Not a deletion of `execute_numeric`. Not decimal support unless step 4 turns out easy. The +output is an answer to "does this fit cleanly", plus whatever the answer implies for #9129's API. + +## Outcome + +It fits, with no change to the `RowFn` API and one change to the machinery. + +Steps 1 through 3 landed together rather than in sequence: once the sink existed, widening it through +`match_each_native_ptype!` and adding the other three operators was the same code. Step 4 leaves +decimal on `execute_numeric_decimal`, which the delegation makes easy since `execute_numeric` still +owns the dtype split. Step 5 deleted the replaced primitive execution, which the measurements below +justify. + +### What the design turned out to be + +`Binary::execute` is untouched. `execute_numeric` keeps its validation, its error messages, its empty +short circuit, and its primitive/decimal split, and only `execute_numeric_primitive` changed: it +builds a `VecExecutionArgs` and calls `ScalarFnVTable::execute(&NumericBinary, &op, ..)`. Everything +the old implementation did around the arithmetic (decoding, the constant-operand collapse, the +all-constant fold, the null-constant short circuit, output allocation, nullability widening, masking, +and the valid-row retry after an overflow behind a null) is now the lifting's. + +`NumericOperator` became the options type. `NumericBinary` is unregistered and deliberately has no +serialization implementation. Persistence now belongs to each `RowFn`, so reusing an options type +does not silently assign the helper a wire contract. `Binary` retains its existing ID and options +serialization, and only primitive execution delegates to `NumericBinary`. + +Three things the old code carried are gone because the row framework removes the distinction they +existed for: + +- `CHECKED_VALUE_LOOP` and `DIV_CHECKS_IN_VALUE_LOOP` chose between a split value/error scan and a + one-pass early-exit kernel, because for integer division the split loop only added a second scan. + A row kernel produces the value and the error bit in the same pass, so there is one loop shape and + no choice to make. `div_i64` got 1.11x faster. +- `checked_apply_lanes` had no caller left. `checked_lanes` stays for decimal. +- `PrimitiveOperand` moved to `compare/primitive.rs`, its only remaining user. + +### The machinery change: the reduction is a word the kernel chooses + +`SinkResult` gained `Accumulated`, the word the executor OR-reduces in a loop-local. Two properties +of that reduction are load-bearing, and each was got wrong once before the numbers made it obvious. + +- **Width no greater than the element.** `DeferredError` held an `i64`, which bounds how many rows a + vector of the reduction covers whatever the element width. That cost `Mul` 3.5x at `i8`, 2.05x at + `i16` and 1.28x at `i32`, and nothing at `i64` where the widths already agree. +- **It lives in a loop-local, not in the sink.** Holding the accumulator as a sink field, reached + through a `&mut` for every row, is a loop-carried memory dependence. It cost the boolean kernels + 2.5x to 10x while leaving the three unsigned multiply kernels untouched. + +Naming the word is also what lets multiplication report the discarded high half of its product +rather than a comparison, which is what recovers its vectorization. `OutputSink` is unchanged and no +sink names the word. + +### Results + +Against the hand-written kernels, divan medians, best of two runs, 65536 rows, Apple M4 Max, with +the decimal, boolean and comparison benchmarks held as controls and moving under 2%: + +| benchmark | hand-written | row framework | | +| --- | --- | --- | --- | +| `mul_u8_nonnull` | 22.91 us | 1.854 us | 12.4x faster | +| `mul_u16_nonnull` | 22.20 us | 3.791 us | 5.9x faster | +| `mul_u32_nonnull` | 24.62 us | 7.124 us | 3.5x faster | +| `div_i64_nonnull` | 40.41 us | 34.83 us | 1.16x faster | +| `mul_i64_nonnull` | 27.37 us | 28.66 us | 1.05x slower | +| `mul_i32_constant` | 7.583 us | 8.041 us | 1.06x slower | + +Everything else lands within 3%, which is inside this host's drift between sessions. The unsigned +multiply win is not attributable to the port: the same defect exists in the hand-written kernels and +is fixed for `develop` separately in vortex-data/vortex#9210, stacked on vortex-data/vortex#9211. +Re-measure the port against `develop` once that lands, because the comparison above flatters it. + +### Measured dead ends + +Recorded so they are not retried. All of these are in vortex-data/vortex#9130 as well. + +- Bounds-check elimination in the row loop is not available. Narrowing the varying view to the row + count buys nothing, and `get_unchecked` is not uniformly a win: about 10% on `mul_u16` and + `mul_u32`, and 22% slower on `mul_u8`. +- A per-argument row source that keeps the `Varying` view when another argument is batch-constant is + 4x slower than the `ArgColumn` branch it replaces, which already vectorizes. +- A batch-constant operand therefore still demotes its neighbours off the slice path. Closing that + needs the row loop monomorphized over which arguments are constant. Revisit when `Compare` moves + onto `RowFn`, since `col < literal` is exactly this shape. + +`mul_i32_constant` is the one regression that survives, and it is inside this host's drift. Let +CodSpeed settle whether it is real. + +### What this implies for #9129 and #9130 + +- The `RowFn` API needed nothing. No new visit method, no options-aware `sink_dtype`, no return + witness. `NumericBinary::FALLIBLE = true` is conservative for every dispatch arm, and each + concrete result type supplies the precise loop behavior. +- `SinkResult::Accumulated` and its two constraints belong in #9130, and are recorded there. +- On kernels this close to the vectorizer's decision boundary, the emitted IR is the reliable gate + and wall clock on one host is not. Two separate interventions here moved a benchmark the wrong + way, and host drift between sessions exceeded the effects under measurement. + +### Final API cleanup and generated code + +The later simplification did not add numeric-specific surface: + +- `NumericBinary` declares `ARG_NAMES = &["lhs", "rhs"]` instead of repeating an argument witness. +- Its `Options = NumericOperator` has no persistence bound or implementation. The registered + `Binary` function remains the sole owner of the serialized `vortex.binary` contract. +- The selected input tuple carries arity, dense-safety, decode fallibility, and filtered-decode + cost. The selected sink and `SinkResult` carry output and deferred-error facts. +- `SinkResult` is sealed, but a numeric function does not need to implement it. It chooses the + supplied unsigned evidence width that matches the primitive element width. +- `OutputSink` remains one abstraction. A later numeric function with multiple logical outputs + should put both builders in one sink rather than add a pair-of-sinks framework type. + +The final cleanup was checked against its parent by cross-compiling the optimized +`row_fn_executor` benchmark for `x86_64-apple-darwin` with `target-cpu=x86-64-v3`. After normalizing +symbol names and metadata, the vector/reduction block for checked `i64` add matched exactly. It +retains `<4 x i64>` loads and adds, vector overflow detection through xor/and/compare operations, +`<4 x i1>` OR accumulation, and a reduction after the loop. The vector body has no call or panic +path, and the scalar tail is unchanged. + +The ordinary `ElementSink` and custom-sink wrapping-add monomorphs also matched exactly. Native +Apple M4 Max measurements over 65,536 rows found RowFn median changes between 1.11% faster and 0.94% +slower, with fastest changes within about 0.17%. Specialized controls drifted more than the RowFn +arms, so there is no measurable native regression from the cleanup. + +This is not an x86 runtime result. It proves that the API edits preserved the optimized x86_64-v3 +loop shape. Runtime confirmation for numeric changes should use the stable public benchmark names +from #9136 on the target host. + +The next session will run on x86 and must perform that confirmation. The #9136 `binary_ops` +benchmark is on `develop` at `9a482c0230`, so compare this branch with the latest +`origin/develop` using the same public benchmark names. Record both exact commits and run each +revision at least twice in alternating order. If possible, pin one core. Report fastest and median +values with the CPU and timer configuration. If a stable case regresses, compare its optimized LLVM +IR before changing the row API or restoring hand-written execution. + +### Verification + +The whole of `binary/numeric/tests.rs` passed unchanged, including +`test_decimal_overflow_on_null_lane_ignored` and the integer-error tests that pin the valid-row +retry. Decimal is untouched and stays on `execute_numeric_decimal`. The final API state also +recorded 67 focused RowFn tests, 179 tensor tests, 230 geo tests, nightly formatting, and full +workspace clippy. Clippy needed `PYO3_NO_PYTHON=1 PYO3_BUILD_EXTENSION_MODULE=1` because the host +Python is 3.9 while the workspace targets the Python 3.11 stable ABI. diff --git a/SCALAR_FN_HANDOFF.md b/SCALAR_FN_HANDOFF.md new file mode 100644 index 00000000000..f99f4cb0eaa --- /dev/null +++ b/SCALAR_FN_HANDOFF.md @@ -0,0 +1,443 @@ + + + +# Handoff: the row scalar-function framework + +This is the concise source of truth for the branch. `STRICT_SCALAR_FN_RESEARCH.md` keeps the full +design history, rejected alternatives, measurements, and generated-code evidence. +`NUMERIC_ROWFN_PLAN.md` records the numeric-binary migration and its narrower performance boundary. +All three are branch-only working notes for agents. They are not intended to land with the API. + +The public design lives in these tracking issues, which now match the implementation: + +- [#9128, Row-oriented scalar functions](https://github.com/vortex-data/vortex/issues/9128) +- [#9129, Define the `RowFn` API](https://github.com/vortex-data/vortex/issues/9129) +- [#9130, Execute `RowFn` over Vortex arrays](https://github.com/vortex-data/vortex/issues/9130) + +The branch is `ct/row-fn`. It is publicly linked from #9128, so do +not rewrite or delete its history. Commit `4becc863ae` contains the final API simplification. Push +only when explicitly requested. + +## Next action: rerun the benchmarks on x86 + +The next session will run on an x86 machine. Rerun the performance comparison there before treating +the implementation as complete. Do not reuse the Apple timings as the final runtime result. + +The production benchmark baseline from #9136 is on `develop` at `9a482c0230`. Fetch the latest +`origin/develop`, record the exact baseline and candidate commits, and run the same public benchmark +binaries at both revisions: + +```bash +cargo bench -p vortex-array --bench binary_ops +cargo bench -p vortex-array --bench like +cargo bench -p vortex-tensor --bench l2_norm +cargo bench -p vortex-tensor --bench inner_product +cargo bench -p vortex-tensor --bench cosine_similarity +cargo bench -p vortex-tensor --bench normalized +cargo bench -p vortex-geo --bench binary_predicates +cargo bench -p vortex-geo --bench distance +cargo bench -p vortex-geo --bench envelope +cargo bench -p vortex-geo --bench predicate_bbox +``` + +Run each revision at least twice in alternating order. If the host allows it, pin the process to one +core. Record the timer and CPU configuration, and compare both fastest and median values. The +benchmark binaries and public names are now shared with `develop`, so the comparison no longer +needs a frozen benchmark-local implementation as its primary control. + +Also run the branch-only `vortex-geo` `null_strategies` diagnostic. It forces branch-and-skip and +filter-and-scatter for the measured nullable geometry shapes. Confirm that automatic selection uses +the faster mechanism for one costly decode at 50% survivors and for two costly decodes at about 81% +survivors. This is the x86 runtime check that remains after the LLVM comparison. + +```bash +cargo bench -p vortex-geo --bench null_strategies +``` + +If a stable benchmark regresses, inspect optimized LLVM IR again. The previous cross-compile proves +that the API cleanup preserved the x86_64-v3 loop shape. The x86 run must confirm runtime effects +from the revised null selector and the target CPU's vectorizer and branch predictor. + +## The API in one screen + +`RowFn` is the author-facing function trait. A function gives the framework its exact argument +names, a conservative fallibility declaration, function-owned persistence, and a value-blind +dispatch over concrete input and sink types: + +```rust +impl RowFn for Example { + type Options = ExampleOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.example"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(encode(options)?)) + } + + fn deserialize( + &self, + metadata: &[u8], + session: &VortexSession, + ) -> VortexResult { + decode(metadata, session) + } + + fn dispatch( + &self, + options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + validate_options(options, args)?; + visitor.visit_prepared_into::<(InputA, InputB), ElementSink, _, _>( + |_| (), + |&(), (lhs, rhs), output| { + *output = compute(lhs, rhs); + }, + ) + } +} +``` + +There are no argument or return witness types. The dispatched tuple is the argument declaration, +the sink owns the output representation, and the row result names the error behavior. Planning +runs the same dispatch as execution and checks the selected types against the function constants. + +## The extension boundary + +The framework is deliberately not sealed wholesale. Function authors need to add decode and output +primitives for their own scalar functions. Only the executor mechanics are closed. + +| API | Boundary | Why | +| --- | --- | --- | +| `RowFn` | open | Defines a scalar function and selects concrete execution types. | +| `InputElement` | open | Adds a new scalar decode primitive, including crate-local domain types. | +| `OutputElement` | open | Adds an ordinary one-value-per-row output primitive. | +| `OutputSink` | open | Adds a custom output representation or builder. | +| `RowVisitor` | sealed | Executor-owned dispatch mechanism with one supported implementation. | +| `ElementTuple` | sealed | Executor-owned tuple recursion, with built-ins through arity 12. | +| `SinkResult` | sealed | Executor-owned loop and error facts trusted by the blanket vtable. | + +`ElementTuple` being sealed does not prevent a function from adding a decode primitive. Implement +`InputElement` and use it inside one of the supplied tuples. Likewise, a function with two logical +outputs should define one `OutputSink` whose state has two fields. The framework does not need a +second tuple or composite-sink abstraction. + +The supplied `SinkResult` forms are: + +- `()` for infallible rows; +- `VortexResult<()>` for an error that must stop immediately; and +- `bool`, `u8`, `u16`, `u32`, or `u64` for error evidence OR-reduced after the loop. + +The unsigned evidence widths let each kernel choose a word no wider than its element type. That is +load-bearing for vectorization, particularly for checked unsigned multiplication. + +## Function-owned persistence + +Persistence belongs to the function ID, not to the Rust options type. `RowFn::Options` has no +serialization supertrait. The `RowFn::serialize` and `RowFn::deserialize` hooks have conservative +defaults, and registered functions override them when their existing wire contract requires it. + +This has three useful consequences: + +- two functions may reuse an options type while choosing different formats; +- a function may deliberately be nonserializable even if another function serializes the same + options type; and +- an unregistered helper such as `NumericBinary` needs no dummy persistence implementation. + +Tensor and geo functions keep their explicit existing formats. Do not introduce a blanket options +wire format or infer serializability from `Options`. + +## One sink abstraction + +`OutputSink` is the complete output contract. It owns the output dtype, allocation, row storage, +row lookup, length proof, and final array construction. `ElementSink` covers the common case. Its +row type is `&mut T`, so the closure writes with ordinary assignment. + +Custom sinks remain available for a real output shape that cannot use `ElementSink`. The unused +public `TensorSink` was removed. No current tensor row function returns tensor-valued rows, and a +90-line public runtime-shaped sink was not justified without a user. Add a custom sink when a real +function needs one, using one sink struct even when it owns several builders. + +Every current sink produces an all-valid child column. The blanket vtable can therefore derive the +function result validity from the input validities. Nullable row outputs remain out of scope. A +sink that emits its own nulls must change that derivation in the same change. + +`OutputSink::sink_dtype` must return a non-nullable dtype. `SUPPORTS_SKIPPED_ROWS` says whether +branch-and-skip may leave placeholder rows behind the result validity. `ERRORS_ARE_DEFERRED` says +whether the sink accepts accumulated error evidence at `finish`. + +## Dispatch and fallibility + +`dispatch` must be pure in `(options, args)`. It sees dtypes, not array values. Planning and +execution both call it, so value-dependent preparation belongs inside `visit_prepared_into`. + +The executor statically checks each dispatched visit: + +- the tuple arity equals `ARG_NAMES.len()`; +- a fallible decoder, early-error result, or deferred result implies `RowFn::FALLIBLE`; +- deferred evidence requires both `RowFn::FALLIBLE` and a sink with + `ERRORS_ARE_DEFERRED = true`; and +- the sink and result agree about their error contract. + +The implications are intentionally one-way. `FALLIBLE = true` is a conservative function-level +claim, while a particular dtype dispatch arm may be infallible. + +`prepare` must not be load-bearing for validation. Empty batches may bypass value preparation, and +the executor needs its safety and fallibility facts before it runs the closure. + +## Null execution policy + +The old public `NullHandling` enum and argument witness were removed. Authors do not select an +execution mechanism. The executor derives a private row policy from the dispatched input and result +types: + +- `Dense` may execute over garbage behind nulls and masks afterward; +- `DenseWithRetry` may execute densely, then retry valid rows when deferred evidence reports an + error; and +- `ValidOnly { filtered_decode_cost }` guarantees that the row closure sees only valid rows. + +An early-failing row or a decoder that is not dense-safe must use valid-only execution. A deferred +kernel may use dense execution because it writes a legal provisional value for every row. If only +garbage behind nulls reports an error, the valid-row retry discards it. + +Valid-only execution has two mechanisms. Filter-and-scatter shrinks inputs before decoding. +Branch-and-skip decodes the original batch and visits set bits from the conjoined validity mask. A +sink that does not support skipped rows automatically falls back to filter-and-scatter. + +The selector needs more than a boolean "decode shrinks" flag. Every `InputElement` declares an +additive `FILTERED_DECODE_COST`, defaulting to zero. `ElementTuple` sums the costs across arguments: + +- cost 0 always prefers branch-and-skip; +- cost 1 prefers branch-and-skip at 50% or more surviving rows; and +- cost 2 or greater prefers branch-and-skip at 85% or more surviving rows. + +This distinction comes from the x86 measurement in #9128. One nullable geometry input at 50% nulls +favored branching, while two independently nullable geometry inputs at 10% nulls each, about 81% +survivors, favored filtering. OR-ing a per-argument flag loses exactly that distinction. + +The values are still a coarse heuristic. There is no evidence yet to separate cost 2 from cost 3, +and the batch-size crossover has not been measured. `NullStrategy` remains only as a test-harness +seam for forcing a mechanism. Do not expose the private row policy as an author contract. + +## Performance and generated-code evidence + +The older Ryzen 9 7950X AVX-512 measurements remain the production-performance record in the +[#9128 follow-up](https://github.com/vortex-data/vortex/issues/9128#issuecomment-5151831802). They +also supplied the per-argument null-selection evidence above. + +The final API cleanup was checked separately against its parent, `53c51d803c`, by cross-compiling +the optimized `row_fn_executor` benchmark for `x86_64-apple-darwin` with `target-cpu=x86-64-v3`. +After normalizing symbol names and metadata, the vector and reduction blocks were identical for all +three executor shapes: + +- ordinary wrapping add through `ElementSink`; +- checked add with deferred evidence; and +- wrapping add through a custom sink. + +The wrapping loops retain 256-bit `<4 x i64>` loads, adds, and stores. The checked loop retains the +same vector loads and adds, derives overflow with vector xor/and/compare operations, accumulates +`<4 x i1>` with vector OR, and reduces after the loop. None of the vector bodies contains a call or +panic path. Scalar tails are unchanged. + +The production tensor benchmarks were also cross-compiled before and after the cleanup. Normalized +arithmetic sequences and counts match for `l2_norm`, inner product, and cosine similarity. Their +ordered floating-point reductions are scalar-unrolled in both revisions because LLVM preserves the +strict reduction order. The cleanup did not remove vectorization because those reductions were not +vectorized before it. + +Native Apple M4 Max timings used 65,536 rows, two alternating before/after runs, 100 samples, and a +0.5-second minimum per arm. RowFn median deltas ranged from 1.11% faster to 0.94% slower. Fastest +deltas stayed within about 0.17%, while specialized controls drifted by as much as 3.7% in their +medians. There is no measurable native regression from the API cleanup. + +This does not replace the required x86 runtime run above. Cross-target IR proves that the hot loop +shape survived, not that the revised null selector has the expected branch-predictor behavior on +x86. + +## Current implementation and checks + +The implementation includes production users in `vortex-array`, `vortex-tensor`, and `vortex-geo`. +`NumericBinary` is an unregistered `RowFn` used only for primitive arithmetic execution. Decimal +arithmetic keeps its existing path. The stable public-path benchmark baseline landed as #9136. + +The checks recorded for the final API state are: + +- 67 focused RowFn tests; +- 179 `vortex-tensor` tests; +- 230 `vortex-geo` tests; +- `cargo +nightly fmt --all`; and +- full workspace clippy, with `PYO3_NO_PYTHON=1 PYO3_BUILD_EXTENSION_MODULE=1` because the host + `/usr/bin/python3` is 3.9 while the workspace requires the Python 3.11 stable ABI. + +The generated-code comparison and native timing evidence are described above and in the final +section of `STRICT_SCALAR_FN_RESEARCH.md`. + +## Review pass: what changed and what was deliberately left + +A review of the three parts (API, execution, implementations). **The author-facing API is +unchanged**: every proposal that would have altered it was backed out, for the reasons below, and +what landed is cleanup, corrected documentation, and test coverage. The emitted IR of every +`visit_prepared_into` monomorph is identical to the pre-review commit. + +API: + +- `InputElement::decode_null_tolerant` overrides that only restated the default were deleted from + the primitive, bool and `TensorRow` elements. `GeometryRow`'s override is the only real one. The + doc now says a dense-safe element should *not* override. +- `ElementTuple` now records why it carries arities past the widest function in tree: it is sealed, + so a downstream crate cannot add the one it needs, and an uninstantiated arity costs only its own + macro expansion. + +Execution: + +- `execute_filtered` and the forced-strategy test seam now share `resolve_validity`, so the mask + materialization and the all-true/all-false shortcuts cannot drift apart between them. +- The dense-retry path's comment was wrong and is corrected. It filters unconditionally because + `execute_dense` is not handed the `branch` closure, **not** because a deferred sink cannot skip + rows: `ERRORS_ARE_DEFERRED` and `SUPPORTS_SKIPPED_ROWS` are independent consts and a sink may + legally set both. + +Implementations: + +- `l2_norm_row` had two copies, in `l2_norm.rs` and `cosine_similarity.rs`. Cosine's prepared and + per-row arms must agree bit for bit, which only holds while both accumulate in the same order, so + the duplicate was an invitation to break exactly the property the comments defend. One copy now + lives in `utils.rs` beside the other shared tensor helpers. +- `CosineSimilarity::reduce_encoded` zips its three slices instead of indexing `0..len` three times + per row, and documents why it materializes where `InnerProduct::reduce_encoded` stays lazy (the + zero-norm guard is a conditional, not an arithmetic factor). +- `IndexedSourceExt::map_checked_into` was deleted from vortex-compute. `CheckedSink` replaced the + split value/evidence pass it served, and it had no caller left. +- `contains_route` and the workspace `geo` dependency both record that the table transcribes geo's + `impl_contains_from_relate!` and must be re-verified on a version bump. `geo` is pinned to + `=0.31.0`: a caret requirement would admit 0.31.x patches, which `cargo update` (or automated + lockfile maintenance) takes with no diff to review, and a patch is free to reshuffle the dispatch + without any API change. The agreement tests stay green wherever relate and the direct algorithm + agree, so the pin, not the suite, is what makes the coupling break only deliberately. + +Split out onto `develop` instead of landing here: + +- **The checked-arithmetic macro collapse.** `primitive.rs` on this branch and on `develop` both + carry four near-identical `CheckedArithmetic` bodies that differ only in `mul_failure`, so the + collapse into one `impl_checked_integer!` belongs on `develop` where every caller benefits. It is + on `claude/collapse-checked-arith-macros`. This branch's `primitive.rs` keeps its four bodies + until `develop` is merged, at which point the collapse arrives with it and the merge conflict is + a member deletion rather than two competing macro structures. +- **The `mul_failure` kernel tests.** The exhaustive 8-bit sweep and the 64-bit probe grid already + exist on `develop` from vortex-data/vortex#9210 and arrive with the same merge. + +Deliberately **not** done: + +- **No `DeferredElementSink`.** `CheckedSink` exists largely because `ElementSink` cannot name an + error at `finish`. A framework sink combining an element output with a type-level message would + remove ~100 lines per function, but there is exactly one deferred-error function. Build it when a + second appears, rather than copying `CheckedSink`. +- **No change to `reduce_encoded`'s probe semantics.** Hoisting the probe out of the strategy paths + and masking a full-length result looks like a simplification and is not one: + `normalized_readthrough_survives_null_rows` pins that a filtered input is no longer `Normalized`, + so which arrays reach `reduce_encoded` is load-bearing and differs per strategy. +- **No PR split.** Recommended landing order, each step individually revertible and separately + benchmarkable: (1) API + lifting with dense/filter only; (2) branch-and-skip + adaptive selection + + its benchmarks; (3) `NumericBinary`; (4) tensor; (5) geo. The seam already supports this split + and no API changes between steps. + +### Three API changes proposed, and why none of them landed + +All three were implemented, run against the suite, and backed out. None prevents a bug, and this +branch's open work is *settling* the API rather than churning it, so they belong in #9129 as +questions decided alongside the rest of the surface: + +- **Should `reduce_encoded` take an explicit `row_count`?** The filtered-count requirement is real + and easy to miss, but `args` are filtered to match, so `args[0].len()` is already both the natural + thing to write and correct. The parameter is documentation, and it costs every implementor a + signature change. What survived is the test: + `reduce_encoded_is_probed_before_and_after_filtering` pins that the rewrite is offered the + original arrays at full length and then the filtered ones at the surviving count. +- **Should `OutputSink::row_count_matches` become `rows_len`?** A length reads cleaner and lets the + executor name what it found. Against that, `row_count_matches` lets a sink fold in its own + invariants, which `SpreadSink` uses for its width check; narrowing it turns that into a panic. + Neither spelling prevents a bug. +- **Should the nullary path go?** A function with no inputs has no validity to lift, which is the + lifting's whole job. But `RowFn` would still give it sink allocation and dtype derivation, so + `random()` or `now()` is not obviously better hand-written, and the path is ~70 lines and tested. + +Trimming `ElementTuple` to arity four was proposed on the same reasoning and backed out for a +stronger one: the trait is sealed, so the arities are the only ones a downstream crate can ever +have. + +### Two changes this pass made and then reverted + +Both were proposed, implemented, reviewed, and backed out on evidence. They are recorded because +each is an attractive idea that a later reader will have again. + +**Making `CheckedSink` safe with `BufferMut::zeroed` costs 1.65 to 1.71x.** Replacing the +`MaybeUninit` storage removes an `unsafe set_len` and reads as a clear win, and `ElementSink`'s own +comment appears to bless it by routing a zeroable placeholder to `alloc_zeroed`. Measured, it is +not: allocate-zeroed-then-fill against allocate-then-fill, interleaved in one process over `u64` +outputs, ran **1.221x** slower at 8 KiB, **1.71x** at 64 KiB, **1.66x** at 512 KiB and **1.71x** at +2 MiB, stable to within 2% across two runs. `alloc_zeroed` does not avoid the write: below glibc's +mmap threshold `calloc` recycles a dirty chunk and memsets it, and above it every fresh page faults +on first touch. The row loop overwrites every slot regardless, so this is a duplicated pass over +the output of the hottest kernel in the system. + +Note the corollary, which is a real optimization nobody has taken: `ElementSink::with_capacity` +pays exactly this on every batch, and only branch-and-skip ever reads a placeholder back. A sink +that allocated uninitialized on the dense and filter paths would recover it. + +**Hoisting `OutputSink::SUPPORTS_SKIPPED_ROWS` into the plan is not sound as an optimization.** +#9130 records "avoid probing `reduce_encoded` twice when branch execution is unsupported" as a +follow-up. It reads as free, and is not, because the branch path probes `reduce_encoded` against +the _original_ arrays before it consults the sink, and that is the only probe that ever sees them +still encoded. Skipping the path early leaves such a function with only the filtered probe, whose +canonical arrays match no encoding fast path. For a function whose reduction is _defined_ to answer +differently from its row loop, which is exactly what `L2Norm` over `Normalized` is, that is a wrong +answer rather than a slow one. Nothing in tree is reachable today only because every `ValidOnly` +dispatch happens to use `ElementSink`. **#9130's follow-up should be struck, not implemented.** +`reduce_encoded_is_probed_before_and_after_filtering` now pins the two probes and their row +counts. + +### On measurement, and what the IR gate does and does not cover + +Wall-clock benchmarking of the row loops was attempted first and abandoned on evidence. Two runs of +the *same* baseline binary, pinned with `taskset -c 2`, 100 samples, disagreed by up to 4x +(`row_wrapping_add_nullable`: 198.8 us then 52.9 us median; `specialized_checked_add`: 185.5 us then +34.4 us). The 4-vCPU shared VM drifts more within a session than any effect being measured, which is +the same conclusion this branch already reached on a dedicated 7950X. + +The gate used instead is the emitted optimized IR of every `visit_prepared_into` monomorph in +`vortex-array`, profiled by vector width, reduction count, overflow-intrinsic survival and bounds +checks, then compared as a multiset before and after. Reproduce with: + +```bash +RUSTFLAGS="--emit=llvm-ir -C codegen-units=1" cargo rustc -p vortex-array --release --lib +``` + +**Its blind spot is worth stating, because it nearly landed a regression.** The IR of a row loop +cannot show an allocator call outside it, so the `BufferMut::zeroed` substitution above passed this +gate cleanly while costing 1.7x. An allocation-strategy change needs its own targeted A/B, which is +cheap to write and immune to the host drift above because both arms run interleaved in one process. +Use the IR gate for loop shape and a focused microbenchmark for anything the loop does not contain. + +## Remaining boundaries + +- Complete the required x86 production and forced-null-strategy benchmark run above before treating + the thresholds or overall performance as settled. +- Keep nullable outputs separate until the first real function can define the validity contract. +- Do not add another sink composition abstraction. Put multiple builders in one custom sink. +- Do not add a general runtime-shaped sink until a production function needs one. +- Keep pattern compilation and other state shared across rows outside `RowFn` when it cannot be + represented as batch preparation. +- Use emitted optimized IR as a gate for numeric changes near LLVM's vectorization boundary, then + use the stable #9136 benchmark names for runtime confirmation. + +## Repository rules for the next agent + +Follow `AGENTS.md`. Keep public APIs small, run narrow checks before workspace-wide checks, and +report blocked checks separately from passing ones. Preserve unrelated working-tree and staging +state. Every commit must include the required `Signed-off-by` trailer. diff --git a/STRICT_SCALAR_FN_RESEARCH.md b/STRICT_SCALAR_FN_RESEARCH.md new file mode 100644 index 00000000000..1c873108345 --- /dev/null +++ b/STRICT_SCALAR_FN_RESEARCH.md @@ -0,0 +1,1790 @@ + + + +# A layered authoring API for strict scalar functions + +**Status: historical design record, with the final API review recorded at the end.** This document +keeps the experiments in the order they happened, including APIs and ports that were later removed. +The current architecture is one `RowFn` authoring trait, private lifting, one sink-backed +`RowVisitor::visit_prepared_into` primitive, and a deliberately open input/output vocabulary. Read +[`SCALAR_FN_HANDOFF.md`](SCALAR_FN_HANDOFF.md) for orientation, then the final section here before +using an earlier sketch. + +> **Later architecture decisions:** `StrictScalarFnVTable`, the columnar ports, returning visits, +> both witness types, `PersistableOptions`, the public `NullHandling`, the aggregate decode-shrinks +> flag, and the unused `TensorSink` were deleted. Framework-only visitor, tuple, and result traits +> are sealed. `InputElement`, `OutputElement`, and `OutputSink` remain open so functions can add +> their own decode and output primitives. Sections below remain the evidence that led to those +> decisions, not the API to implement. + +--- + +## Current benchmark and codegen record + +The authoritative current comparison is the +[x86 AVX-512 re-measurement on issue #9128](https://github.com/vortex-data/vortex/issues/9128#issuecomment-5151831802). +It records the machine, exact refs, stabilized governor, two-run fastest and median results, control +limitations, geo fix, adaptive-null diagnostics, and native LLVM IR/assembly in folded sections. +It supersedes every older shared-VM or pre-#9076 figure in these notes for claims about the current +branch versus `develop`. + +The run used candidate `d293d3cdd59e` plus the recorded geo bbox widening, baseline +`876996fe7846`, and a Ryzen 9 7950X pinned to CPU 4 with the TSC timer and performance governor. +The conclusions that survive into the implementation plan are: + +- sink-only checked add is 1.018-1.226x faster by median than its benchmark-local specialized + control, depending on constant and null shape; +- cosine is 1.40-30.13x faster than current develop; +- prepared overlapping `contains` is 8.60-8.77x faster by median, while the widened bbox gate + restores disjoint polygons to parity with #9076; +- point/constant geo still has real 8.6-14.2% and 10.9-13.2% median regressions; +- `BytesLen` is 1.410-1.411x faster by median on long strings and 1.097x on short strings; +- the global 75% survivor threshold mispredicts both one-input/50%-null and + two-input/10%-null geo cases, so adaptive selection needs element/arity-aware cost data; +- checked-add codegen has AVX-512 error-word accumulation and post-loop vector reduction with no + per-row error branch; current `l2_norm` remains a strict-order scalar reduction. + +The stabilized cosine median ratios preserve the shape and width dependence instead of collapsing +the result into one headline range: + +| shape | width 2 | width 32 | width 256 | +| --- | ---: | ---: | ---: | +| column x column | 5.77-5.79x | 1.93-2.19x | 2.54-2.58x | +| column x constant | 12.44-12.48x | 28.87-28.96x | 30.08-30.13x | +| column x extension constant | 3.05x | 1.40-1.41x | 1.77-1.78x | + +The final geo median ratios, including the bbox patch, are: + +| predicate and shape | develop / row branch | +| --- | ---: | +| contains, column x column points | 0.951-0.964x | +| contains, column x column polygons | 0.999-1.003x | +| contains, constant x points | 0.876-0.921x | +| contains, disjoint polygons | 0.993-0.997x | +| contains, overlapping polygons | 8.60-8.77x | +| intersects, column x column polygons | 0.983-0.995x | +| intersects, points x constant | 0.884-0.902x | +| intersects, disjoint polygons | 1.011-1.019x | +| intersects, overlapping polygons | 1.011-1.023x | + +Here, as above, ratios greater than 1x favor the row branch. The issue comment contains the paired +fastest and median observations rather than only these compact ranges. + +The historical measurements below remain because they explain design decisions and experiments made +while building the prototype; they are not the current before/after performance record. + +--- + +## The design in one screen + +```text +RowFn ──────────blanket──▶ StrictScalarFnVTable ──────blanket──▶ ScalarFnVTable +(row at a time, types (null / constant / validity (full control) + chosen per batch) lifting for a columnar kernel) +``` + +Two authoring traits, one for each axis a strict function actually varies on, plus a third axis (*how +a row is typed, and how its output is delivered*) factored into an open element and sink vocabulary that +neither trait mentions. + +### `StrictScalarFnVTable`, the null/validity lifting + +Write the structural metadata plus one **columnar** kernel that ignores validity. A blanket impl +derives: + +- `is_strict = true`, and a mirrored `validity` a kernel can answer with the conjunction of its child + validities when it never turns a wholly non-null row into a null (see + [Strictness is not totality](#strictness-is-not-totality)), so the planner knows which rows are null + without executing the function. +- `return_dtype` = `return_element_dtype` widened to nullable iff any input is nullable, so the + strictness dtype contract holds by construction rather than per function. +- `execute` = the shared cases before the kernel runs: a null-constant input short-circuits to an + all-null constant, all-constant inputs evaluate one row and broadcast, and partially-null inputs + are handled per `NullHandling` (`Dense` masks after a full pass, `Filter` filters then scatters). +- Options serde, from `PersistableOptions` on the options type. + +This is the layer for a function whose kernel is columnar rather than row-at-a-time: `not` (one `!` +per 64-bit word), `list_length` (a difference of offset buffers), `list_sum` (a grouped accumulator over +the elements child). See [Why three concepts and not fewer](#why-three-concepts-and-not-fewer) for why it +cannot be folded away. + +### `RowFn`, one row with element types chosen per batch + +Name a witness argument tuple and return type, then in `dispatch` pick the concrete element types for +a batch and hand the framework a row closure through a rank-2 visitor. A blanket impl derives the +whole `StrictScalarFnVTable` from it. When the element types are fixed, `dispatch` is a single +`visit` at those types. When one ID spans several widths (`l2_norm` accepts f16/f32/f64), `dispatch` +matches on the input dtypes and visits at the chosen width. + +Everything structural follows from the argument tuple and return type: arity, per-argument dtype +validation, the output dtype, null handling, and fallibility. There is nothing for an implementor to +declare twice or get wrong, because the framework reads it off the types (see +[Properties, not conventions](#properties-not-conventions)). A constant operand is decoded once and +read at stride 0, so a broadcast argument costs one decode rather than one per row. + +Output takes one of two forms, chosen per visit. `visit` takes a closure that **returns** an +`OutputElement`, one owned value per row whose dtype is a property of its Rust type. `visit_into` takes +one that **writes** into an `OutputSink`, allocated once per batch knowing the output dtype and handing +out a place to write. Orthogonally, `visit_prepared` runs a once-per-batch prepare step over the +element values of whichever operands are constant for the batch, and threads its result to every row +by shared reference; plain `visit` is that with unit state (see +[Constant compute](#constant-compute-the-last-quadrant-of-the-lifting)). The sink carries what an owned per-row value cannot: `l2_denorm` writes each row +into a slice of one flat buffer, so its output width comes from the arguments and it allocates once +rather than per row. The executor holds the sink and passes the handle in, so the closure stays `Fn` +and the returning path pays nothing. + +Note that `RowFn` does not *require* totality, it just cannot currently express its absence: both output +forms build an all-valid column, so a row kernel has no way to say "this row is null". An +`impl OutputElement for Option`, or a sink that can push a null, would lift that, at the cost of +revisiting the `validity` law that reads the output validity off the inputs. No function needs it yet, so +it is not there. + +### The element vocabulary, how a row is typed + +`InputElement`, `OutputElement` and `OutputSink` are open traits. A `NativePType`, `bool`, `Bytes` (a +resolved `&[u8]`), and `BytesLen` (a length read from a view without resolving it) ship in the framework, +and `vortex-tensor` adds `TensorRow`, reaching through the extension wrapper into flat storage, plus +`TensorSink` on the output side, in its own crate. Adding `&str`, decimals, or a list row is one impl +that every row function gains, with no framework change. + +--- + +## Why three concepts and not fewer + +The standard applied here: every trait, and every member of every trait, has to have a purpose +nothing else can provide. Testing each against that standard is what the bulk of this research was. + +### `RowFn` and the witnesses are forced, not chosen + +A scalar function's *signature*, meaning its arity and fallibility, is a property of +`(function, options)` with **no input dtypes**: `ScalarFnVTable::arity(&self, options)` and +`is_fallible(&self, options)`, and `ScalarFnSignature` above them, take none. So any framework that +derives arity and fallibility from element types has to be able to name element types *without seeing +dtypes*, which is exactly what `ArgsWitness` / `RetWitness` are. Because `dispatch` *does* see dtypes +and could choose otherwise, some check has to tie the two together, which is the compile-time witness +check below. This cost is not a consequence of the rank-2 encoding: **any** design that derives a +dtype-free signature from per-batch types pays it. + +A previous iteration made the width choice a generic-associated-type family generated by a +`row_family!` macro. Rust cannot abstract over a GAT's bound (`type Args` is +rejected), so that approach needed a trait *and* an adapter per width class, hand-written or +macro-stamped. The rank-2 visitor sidesteps the limit rather than writing around it: the kernel owns +the width `match`, where `T: Float` appears literally inside a `match_each_*_ptype!` arm, and the +framework method `RowVisitor::visit` is generic only over bounds it +owns. The macro, its family traits, and its generated adapters are all deleted. Note that `dispatch` +is not even per-*width*: it can pick different element *kinds* per dtype, which no +bound-parameterized family could. + +### `ElementwiseFn` was not forced, so it is gone + +An earlier revision had a third trait, `ElementwiseFn`, for the fixed-element-type case: name `Args` +and `Ret`, write `apply`. It read cleanly, but it failed the standard. `RowFn` already covers the +fixed case (the dispatch is a single constant `visit`), so `ElementwiseFn` bought roughly seven lines +on exactly one production function (`byte_length`) at the cost of 114 framework lines and a third +link in the blanket-impl chain. The probes settled it: of the functions examined, `not` and `list_sum` +turned out not to be row functions at all, and `list_length` needed the encoding-aware +`reduce_encoded` hook that `ElementwiseFn` never exposed. So the constituency I expected it to have +never materialized, and it is deleted. `byte_length` writes a two-line `dispatch` instead. + +The one-trait-with-defaults alternative (a single `RowFn` with `dispatch` defaulted to visit the +witnesses and `apply` defaulted to `unimplemented!()`) was rejected because it converts a compile +error into a runtime panic: a type implementing neither method compiles, registers, and answers +signature queries with a plausible shape, then panics on first execution. `dispatch` is therefore +required. + +### `StrictScalarFnVTable` cannot be folded into `RowFn` + +`RowFn`'s type surface is *closed*. The output dtype is `OutputElement::element_dtype()`, drawn from +the finite set of `OutputElement` impls, `ElementTuple` exists only for arities 1 to 3, and the loop +is one `apply` per row. Three whole classes of strict function are therefore inexpressible as a +`RowFn` at any cost: + +- **Output dtype outside the element set.** `ext_storage`'s output is an extension array's storage + dtype, so `vortex.geo.box` is a struct and `vortex.uuid` is a `FixedSizeList(u8,16)`. `vortex-geo`'s + zone-map pruning calls `ext_storage` on a `geo.box` statistic, and a row-function port breaks it at + plan time. +- **Variadic arity.** `merge` and `select` take an unbounded number of children, while `RowFn` fixes + `Arity::Exact(n <= 3)`. +- **Sub-row-granular kernels.** `not` negates one 64-bit word at a time, so a row loop over `bool` is + ~64x the memory traffic and, measured, 406x slower at a 64Ki batch (see + [Measurements](#measurements)). + +So the middle layer has a genuine, disjoint constituency: `not`, `list_length`, `list_sum`, and +prospectively `select`, `merge`, `json_to_variant`. "Just a visitor" collapses three concepts to two +rather than to one. + +### Every remaining member earns its place + +A member-by-member audit, with call sites found by grep rather than by guess, turned up nothing +deletable. The non-obvious cases are worth recording: + +- **`RowVisitor::Out`** is what lets one `dispatch` `match` serve both plan time (`Out = DType`, + validate and name the output dtype) and run time (`Out = ArrayRef`, decode and run the loop). The + alternatives, a `{DType, ArrayRef}` enum unwrapped at each site or two separate dispatch hooks, + either add unwrap-panics or duplicate the width `match` in every width-polymorphic function with no + compiler check that the two copies agree. +- **A plan-time visit is unavoidable.** `l2_norm` declares `RetWitness = f64` but dispatches over + f16/f32/f64, so the output dtype read off the witness would be wrong for two of three widths. Also + `TensorRow::validate` rejects an `f32` column against an `f64` witness, and the visit is what + gives cross-argument uniformity for free (`int_max(i16_col, i64_col)` is rejected by + `(T, T)::validate`, not by any `dispatch` body, which only inspects `args[0]`). +- **`ApplyResult` distinct from `OutputElement`** is what lets one trait serve both infallible + (`Ret = f64`) and fallible (`Ret = VortexResult`) kernels without a wrapper. `f64` cannot be + simultaneously fallible and infallible, so the fallibility bit lives on the return *shape* rather + than on the element. + +--- + +## Properties, not conventions + +The framework's real value beyond line count is that two invariants an implementor used to have to +get right are now derived from the types, so an unsound combination cannot be written. + +### Null handling follows from the arguments and the return type + +`NullHandling::Dense` runs the kernel over every row including those behind nulls, then masks. It is +cheaper than filtering and the only option that leaves inputs at their original encoding, so it is +right whenever it is sound. Soundness needs two things, every argument readable behind a null row and +an infallible computation, and both are already visible in the types: + +```rust +const fn row_null_handling() -> NullHandling { + if A::DENSE_SAFE && !row_is_fallible::() { NullHandling::Dense } else { NullHandling::Filter } +} +``` + +Whether a dense read is safe is a property of the *element*, not of the function: reading a whole +value out of a flat buffer is safe (`NativePType`, `bool`, `TensorRow`, `BytesLen`), while following a +stored offset into a data buffer is not (`Bytes`), because arrays only validate the views of their +*valid* rows. This caught a real bug in this branch's own `byte_length`, see +[Problems to extract](#problems-to-extract-onto-develop). + +### Fallibility comes from the return type *and* the element decode + +A function is fallible if its computation can fail (`Ret = VortexResult`) **or** if decoding an +argument can fail on legal data. The second source is real and was missing: `geo_distance`'s row +computation cannot fail, but parsing WKB bytes into a geometry can, for a *valid* row holding +malformed bytes. So `InputElement` carries `DECODE_FALLIBLE`, and fallibility is the disjunction: + +```rust +const fn row_is_fallible() -> bool { A::DECODE_FALLIBLE || R::FALLIBLE } +``` + +`is_fallible` gates dict-value pushdown (`arrays/dict/compute/rules.rs`), which speculatively +evaluates a function over *unreferenced* dictionary values, so a function that under-reports +fallibility fails a query on rows it never needed. + +### The witness is checked at compile time + +Arity, dense-safety and fallibility must not vary between the choices `dispatch` makes, because the +framework acts on them before dispatching. Since (with `ElementwiseFn` gone) *every* function names +its element tuple twice, once as `ArgsWitness` and once in the `visit`, the check that the two agree +is load-bearing, and it is a compile-time `const` assert inside each visit: + +```rust +const fn assert_witness_agrees() { + assert!(A::ARITY == ::ARITY, "…"); + assert!(A::DENSE_SAFE == ::DENSE_SAFE, "…"); + assert!(row_is_fallible::() == row_is_fallible::(), "…"); +} +``` + +Monomorphizing any dispatch arm evaluates it, so even a `match` arm that never runs at a given width +is checked, and a disagreement fails the build pointing at the exact `visit::<…>` call. It compares +the raw arity/dense-safety/fallibility rather than the derived `NullHandling`, which collapses +dense-safety and fallibility together and would miss an arm that flipped both. A `compile_fail` +doctest pins that a lying witness does not compile. This replaced a runtime check that ran three +times per array (plan, execute, deserialize). + +--- + +## Strictness is not totality + +This is the finding that decides what the middle layer may derive. Note that +[#9033](https://github.com/vortex-data/vortex/pull/9033) reached the same conclusion independently and +has since landed, so this section is no longer the argument for the finding, only for the API that +follows from it. + +Before #9033, the `is_strict` documentation stated the validity-equivariance law, +`f(…, mask(aⱼ, m), …) == mask(f(…, aⱼ, …), m)`, and then asserted as "consequence 1" that output +validity is the conjunction of input validities. **Consequence 1 does not follow from the law.** It +needs an extra premise: that the kernel never turns a wholly non-null row into a null. #9033 replaced +that equality with a one-sided bound, `valid(f(a₁, …, aₖ)) ⊆ valid(a₁) ∧ … ∧ valid(aₖ)`, which is the +vocabulary this branch uses. `docs/strictness-and-validity-pushdown.typ` proves the law and the +null-propagation reading are the same property, and separates what does not follow from either. + +`list_sum` is the counterexample. Summing a valid *empty* list yields null. It still satisfies the law +(a null it introduces at a valid row appears identically on both sides of the equation and cancels), +so it is genuinely strict, but its output validity is *narrower* than its input validity. + +Two properties, then, not one: + +| property | what needs it | +| --- | --- | +| **strict** (null propagation, equivalently validity equivariance) | every validity push-down, the thing we actually want | +| **total** (non-null in implies non-null out) | upgrading the `⊆` bound to `=`, so validity is precomputable | + +The old blanket impl derived `validity = union_child_validities` for *every* implementor, which needs +totality while the trait only requires strictness. Every current implementor happens to be total, so +nothing was broken, but a partial function joining the layer would get a `validity` that contradicts +what it computes: `arr.validity()` would report all-valid while `arr.execute()` yields the null, since +`ValidityVTable::validity` evaluates the derived expression. `list_sum` was about to be +exactly that, and is now ported onto the layer as the first non-total member. + +#9033 says a function satisfying the stronger equality "can advertise that through +`ScalarFnVTable::validity`". That is the same idea as `is_total`, moved from a hand-written method to a +boolean, because a blanket impl cannot hand-write `validity` per function: it needs the property as +data in order to decide whether to derive one. + +The fix needs no new property. `validity` is mirrored on `StrictScalarFnVTable` alongside `reduce`, +defaulting to `None`, and a kernel that satisfies the equality answers it with +`union_child_validities`. The unsound direction is the one that now takes work, and the safe default is +what a function gets for free. + +An earlier revision of this branch instead added an `is_total` method and derived `validity` from it. +That was strictly worse: it introduced a concept the codebase did not have, in order to compute +something a function can just say directly. It is gone. The `RowFn` blanket impl answers `validity` +for every row function, justified by its own output vocabulary (no `OutputElement` is nullable, so no +row kernel can introduce a null), which keeps the row layer at zero boilerplate. + +Note that strictness rather than totality gates membership either way: `is_null` is total but +disqualified, because it inspects validity and so does not propagate nulls. That is also why the trait +is not called `TotalFnVTable`. + +> **A related latent issue, deliberately not fixed here.** Four functions declare `is_strict = true` +> and are strict-but-not-total: `get_item` (a nullable field under a non-null struct), `mask`, +> `variant_get`, `geo_envelope`. None is broken today, since `get_item` leaves `validity` at the +> default and `mask` overrides it correctly, but any that grows a conjunction-shaped `validity` +> derivation would be wrong. This predates the branch and belongs in its own investigation. + +--- + +## Problems to extract onto develop + +The framework surfaced three problems that are not really about the framework. Each is filed +separately and I think each should land as its own PR rather than riding in on this one. Note that +none of them is a live miscompute on `develop` today, which is worth saying plainly, because the +branch's own commit messages describe fixes to *this branch's* code. + +1. **Strict-but-non-total validity derivation ([#9091]).** The `is_strict` documentation presents + totality as a consequence of strictness when it is an independent premise (see above). Nothing + derives validity from `is_strict` automatically, so nothing is wrong today, but the doc invites the + next strict-but-partial function to write `validity: union_child_validities` and be silently wrong. + **Superseded by [#9033], which lands the documentation correction on `develop`.** This branch needs + nothing beyond that, since it now mirrors `validity` rather than deriving it from a property. + +2. **Views behind null rows are unvalidated ([#9090]).** `VarBinViewArray::validate_views` only + validates the views of *valid* rows, so a legal array can hold a view behind a null row naming a + buffer that does not exist, and resolving it densely panics (`index out of bounds: the len is 1 but + the index is 9`). On this branch, expressing byte length as "a function of the row's bytes" quietly + changed *what gets decoded* and hit that panic. The fix here reads the length out of the view + (`BytesLen`) and never resolves the row, and + `test_byte_length_ignores_unresolvable_views_behind_nulls` pins it (verified to panic without the + fix). `develop`'s `byte_length` was already immune, since it also read `view.len()`, so the + extraction is that regression test rather than a code change. The doc half is also covered by + [#9033], which deletes the dense-evaluation "consequence 2" outright rather than narrowing it. That + leaves `InputElement::DENSE_SAFE` as the only place the licence is written down, per element rather + than as a blanket claim, which is where it belongs. + +3. **Bit-at-a-time bool packing ([#9092]).** `OutputElement for bool` used `BitBuffer::from_iter`, + where the `Vec` is already owned and contiguous so `BitBuffer::from` routes to the + multiversioned SIMD packer. Measured **6.6 to 7.9x faster** on the packing step, for every + bool-returning row function. Note that `OutputElement` only exists on this branch, so the + develop-side instance of the same pattern is a different call site: + `encodings/sequence/src/compute/compare.rs` builds an n-bit result with a per-row predicate when it + already knows the single set index. I have not benchmarked that site. + +[#9033]: https://github.com/vortex-data/vortex/pull/9033 +[#9090]: https://github.com/vortex-data/vortex/issues/9090 +[#9091]: https://github.com/vortex-data/vortex/issues/9091 +[#9092]: https://github.com/vortex-data/vortex/issues/9092 + +--- + +## Audit: can the four `StrictScalarFnVTable` impls really not be `RowFn`? + +There were exactly four in production when this audit ran. Auditing each against the two questions that +matter, rather than repeating the earlier verdicts, **not one of them was structurally impossible**. Every +"cannot" in this document was really "cannot with the trait signed as it is today". One of the four, +`l2_denorm`, has since moved onto `RowFn`, so three remain. Recording the distinction because it is the +difference between a limit and a decision. + +| function | signature expressible? | kernel row-shaped? | what it would take | +| --- | --- | --- | --- | +| `not` | **yes**, `(bool,) -> bool`, both elements exist | **no** | nothing. It can be a `RowFn` today and should not be: `!bits` is one `!` per 64-bit word, in place when unshared, against 16k closure calls and a `Vec` repack | +| `list_length` | output is a fixed `U64`; input needs a `ListLen` element | **no** | one new element. Still should not: the answer is a child array or one constant | +| `list_sum` | output is one number per row, so nearly: only the *nullability* is unexpressible | **no** | `impl OutputElement for Option` and a list element, but the kernel is the real blocker | +| `l2_denorm` | **yes, now**: an `OutputSink` names its dtype from the arguments | yes, per-row scaling | **done**, see below | + +**A varying output dtype was already supported, and listing it as a blocker was wrong.** `dispatch` +chooses element types per batch and `return_element_dtype` routes through it, so `R::Out::element_dtype()` +is already answered per dispatch arm. `l2_norm` relies on this today, visiting `::<(TensorRow,), T>` +with `T` ranging over the float widths. The compile-time witness check pins only arity, dense-safety and +fallibility, deliberately leaving the output type free to vary. What `l2_denorm` needed was different and +narrower: its output dtype depends on the input *dtype* in a way no choice of element type can express, +because the extension dtype carries a shape. That is what `OutputSink::sink_dtype(args)` supplies. + +**`list_sum`'s output side is the easy part; its kernel is not.** One number per row means it needs only +a nullable output element, no write-into-buffer machinery. But `execute_strict` is not a per-row sum: it +builds a `GroupedAccumulator` over `Sum`, calls `accumulate_list`, and then `mask_empty_lists` computes +per-group emptiness with `count_range` popcounts, with all-true and all-none fast paths and an early +return when nothing needs masking. Porting it to a row loop would hand-roll the shared aggregate +framework, lose the overflow modes that `NumericalAggregateOpts` selects, and trade SIMD popcounts for +per-row checks. That puts it in the same category as `not`: expressible, and worse. + +So `l2_denorm` was the only one of the four whose kernel actually wants to be a row loop, which is why it +was the right first target despite needing the larger output-side change. + +Two readings follow. + +**The honest framing is "can, and here is whether it is worth it."** For `not` and `list_length` the +answer is a flat no on performance grounds, and those are settled. For `list_sum` the answer is +yes-with-changes, and the change it wants is a nullable output, which the sink could supply but which the +`validity` law argues against (see below). + +**`l2_denorm` was the one worth doing, and it is done.** Its kernel genuinely is per-row scaling, and +it carried the `unsafe` the other three tensor ports removed. What it needed was a second visit method +whose closure *writes* its row instead of returning it, generalized to an `OutputSink` rather than +hardcoding `&mut [T]`, because the same mechanism covers three gaps recorded separately in these notes: + +- **runtime-shaped output**: the sink is a preallocated flat buffer and the per-row handle a + `&mut [T]` slice of it, so `l2_denorm` allocates once per batch rather than once per row. This is what + shipped. +- **`str -> str` without the double copy**: the sink is one growing byte buffer plus views, and + `upper`/`lower`/`replace` push into it. Strictly better than the `Cow` output element considered + above, which still copies each row once. Not built, but the trait admits it unchanged. +- **nullable output**: a sink *could* push a null, which would remove the need for + `impl OutputElement for Option` as a separate patch. Deliberately **not** taken: both output forms + build an all-valid column today, and that is exactly what lets the blanket `validity` return + `union_child_validities`. Adding nulls to either form has to come with that law being revisited. + +### What shipped + +```rust +pub trait OutputSink: 'static + Sized { + type Row<'a> where Self: 'a; + fn sink_dtype(args: &[DType]) -> VortexResult; + fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; + fn row(&mut self, index: usize) -> Self::Row<'_>; + fn finish(self) -> VortexResult; +} + +fn visit_into( + self, + apply: impl Fn(A::Elems<'_>, S::Row<'_>) -> R, +) -> VortexResult; +``` + +**The executor threads the sink, not the closure**, so `apply` stays `Fn` and the existing `visit` pays +nothing. That was the design constraint, not an accident: relaxing `visit` itself to `FnMut` measured at +8 to 11% (see the `like` discussion), and a handle passed in per row avoids captured mutable state +entirely. Measured after the fact, `l2_norm` is unchanged at 69.05 µs against the 69.44 µs recorded +before the sink landed. + +**Step 1 of the earlier plan turned out to be unnecessary.** The plan called for widening +`OutputElement::element_dtype()` to take `args`. It never happened, because `sink_dtype(args)` puts the +argument-dependence on the *sink* instead, leaving all three existing `OutputElement` impls untouched. +That is the better split: an element's dtype genuinely is a property of its Rust type, and only the +thing that needs the arguments asks for them. + +**The `RetWitness` split resolved as predicted.** It carried two roles, *what dtype* and *is it +fallible*, and only the second is readable before `dispatch` picks a form. So `RowResult` now holds just +`const FALLIBLE`, with `ApplyResult: RowResult` adding the output element and `SinkResult: RowResult` +adding nothing but the error, and `RowFn::RetWitness` is bounded by `RowResult`. A returning dispatch +names `f64` or `VortexResult`; a writing one names `()` or `VortexResult<()>`. Coherence permits +this: `impl RowResult for ()` does not overlap `impl RowResult for T` because +`(): OutputElement` does not hold and no downstream crate can make it hold, the same negative reasoning +the pre-existing `ApplyResult` impls already relied on. + +**A new limit, worth naming.** `sink_dtype` sees the input dtypes but **not** the function's options, +because `OutputSink` does not know the `RowFn`'s `Options` type. A function whose output dtype depends +on an option value therefore still drops to `StrictScalarFnVTable`, whose `return_element_dtype` sees +both. Nothing in the repository needs it, and threading options through later is additive. + +### Results + +`unsafe` in `l2_denorm.rs` went from 8 blocks to 6. The two removed are the memory-safety ones on the +kernel path: `FixedSizeListArray::new_unchecked` in the constant-norms path, now `try_new` (the norm is +cast to the element dtype first, so the product stays non-nullable and the check passes), and +`PrimitiveArray::new_unchecked` in `build_tensor_array`, now `new`. That second one is an independent +cleanup rather than something the port forced. + +The 6 remaining are not of that kind and are not the row layer's business: four are calls to +`L2Denorm::new_array_unchecked`, an `unsafe fn` whose contract is the *semantic* unit-norm invariant and +not memory safety, and two are buffer pushes inside `normalize_as_l2_denorm`, a helper that builds the +normalized child and is not a scalar function at all. + +**Performance: the sink is faster than the kernel it replaced**, which was not the expected outcome. +`vortex-tensor/benches/l2_denorm.rs`, `fastest` column, both configurations run twice, 16384 rows, +non-nullable. The control implements `StrictScalarFnVTable` with the pre-port body, so it shares the +strict lifting and the gap is the row layer alone: + +| width | sink | pre-port kernel | ratio | +| --- | --- | --- | --- | +| 2 | 88.02 / 88.16 µs | 60.19 / 60.45 µs | sink 1.46x slower *(since fixed, see below)* | +| 32 | 482.0 / 515.5 µs | 1.175 / 1.014 ms | sink **2.1x faster** | +| 256 | 10.23 / 10.43 ms | 20.41 / 22.48 ms | sink **2.0x faster** | + +The likely cause of the win is that the pre-port kernel collected a `flat_map` over rows into a fresh +`Buffer`, and `flat_map` is not `TrustedLen`, so that `collect` grew the buffer with a capacity check +per element. The sink allocates once with `BufferMut::zeroed` and each row writes a slice of it, which +vectorizes. The zeroing is not a separate pass at these sizes, since large allocations come back zeroed +from the allocator. This is a hypothesis consistent with the width scaling rather than something +profiled. + +Width 2 showed the same regression as `l2_norm`'s, and for the same reason: both read tensor rows through +`TensorRow`, whose `get` re-derived a typed slice per row. Typing the column at decode time took +`l2_denorm` from 88.0 µs to **48.9 µs** at width 2, ahead of this control rather than behind it. See +[the like-for-like comparison](#the-like-for-like-comparison-and-the-per-row-cost-that-was-hiding-in-it) +for the measurement and for the wrong diagnosis it corrects. + +The constant-norms fast path moved to `reduce_encoded`, which sees the argument arrays before the row +loop. It keeps both of its cases (unit norms return the normalized child untouched, any other constant +rewrites the storage elements through one multiply), and it still fires for a filtered batch because +filtering a constant yields a constant. + +**Two visit methods do not cover everything, and it is worth being precise about the residue.** They +cover every function whose output is *computed* per row, returned or written. What stays columnar is +output that *aliases* its input, since `trim` and `substring` want to keep the input's data buffer and +rewrite only views, copying nothing, and a sink still copies bytes into itself. Likewise kernels whose +natural unit is not a row (`not`'s word-at-a-time negation, `binary`'s slice kernels) gain nothing. + +The sink is also what a `str -> str` string library needs. After reclassifying `L2Denorm` as an +encoding, that string library becomes the prospective first production user rather than a second +one. The experiment still demonstrates that the generic sink can carry runtime-shaped and +builder-backed outputs without making the returning path pay, but it should not be stabilized from +the tensor experiment alone. + +--- + +## Constant compute: the last quadrant of the lifting + +The lifting's constant handling was complete on the data side and absent on the compute side. A +null-constant input short-circuits, all-constant inputs fold to one row, and a constant operand is +decoded once and read at stride 0. What nothing owned was kernel computation that depends only on a +constant argument: `cosine_similarity(rows, query)` with a broadcast query re-accumulated +`norm(query)`, an O(width) pass plus a sqrt, once per row, and the geo predicates rebuilt the +constant side's topology graph, R-tree, or bounding box once per row. `cosine_similarity` escaped +partially by hand-writing a `reduce_encoded` rewrite, and the survey found that rewrite already +wrong for the literal shape, which is the argument for framework ownership stated as a correctness +fact: one hand-written constant path per function is one place per function to rot on +encoding-normalization details. + +### Where the hook can live, and where it cannot + +The hoist needs three things at once: knowing which arguments are constant, having their decoded +values, and a typed place for the function to compute from them. Constness is a per-batch value +fact (a RunEnd slice landing inside one run, a per-chunk compression decision), so: + +- **`dispatch` cannot see it.** It runs at plan time and run time and must choose identical element + types at both; values do not exist at plan time. +- **Element types cannot encode it.** A `Const` wrapper element would need value-aware dispatch + to be chosen, splitting plan/run monomorphizations in exactly the way the witness deliberately + does not pin, and costing 2^arity dispatch arms. The salvageable half of the idea, + framework-internal value-driven specialization, already exists as the stride-0 `ArgColumn`. +- **The closure cannot memoize it.** An `unsync::OnceCell` capture compiles under `Fn`, but without + constness information it is wrong (it would cache row 0 of a varying operand), and with that + information it saves nothing over a prepare step while planting an unhoistable load inside the + loop. + +That leaves one point: inside the visit, after decode, where `ArgColumn` already knows each +column's stride. `ElementTuple` gains `ConstElems<'a>`, the element tuple with every slot wrapped +in `Option` (`Some` iff that operand is batch-constant), and the visitor gains: + +```rust +fn visit_prepared( + self, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>) -> R, +) -> VortexResult; +``` + +`prepare` runs once per batch; its result reaches every row by `&P`, so `apply` stays `Fn` and the +loop keeps the shape the FnMut measurement forbids changing. `P` names no column lifetime, so +prepared state provably cannot alias the columns the loop reads. Plain `visit` is now a *provided* +method, `visit_prepared` with unit state: the ZST erases under monomorphization (measured, l2_norm +non_nullable at 33.38 us against the 32.83 us hand-written control, parity), the duplicate row loop +is deleted, and the visitor's method count grows with genuine axes (how output is delivered) rather +than with feature combinations. + +`prepare` is infallible in v1: it refines values the row loop could compute itself, and fallibility +is read off the witnesses before dispatch, so a failing prepare would have nowhere to be declared. +The extension (prepare returning `VortexResult

`, riding the existing fallibility axis) is +documented next to the method and deliberately unbuilt, because no adopter needs it. + +Three boundary facts worth stating because they will bite someone: + +- **Prepare must never be load-bearing for validation.** An empty batch decodes every operand as + non-constant (there is no row 0 to slice), so a prepare that validated its constant would + silently not run. Validation belongs to `validate` and the dtype rules. +- **What counts as a batch constant is wider than the constant encoding.** The stride-0 decode sees + one level through two wrappers that spell "the same value in every row" without being it: + `MaskedArray(ConstantArray)`, how the compressor spells an all-same-with-nulls chunk (sound + because the lifting owns validity entirely, so the value the loop reads behind a null row is + unobservable), and `Extension` over constant storage, the shape extension builders produce before + `ExtensionConstantRule` normalizes it. +- **`P` having no `Send`/`Sync` bound is load-bearing.** geo's `PreparedGeometry` carries + `Rc`/`RefCell` and could not be prepared state otherwise. The flip side, recorded so it is a + decision rather than a surprise: adding such bounds later (a parallel row loop, say) is a + breaking change to real adopters, not a relaxation. + +### What it bought, measured + +**cosine_similarity, and a lesson in ILP.** The closure accumulated the rhs norm per element and +sqrt'd it per row, a third of the arithmetic plus one of two sqrts. Hoisting it moved the benchmark +by only ~5% at width 32 and ~3% at 256 (16384 rows, fastest column), far under the flop count, +because the loop is latency-bound on the serial dot-product FMA chain (FP reassociation is illegal) +and the removed accumulation was executing in the chain's spare ILP slots. The measurable saving is +the hoisted sqrt. The row is bit-identical either way, each arm accumulating in the same order as +the unprepared kernel. + +The lesson generalizes and is the honest scoping of the feature: **"removes an O(width) pass per +row" is not "saves time" when that pass rides in ILP slack.** The work that collects the full +saving is work that extends the dependency chain: parses, tree builds, prepared structures. Which +is exactly what the geo numbers then showed. + +**The geo predicates, where the win lives.** `contains` substitutes an owned +`PreparedGeometry<'static>` of the constant operand (r-tree plus self-noded topology, built lazily +inside `P` through a `OnceCell` so point-row batches never pay for it) into relate exactly where +geo routes `Contains` through relate, argument order preserved including the `MultiPolygon` +reversal; direct pairings keep geo's own algorithms untouched. `intersects` hoists the constant +side's `bounding_rect` and replays geo's own disjoint-bboxes early-out, gated to fire only where +geo makes exactly that comparison first. `distance` was investigated and left alone: geo builds +R-trees for both sides inside a private helper on every call, so there is no seam to reuse one, and +the finding is recorded as a doc comment on its dispatch. 16384 rows, fastest column, two runs: + +| arm | before | after | change | +| --- | --- | --- | --- | +| contains, constant x polygons, overlapping | 457.5 / 458.0 ms | 50.88 / 50.00 ms | **9.1x** | +| contains, constant x polygons, disjoint | 7.04 / 7.05 ms | 3.97 / 3.74 ms | **1.9x** | +| contains, constant x points (direct route) | 3.15 / 3.08 ms | 3.22 / 3.15 ms | unchanged | +| contains, column x column | 3.56 / 6.29 ms | 3.68 / 6.40 ms | unchanged | +| intersects, polygons disjoint x constant | 6.81 / 6.72 ms | 3.20 / 3.14 ms | **2.1x** | +| intersects, polygons overlapping x constant | 9.57 / 9.48 ms | 9.87 / 9.63 ms | 1-3% slower, accepted | +| intersects, points and column x column arms | 3.20 / 5.98 ms | 3.21 / 5.92 ms | unchanged | + +The overlapping-intersects arm is the disclosed tradeoff: the hoisted bbox check is an early-out, +so where it rarely fires the row pays for it. The port was an out-of-sample test of the API and +passed it: **zero framework changes were needed**, matching the element vocabulary's earlier record +(`TensorRow`, `GeometryRow`, `TensorSink`, each added in its own crate). + +**Deleting the hand-written path made its shape faster.** With `Extension`-over-constant visible to +the stride-0 decode, cosine's `reduce_encoded` constant routing (manufacture an `L2Denorm` from a +constant operand, answer through the denorm paths) became deletable. Its shape then sped up: + +| width | through the deleted rewrite | through the row loop + prepare | +| --- | --- | --- | +| 2 | 118.8 us | **63.08 us** | +| 32 | 554.0 us | **377.9 us** | +| 256 | 5.159 ms | **3.007 ms** | + +Both constant spellings now measure identically (63.08 vs 62.72 us at width 2). The hand-written +fast path was 1.5-1.9x slower than the framework path that replaced it, on top of having missed the +literal shape entirely. That is the dedup argument in its strongest form: not fewer lines, but +fewer wrong ones. + +### The one unenforceable thing + +The design's benefit rests on LLVM treating the per-row branch on the prepared `Option` as +loop-invariant. Three outcomes exist per call site: unswitched (intended), if-converted (both arms +computed, the hoist silently evaporates while staying correct), or retained (a branch in a cheap +scalar kernel can block vectorization). For every real adopter the hoisted work is a loop or a +parse, which cannot be speculated, so the worst case degrades to one predicted branch per row, the +same cost class as the bounds check kept over `unsafe`. It is still a hope rather than a contract, +and the convention that polices it is stated in the trait-choice guide: every adopter lands with a +constant/non-constant benchmark pair, and the non-constant arm must not move. + +### Rejected alongside + +- **`Const` wrapper elements**: needs value-aware dispatch; splits plan/run; 2^arity dispatch + arms. Dead on the purity invariant. +- **Closure-internal `OnceCell` memoization**: wrong without constness plumbing, redundant with it. + Distinct from the `OnceCell` *inside `P`* that contains uses, which is constness-aware and only + defers an expensive build. +- **Plan-time currying through `reduce`** (folding a Literal into Options as a compiled variant): + the only design that amortizes across batches, deferred because `PersistableOptions` admits only + the source value, it misses every run-time-only constant, and re-currying bifurcates function + identity, silently detaching encoding kernels keyed on the original function. Revisit only if + per-batch prepare cost ever measures as material. +- **`visit_prepared_into`** (sink plus prepare): no user. `l2_denorm`'s constant case is a bulk + answer in `reduce_encoded`, not a prepared loop. The asymmetry is deliberate and cheap to fix + when a user appears. + +--- + +## Is there anything left to port? + +Asked directly: could the remaining hand-written vtables move onto `RowFn` if the element vocabulary +covered more types? Classifying all ~30 of them says no, and says the vocabulary is not what is +stopping them. + +| blocker | count | members | +| --- | --- | --- | +| **Not strict.** `RowFn` implies strict, so these cannot reach it at all. | 12 | `between`, `case_when`, `cast`, `dynamic`, `fill_null`, `is_null`, `is_not_null`, `list_contains`, `pack`, `stat`, `row_size`, `zip` | +| **The answer already exists in bulk.** Zero-copy child projection, a metadata field, or a vectorized slice kernel. A row loop would be strictly slower. | 12 | `not`, `list_length`, `binary`, `mask`, `ext_storage`, `get_item`, `select`, `merge`, `variant_get`, `geo.envelope`, `json_to_variant`, `row_encode` | +| **No element rows to read.** Zero-arity, or a type-erasure adapter. | 5 | `literal`, `root`, `row_idx`, `row_count`, `ForeignScalarFnVTable` | +| **Output side.** Nullable output, or an output dtype that depends on runtime data. | 2 | `list_sum`, `geo.envelope` | +| **Value-dependent per-batch setup.** | 1 | `like` | + +`geo.envelope` is the one function counted twice: its output is a struct-of-four extension type *and* +its fast paths hand back existing child arrays untouched. + +`binary` deserves a note, since on strictness alone it looks portable: only its Kleene `And`/`Or` are +non-strict, and `is_strict` already varies by operator, so comparison and arithmetic go through the +strict lifting today. What keeps it columnar is the kernel. `collect_zip_bits` and `LaneZip` run over +`as_slice()` pairs as tight vectorizable loops, with a separate constant-operand path +(`collect_bits(lhs, |a| a.is_eq(rhs))`). Routing that through a per-row closure and `ArgColumn::get` +would give up the slice-level vectorization for nothing. + +Three things follow. + +**The porting well is dry.** The eight functions on `RowFn` (`byte_length`, the four tensor kernels, the +three geo kernels) are the complete set in this repository that wants a row loop. Every remaining one is +blocked, and forcing any of them onto `RowFn` would cost performance rather than save lines. `l2_denorm` +was the last one the vocabulary was actually keeping out, and the sink let it in. + +**Missing elements are not the constraint.** Only `list_contains` would need new input vocabulary, and +it is independently blocked by non-strictness, so a list element would not unblock a single function +today. A list *input* element is nonetheless easy (`Bytes` already proves the shape: `Elem<'a>` is a +GAT, so `&'a [T]` works), and `list_length` could even be a `RowFn` given a `ListLen` element in the +style of `BytesLen`. It should not be, because its answer is a child array or one constant. + +**`like` is a new gap, and the sharpest one.** It is strict, infallible, `(Utf8, Utf8) -> Bool`: on +signature alone it is the ideal `RowFn`. Two things block it, and measuring both is what settled where +it belongs. + +Its constant-pattern path is fine. `reduce_encoded` already sees the argument arrays before the row +loop, so compiling the pattern once and evaluating in bulk has a home, and a constant operand stays +constant even through a filtered batch. No new hook needed for that case. + +Its *per-row* pattern path is what blocks it. That path memoizes the compiled pattern across +consecutive rows carrying the same one, and a `RowFn` closure is `impl Fn`, so it can hold no such +state. Defeating the cache costs **5.7x** (`like_per_row_distinct_patterns` 249.1 µs against +`like_per_row_patterns` 44.03 µs, 2048 rows, same matching work in both), which is the same shape of +regression the constant-operand stride fixed for geo. + +Relaxing the closure to `impl FnMut` would restore the cache, and it compiles as a one-word change. +It is not free. Measured on `byte_length_element`, `fastest` column, both configurations run twice: + +| case | `Fn` | `FnMut` | delta | +| --- | --- | --- | --- | +| `long_strings_bytes_len` 4096 | 11.15 µs | 12.08 µs | +8.3% | +| `long_strings_bytes_len` 65536 | 166.4 µs | 181.7 µs | +9.2% | +| `long_strings_bytes_slice` 4096 | 14.75 µs | 15.97 µs | +8.3% | +| `short_strings_bytes_len` 65536 | 166.2 µs | 180.4 µs | +8.5% | +| `short_strings_bytes_slice` 65536 | 180.9 µs | 200.3 µs | +10.7% | + +Capturing the closure by `&mut` inhibits the vectorization the shared capture allows, so `FnMut` +taxes every row function 8 to 11% to enable state that one function wants. Keep `visit` on `Fn`. + +The conclusion is that `like` does not want a row loop at all: its general path needs cross-row state, +and its fast path is bulk. What it wants is to declare `(Utf8, Utf8) -> Bool` through the element +vocabulary and keep its own kernel, which is the missing cell below. A per-batch setup hook would not +have been enough on its own, since the state `like` needs is mutable *across* rows rather than fixed +before them. + +A second, smaller thing blocks `like` too: it renders custom SQL through `fmt_sql`, and neither +`StrictScalarFnVTable` nor `RowFn` forwards that, so today porting any function with bespoke SQL +rendering would silently lose it. + +--- + +## Known gaps and future work + +Found by the porting probes, left unfixed here because each is a larger change with its own review +surface. Recorded so they are decisions rather than surprises. + +- **~~No constant-operand affordance.~~ Fixed twice over.** A partially-constant call used to decode + the constant column in full, so a broadcast operand cost one decode per row (measured: a broadcast + query vector cost the same as a genuine column, 234 ms vs 226 ms at 50k x 256). That was what kept + the geo functions off `RowFn`. Each decoded column now carries a stride, 0 for a constant, and the + geo functions are row functions. Constant *compute* was the remaining half, closed by + `visit_prepared` (see [Constant compute](#constant-compute-the-last-quadrant-of-the-lifting)). +- **`NullHandling::Dense` is chosen on safety alone, with no cost input.** For a fixed-width element + (`TensorRow`) dense is unambiguously cheaper. For an unbounded-width row (a nested list) the garbage + behind a null row need only be *in bounds*, so it can span the whole elements array, which is + pathologically O(nulls x elements). No current function hits this, but the choice should consider + width. +- **`OutputElement::build(Vec)` forces materialization.** A row function's output is always a + freshly built `Vec` turned into a `PrimitiveArray`, so it cannot return a `ConstantArray` or a lazy + child. This is why `list_length` is a columnar `StrictScalarFnVTable` rather than a `RowFn`, since a + row port would materialize one `u64` per row and lose the `FixedSizeList` constant. A columnar output + escape that stays inside the framework ("given the decoded columns, can you produce the whole output + at once?") would let `list_length`, `byte_length` and `not` share one abstraction. +- **The missing cell.** The two authoring traits cover *declare-signature-once + row-loop* (`RowFn`) + and *hand-write-signature + own-kernel* (`StrictScalarFnVTable`). The cell for + *declare-signature-once + own-kernel* is empty, so a columnar function hand-writes five signature + methods (`arity`, `child_name`, `return_element_dtype`, `null_handling`, `is_fallible`) that are all + mechanically derivable from an element tuple. + + **It is buildable.** The obvious worry is coherence, since `RowFn` already blanket-impls + `StrictScalarFnVTable` and a second blanket impl of the same trait is a hard E0119 conflict. The way + through is to layer rather than branch, putting the new trait *between* the two: + + ```text + StrictScalarFnVTable <-blanket- StrictSignature <-blanket- RowFn + ``` + + One blanket impl per edge, so nothing overlaps, and a columnar function hand-writes `StrictSignature` + while a row function reaches it through `RowFn`. Compiling the shape confirms a hand-written impl + coexists with the blanket one, including from a *downstream* crate, because within the crate that owns + the type rustc can see the blanket impl's bound does not hold. This is not a new trick here: + `impl ScalarFnVTable for V` already coexists with `Like`'s and `Between`'s + hand-written `ScalarFnVTable` impls the same way. + + **The user count is 3, not 12, and 2 of those need an element first.** Being in the columnar category + is not enough: the function's *signature* has to be expressible in the vocabulary, and + `element_dtype()` taking no arguments rules out every function whose return dtype is derived from its + input at runtime. That is most of them: `mask` returns `arg_dtypes[0].as_nullable()`, `ext_storage` + returns `ext_dtype.storage_dtype()`, `get_item` and `select` a projection of the input struct, + `variant_get` an options-derived dtype, `binary` a width negotiated between operands. What is left is + `not` (`(bool,) -> bool`, usable today), `like` (`(Bytes, Bytes) -> bool`, usable today once `fmt_sql` + forwards), and `list_length` (needs a `ListLen` element in the style of `BytesLen`). + + So this is worth building *after* the elements that give it a third user, not before. Against ~140 + lines of new trait and blanket impl it would save roughly 20 lines per function, which at one usable + caller is a wrapper with one impl. The cheap interim is to make `validate_row_args`, + `row_null_handling` and `row_is_fallible` public, which turns each hand-written signature method into + a one-liner and removes the *logic* duplication (each function currently rolling its own dtype check + and asserting rather than deriving its null handling) without adding a layer. +- **No nullable output element, so no non-total `RowFn`.** `OutputElement::build` always produces an + all-valid column, so a row kernel cannot return a null from a valid row. `impl OutputElement for + Option` is the whole fix. Left out because nothing needs it *yet*: `list_sum` would need it, but + is columnar for independent reasons too (the grouped-accumulator path and the `FixedSizeList` + constant). +- **No borrowed output element, so no zero-copy row function.** A row closure returns an + `ApplyResult`, which is `'static`, so its result cannot borrow from the input columns. Note the + asymmetry with the input side, where `InputElement::Elem<'a>` is a GAT and borrows freely. Every + `str -> str` function therefore copies: `OutputElement for String` allocates one `String` per row + and then rebuilds views from them. A string library would hit this on its first `upper`. Two + distinct fixes, of increasing scope: + - `upper`, `lower` and `replace` genuinely allocate, and want a `Cow<'a, str>` output element. That + needs `OutputElement` to grow its own lifetime GAT and `build` to take an iterator rather than a + `Vec`, so a borrowed row passes through without a copy and an owned one is built in place. + - `trim`, `substring`, `left` and `right` want more than a `Cow` can give. Their result is a + *slice* of the input, so the right kernel keeps the input's data buffer entirely and rewrites + only the views, copying no bytes. That stays columnar whatever the output element can express. + + Predicates and measurements (`starts_with`, `contains`, `byte_length`) have none of this problem + and are already the best case for `RowFn`, so the split for a string library falls along the return + type rather than the argument type. + + **A plain higher-ranked bound does not get there,** which is worth recording because it looks like + it should. Writing the visit as `impl for<'a> Fn(A::Elems<'a>) -> R::Elem<'a>` fails with + [E0582]: the `Fn` sugar puts `R::Elem<'a>` in an `Output` binding, and rustc requires the bound + lifetime to appear *structurally* in the trait's input types before a binding may reference it. An + opaque projection `A::Elems<'a>` does not count, even though it plainly mentions `'a`. Three routes + around it, measured by compiling each: + + | route | works | cost | + | --- | --- | --- | + | concrete input type instead of `A::Elems<'a>` | yes | gives up the element abstraction | + | custom callable trait with a generic `apply` method | yes | callers write a struct per kernel, not a closure, and the impl must spell `::Elem<'a>` rather than `&'a str`, or hit [E0195] | + | pass a zero-sized `Row<'a>(PhantomData<&'a ()>)` token beside the row | yes | closures survive, but every row closure grows an ignored parameter | + + The third is the one to build on: the token makes `'a` appear structurally in the `Fn`'s inputs, + which satisfies E0582 and lets the `Output` binding reference it, and plain closures still infer. + The ignored parameter is a tax on *every* row function though, so the shape to prefer is a second + visit method for lending kernels, leaving today's `visit` untouched for the `'static` majority. + + **Still open, and not what `visit_into` is.** The sink method added since is a second visit method, but + for a closure that *writes* rather than one that *lends*: its output is owned by the sink, not borrowed + from the row. A lending visit would still need the `Row<'a>` token. The precedent it sets is that + adding a third visit method costs the existing ones nothing, which is the same additive shape. + + [E0582]: https://doc.rust-lang.org/error_codes/E0582.html + [E0195]: https://doc.rust-lang.org/error_codes/E0195.html +- **~~`OutputElement::element_dtype()` takes no arguments,~~ Resolved, and not the way this predicted.** + An element's output dtype is a property of its Rust type and cannot depend on runtime data, which is + what kept `l2_denorm` columnar: it returns whole tensor rows, and a tensor's dtype carries its shape. + + Calling that a law was wrong, and the fix was recorded here as "widen `element_dtype` to take `args`". + That is *not* what shipped, and the shipped version is better. `OutputSink::sink_dtype(args)` puts the + argument-dependence on the sink, so all three `OutputElement` impls keep their no-argument + `element_dtype()` and only the thing that needs the arguments asks for them. + + This gap also named the real blocker correctly: `build(values: Vec)` with `Self = Vec` means + one heap allocation per row and then a flatten, against a columnar kernel that scales the flat storage + buffer in a single pass. At 16k rows that is 16k allocations versus zero, and no amount of dtype + plumbing fixes it. The prescription it drew, "an output element that writes into a preallocated flat + buffer (`fn apply(row, out: &mut [T])`)", is exactly what `OutputSink` is, generalized past `&mut [T]` + so a byte buffer works too. See + [the audit](#audit-can-the-four-strictscalarfnvtable-impls-really-not-be-rowfn) for what it cost and + bought. + + Note also what *not* to do on the input side: replacing the generic `TensorRow` with a + non-generic element whose `Elem<'a>` is an enum over `f16`/`f32`/`f64` would move the width choice + from monomorphization into a branch inside the row loop. That is precisely what + `match_each_float_ptype!` plus a generic element exists to avoid, so it would cost every tensor + kernel its inner-loop specialization. +- **~~The witness carries four scalars through two associated types.~~ Not a gap.** This looked like + the framework's weakest joint, since `ArgsWitness` and `RetWitness` are read *only* for `ARITY`, + `DENSE_SAFE`, `DECODE_FALLIBLE` and `FALLIBLE`, and for a multi-dispatch function the witness names + an arbitrary representative (`L2Norm` says `f64` for no reason a reader can see). The plan was to + collapse them into three consts. + + Checking the signatures says no. `arity`, `null_handling` and `is_fallible` on + `StrictScalarFnVTable` all take *only* the options, with no input dtypes, while `dispatch` needs + dtypes to choose. So those three answers **must** be dtype-independent, which means they cannot be + read off whatever element types a batch picks, which is exactly why a separate declaration has to + exist. The witness is not redundant bookkeeping; it is the only place those facts can live. + + Given that, types beat consts. With types, dense-safety and fallibility are *derived* from the + element types, so the only available mistake is a witness that disagrees with the dispatch, and that + is a build error. With three hand-written consts an implementor could state a fact wrongly *and* + visit consistently with their mistake. Converting would be a notation change that removes a + derivation, not a fragility fix. Left alone, with the reason now recorded on `ArgsWitness` so the + next reader does not re-open it. + + What is left of the original complaint is presentational: the arbitrary representative reads oddly. + A doc line on each multi-dispatch implementor saying why the width shown is arbitrary is the whole + fix. +- **`InputElement` is an open trait with required consts.** Adding `DECODE_FALLIBLE` broke every + out-of-crate element (`TensorRow`) until updated. If elements are a real extension point for other + crates, `DENSE_SAFE` / `DECODE_FALLIBLE` should carry conservative defaults. +- **`DENSE_SAFE`'s doc guidance is subtly wrong for lists.** It says `false` for "any element that + follows an offset," but a list element *is* dense-safe, because list arrays validate + `offsets[i] + sizes[i] <= elements.len()` for every row including nulls. Following the doc literally + would put `list_length` on `Filter` and lose its encoding fast paths. + +--- + +## What the ports bought + +**Not line count.** That was the first justification I reached for and it does not hold up: `row/` is +514 code lines and `strict/` is 269, against roughly 470 lines saved across six kernels. Near +break-even. Nor is it bug fixes, since none of the three extracted problems is a live miscompute on +`develop`. + +**It is `unsafe`.** Every hand-written kernel in `vortex-tensor` ended the same way: + +```rust +// SAFETY: The buffer length equals `len`, which matches the source validity length. +Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) +``` + +A kernel that computes its own values *and* carries its input's validity has to assert that the two +lengths agree, and the only tool for that is `new_unchecked`. The framework never pairs them: +[`OutputElement::build`] returns a non-nullable column, and the strict lifting applies validity +afterwards by masking. The invariant stops being asserted and becomes unrepresentable. + +Counting production `unsafe` blocks, test modules excluded: + +| function | layer it moved to | `unsafe` on `develop` | `unsafe` now | +| --- | --- | --- | --- | +| `l2_norm` | `RowFn` | 1 | 0 | +| `inner_product` | `RowFn` | 3 | 0 | +| `cosine_similarity` | `RowFn` | 3 | 0 | +| `l2_denorm` | `RowFn` (was `StrictScalarFnVTable`) | 8 | 6 | + +**This started as a controlled experiment and the control has since been ported, so read it in two +stages.** For most of this branch's life `l2_denorm` stayed on `StrictScalarFnVTable` and held all 8 of +its blocks while the three functions that moved onto the row layer lost all of theirs. Same crate, same +reviewers, same standards, so the row layer was what removed them rather than the strict lifting or the +port itself. That is the inference the control bought, and it is still the argument. + +`l2_denorm` then moved onto the row layer too, via `OutputSink`, and dropped to 6. The two it lost are +exactly the memory-safety ones on its kernel path, which is the pattern the other three showed. Of those +two, one (`FixedSizeListArray::new_unchecked` in the constant-norms path) is attributable to the port and +one (`PrimitiveArray::new_unchecked` in `build_tensor_array`) is an independent cleanup noticed along the +way. Its 6 remaining blocks are a different kind and are not the row layer's business: four call +`L2Denorm::new_array_unchecked`, an `unsafe fn` guarding the *semantic* unit-norm invariant rather than +memory safety, and two are buffer pushes in `normalize_as_l2_denorm`, a helper that is not a scalar +function. + +`develop`'s `l2_norm` also hand-rolled a 25-line constant-array fast path that the strict lifting now +does generically for every function, and computed its output nullability by hand. + +This is the justification to carry onto a clean branch. It also bounds the claim: a `vortex-tensor` +local helper owning the same invariant would remove the same `unsafe`, so what earns the *generic* +placement in `vortex-array` is that `vortex-geo`'s three predicates and `byte_length` use it too, +over three different element types. Two downstream crates plus core is the second-caller test met, not +anticipated. + +### What it costs + +Removing that `unsafe` is not free, because `new_unchecked` was buying something: the old kernel paired +its freshly built buffer with the input's validity in one step, so a nullable input cost it nothing +extra. The framework builds a non-nullable column and the lifting applies validity afterwards, which +for `Validity::Array` means materializing a mask and running a separate pass. + +That pass is `O(rows)` while the kernel is `O(rows * width)`, so width amortizes it. Measured on +`vortex-tensor/benches/l2_norm.rs`, 16384 rows, `fastest` column: + +| width | non-nullable | nullable | cost of the extra pass | +| --- | --- | --- | --- | +| 2 | 68.87 µs | 70.44 µs | +2.3% | +| 32 | 241.4 µs | 243.9 µs | +1.0% | +| 256 | 2.513 ms | 2.529 ms | +0.6% | + +So 1 to 2% on nullable input, worst at the narrowest vector anyone would store, and nothing at all on +non-nullable input where no mask is applied. Trading that for eight memory-safety `unsafe` blocks is the +right side of the deal. + +These figures are near this machine's noise floor and should be re-confirmed on quieter hardware before +being quoted. The larger measurements in these notes (the 5.7x `like` cache loss, the 8 to 11% `FnMut` +tax, the 2x width-2 per-row cost and its removal, the 2x `l2_denorm` sink win) are well clear of it. + +### The like-for-like comparison, and the per-row cost that was hiding in it + +The table above compares the framework against itself, so it isolates the masking pass but says nothing +about the rest of the machinery. `PrePortL2Norm` in the same benchmark closes that: a bench-local +`ScalarFnVTable` running the identical arithmetic, indexing the flat slice directly into a `Buffer` and +attaching validity in one step. + +This measurement found a real defect in the tensor element, and the diagnosis recorded here first was +wrong in a way worth keeping visible. + +**What was measured, and the wrong inference.** `fastest` column, non-nullable, 16384 rows: + +| width | framework | pre-port | delta | +| --- | --- | --- | --- | +| 2 | 68.85 µs | 32.85 µs | **2.10x slower** | +| 32 | 266.6 µs | 255.5 µs | +4% | +| 256 | 2.564 ms | 2.512 ms | +2% | + +The gap in absolute terms is 36 µs at width 2 and 11 µs at 32, and the conclusion drawn was "a cost that +shrinks as total work grows is a constant being amortized, so the framework carries tens of microseconds +of fixed per-batch setup." That reasoning does not hold. 36 µs over 16384 rows is 2.2 ns/row, which is a +*per-row* cost; it stops showing at width 32 because the kernel there is memory-bound and absorbs extra +CPU work in its stalls. Reading "shrinks with width" as "fixed per batch" skipped dividing by the row +count. + +**The actual cause was one per-row accessor, in the tensor element.** `TensorRow::get` called +`FlatElements::row::(i)`, which per row re-derived its typed slice: a ptype comparison against the +stored `PType`, a host-buffer downcast out of the buffer handle, a length division, and then two range +indexings with a bounds check each. All of it loop-invariant except the offset. This is exactly the +hidden-cost-accessor pattern the repository guidelines warn about, and it was written into the element +rather than found in the framework. + +The fix types the column at decode time instead of per row. `TensorRow` is already generic over `T`, +so its `Column` can be a `Buffer` plus a stride, and `get` becomes one multiply and one range index +into a typed slice. `FlatElements` keeps its untyped `row` for the callers that read a handful of rows. + +**After, same bench, same run:** + +| width | framework | pre-port | delta | +| --- | --- | --- | --- | +| 2 | **33.32 µs** | 32.83 µs | **parity, 1.01x** | +| 32 | **227.4 µs** | 258.9 µs | framework **1.14x faster** | +| 256 | **2.422 ms** | 2.522 ms | framework **1.04x faster** | + +The pre-port column is stable across both runs (32.85 then 32.83 µs at width 2), which is what makes +this comparison trustworthy; only the framework side moved. `l2_denorm` gained the same way, from +88.0 µs to 48.9 µs at width 2, since it reads its tensor argument through the same element. + +Three things follow. + +**The row layer was never the cost.** The 2x was one accessor in one element implementation, and the +generic machinery around it (the visitor, the witness, the strict lifting's bookkeeping, `reduce_encoded`'s +probe, the dispatch width match) does not measurably show up at 16384 rows. The planned decomposition +into "strict lifting versus row layer" is moot: neither was it. + +**An element is a performance-critical surface, and nothing in the framework says so.** `InputElement::get` +is documented as needing to be `O(1)`, which `FlatElements::row` technically was. `O(1)` is the wrong +contract; the right one is that `get` must not repeat work that is constant across the batch, because it +is the one function called once per row. `decode` exists precisely to hold that work, and the element +vocabulary's whole promise (anyone can add an element in their own crate) means this trap is now +available to every future implementor. + +**The framework being generic is what let one fix pay out twice.** `l2_norm`, `inner_product`, +`cosine_similarity` and `l2_denorm` all read tensor rows through this element, so a single change moved +all four. That is the case for the shared layer stated in performance terms rather than in line counts. + +### What the harness actually costs, from the optimized IR + +The measurements above say the harness is free at 16384 rows. Reading the post-optimization LLVM IR says +*why*, and settles whether more `#[inline]` would buy anything. Emitted with +`cargo rustc --release -p vortex-tensor --lib -- --emit=llvm-ir -Cdebuginfo=0`, reading the `l2_norm` f64 +arm. + +**The whole stack is already one function.** `execute_row_loop`, `ElementTuple::get` and the row closure +have no `define` of their own anywhere in the module. They survive only as basic-block *labels* carrying +`.exit.i.i.i…` suffixes about sixteen `.i` deep, which is inline-depth notation: the engine's +`ScalarFnVTable::execute`, `execute_dense`, `execute_strict`, `dispatch`, `RowVisitor::visit`, +`execute_row_loop`, `A::get` and the closure are all inlined into a single body. Adding `#[inline]` +anywhere on that path cannot help, because nothing on it is still a call. + +**Per batch the harness leaves five calls**, each correctly placed outside the loop: one +`ArgColumn::decode` per argument, one `tensor_element_ptype` for the width match, one `reduce_encoded`, +one `OutputElement::build` after the loop exits, and the output allocation. + +**Per row it leaves this, and nothing else:** + +```llvm +%row = phi i64 [ 0, %preheader ], [ %next, %loop_latch ] +%next = add nuw i64 %row, 1 +%start = mul i64 %row, %stride ; ArgColumn's stride, fused with list_size +%end = add i64 %start, %list_size +%ovf = icmp ult i64 %end, %start ; the two halves of one slice range check +%oob = icmp ugt i64 %end, %len +br i1 (or %ovf, %oob), label %slice_index_fail, label %body ; cold side out of line +%rowp = getelementptr inbounds nuw double, ptr %elements, i64 %start +%endp = getelementptr inbounds nuw i8, ptr %rowp, i64 %list_size_bytes +... ; element loop, 8x unrolled +%out = getelementptr inbounds nuw double, ptr %values, i64 %row +store double %result, ptr %out +``` + +About ten integer ops and one always-taken branch. The element loop underneath is 8x unrolled with a +serial `fadd` chain (LLVM correctly refuses to reassociate the float sum) terminating on `icmp eq ptr` +against `%endp`, which is what a hand-written `iter().map(|x| x * x).sum().sqrt()` compiles to: the +`Elem<'a> = &'a [T]` GAT is fully scalar-replaced, and the slice iterator becomes pointer bumping at +fixed byte offsets. + +**The one removable cost is not worth removing.** The surviving per-row branch is the range check on +`&elements.as_slice()[start..start + list_size]`. LLVM cannot hoist it because nothing tells it +`len == rows * list_size`. Eliminating it means `get_unchecked`, and this framework's stated value is +removing `unsafe` from kernels, so buying back a perfectly-predicted branch with an unchecked index is +the wrong direction. It is also already hidden: at width 2 the row's `sqrt` alone has longer latency than +the whole index computation. + +LLVM also unswitched the row loop on `list_size == 0` and emitted a zero-width specialization that stores +`0.0` per row. Harmless, and a sign the loop was simple enough to reason about completely. + + + +[`OutputElement::build`]: vortex-array/src/scalar_fn/row/element/mod.rs + +Production lines, before and after: + +| function | layer | before | after | +| --- | --- | --- | --- | +| `byte_length` | `RowFn` (fixed) | n/a | 23 (impl) | +| `list_length` | `StrictScalarFnVTable` | 189 | 143 | +| `not` | `StrictScalarFnVTable` | 76 (impl) | 53 (impl) | +| `list_sum` | `StrictScalarFnVTable` | 78 (impl) | 56 (impl) | +| `l2_norm` | `RowFn` (width) | 254 | 96 | +| `inner_product` | `RowFn` (width) | 277 | 112 | +| `cosine_similarity` | `RowFn` (width) | 309 | 203 | +| `l2_denorm` | `RowFn` (width, sink) | 731 | 618 | +| geo x 3 | `RowFn` (fixed) | 51 each (impl) | 15 each (impl), plus one shared element | + +Nothing outside the functions' own crates changed: the `L2DenormScheme` compressor and every +`ExactScalarFn` matcher are untouched, because the encoding-aware push-downs key off the function +*type* rather than its vtable layer. + +The line-count case does not close on its own. The framework is ~1670 production lines (up from ~1510 +before the sink, which added `result.rs`, `sink.rs` and a second visit path) and removes ~870 across the +ported functions, so **net this branch adds lines**, amortizing around the fourteenth function against +~20 strict candidates in the tree. To be honest, the case for merging is the marginal +cost of the *next* function (~15 lines, and the invariants above enforced rather than reviewed), plus +the correctness the type-derived properties buy, rather than the diff. + +--- + +## Measurements + +`vortex-array/benches/byte_length_element.rs`, element choice for `byte_length`, whole-execution +medians: + +| input | `BytesLen` | `Bytes` | | +| --- | --- | --- | --- | +| 64Ki non-inlined rows | **206 µs** | 256 µs | 24% faster | +| 64Ki inlined rows | **207 µs** | 215 µs | 4% faster | + +`vortex-array/benches/strict_validity.rs`, how the `Dense` path applies validity, same kernel in both +arms: + +| | `lazy` | `eager` | | +| --- | --- | --- | --- | +| 64Ki, one call | **9.0 µs** | 75.3 µs | 8.3x faster | +| 1Mi, one call | 1.357 ms | 1.357 ms | parity | +| 64Ki, chain of 3 | **28.3 µs** | 30.6 µs | 7% faster | + +`Validity::and` is already lazy, so the conjunction is never materialized to be applied. Only +`NullHandling::Filter` needs positions, and only it pays for them. + +`not`, word-wise kernel against the row loop it would have if it were a `RowFn` (release, identical +outputs asserted): + +| len | word-wise `!` | row loop + `bool::build` | +| --- | --- | --- | +| 64Ki | 927 ns | 376 µs (**406x**) | +| 1Mi | 10.3 µs | 5.83 ms (**569x**) | + +This is why `not` is a columnar `StrictScalarFnVTable` rather than a row function. + +--- + +## Rejected alternatives + +- **A wrapper type instead of a blanket impl** (`Strict`): forces churn at every call site, + meaning matchers, kernel registrations, and expression constructors. The blanket impl means a port + edits only the function's own impl block. +- **A `row_family!` macro, a per-crate GAT family, or a framework GAT family**: three encodings of + "element types as a function of the width," all paying for the same limit (the width bound has to + appear literally in a GAT), so each width class needed its own trait *and* adapter. The rank-2 + visitor replaces the whole lineage with one non-generic trait method and no generated code. +- **`ElementwiseFn` as a third trait**: subsumed by `RowFn` with a constant dispatch, see above. +- **One `RowFn` with defaulted `dispatch` and `apply`**: converts "define nothing" from a compile + error into a runtime panic. +- **Renaming `StrictScalarFnVTable` to `TotalFnVTable`**: the trait admits non-total members on + purpose, so the name would be wrong. +- **An `is_total` method feeding a derived `validity`**: a new concept to compute what a function can + state directly. Mirroring `validity` with a `None` default makes the unsound answer the one that + takes work. +- **Macro-generated per-type constructors**: a bespoke API per function, where the general + `ScalarFnFactoryExt::try_new_array` is what every other scalar function already uses. +- **A separate `FallibleElementwiseFn`**: an associated return type (`ApplyResult`) costs one line per + function instead of a whole trait and a spent coherence slot. + +## Null strategies and the non-strict frontier + +The question that opened this chapter: with the strict trait retiring into a private lifting under +`RowFn`, could the row framework also serve non-strict functions, where the kernel sees each input +as an `Option` and owns null semantics itself? The prior expectation was "probably not useful or +performant, but worth establishing why." The answer splits into three verdicts, one per axis, and +the investigation surfaced a fourth result nobody asked for that is worth more than the question. + +Method: a survey of every non-strict `ScalarFnVTable` impl in the workspace plus every consumer of +`is_strict` and `validity()`, and a working prototype (worktree branch `proto/null-strategies`, +2,034-line diff, not for merging) that implemented both a branch-and-skip execution strategy and a +`Nullable` input element, benchmarked on 65,536-row batches at null densities from 0% to 90%. +All 435 vortex-array scalar_fn tests and 223 vortex-geo tests pass with the prototype strategy both +off and on, including new hostile tests (out-of-bounds views and poison divisors behind null rows) +proving the kernel never runs behind a null. + +### Verdict 1: null-visible inputs have no customer, and now we know the price + +The survey found 15 non-strict functions. Thirteen are cheap columnar mask algebra or pure +structure. The canonical case is Kleene `AND`: a fused kernel computing values and validity +together at roughly six bitwise ops per 64 rows, with validity `(lv & rv) | (lv & !l) | (rv & !r)`. +The prototype measured a row-function Kleene `AND` over `(Nullable, Nullable)` against +it: **250x to 1,030x slower** depending on density. That is the honest price of spelling bitwise +logic one row at a time, and no framework design recovers it. + +The remaining two, `RowEncode` and `RowSize` in vortex-row, are the only genuinely expensive +null-visible per-row kernels in the tree, and they are excluded by something the Option tier does +not touch: they are variadic over heterogeneous column types with a shared per-row write cursor, +which the fixed-arity tuple witness cannot express. Null-visible inputs alone unlock nothing. + +Four functions (Kleene `AND`/`OR`, `zip`, `case_when`, `list_contains`) have **value-dependent +output validity**: `false AND null` is a *valid* `false`. For these no validity expression over +child validities exists even in principle, so the lifting's derivations (validity expression, mask +motion, dictionary push-down eligibility) are unavailable by definition rather than by +implementation gap. Any future Option-input tier must let the kernel author value and validity +together, which is to say it must be a different trait, not a mode of this one. + +What `is_strict = false` forfeits is exactly enumerable: the dictionary values push-down +(`arrays/dict/compute/rules.rs`), the dict-layout below-decode push-down +(`vortex-layout/src/layouts/dict/reader.rs`), and, when `validity()` is also `None`, lazy validity +on an unexecuted `ScalarFnArray` degrades to executing the kernel to read its nulls. Nothing in +vortex-scan, vortex-file, or the engine integrations consumes strictness. + +Mechanically, `Nullable` works exactly as sketched: `Elem<'a> = Option>`, decode +materializes the validity mask once, `get(i)` consults it, `DENSE_SAFE = true` by construction. +Niche packing is free for every by-reference element (`Option<&[u8]>`, `Option<&str>`, +`Option<&[T]>`, `Option<&Geometry>`, `Option` all compile-time asserted same-size) and +doubles every by-value primitive, which are precisely the elements that were already dense-safe +and never needed a strategy. The prototype's geo `contains` over `(Nullable, const)` +tracked branch-and-skip within 2-8%, so the shape is viable for a kernel that wants null +visibility for semantic reasons. Nothing in the tree does. **Do not build it; keep the survey's +constraint list for whenever a real variadic or null-visible demand shows up.** + +### Verdict 2: Option outputs inside the strict tier are the real demand + +Strictness is a subset bound, `valid(out) ⊆ valid(in)`, so a kernel that turns a valid row into a +null is still strict, and the strict lifting already keeps kernel-produced nulls, unioned with the +lifted ones. What excludes such functions from `RowFn` today is only the all-valid-output rule on +`OutputElement`. Two in-tree functions are shaped exactly like this: `list_sum` (a valid empty +list sums to null; the module doc names it as the canonical exclusion) and `variant_get` +(expensive per-row path traversal where a missing path yields null). The extension is small and +local: an `Option` output form whose element dtype is nullable and whose build sets validity, +`RetWitness` gaining a nullability bit alongside `FALLIBLE`, and the derived `validity()` moving +from `union_child_validities` to `None` for such functions, which costs them lazy validity but is +already the status quo for both named candidates. `is_strict` stays `true`. **This is the piece +worth building.** + +### Verdict 3: branch-and-skip, the result nobody asked for + +Today the derived null handling is binary: `Dense` (run over garbage, mask after) when every +element is dense-safe and the kernel infallible, else `Filter` (filter every input to the +conjoined-valid rows, run, scatter back). The prototype added the missing third strategy: +materialize the conjoined mask once, run over the *unfiltered* inputs visiting only set rows +word-at-a-time (`BitBuffer::for_each_set_index`), pre-fill the output with garbage, mask exactly +as Dense does. Fallible kernels stay sound because apply never runs behind a null. + +Measured against Filter at 65,536 rows (divan fastest, two runs): + +| workload | 1% nulls | 10% | 25% | 50% | 90% | +| --- | --- | --- | --- | --- | --- | +| `byte_length` at `Bytes` (cheap kernel) | branch 1.8x | 2.6x | 3.8x | 4.7x | **5.9x** | +| geo `contains`, one nullable operand | branch 1.07x | 1.11x | 1.18x | 1.11x | filter 1.38x | +| geo `contains`, two nullable operands | branch 1.06x | even | filter 1.2x | filter 1.9x | filter 11.3x | + +For the cheap kernel Filter never wins: at even 1% nulls, filtering the input plus scattering the +output costs more than the entire branch-side loop. For the expensive kernel the governing +quantity is the **surviving-row fraction**: branch pays O(n) decode regardless, Filter pays +O(survivors) decode plus filter and scatter. Geo's ablation makes the mechanism explicit: filter +plus scatter are under 4% of `contains`' total, so Filter's entire advantage at sparse validity is +the shrunken arrow-export-and-parse, while for `byte_length` those same two steps are 20-40% of +Filter's total and pure waste. Crossover lands near 50-75% surviving rows for one nullable operand +and lower with two (the conjoined fraction shrinks quadratically). + +The strategy is invisible to function authors: it slots under the existing derived null handling, +selectable per batch from `Mask::true_count`, with Filter kept for the sparse tail. **This is now +implemented on this branch** (see "Adaptive null strategy, as shipped" below); the rest of this +section records the prototype evidence that justified it. The prototype +also validated the two supporting pieces: a null-tolerant `decode_branch` on `InputElement` +(defaulting to plain decode, correct for bulk canonicalizations) and `OutputElement::garbage()` +for pre-fill. Production caveats recorded in the prototype report: `reduce_encoded` is not +consulted on the branch path, sinks fall back to Filter, the toggle must become per-execution and +cost-based, and geo's null-tolerant decode covered Point and Polygon only, still paying a +full-length arrow export that a run-slicing decode would shrink. The prototype's conclusion, since borne out: `Bytes`-element functions were paying the +Filter tax on every nullable batch, and most of it is recoverable. + +### Adjacent findings, recorded so they are not relearned + +- `Between::validity` declares the strict three-way conjunction while its fallback execute path + joins two comparisons with Kleene `AND`; with per-row nullable bounds the lazy validity and the + executed result disagree (a valid `false` reported as null). Pre-existing on develop, + independent of this work, slated-for-removal expression; deserves an issue. +- `not` is already at the optimum reachable through the current ownership model: `to_bit_buffer()` + is a handle clone, the source array keeps the buffer shared, so in-place negation (a real 19% on + uniquely owned buffers) is unreachable without redesigning `ExecutionArgs` ownership. Encoded + NOT flows through `NotReduce` (Constant, Sparse) and generic per-encoding push-down (Dictionary, + RunEnd) at 13-24x below canonical cost; `NotKernel` has no implementations and looks like dead + code. The three columnar ports of the retired strict trait revert entirely. +- The strict lifting's small-batch overhead is generic prelude bookkeeping (collect inputs, + compute the declared dtype, conjoin validity), not any single avoidable allocation; ablations + including SmallVec found nothing independently beneficial, and the earlier -10%-at-100-rows + reading did not reproduce uniformly. The row layer can eventually monomorphize the prelude over + its compile-time arity (`[ArrayRef; N]` via the tuple witness), which is the only structural + answer if small batches ever matter. + +## Adaptive null strategy, as shipped + +Branch-and-skip is implemented as a third null strategy, chosen per batch by the lifting. Nothing +about a function's definition changes: the row layer already derived `Dense` or `Filter` from the +element types, and `Filter` now names a *contract* (the kernel never sees a row null in any input) +rather than a mechanism. Two mechanisms satisfy that contract, and the lifting picks between them +where the conjoined mask is materialized. + +The selection rule needs one fact the framework cannot infer, so elements state it: +`InputElement::DECODE_SHRINKS_WHEN_FILTERED`, defaulted `false`, is `true` for an element whose +decode parses every row (geometry from coordinate storage) and `false` for a bulk canonicalization +(bytes, bools, primitives). Getting it wrong is a performance bug, never a correctness bug. +`ElementTuple` ORs it across arguments, the witness check pins it like dense-safety and +fallibility, and the rule is: + +```text +branch-and-skip, UNLESS some argument's decode shrinks when filtered + AND fewer than BRANCH_MIN_SURVIVING_FRACTION (0.75) of rows survive +``` + +Two supporting hooks: `InputElement::decode_null_tolerant` (defaults to the ordinary decode, sound +because the branch loop never resolves an unset row, so hostile bytes behind a null are never +touched) and `OutputElement::placeholder` (the pre-fill written behind nulls, masked before anyone +observes it). Geo overrides the decode for Point and Polygon; other geometry types report +unsupported and the selection falls back to Filter, which is tested rather than asserted in a +comment. Sinks stay on Dense/Filter, documented at the visitor. `reduce_encoded` runs on the +branch path over the *original* encodings, which is strictly better for encoding fast paths than +Filter's canonical copies, and its contract doc now states the row count differs per strategy. + +The original forced-filter, forced-branch and auto measurements used 65,536 rows on a shared 4-vCPU +VM: + +| workload | 1% | 5% | 10% | 25% | 50% | 90% | +| --- | --- | --- | --- | --- | --- | --- | +| `byte_length` at `Bytes`, auto over filter | 5.0x | 5.3x | 5.8x | 4.0x | 4.5x | 6.3x | +| geo `contains` x const, auto picks | branch | branch | branch | branch | filter | filter | +| geo `contains` x column, auto picks | branch | branch | branch | filter | filter | filter | + +Those historical rows justified shipping branch-and-skip, but they no longer calibrate the global +threshold. The controlled x86 AVX-512 rerun used a Ryzen 9 7950X pinned to CPU 4, TSC timing, a +performance governor, 60 samples for 2-4 seconds per arm, and two runs. Its representative medians +were: + +| workload | auto | branch | filter | verdict | +| --- | ---: | ---: | ---: | --- | +| one nullable, 50% nulls | 5.999-6.050 ms | 5.560-5.642 ms | 6.026-6.049 ms | auto filters, branch is 6-8% lower latency | +| two nullable, 10% nulls | 10.40-10.48 ms | 10.49-10.60 ms | 10.20-10.34 ms | auto branches, filter is 2.5-2.8% lower latency | +| two nullable, 25% nulls | 7.502-7.678 ms | 9.156-9.285 ms | 7.588-7.749 ms | auto correctly filters; filter is 1.21-1.22x faster than branch | +| two nullable, 90% nulls | 277.1-277.4 us | 3.232-3.253 ms | 277.7-278.5 us | auto matches filter; filter is about 11.6x faster than branch | + +The two misses point in opposite directions. A 50% surviving one-element decode still favors +branch, while an approximately 81% surviving two-element decode already favors filter. A single +threshold against the conjoined survivor fraction therefore cannot represent both decode cost and +arity. Replace it with per-element/arity inputs or a small estimated-cost comparison when this work +moves onto production branches. Batch size remains an unmeasured input to that model. + +Verified independently of the implementing agent: 3,441 tests pass across vortex-array and +vortex-geo (17 new: hostile out-of-bounds views behind nulls, a fallible kernel with poison +divisors behind nulls in one and both operands, conjoined-mask honoring, constant operands, real +errors still propagating, geo filter-versus-branch agreement, the unsupported-geometry fallback, +and six selection-rule cases), vortex-tensor's 164 pass unchanged, clippy `--all-targets +--all-features` is silent on both crates, and fmt and whitespace are clean. + +Open items, none blocking: the branch fallback probes `reduce_encoded` twice when the dispatch +turns out unsupported (cheap encoding check, no in-tree function affected since every +`reduce_encoded` implementor is a dense-path tensor function); geo's null-tolerant decode still +arrow-exports the full column, and slicing runs of valid rows would blunt Filter's sparse-validity +advantage enough to retire the threshold for geo; the fallible branch loop pays one `is_none` +check per set row after the first error because `for_each_set_index` cannot early-return. + +## The strict trait, deleted + +`StrictScalarFnVTable` is gone. Not made private: deleted, with its lifting kept as private +machinery under `vortex-array/src/scalar_fn/row/lift.rs`. The chain is now `RowFn` -> +`ScalarFnVTable`, one blanket impl, no intermediate trait. + +Three things converged on that. First, reverting the columnar ports left the trait with exactly one +implementor, the blanket impl over `RowFn`, and a trait with one impl is indirection rather than +abstraction. Second, the mirroring tax existed only because that blanket impl occupied the +`ScalarFnVTable` slot: `reduce` and `validity` were forwarded so a strict function could override +them despite being unable to implement `ScalarFnVTable` itself. `RowFn` keeps `validity`, because +all-valid outputs make it the child conjunction, and the `reduce` mirror went with the trait since +no adopter ever used it. Third, the naming objection a local review raised was real and is now +moot: `is_strict` names the semantic property `valid(f(x)) ⊆ valid(x)` that pushdown consumes, +while the trait demanded the *operational* property that a kernel may run over the garbage behind a +null row or over a filtered copy. Those are independent, `Bytes` being strict and not dense-safe, +so the trait was named for the wrong one of the two. + +What replaced each member: `execute_strict` and `execute_strict_branch` are the two closures +`Batch::execute` takes, `decode_shrinks_when_filtered` is a `Batch` field read off +`ElementTuple::DECODE_SHRINKS_WHEN_FILTERED`, `return_element_dtype` is what a visit returns before +`ScalarFnVTable::return_dtype` widens it, `null_handling` is `row_null_handling` over the witnesses, +and options serde is `RowFn::Options: PersistableOptions` delegated from the blanket impl. +`Batch` carries one batch's facts (id, arguments, collected inputs, conjoined validity, declared +return dtype, null handling, and the decode-shrinks flag) and takes the kernel as closures rather +than through a trait, which is the point: there is no second implementor to name. + +The one behaviour deliberately dropped is the runtime rejection of `Dense` paired with a fallible +kernel. `row_null_handling` derives the pairing from the same witnesses `is_fallible` reads, so the +combination cannot be constructed, and the requirement now lives in `NullHandling::Dense`'s doc +pointing at the derivation. Four tests went with the trait: three described a strict kernel that +returns nulls of its own (`list_sum`'s shape), which no `RowFn` can be until the `Option` output +form of open item 3 exists, and one pinned the `reduce` mirror. + +`PersistableOptions` survives with `EmptyOptions` as its only implementor, since every row function +in tree uses it. That is a bound on `RowFn::Options` rather than a speculative trait, and the +reverted `list_sum` port is what removed its second implementor. + +If a non-row columnar kernel ever wants the lifting, extract the trait then, named for the lifting +contract rather than for strictness, with that kernel as its first user. + +## Sink-only execution, the final prototype + +The last executor revision collapses every row function onto one primitive: + +```rust +visitor.visit_prepared_into::( + |constant_args| prepare(constant_args), + |state, args, output| write_one_row(state, args, output), +) +``` + +The ordinary case uses unit preparation and `ElementSink`. A tensor uses `TensorSink` so the +input dtype can determine the runtime row width. A future string transform can own one batch-wide +builder. These are not different executor modes, so the API no longer gives them different visit +methods. + +### Why the return witness disappeared + +A returning row closure needed a return witness before dispatch so `return_dtype` and fallibility +could be derived without knowing which dtype arm dispatch would select. Once every closure writes +through a sink, the sink already answers the output question: + +- `sink_dtype(args)` supplies the non-nullable element or runtime-shaped dtype. +- `with_capacity` allocates once for the batch. +- `rows` borrows the loop-local storage once. +- `row_count_matches` proves the output bound once. +- `row` hands one slot into the closure. +- `finish` builds the column and interprets any deferred error. + +`RowFn::ArgsWitness` remains load-bearing because arity and input decode properties are needed +before dispatch. `RowFn::FALLIBLE` remains because `ScalarFnVTable::is_fallible` is queried without +input dtypes. There is no analogous need for a return witness. + +The closure stays `Fn`, not `FnMut`. An earlier sink design captured `&mut Sink` in the closure and +measured 8 to 11% slower because the mutable capture blocked loop vectorization. The executor now +owns the sink, borrows its rows once, and passes a row slot as an ordinary argument. + +### Errors without a per-row result branch + +`SinkResult` has three implementations: + +- `()` for an infallible write. +- `VortexResult<()>` for an error that must exit immediately. +- `DeferredError` for a row that can write a legal provisional value and report failure after the + loop. + +Checked integer addition is the motivating deferred case. Its sink writes the wrapping sum, each +row returns a word whose sign bit means overflow, and the executor OR-reduces those words. `finish` +returns the overflow error only when the final word has its sign bit set. No `Result` discriminant +or conditional error branch is required per row. + +Nullable dense execution needs one extra rule. Garbage behind a null may overflow even when every +valid row succeeds. When dense execution finishes with a deferred error, the lifting materializes +the conjoined validity and retries only valid rows. A successful retry proves the first error came +only from discarded rows; a second deferred error is real. This preserves strict null propagation +without giving up the dense vector loop on the common path. + +This is deliberately narrow. Parsing, allocation, and any computation that cannot produce a legal +provisional row still returns `VortexResult<()>` and receives valid-row-only execution. + +### Skipped rows are a sink property + +`OutputSink::SUPPORTS_SKIPPED_ROWS` replaces the earlier blanket statement that sinks cannot use +branch-and-skip. `ElementSink` pre-fills `OutputElement::placeholder` and supports skipped rows. +A custom sink may do the same, or decline and let the lifting filter and scatter. The semantic +contract remains that skipped values are legal but arbitrary and are masked before the result +escapes. + +### Final executor measurements and IR + +The authoritative `row_fn_executor` run used 65,536 `i64` rows, 100 samples, a one-second minimum +per arm, TSC timing, CPU 4, and a performance governor on the Ryzen 9 7950X. Each cell is the range +across two runs as fastest / median: + +| workload | specialized | sink-only `RowFn` | specialized / `RowFn` | +| --- | ---: | ---: | ---: | +| checked add, two columns | 131.5-132.3 / 132.3-133.3 us | 128.4-129.6 / 129.5-130.9 us | 1.021-1.024x / 1.018-1.022x | +| checked add, column and constant | 16.90-16.93 / 17.10-17.21 us | 13.82-13.85 / 14.04 us | 1.222x / 1.218-1.226x | +| checked add, nullable columns | 133.8-134.8 / 136.1 us | 128.4-128.7 / 130.6-131.8 us | 1.042-1.047x / 1.033-1.042x | + +The native release IR has `<8 x i64>` vector error-word accumulators and +`llvm.vector.reduce.or.v8i64`. The two-column assembly is four-way unrolled over AVX-512 `zmm` +registers, producing 32 `i64` rows per iteration with four `vpaddq` instructions. Overflow bits +accumulate through vector xor/ternary-OR operations and reduce after the loop; there is no per-row +result discriminant or error branch. The specialized arm remains benchmark-local, and no production +deferred-error user exists yet. + +Other final diagnostic medians: + +- `strict_validity` lazy versus eager stayed within 2% across 65,536 and 1,048,576 rows, including + a chain of three calls. +- `byte_length_element` found `BytesLen` 1.410-1.411x faster by median than resolving a byte slice + for long strings and 1.097x for short/inlined strings at 65,536 rows. This justifies the element + choice but is not a production benchmark. +- `null_strategy_bytes` auto matched branch-and-skip; at 90% nulls it took 24.95 us against + 175.4 us for filter-and-scatter. +- Geo auto broadly tracks branch at dense validity and filter at sparse validity, but the controlled + x86 run found the two threshold misses recorded above. The full forced-strategy matrix remains an + implementation diagnostic, not permanent CodSpeed coverage. +- Distinct per-row LIKE patterns took 126.4 us against 26.87 us for a repeated pattern, 4.7x + slower. That is the measured reason LIKE remains a stateful columnar implementation. + +### Durable benchmark boundary + +Draft PR [#9136](https://github.com/vortex-data/vortex/pull/9136) now owns the stable production +benchmark names. At `bf814bbe02cb` it covers public-path byte length; signed and unsigned add, +including constant and nullable inputs; repeated and distinct LIKE patterns; tensor functions and +the `Normalized` encoding; and geo contains, intersects, and distance with constant and nullable +shapes. It also reduces the expensive overlapping-contains simulation to 1,024 rows and uses +vendored `mimalloc` in allocating binaries. + +Do not merge the research harnesses above into that permanent suite. They compare internal +strategies or frozen controls that do not exist on develop. Land #9136 first, then use its identical +benchmark names to gate each production implementation PR through CodSpeed's compiled amd64/AVX2 +simulation. Keep local Divan for real wall-clock diagnosis and generated IR for explaining a +regression. + +### Final API consequence + +Issue 9129's current sketch is obsolete: it still has `RetWitness`, `visit`, `visit_prepared`, and +`visit_into`. Issue 9130 still says sink-backed execution cannot branch-and-skip. Update both before +using their checklists to cut the implementation stack. The prototype to carry forward is: + +```text +RowFn + -> dispatches Args + OutputSink through visit_prepared_into + -> private Batch lifting chooses dense, branch-and-skip, or filter-and-scatter + -> ElementSink covers ordinary output + -> custom sinks cover runtime shape and deferred errors + -> ScalarFnVTable blanket impl exposes the function +``` + +Nullable outputs remain separate. A sink can build values plus validity, but doing so invalidates +the unconditional `validity() = union_child_validities` derivation. That semantic change should +land with its first strict non-total user, not inside the initial sink executor. + +--- + +## Final API simplification review + +This section supersedes every earlier API sketch in this document. In particular, do not carry +forward `ArgsWitness`, `RetWitness`, `PersistableOptions`, public `NullHandling`, +`DECODE_SHRINKS_WHEN_FILTERED`, or `TensorSink`. + +The review started from two constraints. The public API should expose only decisions a function +author can meaningfully make, and the executor should not trust facts fabricated by downstream +implementations. Applying both constraints removed more framework surface without preventing a +function from defining domain-specific rows. + +### The final extension boundary + +The framework is selectively sealed: + +- `RowFn` remains open. It names the function, options, argument names, fallibility, persistence, + and dtype-based dispatch. +- `InputElement` remains open. This is how a crate adds a new decoder for a geometry, tensor view, + byte view, or another domain scalar. +- `OutputElement` remains open for ordinary one-value-per-row outputs. +- `OutputSink` remains open for output representations that need their own builder or row state. +- `RowVisitor`, `ElementTuple`, and `SinkResult` are sealed because their implementations assert + executor facts used by the blanket vtable. + +Sealing `ElementTuple` does not seal decoding. The framework supplies tuple recursion for arities 0 +through 12, and a function places any open `InputElement` implementation inside those tuples. +Sealing `SinkResult` likewise does not seal output representation. A custom `OutputSink` selects one +of the supplied result behaviors. + +This keeps the author vocabulary extensible while avoiding public implementations that can lie +about arity, dense safety, result fallibility, deferred errors, or skipped-row support. + +### Dispatch contains its own evidence + +`RowFn` no longer has argument or return witnesses. `ARG_NAMES.len()` is the exact arity. The types +selected by `dispatch` carry the remaining evidence: + +```text +(InputElement, ...) + OutputSink + SinkResult + -> arity and decode properties + -> output representation and dtype + -> row fallibility and deferred-error word +``` + +The visitor asserts at compile time that the dispatched tuple arity matches `ARG_NAMES`, a +fallible decoder or result implies `RowFn::FALLIBLE`, and deferred evidence is accepted by the +selected sink. These are implications rather than equalities. A function may conservatively +declare `FALLIBLE = true` while selecting an infallible arm for some dtypes. + +This is enough for planning because dispatch is pure in `(options, args)`. It is also simpler than +duplicating the same tuple in a witness and every dispatch arm, then proving that the declarations +agree. + +### Persistence follows the function ID + +`Options: PersistableOptions` assigned one wire contract to a Rust type. That was the wrong owner. +Two functions may reuse an options type while choosing different encodings or serializability, and +an unregistered function should not invent persistence merely because its options type supports it. + +The final `RowFn` therefore owns `serialize` and `deserialize` hooks. Serialization defaults to +`Ok(None)`, and deserialization defaults to an error. Registered tensor and geo functions preserve +their explicit existing formats. The unregistered `NumericBinary` needs no otherwise-unused +serialization implementation for `NumericOperator`. + +### One custom sink is enough + +`OutputSink` already permits arbitrary internal state. A function that needs two builders defines +one sink with two fields rather than asking the executor to understand pairs of sinks. The same +rule applies to other composite or runtime-shaped results: express the shape inside one sink and +add framework abstraction only after two real users expose shared mechanics. + +The public `TensorSink` had no user after `l2_denorm` became the `Normalized` encoding. `l2_norm`, +inner product, and cosine similarity all return scalar rows through `ElementSink`. Removing +`TensorSink` avoids stabilizing roughly 90 lines of runtime-shaped row behavior without preventing +a future tensor-valued function from defining a private sink. + +`ElementSink` also no longer needs an `ElementRow` wrapper. Its row is `&mut T`, and a closure +writes with `*output = value`. The sink still pre-fills legal placeholders so branch-and-skip may +leave masked rows untouched. + +### Per-argument filtered-decode cost + +The aggregate `DECODE_SHRINKS_WHEN_FILTERED` flag was measurably lossy. OR-ing the flag made one +expensive decode indistinguishable from two, even though the x86 data selected opposite mechanisms: + +- one nullable geometry argument at 50% nulls favored branch-and-skip; and +- two independently nullable geometry arguments at 10% nulls, about 81% surviving rows, favored + filter-and-scatter. + +`InputElement::FILTERED_DECODE_COST` now defaults to zero, and each tuple adds the costs of all its +arguments. The batch selector uses the following coarse policy: + +- cost 0 always branches; +- cost 1 branches at 50% or more survivors; and +- cost 2 or greater branches at 85% or more survivors. + +The exact values come from the measured cases rather than a general cost model. There is not yet +enough evidence to distinguish two costly arguments from three, or to make the crossover depend on +batch size. Keep the value additive so a later selector can use that information without another +author-facing API change. + +The old public `NullHandling` enum is gone. The executor privately derives `Dense`, +`DenseWithRetry`, or `ValidOnly { filtered_decode_cost }`. Authors declare local safety and cost on +their input/result types, not a global mechanism. `NullStrategy` survives only in the test harness +to force branch-and-skip or filter-and-scatter. + +### Deferred errors stay in a loop-local word + +The numeric migration confirmed two constraints on deferred error evidence: + +- the accumulated word must be no wider than the element type; and +- the accumulator must live in the generated loop, not behind a mutable sink reference. + +The sealed `SinkResult` implementations for `bool`, `u8`, `u16`, `u32`, and `u64` preserve both. +Checked multiplication can report discarded high bits directly, LLVM can accumulate those words in +vectors, and `finish` turns the final evidence into the function error. `VortexResult<()>` remains +the separate early-exit form for a row that cannot write a legal provisional value. + +### Code generation after the simplification + +The final cleanup at `4becc863ae` was compared with parent `53c51d803c` using rustc 1.91.0 and LLVM +21.1.2. Both revisions were cross-compiled with: + +```bash +cargo rustc -p vortex-array --bench row_fn_executor --profile bench \ + --target x86_64-apple-darwin -- \ + --emit=llvm-ir -C codegen-units=1 -C target-cpu=x86-64-v3 +``` + +The optimized executor monomorphs were normalized to remove revision-specific symbol names and +metadata. Their vector/reduction block hashes matched exactly for wrapping add through +`ElementSink`, checked add with deferred evidence, and wrapping add through the custom `I64Sink`. + +The two wrapping paths retain 256-bit `<4 x i64>` loads, adds, and stores across six vector loop +bodies covering constant and varying inputs. Checked add retains `<4 x i64>` arithmetic, derives +overflow with vector xor/and/compare operations, ORs `<4 x i1>` evidence in the vector loop, and +reduces after the loop. The vector bodies have no calls or panic references. Scalar tails are +present in both revisions. + +The production tensor benchmark IR was checked separately for `l2_norm`, inner product, and cosine +similarity. After normalizing SSA and metadata, arithmetic sequences and instruction counts matched +between revisions for both `f32` and `f64`. Their ordered floating-point reductions remain +eightfold scalar-unrolled in both revisions. They were not vectorized before the cleanup, so the +API change did not cause that property. + +Native Apple M4 Max `row_fn_executor` timings used 65,536 rows, two alternating revisions, 100 +samples, and a 0.5-second minimum per arm. RowFn median deltas ranged from 1.11% faster to 0.94% +slower. Fastest deltas stayed within approximately 0.17%, while specialized controls had median +drift as high as 3.7%. That is no measurable native regression. + +This evidence is deliberately bounded. Cross-target optimized IR shows that the API cleanup did +not change the x86_64-v3 hot loops. It cannot establish the runtime effect of the new null selector +on an x86 branch predictor. Re-run the measured null shapes on x86 before changing or declaring the +50% and 85% thresholds settled. + +### Required x86 rerun + +The next session will run on an x86 machine. It must rerun the production comparison before this +performance record is considered complete. The #9136 benchmark baseline is now on `develop` at +`9a482c0230`, including the public binary, tensor, and geo benchmark binaries used by this work. +Fetch the latest `origin/develop`, record both exact revisions, and compare the branch against +`develop` with the same benchmark names. + +Run `binary_ops` and `like` from `vortex-array`. Run `l2_norm`, `inner_product`, +`cosine_similarity`, and `normalized` from `vortex-tensor`. Run `binary_predicates`, `distance`, +`envelope`, and `predicate_bbox` from `vortex-geo`. Use at least two alternating runs per revision. +If the host permits it, pin one core. Report both fastest and median values with the CPU, timer, and +governor configuration. + +The stable production binaries are the cross-revision gate because they now exist on `develop`. +The branch-only `vortex-geo` `null_strategies` benchmark remains the forced-policy diagnostic. Run +it on the same x86 host to verify both measured selector decisions: one costly decode at 50% +survivors must select the faster mechanism, and two costly decodes at about 81% survivors must do +the same. Inspect optimized LLVM IR again for any stable regression before changing the API or the +selector. + +### Final verification state + +The final API state recorded 67 focused RowFn tests, 179 tensor tests, and 230 geo tests. Nightly +formatting passed. Full workspace clippy passed with +`PYO3_NO_PYTHON=1 PYO3_BUILD_EXTENSION_MODULE=1`, required because the host `/usr/bin/python3` is +3.9 while the workspace targets the Python 3.11 stable ABI. + +Issues #9128, #9129, and #9130 were updated to this API. The durable public-path benchmark baseline +from #9136 is now in the repository. Earlier statements in this document that those issues or the +baseline still need updating are historical only. diff --git a/docs/strictness-and-validity-pushdown.typ b/docs/strictness-and-validity-pushdown.typ new file mode 100644 index 00000000000..d45d87a37c9 --- /dev/null +++ b/docs/strictness-and-validity-pushdown.typ @@ -0,0 +1,243 @@ +#set page(paper: "a4", margin: 2.2cm, numbering: "1 / 1") +#set text(font: "Libertinus Serif", size: 10.5pt) +#set par(justify: true, leading: 0.62em) +#set heading(numbering: "1.") +#show heading: it => block(above: 1.4em, below: 0.8em, it) +#show raw: it => text(font: "Noto Sans Mono", size: 0.88em, it) +#set table(stroke: 0.4pt + luma(65%), inset: 5pt) + +#let mask = math.op("mask") +#let valid = math.op("valid") +#let N = text(fill: rgb("#b03a2e"), weight: "bold", [NULL]) + +#let node(body, fill: luma(96%)) = box( + inset: (x: 7pt, y: 5pt), radius: 3pt, stroke: 0.5pt + luma(55%), fill: fill, body, +) + +#let lead(body) = block( + inset: (x: 10pt, y: 8pt), radius: 3pt, fill: luma(97%), + stroke: (left: 2pt + rgb("#2c3e50")), width: 100%, body, +) + +#align(center)[ + #text(size: 17pt, weight: "bold")[Strictness and validity push-down] + #v(-0.4em) + #text(size: 12pt)[the same value law, once partiality is accounted for] +] + +#v(1em) + +#lead[ + *Summary.* A row-local function may be pushed through an input's validity exactly when it is strict + in that argument *and* remains defined after validity masks that argument. The first condition is the + usual null-propagation meaning of `is_strict`; the second matters only for partial functions. It is + automatic for an infallible function. Return-dtype representability, totality, speculative errors, + and `Dense` safety remain separate concerns. +] + += Model + +Scalar functions are *row-local*: output row $i$ depends only on input rows $i$. They are also assumed +deterministic and insensitive to the bytes behind nulls. Equality below is therefore *logical equality* +$eq.triple$: equal length, equal validity, and equal values at valid rows. + +A mask is a non-nullable boolean column. It applies validity without changing valid values: + +$ mask(a, m)[i] = cases(#N &"if" not m[i], a[i] &"otherwise") $ + +For example, masking does not distinguish a newly nulled row from one that was already null: + +#figure( + table( + columns: 4, + align: center, + table.header([$i$], [$a$], [$m$], [$mask(a, m)$]), + [0], [10], [`true`], [10], + [1], [20], [`false`], N, + [2], N, [true], N, + ), + caption: [Rows 1 and 2 are both null after masking, for different reasons.], +) + +The function $f$ may be partial: an evaluation can error instead of returning a column. Statements +about its result are quantified only where that evaluation succeeds. + += The law and its missing premise + +Fix an argument position $j$. + +#lead[ + *$(S_j)$ Strictness.* If $f(a_1, ..., a_k)$ succeeds and $a_j[i] = #N$, its output at $i$ is #N. + + *$(C_j)$ Mask closure.* If $f(a_1, ..., a_k)$ succeeds, then + $f(a_1, ..., mask(a_j, m), ..., a_k)$ succeeds for every mask $m$. + + *$(M_j)$ Validity equivariance.* Whenever $f(a_1, ..., a_k)$ succeeds, the masked evaluation also + succeeds and + $ f(a_1, ..., mask(a_j, m), ..., a_k) eq.triple mask(f(a_1, ..., a_k), m). $ +] + +$(M_j)$ is the law used by a validity push-down: compute after masking one argument, or compute first +and mask the result. It includes definedness of both sides, rather than treating an error as a value. + +#pagebreak() + +For an ordinary addition, $(M_1)$ says the following two columns agree. The evaluation after masking +is defined, and strictness makes its second row null. + +#figure( + table( + columns: 6, + align: center, + table.header( + [$i$], [$a_1$], [$a_2$], [$m$], + [mask first, then add], [add first, then mask], + ), + [0], [1], [10], [`true`], [11], [11], + [1], [2], [20], [`false`], N, N, + [2], [3], [30], [`false`], N, N, + ), + caption: [The two orders differ only in the unobserved bytes behind null rows.], +) + +#lead[ + *Theorem.* For a row-local deterministic function, + $ (S_j) " and " (C_j) quad arrow.l.r quad (M_j). $ + Consequently, full strictness plus mask closure in every argument is exactly what licenses every + per-argument validity push-down. +] + +== Forward: strictness and closure imply the law + +Assume $(S_j)$ and $(C_j)$, and start with any successful evaluation +$f(a_1, ..., a_k)$. By closure, the left side below also succeeds. Fix a row $i$; row-locality means +there are only two cases to check: + +#figure( + table( + columns: (auto, 1fr, 1fr), + align: (center, left, left), + table.header([mask bit], [left: compute after masking], [right: mask after computing]), + [$m[i] = $ `true`], + [the input at row $i$ is unchanged, so this is $f(a_1, ..., a_k)[i]$], + [masking preserves $f(a_1, ..., a_k)[i]$], + [$m[i] = $ `false`], + [argument $j$ is #N; the successful left evaluation is #N by $(S_j)$], + [the mask makes the result #N by definition], + ), + caption: [Each row agrees, so the columns are logically equal.], +) + +This proves $(M_j)$. Notice the distinct jobs of the two premises: closure establishes that the left +evaluation exists; strictness establishes its value at masked rows. + +== Reverse (by contrapositive): the law implies strictness and closure + +$(M_j)$ explicitly includes $(C_j)$. To obtain $(S_j)$, use its contrapositive: suppose a successful +input $b$ has a null in argument $j$ at row $i$, but gives a non-null result $v$ there. This is exactly +the negation of $(S_j)$, and we will derive a contradiction with $(M_j)$. + +Choose a mask $m$ that is false only at $i$, and write +$b'_j = mask(b_j, m)$. At row $i$, $b_j[i]$ was already #N; at every other row, $m$ is true. Thus +$b'_j eq.triple b_j$. Replacing $b_j$ by $b'_j$ changes no logical input value, including at the one +row we care about. + +Now apply $(M_j)$ to the successful input $b$. Its left-hand side is precisely the evaluation with +$b'_j = mask(b_j, m)$, and it guarantees that evaluation succeeds. At row $i$, the common left-hand +side has these two incompatible values: + +$ f(b_1, ..., mask(b_j, m), ..., b_k)[i] + = f(b_1, ..., b'_j, ..., b_k)[i] + = f(b_1, ..., b_j, ..., b_k)[i] = v != #N. $ + +But $(M_j)$ also says + +$ f(b_1, ..., mask(b_j, m), ..., b_k)[i] + = mask(f(b_1, ..., b_j, ..., b_k), m)[i] = #N. $ + +The first line uses the definition of $b'_j$, then row-locality and $b'_j eq.triple b_j$; the second is +$(M_j)$ and $m[i] = $ `false`. We do not use $(S_j)$ here --- it is the fact being proved. One successful +evaluation cannot be both $v$ and #N, so the assumed counterexample cannot exist. Therefore $(M_j)$ +implies $(S_j)$. $square.stroked$ + +#pagebreak() + +The closure premise is necessary. A binary function that succeeds on $(0, 1)$, errors on $(#N, 1)$, +and otherwise returns null whenever it does evaluate with a null first argument satisfies $(S_1)$ under +the partiality convention, but not $(M_1)$: masking the first input turns a successful evaluation into +an error. Defining strictness to require a *successful* null result on every null input is an equivalent +way to build this premise into $(S_j)$. + +#figure( + table( + columns: 4, + align: center, + table.header([input], [$f$], [after masking argument 1], [$f$ after masking]), + [$(0, 1)$], [0], [$(#N, 1)$], [*error*], + ), + caption: [The function is vacuously strict at $(#N, 1)$ because it does not return a non-null value; + nevertheless, it cannot satisfy the masked-evaluation law.], +) + += What the optimizer uses + +The dictionary rule has the shape + +#align(center)[ + #grid( + columns: 3, column-gutter: 1.2em, align: horizon, + node[`f(dict(codes, values), c)`], + text(size: 13pt)[$arrow.r.long$], + node(fill: rgb("#eafaf1"))[`dict(codes, f(values, c))`], + ) +] + +A null code masks only the dictionary argument while $c$ stays live, so this requires $(M_j)$ for that +argument, not a weaker law that masks all arguments together. Kleene `AND` illustrates the difference: +`false AND NULL` is `false`, so masking only its second argument is not equivariant. + +#table( + columns: 6, + align: center, + table.header( + [$a_1$], [$a_2$], [$m$], [mask $a_2$, then `AND`], [`AND`, then mask], [result], + ), + [`false`], [`true`], [`false`], [`false`], N, [not $(M_2)$], +) + +Value equivalence is not enough for this rewrite when $f$ is fallible. It evaluates *every* dictionary +value, including values with no live code; `div(100, 0)` can then error on the rewritten side although +the original never evaluated it. Thus the dictionary rule also needs its existing no-speculative-error +condition (normally `!is_fallible`). Mask closure addresses masked input rows; it does not make dead +dictionary values safe to evaluate. + += Independent obligations + +#table( + columns: (auto, 1fr, 1fr), + align: (left, left, left), + table.header([property], [statement], [what it enables]), + [strict + mask-closed], [null inputs produce null outputs and remain evaluable], + [validity push-down], + [representable], [the declared return dtype admits required nulls], + [advertising `is_strict`], + [total], [valid inputs never produce null], + [precomputing output validity], + [infallible], [no legal evaluation errors], + [speculative evaluation], + [dense-safe], [bytes behind nulls may be read safely], + [`NullHandling::Dense`], +) + +Representability is a type-level obligation: a strict `cast` with a pinned non-nullable return type +cannot represent the null its value semantics demand. Totality is different again. A strict `list_sum` +may return null for a valid empty list, so strictness only gives + +$ valid(f(a_1, ..., a_k)) subset.eq valid(a_1) " and " dots " and " valid(a_k). $ + +Equality, and hence a precomputed output-validity mask, additionally needs totality. + +`RowFn` supplies strictness structurally. Its `Filter` path evaluates only rows valid in every input and +scatters nulls back; its `Dense` path evaluates all rows then applies that combined validity. The latter +still needs `InputElement::DENSE_SAFE`, because an invalid string view may hold unsafe bytes. That is an +operational property of an element representation, not a consequence of strictness. From ea58061b5d2fe663642144a3f03ea1202e07baed Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 13:35:38 -0400 Subject: [PATCH 007/160] Fix RowFn documentation checks Signed-off-by: Connor Tsui --- SCALAR_FN_HANDOFF.md | 17 ++++++------- STRICT_SCALAR_FN_RESEARCH.md | 25 ++++++++++---------- vortex-array/src/scalar_fn/row/sink.rs | 2 +- vortex-array/src/scalar_fn/row/tests/sink.rs | 2 +- vortex-spatial/benches/null_strategies.rs | 6 ++--- 5 files changed, 27 insertions(+), 25 deletions(-) diff --git a/SCALAR_FN_HANDOFF.md b/SCALAR_FN_HANDOFF.md index f99f4cb0eaa..bff6cf82347 100644 --- a/SCALAR_FN_HANDOFF.md +++ b/SCALAR_FN_HANDOFF.md @@ -34,10 +34,10 @@ cargo bench -p vortex-tensor --bench l2_norm cargo bench -p vortex-tensor --bench inner_product cargo bench -p vortex-tensor --bench cosine_similarity cargo bench -p vortex-tensor --bench normalized -cargo bench -p vortex-geo --bench binary_predicates -cargo bench -p vortex-geo --bench distance -cargo bench -p vortex-geo --bench envelope -cargo bench -p vortex-geo --bench predicate_bbox +cargo bench -p vortex-spatial --bench binary_predicates +cargo bench -p vortex-spatial --bench distance +cargo bench -p vortex-spatial --bench envelope +cargo bench -p vortex-spatial --bench predicate_bbox ``` Run each revision at least twice in alternating order. If the host allows it, pin the process to one @@ -45,13 +45,13 @@ core. Record the timer and CPU configuration, and compare both fastest and media benchmark binaries and public names are now shared with `develop`, so the comparison no longer needs a frozen benchmark-local implementation as its primary control. -Also run the branch-only `vortex-geo` `null_strategies` diagnostic. It forces branch-and-skip and +Also run the branch-only `vortex-spatial` `null_strategies` diagnostic. It forces branch-and-skip and filter-and-scatter for the measured nullable geometry shapes. Confirm that automatic selection uses the faster mechanism for one costly decode at 50% survivors and for two costly decodes at about 81% survivors. This is the x86 runtime check that remains after the LLVM comparison. ```bash -cargo bench -p vortex-geo --bench null_strategies +cargo bench -p vortex-spatial --bench null_strategies ``` If a stable benchmark regresses, inspect optimized LLVM IR again. The previous cross-compile proves @@ -262,7 +262,8 @@ x86. ## Current implementation and checks -The implementation includes production users in `vortex-array`, `vortex-tensor`, and `vortex-geo`. +The implementation includes production users in `vortex-array`, `vortex-tensor`, and +`vortex-spatial`. `NumericBinary` is an unregistered `RowFn` used only for primitive arithmetic execution. Decimal arithmetic keeps its existing path. The stable public-path benchmark baseline landed as #9136. @@ -270,7 +271,7 @@ The checks recorded for the final API state are: - 67 focused RowFn tests; - 179 `vortex-tensor` tests; -- 230 `vortex-geo` tests; +- 230 `vortex-spatial` tests; - `cargo +nightly fmt --all`; and - full workspace clippy, with `PYO3_NO_PYTHON=1 PYO3_BUILD_EXTENSION_MODULE=1` because the host `/usr/bin/python3` is 3.9 while the workspace requires the Python 3.11 stable ABI. diff --git a/STRICT_SCALAR_FN_RESEARCH.md b/STRICT_SCALAR_FN_RESEARCH.md index 1c873108345..e0d5c8ad1d5 100644 --- a/STRICT_SCALAR_FN_RESEARCH.md +++ b/STRICT_SCALAR_FN_RESEARCH.md @@ -200,9 +200,9 @@ is one `apply` per row. Three whole classes of strict function are therefore ine `RowFn` at any cost: - **Output dtype outside the element set.** `ext_storage`'s output is an extension array's storage - dtype, so `vortex.geo.box` is a struct and `vortex.uuid` is a `FixedSizeList(u8,16)`. `vortex-geo`'s - zone-map pruning calls `ext_storage` on a `geo.box` statistic, and a row-function port breaks it at - plan time. + dtype, so `vortex.st.box` is a struct and `vortex.uuid` is a `FixedSizeList(u8,16)`. + Zone-map pruning in `vortex-spatial` calls `ext_storage` on an `st.box` statistic, and a + row-function port breaks it at plan time. - **Variadic arity.** `merge` and `select` take an unbounded number of children, while `RowFn` fixes `Arity::Exact(n <= 3)`. - **Sub-row-granular kernels.** `not` negates one 64-bit word at a time, so a row loop over `bool` is @@ -721,12 +721,12 @@ stopping them. | blocker | count | members | | --- | --- | --- | | **Not strict.** `RowFn` implies strict, so these cannot reach it at all. | 12 | `between`, `case_when`, `cast`, `dynamic`, `fill_null`, `is_null`, `is_not_null`, `list_contains`, `pack`, `stat`, `row_size`, `zip` | -| **The answer already exists in bulk.** Zero-copy child projection, a metadata field, or a vectorized slice kernel. A row loop would be strictly slower. | 12 | `not`, `list_length`, `binary`, `mask`, `ext_storage`, `get_item`, `select`, `merge`, `variant_get`, `geo.envelope`, `json_to_variant`, `row_encode` | +| **The answer already exists in bulk.** Zero-copy child projection, a metadata field, or a vectorized slice kernel. A row loop would be strictly slower. | 12 | `not`, `list_length`, `binary`, `mask`, `ext_storage`, `get_item`, `select`, `merge`, `variant_get`, `spatial.envelope`, `json_to_variant`, `row_encode` | | **No element rows to read.** Zero-arity, or a type-erasure adapter. | 5 | `literal`, `root`, `row_idx`, `row_count`, `ForeignScalarFnVTable` | -| **Output side.** Nullable output, or an output dtype that depends on runtime data. | 2 | `list_sum`, `geo.envelope` | +| **Output side.** Nullable output, or an output dtype that depends on runtime data. | 2 | `list_sum`, `spatial.envelope` | | **Value-dependent per-batch setup.** | 1 | `like` | -`geo.envelope` is the one function counted twice: its output is a struct-of-four extension type *and* +`spatial.envelope` is the one function counted twice: its output is a struct-of-four extension type *and* its fast paths hand back existing child arrays untouched. `binary` deserves a note, since on strictness alone it looks portable: only its Kleene `And`/`Or` are @@ -997,7 +997,7 @@ does generically for every function, and computed its output nullability by hand This is the justification to carry onto a clean branch. It also bounds the claim: a `vortex-tensor` local helper owning the same invariant would remove the same `unsafe`, so what earns the *generic* -placement in `vortex-array` is that `vortex-geo`'s three predicates and `byte_length` use it too, +placement in `vortex-array` is that `vortex-spatial`'s three predicates and `byte_length` use it too, over three different element types. Two downstream crates plus core is the second-caller test met, not anticipated. @@ -1241,7 +1241,7 @@ Method: a survey of every non-strict `ScalarFnVTable` impl in the workspace plus `is_strict` and `validity()`, and a working prototype (worktree branch `proto/null-strategies`, 2,034-line diff, not for merging) that implemented both a branch-and-skip execution strategy and a `Nullable` input element, benchmarked on 65,536-row batches at null densities from 0% to 90%. -All 435 vortex-array scalar_fn tests and 223 vortex-geo tests pass with the prototype strategy both +All 435 vortex-array scalar_fn tests and 223 vortex-spatial tests pass with the prototype strategy both off and on, including new hostile tests (out-of-bounds views and poison divisors behind null rows) proving the kernel never runs behind a null. @@ -1327,7 +1327,7 @@ selectable per batch from `Mask::true_count`, with Filter kept for the sparse ta implemented on this branch** (see "Adaptive null strategy, as shipped" below); the rest of this section records the prototype evidence that justified it. The prototype also validated the two supporting pieces: a null-tolerant `decode_branch` on `InputElement` -(defaulting to plain decode, correct for bulk canonicalizations) and `OutputElement::garbage()` +(defaulting to plain decode, correct for bulk canonicalization) and `OutputElement::garbage()` for pre-fill. Production caveats recorded in the prototype report: `reduce_encoded` is not consulted on the branch path, sinks fall back to Filter, the toggle must become per-execution and cost-based, and geo's null-tolerant decode covered Point and Polygon only, still paying a @@ -1410,7 +1410,7 @@ arity. Replace it with per-element/arity inputs or a small estimated-cost compar moves onto production branches. Batch size remains an unmeasured input to that model. Verified independently of the implementing agent: 3,441 tests pass across vortex-array and -vortex-geo (17 new: hostile out-of-bounds views behind nulls, a fallible kernel with poison +vortex-spatial (17 new: hostile out-of-bounds views behind nulls, a fallible kernel with poison divisors behind nulls in one and both operands, conjoined-mask honoring, constant operands, real errors still propagating, geo filter-versus-branch agreement, the unsupported-geometry fallback, and six selection-rule cases), vortex-tensor's 164 pass unchanged, clippy `--all-targets @@ -1767,12 +1767,13 @@ Fetch the latest `origin/develop`, record both exact revisions, and compare the Run `binary_ops` and `like` from `vortex-array`. Run `l2_norm`, `inner_product`, `cosine_similarity`, and `normalized` from `vortex-tensor`. Run `binary_predicates`, `distance`, -`envelope`, and `predicate_bbox` from `vortex-geo`. Use at least two alternating runs per revision. +`envelope`, and `predicate_bbox` from `vortex-spatial`. Use at least two alternating runs per +revision. If the host permits it, pin one core. Report both fastest and median values with the CPU, timer, and governor configuration. The stable production binaries are the cross-revision gate because they now exist on `develop`. -The branch-only `vortex-geo` `null_strategies` benchmark remains the forced-policy diagnostic. Run +The branch-only `vortex-spatial` `null_strategies` benchmark remains the forced-policy diagnostic. Run it on the same x86 host to verify both measured selector decisions: one costly decode at 50% survivors must select the faster mechanism, and two costly decodes at about 81% survivors must do the same. Inspect optimized LLVM IR again for any stable regression before changing the API or the diff --git a/vortex-array/src/scalar_fn/row/sink.rs b/vortex-array/src/scalar_fn/row/sink.rs index 50d07d1b7ff..886f9f5f9e1 100644 --- a/vortex-array/src/scalar_fn/row/sink.rs +++ b/vortex-array/src/scalar_fn/row/sink.rs @@ -23,7 +23,7 @@ use crate::scalar_fn::OutputElement; /// Relaxing the row closure to `FnMut` instead was measured at 8 to 11%, because a captured `&mut` /// inhibits vectorization of the loop. /// - **[`sink_dtype`](Self::sink_dtype) sees the input dtypes**, unlike -/// [`OutputElement::element_dtype`](crate::scalar_fn::OutputElement::element_dtype), which takes +/// [`OutputElement::element_dtype`], which takes /// none. That is the whole reason a runtime-shaped output fits here: the width comes out of the /// arguments. /// diff --git a/vortex-array/src/scalar_fn/row/tests/sink.rs b/vortex-array/src/scalar_fn/row/tests/sink.rs index 95b5e2df05f..45b91ed191d 100644 --- a/vortex-array/src/scalar_fn/row/tests/sink.rs +++ b/vortex-array/src/scalar_fn/row/tests/sink.rs @@ -355,7 +355,7 @@ fn a_failing_row_is_never_reached_behind_a_null() -> VortexResult<()> { } /// The sink names its output dtype from the input, so a wrong input dtype is rejected at plan -/// time rather than producing a mis-typed column. +/// time rather than producing a mistyped column. #[test] fn the_sink_dtype_validates_its_input() { let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); diff --git a/vortex-spatial/benches/null_strategies.rs b/vortex-spatial/benches/null_strategies.rs index 1ef453d645b..89f43f97429 100644 --- a/vortex-spatial/benches/null_strategies.rs +++ b/vortex-spatial/benches/null_strategies.rs @@ -32,13 +32,13 @@ use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::NullStrategy; use vortex_array::scalar_fn::execute_row_fn_with_strategy; use vortex_array::validity::Validity; +use vortex_session::VortexSession; use vortex_spatial::scalar_fn::contains::SpatialContains; -use vortex_spatial::test_harness::geo_session; use vortex_spatial::test_harness::point_column; use vortex_spatial::test_harness::polygon_column; -use vortex_session::VortexSession; +use vortex_spatial::test_harness::spatial_session; -static SESSION: LazyLock = LazyLock::new(geo_session); +static SESSION: LazyLock = LazyLock::new(spatial_session); fn main() { LazyLock::force(&SESSION); From 66b874290ecd5884c88b6ced8438fd2ef7a8df65 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 14:49:20 -0400 Subject: [PATCH 008/160] Record RowFn regression and landing plan Signed-off-by: Connor Tsui --- AGENTS.md | 6 ++++ NUMERIC_ROWFN_PLAN.md | 28 ++++++++++++---- SCALAR_FN_HANDOFF.md | 63 ++++++++++++++++++++++-------------- STRICT_SCALAR_FN_RESEARCH.md | 39 ++++++++++++++++++++++ 4 files changed, 104 insertions(+), 32 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index a0e4c6558cd..30d67e18cb5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -132,6 +132,12 @@ Notes: Avoid hidden-cost per-element accessors in hot loops, follow the performance guidance in `STYLE.md`, and benchmark changes to hot paths. +Treat branchless indexing as a code-generation hypothesis, not as an optimization by itself. A +runtime expression such as `index & mask` can make a slice index non-affine, retain bounds checks, +and block vectorization. Inspect generated code before replacing a loop-invariant enum match because +LLVM can unswitch the match into specialized loops. For binary kernels, benchmark varying x varying, +varying x constant, constant x varying, and nullable constant shapes separately. + ## Tests - Strongly consider `rstest` cases when parameterizing repetitive test logic. diff --git a/NUMERIC_ROWFN_PLAN.md b/NUMERIC_ROWFN_PLAN.md index 2f69b708b01..5a304e88102 100644 --- a/NUMERIC_ROWFN_PLAN.md +++ b/NUMERIC_ROWFN_PLAN.md @@ -185,19 +185,33 @@ Re-measure the port against `develop` once that lands, because the comparison ab ### Measured dead ends -Recorded so they are not retried. All of these are in vortex-data/vortex#9130 as well. +Recorded so they are not retried. The entries that predate the broadcast-index-mask experiment are +also in vortex-data/vortex#9130. - Bounds-check elimination in the row loop is not available. Narrowing the varying view to the row count buys nothing, and `get_unchecked` is not uniformly a win: about 10% on `mul_u16` and `mul_u32`, and 22% slower on `mul_u8`. - A per-argument row source that keeps the `Varying` view when another argument is batch-constant is 4x slower than the `ArgColumn` branch it replaces, which already vectorizes. -- A batch-constant operand therefore still demotes its neighbours off the slice path. Closing that - needs the row loop monomorphized over which arguments are constant. Revisit when `Compare` moves - onto `RowFn`, since `col < literal` is exactly this shape. - -`mul_i32_constant` is the one regression that survives, and it is inside this host's drift. Let -CodSpeed settle whether it is real. +- Pairing each varying view with a runtime index mask also fails. Commit `ad24700088` used + `index & usize::MAX` for varying inputs and `index & 0` for constants. On x86 with + `RUSTFLAGS="-C target-feature=+avx2"`, the constant numeric cases became approximately 4x to 7x + slower in wall time while non-constant cases stayed at parity. CodSpeed reported smaller but + consistent regressions: `add_i64_constant` 31.34%, `sub_i64_constant` 32.38%, and + `mul_i32_constant` 46.24%. +- The disassembly explains the mask result. Each `index & mask` remained behind a slice bounds + check, so LLVM saw a non-affine index and emitted a scalar loop. The old enum match was + loop-invariant, and LLVM unswitched it into constant-pattern loops with affine varying + accesses. Removing a branch removed information that the vectorizer needed. + +Commit `ad24700088` was removed from `ct/row-fn` history. The clean head after the rewrite is +`ea58061b5d`. A compile-time varying x constant or constant x varying specialization remains a +possible design, but it is not work for the first PR. Implement it only after the clean branch has +a stable mixed-constant regression against the current merge base. + +The earlier `mul_i32_constant` result was within Apple host drift and predates the mask experiment. +It does not establish parity against current `develop`. Rerun the clean candidate and current merge +base on x86 before deleting the hand-written kernels in a mergeable PR. ### What this implies for #9129 and #9130 diff --git a/SCALAR_FN_HANDOFF.md b/SCALAR_FN_HANDOFF.md index bff6cf82347..8b5d6a6d053 100644 --- a/SCALAR_FN_HANDOFF.md +++ b/SCALAR_FN_HANDOFF.md @@ -14,18 +14,37 @@ The public design lives in these tracking issues, which now match the implementa - [#9129, Define the `RowFn` API](https://github.com/vortex-data/vortex/issues/9129) - [#9130, Execute `RowFn` over Vortex arrays](https://github.com/vortex-data/vortex/issues/9130) -The branch is `ct/row-fn`. It is publicly linked from #9128, so do -not rewrite or delete its history. Commit `4becc863ae` contains the final API simplification. Push -only when explicitly requested. +The branch is `ct/row-fn`, and draft PR #9255 remains the integration and research branch. Its +history was rewritten at `ea58061b5d` to remove the regressing broadcast-index-mask experiment. +Do not use the draft PR as the first mergeable change. Cut the first PR from the latest +`origin/develop`, and keep this branch as the source for later tensor and spatial ports. Push or +rewrite either branch only when explicitly requested. -## Next action: rerun the benchmarks on x86 +## Next action: cut the vortex-array PR -The next session will run on an x86 machine. Rerun the performance comparison there before treating -the implementation as complete. Do not reuse the Apple timings as the final runtime result. +The first mergeable PR must stay within `vortex-array` and contain: -The production benchmark baseline from #9136 is on `develop` at `9a482c0230`. Fetch the latest -`origin/develop`, record the exact baseline and candidate commits, and run the same public benchmark -binaries at both revisions: +1. the `RowFn` API, lifting, executor, and focused behavioral tests. +2. the primitive `NumericBinary` port as its production consumer. +3. only the executor and numeric benchmarks needed to support its performance claim. + +Do not include the tensor or spatial ports, these branch-only working notes, the unrelated `like` +benchmark additions, or the fixed-size-list test. `NumericBinary` is the only `RowFn` consumer in +`vortex-array` on this branch. It already exercises varying and constant inputs, all-constant +folding, null constants, nullable execution, deferred overflow evidence, and the valid-row retry. +Do not add another consumer only to make the PR appear broader. + +The numeric commit deletes the now-unused `vortex-compute::lane_kernels::map_into` helper. Leave +that helper in place for a strictly `vortex-array`-only PR, and remove it in a separate cleanup. + +The first PR must establish parity against the latest `origin/develop`, not the integration +branch's old merge base. Run the public `binary_ops` benchmark on x86 with identical build flags at +both revisions. Cover varying x varying, varying x constant, constant x varying, and nullable plus +constant inputs. Run each revision at least twice in alternating order. Record the exact commits, +CPU, timer, pinning, fastest values, and medians. Inspect optimized LLVM IR or assembly for every +stable regression before changing the row API. + +The production benchmark commands across the staged work are: ```bash cargo bench -p vortex-array --bench binary_ops @@ -40,23 +59,17 @@ cargo bench -p vortex-spatial --bench envelope cargo bench -p vortex-spatial --bench predicate_bbox ``` -Run each revision at least twice in alternating order. If the host allows it, pin the process to one -core. Record the timer and CPU configuration, and compare both fastest and median values. The -benchmark binaries and public names are now shared with `develop`, so the comparison no longer -needs a frozen benchmark-local implementation as its primary control. - -Also run the branch-only `vortex-spatial` `null_strategies` diagnostic. It forces branch-and-skip and -filter-and-scatter for the measured nullable geometry shapes. Confirm that automatic selection uses -the faster mechanism for one costly decode at 50% survivors and for two costly decodes at about 81% -survivors. This is the x86 runtime check that remains after the LLVM comparison. +For the spatial PR, also run the branch-only `vortex-spatial` `null_strategies` diagnostic. It +forces branch-and-skip and filter-and-scatter for the measured nullable geometry shapes. Confirm +that automatic selection uses the faster mechanism for one costly decode at 50% survivors and for +two costly decodes at about 81% survivors. ```bash cargo bench -p vortex-spatial --bench null_strategies ``` -If a stable benchmark regresses, inspect optimized LLVM IR again. The previous cross-compile proves -that the API cleanup preserved the x86_64-v3 loop shape. The x86 run must confirm runtime effects -from the revised null selector and the target CPU's vectorizer and branch predictor. +The public benchmark names are shared with `develop`, so cross-revision comparisons do not need a +frozen benchmark-local implementation as their primary control. ## The API in one screen @@ -343,10 +356,10 @@ Deliberately **not** done: and masking a full-length result looks like a simplification and is not one: `normalized_readthrough_survives_null_rows` pins that a filtered input is no longer `Normalized`, so which arrays reach `reduce_encoded` is load-bearing and differs per strategy. -- **No PR split.** Recommended landing order, each step individually revertible and separately - benchmarkable: (1) API + lifting with dense/filter only; (2) branch-and-skip + adaptive selection - + its benchmarks; (3) `NumericBinary`; (4) tensor; (5) geo. The seam already supports this split - and no API changes between steps. +- **No mixed-constant specialization without a failing benchmark.** The broadcast-index-mask + experiment regressed numeric constants by 4x to 7x on x86 and was removed. Keep the current + executor for the first PR. Add a specialized varying x constant or constant x varying loop only + after the clean branch has a stable regression against the current merge base. ### Three API changes proposed, and why none of them landed diff --git a/STRICT_SCALAR_FN_RESEARCH.md b/STRICT_SCALAR_FN_RESEARCH.md index e0d5c8ad1d5..b5125ffc953 100644 --- a/STRICT_SCALAR_FN_RESEARCH.md +++ b/STRICT_SCALAR_FN_RESEARCH.md @@ -73,6 +73,45 @@ fastest and median observations rather than only these compact ranges. The historical measurements below remain because they explain design decisions and experiments made while building the prototype; they are not the current before/after performance record. +### Later broadcast-index-mask experiment + +Commit `ad24700088` tried to preserve a varying neighbor's decoded slice when another input was +constant. Every argument exposed `(Varying, mask)`, where the mask was `usize::MAX` for a varying +column and `0` for a one-row constant, and the fallback loop indexed each input with +`index & mask`. The all-varying loop was unchanged. + +The design regressed the numeric mixed-constant cases. The comparison used baseline `fed7038` and +candidate `edb3953`, with `RUSTFLAGS="-C target-feature=+avx2"` at both revisions: + +| benchmark | `fed7038` | `edb3953` | result | +| --- | ---: | ---: | ---: | +| `add_i64_constant` | 9.5-9.7 us | 37-71 us | approximately 4x to 7x slower | +| `sub_i64_constant` | 9.4-9.7 us | 37-45 us | approximately 4x slower | +| `mul_i32_constant` | 10.7-11.1 us | 42-43 us | approximately 4x slower | +| `add_i64_nonnull` | 11.1 us | 11.2 us | parity | +| `mul_i32_nonnull` | 13.9 us | 13.4 us | parity | + +The x86 report did not record the CPU, timer, pinning, or fastest and median values separately, so +these wall-clock values diagnose the code-generation failure rather than satisfy the release gate. +CodSpeed reported the same direction at a smaller magnitude: `add_i64_constant` 31.34%, +`sub_i64_constant` 32.38%, and `mul_i32_constant` 46.24% slower. + +The generated assembly kept two bounds checks per row and performed scalar loads. The runtime mask +made each varying index non-affine, so LLVM could not prove it in bounds or vectorize the loop. The +previous `ArgColumnKind` match was loop-invariant, which allowed LLVM to unswitch the numeric loop +into constant-pattern variants. The experiment optimized the branch count and discarded the +information that enabled vectorization. + +The commit was removed from `ct/row-fn` history. The clean integration head is `ea58061b5d`. Two +unpinned Divan runs on an Apple M4 Max, with 41 ns timer precision, restored the constant medians to +8.71-8.73 us for `add_i64`, 8.79-9.00 us for `sub_i64`, and 5.71-5.75 us for `mul_i32`. These +values prove that the mask regression is gone. They are not an x86 comparison against current +`develop`. + +Do not reintroduce a runtime mask or another runtime-shaped per-argument source. A later +mixed-constant optimization must monomorphize the loop over the constant pattern and must first be +justified by a stable benchmark against the current merge base. + --- ## The design in one screen From 02beb2d4e55452338e448da68d05b7db4baa4059 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 19:34:47 -0400 Subject: [PATCH 009/160] Port owned RowFn numeric execution Signed-off-by: "Connor Tsui" --- .../scalar_fn/fns/binary/numeric/primitive.rs | 10 +- .../src/scalar_fn/fns/binary/numeric/row.rs | 218 +++--------------- vortex-array/src/scalar_fn/row/element/mod.rs | 8 +- .../src/scalar_fn/row/element/tuple.rs | 32 ++- vortex-array/src/scalar_fn/row/execute.rs | 82 +++++++ vortex-array/src/scalar_fn/row/lift.rs | 39 ++-- vortex-array/src/scalar_fn/row/mod.rs | 9 +- vortex-array/src/scalar_fn/row/row_fn.rs | 30 +++ .../scalar_fn/row/tests/null_strategies.rs | 2 +- vortex-array/src/scalar_fn/row/vtable.rs | 144 +++++++++--- vortex-compute/src/lane_kernels/map_into.rs | 64 +++++ 11 files changed, 394 insertions(+), 244 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 6de544eaaa2..7a6c1de4641 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -8,9 +8,10 @@ //! Keeping them apart is what lets [`row`](super::row) write a value for every row and reduce the //! evidence without a branch, so the loop vectorizes. +use std::ops::BitOrAssign; + use crate::dtype::NativePType; use crate::dtype::half::f16; -use crate::scalar_fn::SinkResult; /// Checked addition, failing on integer overflow. pub(super) struct CheckedAdd; @@ -33,9 +34,9 @@ pub(super) struct CheckedDiv; /// never compares, so the multiply stays a widening vector multiply and the reduction stays a /// vector OR. **The width must not exceed the element's**, or the reduction becomes the loop's /// bottleneck instead of the arithmetic. -pub(super) trait Failure: SinkResult + Copy + Default {} +pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} -impl + Copy + Default> Failure for T {} +impl Failure for T {} /// One arithmetic operator at one width, as a value and its failure evidence. /// @@ -309,7 +310,6 @@ impl_checked_float!(f16, f32, f64); #[cfg(test)] mod tests { use super::CheckedArithmetic; - use crate::scalar_fn::SinkResult; /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change /// pivots, and a spread of magnitudes that straddles the 64-bit split. @@ -336,7 +336,7 @@ mod tests { /// hold each against `checked_mul`, whose `None` is the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { - let failed = ::occurred(lhs.mul_failure(rhs)); + let failed = lhs.mul_failure(rhs) != ::default(); assert_eq!(failed, reference.is_none(), "{lhs:?} * {rhs:?}"); } diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index aab36bf196b..cbbfa9e22da 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -3,27 +3,15 @@ //! The primitive arithmetic operators as a [`RowFn`]. //! -//! [`Binary`] keeps its ID, its options serialization, and its strictness, fallibility and validity -//! contracts, and delegates only the _execution_ of `Add`, `Sub`, `Mul` and `Div` over primitive -//! columns to [`NumericBinary`]. Delegation rather than conversion is what makes the port possible -//! at all: `Binary` also covers Kleene `And`/`Or`, which are not strict, and the six comparisons, -//! which are infallible, so no single [`RowFn`] can stand in for the whole function. +//! [`Binary`] keeps its ID, options serialization, and semantic contracts. It delegates only the +//! execution of primitive `Add`, `Sub`, `Mul`, and `Div` to [`NumericBinary`]. The helper is not +//! registered and appears in no serialized expression. //! -//! [`NumericBinary`] is not registered and appears in no serialized expression. It is reached only -//! through the [`ScalarFnVTable::execute`] that the blanket [`RowFn`] implementation provides, so -//! it needs no rewrite rule, no ID in the registry, and no wire format of its own. -//! -//! Everything the previous hand-written implementation did around the arithmetic itself now comes -//! from the lifting: input decoding, the constant operand collapse, the all-constant fold, the -//! null-constant short circuit, output allocation, nullability widening, and masking. What is left -//! here is the per-type checked operation and the sink that carries its overflow bit. +//! Shared lifting owns decoding, constant handling, output allocation, nullability, validity, and +//! nullable retry. The declaration below contains only type dispatch and the per-row operation. //! //! [`Binary`]: crate::scalar_fn::fns::binary::Binary -use std::marker::PhantomData; -use std::mem::MaybeUninit; - -use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::registry::CachedId; @@ -35,27 +23,18 @@ use super::primitive::CheckedPrimitiveOp; use super::primitive::CheckedSub; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::NativePType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_native_ptype; use crate::scalar::NumericOperator; -use crate::scalar_fn::DeferredError; -use crate::scalar_fn::OutputSink; use crate::scalar_fn::RowFn; use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; -use crate::validity::Validity; /// Execute a numeric operation between two primitive-typed arrays. -/// -/// The caller has already established that both operands are primitive, of the same type, and of -/// the same length. pub(super) fn execute_numeric_primitive( lhs: &ArrayRef, rhs: &ArrayRef, @@ -67,10 +46,7 @@ pub(super) fn execute_numeric_primitive( ScalarFnVTable::execute(&NumericBinary, &op, &args, ctx) } -/// The four arithmetic operators of [`Binary`] over primitive columns, as one row function per -/// operator and width. -/// -/// [`Binary`]: crate::scalar_fn::fns::binary::Binary +/// The primitive arithmetic operators as a row function. #[derive(Clone)] struct NumericBinary; @@ -79,9 +55,8 @@ impl RowFn for NumericBinary { const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; - /// Only the integer widths can overflow, and only integer division can divide by zero, but - /// fallibility is declared without input dtypes. The float widths are therefore covered by the - /// same `true`, which costs them nothing: a deferred error keeps the batch on the dense path. + // Fallibility is declared before dispatch knows the primitive width. The float widths inherit + // this conservative declaration at no execution cost. const FALLIBLE: bool = true; fn id(&self) -> ScalarFnId { @@ -89,29 +64,28 @@ impl RowFn for NumericBinary { *ID } - fn dispatch( + fn dispatch( &self, op: &Self::Options, args: &[DType], - visitor: V, - ) -> VortexResult { + visitor: Visitor, + ) -> VortexResult { let ptype = operand_ptype(args)?; - match_each_native_ptype!(ptype, |T| { + match_each_native_ptype!(ptype, |Primitive| { match op { - NumericOperator::Add => visit_checked::(visitor), - NumericOperator::Sub => visit_checked::(visitor), - NumericOperator::Mul => visit_checked::(visitor), - NumericOperator::Div => visit_checked::(visitor), + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_checked::(visitor), } }) } } -/// The width both operands are read at. +/// Return the primitive width selected by the left operand. /// -/// Only the left operand is inspected. `(T, T)` validates each argument against the chosen width, -/// so a right operand of a different type is rejected by the visit rather than here. +/// The visited `(Primitive, Primitive)` tuple validates both operands against this width. fn operand_ptype(args: &[DType]) -> VortexResult { let lhs = args .first() @@ -120,150 +94,22 @@ fn operand_ptype(args: &[DType]) -> VortexResult { PType::try_from(lhs) } -/// Visit at two `T` columns, applying `Op` per row into the sink that defers its overflow bit. -/// -/// The const block enforces, at monomorphization time, the width rule stated on -/// [`Failure`](super::primitive::Failure): evidence wider than the element would make the -/// OR-reduction rather than the arithmetic decide how many rows fit in a vector. -fn visit_checked(visitor: V) -> VortexResult +/// Visit two primitive columns and defer one OR-reducible failure word per row. +fn visit_checked(visitor: Visitor) -> VortexResult where - T: NativePType, - Op: CheckedPrimitiveOp, - V: RowVisitor, + Primitive: NativePType, + Operator: CheckedPrimitiveOp, + Visitor: RowVisitor, { - const { - assert!( - size_of::() <= size_of::(), - "failure evidence must be no wider than the value, or it bounds the vector width" - ) - }; - - visitor.visit_prepared_into::<(T, T), CheckedSink, _, _>( + visitor.visit_prepared_deferred::<(Primitive, Primitive), Primitive, _, Operator::Failure>( |_| (), - |&(), (lhs, rhs), output| output.write(lhs, rhs), - ) -} - -/// The output column of one checked arithmetic batch, reporting failure once after the row loop. -/// -/// Deferring the failure is what keeps a fallible kernel on the dense path: every row writes a -/// value unconditionally and OR-reduces its failure evidence, so the loop holds no branch and no -/// `Result` discriminant. The lifting retries a nullable batch over only its valid rows if that -/// reduction is non-zero, which is what makes an overflow behind a null row invisible. -/// -/// The reduction lives in the sink rather than in the row closure's return type so that its width -/// is [`Op::Failure`](CheckedPrimitiveOp::Failure), the operator's choice, rather than one bit. That -/// is what lets unsigned multiplication report its discarded high half instead of a comparison, and -/// so stay vectorized. -/// -/// **The storage is deliberately uninitialized, not zeroed.** Substituting `BufferMut::zeroed` to -/// make the sink safe was measured at **1.65 to 1.71x** the cost of allocate-and-fill, stable across -/// two runs and every batch size from 8 KiB to 2 MiB, because `alloc_zeroed` does not avoid the -/// write: below glibc's mmap threshold `calloc` recycles a dirty chunk and memsets it, and above it -/// the first touch of each fresh page faults instead. The row loop overwrites every slot regardless, -/// so that pass is pure duplicate work on the hottest kernel in the system. This is the case the -/// repository's "avoid `unsafe` unless it is necessary" rule leaves room for: the safe spelling -/// exists, and it costs a second pass over the output. -/// -/// Rows are written into uninitialized storage, so this sink cannot finish a batch whose rows were -/// not all visited, and leaves [`OutputSink::SUPPORTS_SKIPPED_ROWS`] at `false`. Nothing is lost: -/// `SUPPORTS_SKIPPED_ROWS` is what makes branch-and-skip unavailable, which is the guard that keeps -/// the uninitialized slots sound. Note this is _not_ implied by the dispatch policy alone: a -/// deferred result still reaches the executor's valid-only policy whenever its arguments are not -/// dense-safe, so the `false` here is load-bearing rather than a restatement. -struct CheckedSink> { - /// The result values, initialized one row at a time up to `row_count`. - values: BufferMut, - - /// The batch length, which is the capacity `values` was allocated with. - row_count: usize, - - /// The operation applied to every row, which names the error reported by - /// [`finish`](OutputSink::finish). - op: PhantomData, -} - -/// The uninitialized output slots of a [`CheckedSink`], borrowed once for the row loop. -struct CheckedRows<'a, T: NativePType, Op: CheckedPrimitiveOp> { - values: &'a mut [MaybeUninit], - op: PhantomData, -} - -/// One output slot of a [`CheckedSink`]. -struct CheckedRow<'a, T: NativePType, Op: CheckedPrimitiveOp> { - value: &'a mut MaybeUninit, - op: PhantomData, -} - -impl> CheckedRow<'_, T, Op> { - /// Apply `Op` to one row, writing its value and handing back its failure evidence. - /// - /// The value is written whether or not the operation failed, since a failing row is either - /// masked away as null or turned into a batch error before it can be read. The evidence is - /// returned rather than reduced here so the executor can keep the reduction in a register, and - /// it is `Op`'s own width so the row never has to compare. - fn write(self, lhs: T, rhs: T) -> Op::Failure { - let (value, failure) = Op::apply(lhs, rhs); - self.value.write(value); - - failure - } -} - -impl> OutputSink for CheckedSink { - const ERRORS_ARE_DEFERRED: bool = true; - - type Rows<'a> - = CheckedRows<'a, T, Op> - where - Self: 'a; - type Row<'a> - = CheckedRow<'a, T, Op> - where - Self: 'a; - - fn sink_dtype(_args: &[DType]) -> VortexResult { - Ok(DType::Primitive(T::PTYPE, Nullability::NonNullable)) - } - - fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { - Ok(Self { - values: BufferMut::with_capacity(rows), - row_count: rows, - op: PhantomData, - }) - } - - fn rows(&mut self) -> Self::Rows<'_> { - let row_count = self.row_count; - CheckedRows { - values: &mut self.values.spare_capacity_mut()[..row_count], - op: PhantomData, - } - } - - fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { - rows.values.len() == row_count - } - - fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { - CheckedRow { - value: &mut rows.values[index], - op: PhantomData, - } - } - - fn finish(mut self, error: DeferredError) -> VortexResult { - if error.occurred() { - return Err(vortex_err!(InvalidArgument: "{}", Op::ERROR)); - } - - // SAFETY: the sink reports `SUPPORTS_SKIPPED_ROWS = false`, so every path that reaches - // `finish` without an error has written all `row_count` slots: dense execution visits - // `0..row_count`, and the valid-row retry runs densely over a sink allocated for exactly - // the filtered rows. - unsafe { self.values.set_len(self.row_count) }; + |&(), (lhs, rhs)| Operator::apply(lhs, rhs), + |failure| { + if failure != ::default() { + return Err(vortex_err!(InvalidArgument: "{}", Operator::ERROR)); + } - Ok(PrimitiveArray::new(self.values.freeze(), Validity::NonNullable).into_array()) - } + Ok(()) + }, + ) } diff --git a/vortex-array/src/scalar_fn/row/element/mod.rs b/vortex-array/src/scalar_fn/row/element/mod.rs index fc88690b6b0..384d8bc3aa7 100644 --- a/vortex-array/src/scalar_fn/row/element/mod.rs +++ b/vortex-array/src/scalar_fn/row/element/mod.rs @@ -8,9 +8,10 @@ //! `vortex-tensor`'s `TensorRow` drills through an extension wrapper into its storage. //! //! The two directions are deliberately asymmetric. [`InputElement::Elem`] is a GAT, so an input row -//! can borrow out of the decoded column, while an [`OutputElement`] is one owned value written into -//! an [`ElementSink`](crate::scalar_fn::ElementSink). Runtime-shaped output uses a custom -//! [`OutputSink`](crate::scalar_fn::OutputSink) instead. +//! can borrow out of the decoded column, while an [`OutputElement`] is one owned value returned by +//! an owned row computation or written through +//! [`ElementSink`](crate::scalar_fn::ElementSink); runtime-shaped output uses a custom +//! [`OutputSink`](crate::scalar_fn::OutputSink). use vortex_error::VortexResult; @@ -29,6 +30,7 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; +pub use tuple::IndexedElementTuple; pub(super) use tuple::batch_constant; /// An element type that can be read row-wise out of an input column. diff --git a/vortex-array/src/scalar_fn/row/element/tuple.rs b/vortex-array/src/scalar_fn/row/element/tuple.rs index a3c31cfeb2e..fad02639e1a 100644 --- a/vortex-array/src/scalar_fn/row/element/tuple.rs +++ b/vortex-array/src/scalar_fn/row/element/tuple.rs @@ -3,6 +3,8 @@ //! Argument lists built from [`InputElement`]s, and the per-argument decode behind them. +use vortex_compute::lane_kernels::IndexedSource; +use vortex_compute::lane_kernels::LaneZip; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; @@ -13,6 +15,7 @@ use crate::arrays::Masked; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::masked::MaskedArraySlotsExt; use crate::dtype::DType; +use crate::dtype::NativePType; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::InputElement; @@ -162,9 +165,8 @@ pub trait ElementTuple: 'static + private::Sealed { /// /// `Some` marks an argument whose operand is constant for the batch and carries the element /// every row reads; `None` marks one that varies by row. This is what - /// [`visit_prepared_into`](crate::scalar_fn::RowVisitor::visit_prepared_into) hands to its prepare - /// closure, so a kernel can hoist work that depends only on a constant argument out of the - /// row loop. + /// [`RowVisitor`](crate::scalar_fn::RowVisitor) hands to a visit's prepare closure, so a kernel + /// can hoist work that depends only on a constant argument out of the row loop. type ConstElems<'a>; /// The number of arguments. @@ -223,6 +225,22 @@ pub trait ElementTuple: 'static + private::Sealed { fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; } +/// An argument tuple that can expose independent indexed reads after one length validation. +/// +/// This trait is sealed through [`ElementTuple`]. Tuples without a natural indexed source continue +/// to use ordinary row access and output sinks. +pub trait IndexedElementTuple: ElementTuple { + /// The source shared execution uses for a dense all-varying loop. + /// + /// Its length must be the common varying-column length. For every valid index it must preserve + /// row order, return the same value as [`ElementTuple::get_varying`], and uphold the unchecked + /// read contract of [`IndexedSource`]. + type Source<'a>: IndexedSource>; + + /// Borrow a source from columns already validated to vary within the batch. + fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a>; +} + impl private::Sealed for () {} impl ElementTuple for () { @@ -368,3 +386,11 @@ element_tuple!(9; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8); element_tuple!(10; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9); element_tuple!(11; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10); element_tuple!(12; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11); + +impl IndexedElementTuple for (A, B) { + type Source<'a> = LaneZip<&'a [A], &'a [B]>; + + fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a> { + LaneZip::new(columns.0, columns.1) + } +} diff --git a/vortex-array/src/scalar_fn/row/execute.rs b/vortex-array/src/scalar_fn/row/execute.rs index 345bed991fe..a92e3e43998 100644 --- a/vortex-array/src/scalar_fn/row/execute.rs +++ b/vortex-array/src/scalar_fn/row/execute.rs @@ -6,6 +6,10 @@ //! These back the blanket impls in [`row_fn`](super::row_fn) and are deliberately not public: //! [`RowFn`](crate::scalar_fn::RowFn) is the abstraction, these are its internals. +use std::mem::needs_drop; +use std::ops::BitOrAssign; + +use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -19,6 +23,8 @@ use crate::dtype::DType; use crate::scalar_fn::DeferredError; use crate::scalar_fn::ElementTuple; use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; use crate::scalar_fn::OutputSink; use crate::scalar_fn::SinkResult; @@ -42,6 +48,19 @@ impl RowExecution { } } +/// Validate the input dtypes of an owned-output row function and return its output dtype. +pub(super) fn validate_row_output( + args: &[DType], +) -> VortexResult { + A::validate(args)?; + let dtype = O::element_dtype(); + vortex_ensure!( + !dtype.is_nullable(), + "row output elements must declare a non-nullable dtype, got {dtype}", + ); + Ok(dtype) +} + /// Validate the input dtypes of a sink-writing row function and return the dtype its sink builds. /// /// The output dtype may be a function of the inputs. A sink can also own a batch-wide builder, such @@ -58,6 +77,69 @@ pub(super) fn validate_row_sink( Ok(dtype) } +/// Decode every input column once, then store owned row outputs and reduce deferred failures. +pub(super) fn execute_row_output_prepared( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>) -> (O, F), + finish_failure: impl FnOnce(F) -> VortexResult<()>, +) -> VortexResult +where + A: IndexedElementTuple, + O: OutputElement, + F: Copy + Default + BitOrAssign, +{ + const { + assert!( + !needs_drop::(), + "owned deferred outputs must not require drop glue" + ) + }; + + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let columns = A::decode(args, ctx)?; + let state = prepare(A::constants(&columns)); + let failed; + + { + let output = &mut values.spare_capacity_mut()[..row_count]; + + if let Some(varying) = A::varying(&columns) { + vortex_ensure!( + A::varying_len_matches(&varying, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + failed = + A::indexed_source(&varying).map_checked_into(output, |elems| apply(&state, elems)); + } else { + vortex_ensure!( + A::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut accumulated = F::default(); + for index in 0..row_count { + let (value, failure) = apply(&state, A::get(&columns, index)); + output[index].write(value); + accumulated |= failure; + } + failed = accumulated; + } + } + + // SAFETY: normal completion of either loop initializes every slot in `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + match finish_failure(failed) { + Ok(()) => Ok(RowExecution::Output(O::build(values))), + Err(error) => Ok(RowExecution::DeferredError(error)), + } +} + /// Decode every input column once, allocate the sink once, then write one row at a time. /// /// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and the loop keeps the diff --git a/vortex-array/src/scalar_fn/row/lift.rs b/vortex-array/src/scalar_fn/row/lift.rs index a842ef5deef..0a44b066826 100644 --- a/vortex-array/src/scalar_fn/row/lift.rs +++ b/vortex-array/src/scalar_fn/row/lift.rs @@ -76,9 +76,9 @@ impl ExecutionArgs for BorrowedExecutionArgs<'_> { /// The arguments handed to one kernel invocation. /// -/// `arrays` may be filtered or sliced, while `dtypes` and `sink_dtype` always describe the original -/// planned batch. Keeping them together prevents an execution path from accidentally pairing an -/// input view with unrelated planning metadata. +/// `arrays` may be filtered or sliced, while `dtypes` and `output_dtype` always describe the +/// original planned batch. Keeping them together prevents an execution path from accidentally +/// pairing an input view with unrelated planning metadata. #[derive(Clone, Copy)] pub(super) struct KernelArgs<'a> { /// The executor-facing view, including the row count for this invocation. @@ -90,14 +90,14 @@ pub(super) struct KernelArgs<'a> { /// The original input dtypes used to select the row implementation. pub(super) dtypes: &'a [DType], - /// The non-nullable dtype allocated by the selected output sink. - pub(super) sink_dtype: &'a DType, + /// The non-nullable dtype built by the selected output capability. + pub(super) output_dtype: &'a DType, } /// The execution policy and output dtype selected by a planning visit. pub(super) struct BatchPlan { - /// The non-nullable dtype built by the selected sink. - pub(super) sink_dtype: DType, + /// The non-nullable dtype built by the selected output capability. + pub(super) output_dtype: DType, /// How this concrete dispatch executes nullable rows. pub(super) policy: RowPolicy, @@ -118,6 +118,17 @@ pub(super) enum RowPolicy { } impl RowPolicy { + /// The policy for an owned output carrying batch-deferred failure evidence. + pub(super) const fn for_deferred_output() -> Self { + if A::DENSE_SAFE && !A::DECODE_FALLIBLE { + Self::DenseWithRetry + } else { + Self::ValidOnly { + filtered_decode_cost: A::FILTERED_DECODE_COST, + } + } + } + /// The policy one concrete dispatch executes nullable rows under. /// /// Note what is deliberately **not** read here: [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. Hoisting @@ -130,7 +141,7 @@ impl RowPolicy { /// answer differently from its row loop, that is a wrong answer rather than a slow one. /// /// [`OutputSink::SUPPORTS_SKIPPED_ROWS`]: crate::scalar_fn::OutputSink::SUPPORTS_SKIPPED_ROWS - pub(super) const fn for_dispatch() -> Self { + pub(super) const fn for_sink() -> Self { if A::DENSE_SAFE && !A::DECODE_FALLIBLE && !R::FALLIBLE { if R::DEFERRED { Self::DenseWithRetry @@ -178,8 +189,8 @@ pub(super) struct Batch<'a> { /// against. Already widened to nullable if any input is nullable. result_dtype: DType, - /// The non-nullable dtype the dispatched sink builds, computed once while planning. - sink_dtype: DType, + /// The non-nullable dtype the dispatched output capability builds, computed while planning. + output_dtype: DType, /// How the concrete dispatch executes nullable rows. policy: RowPolicy, @@ -203,9 +214,9 @@ impl<'a> Batch<'a> { let arg_dtypes: SmallVec<[DType; 4]> = inputs.iter().map(|input| input.dtype().clone()).collect(); let plan = plan(&arg_dtypes)?; - let nullability = plan.sink_dtype.nullability() + let nullability = plan.output_dtype.nullability() | Nullability::from(arg_dtypes.iter().any(DType::is_nullable)); - let result_dtype = plan.sink_dtype.with_nullability(nullability); + let result_dtype = plan.output_dtype.with_nullability(nullability); let mut validity = Validity::NonNullable; for input in &inputs { @@ -219,7 +230,7 @@ impl<'a> Batch<'a> { arg_dtypes, validity, result_dtype, - sink_dtype: plan.sink_dtype, + output_dtype: plan.output_dtype, policy: plan.policy, }) } @@ -507,7 +518,7 @@ impl<'a> Batch<'a> { execution, arrays, dtypes: &self.arg_dtypes, - sink_dtype: &self.sink_dtype, + output_dtype: &self.output_dtype, } } diff --git a/vortex-array/src/scalar_fn/row/mod.rs b/vortex-array/src/scalar_fn/row/mod.rs index f62d8bb8d49..d9799fff678 100644 --- a/vortex-array/src/scalar_fn/row/mod.rs +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -19,9 +19,11 @@ //! [`RowFn`] does not say how a row is _stored_, which is the element's job: `vortex-tensor` adds a //! `TensorRow` [`InputElement`] and writes ordinary kernels over it. //! -//! Output always goes through [`RowVisitor::visit_prepared_into`]. [`ElementSink`] covers one owned -//! [`OutputElement`] per row; custom [`OutputSink`] implementations cover runtime-shaped rows. The -//! prepare closure sees every batch-constant input and returns shared state for the row loop. Pass +//! Output has two capabilities. [`RowVisitor::visit_prepared_deferred`] returns one independent +//! [`OutputElement`] and failure word per row, letting shared execution own the stores and choose an +//! indexed dense source. [`RowVisitor::visit_prepared_into`] writes through an [`OutputSink`] for +//! runtime-shaped rows, shared builders, skip-capable output, and values requiring drop glue. Both +//! prepare closures see every batch-constant input and return shared state for the row loop. Pass //! `|_| ()` when there is nothing to prepare. //! //! A kernel that can safely write a provisional value uses [`DeferredError`] instead of returning @@ -43,6 +45,7 @@ mod element; pub use element::ElementTuple; +pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; #[cfg(any(test, feature = "_test-harness"))] diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index f2a63636b1d..f790d21d8a0 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -6,6 +6,7 @@ use std::fmt::Debug; use std::fmt::Display; use std::hash::Hash; +use std::ops::BitOrAssign; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -15,6 +16,8 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; use crate::scalar_fn::OutputSink; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::SinkResult; @@ -116,6 +119,33 @@ pub trait RowVisitor: private::Sealed { /// What this visit produces. type Out; + /// Visit at indexed argument tuple `A`, returning one independently owned output and one + /// deferred failure word per row. + /// + /// The executor allocates and writes the output column. It reads through + /// [`IndexedElementTuple`] when every argument varies; batches containing a constant use + /// ordinary row access selected once outside the loop. `F::default()` **must** mean success, + /// including for empty execution, and `|=` must combine the evidence from independent rows. + /// `finish_failure` runs once after the loop: it must return `Ok(())` for successful evidence + /// and may report only the operation's row error for failed evidence. That error is deferred, + /// so nullable lifting may retry over only valid rows. + /// + /// `A` **must** have the arity declared by [`RowFn::ARG_NAMES`]. This method requires + /// [`RowFn::FALLIBLE`] to be `true`, and `O` must be no narrower than `F` so failure reduction + /// does not constrain vector width. `O` must not require drop glue. Use + /// [`visit_prepared_into`](Self::visit_prepared_into) for non-indexed tuples, runtime-shaped + /// output, shared builders, and output that requires drop. + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>) -> (O, F), + finish_failure: impl FnOnce(F) -> VortexResult<()>, + ) -> VortexResult + where + A: IndexedElementTuple, + O: OutputElement, + F: 'static + Copy + Default + BitOrAssign; + /// Visit at argument tuple `A`, preparing shared state once and writing every output row into /// sink `S`. /// diff --git a/vortex-array/src/scalar_fn/row/tests/null_strategies.rs b/vortex-array/src/scalar_fn/row/tests/null_strategies.rs index d8a7fe18816..487cb4786f6 100644 --- a/vortex-array/src/scalar_fn/row/tests/null_strategies.rs +++ b/vortex-array/src/scalar_fn/row/tests/null_strategies.rs @@ -493,7 +493,7 @@ mod selection { #[test] fn planning_adds_decode_cost_across_arguments() { assert_eq!( - RowPolicy::for_dispatch::<(TrackedI64<1>, TrackedI64<1>), ()>(), + RowPolicy::for_sink::<(TrackedI64<1>, TrackedI64<1>), ()>(), RowPolicy::ValidOnly { filtered_decode_cost: 2 } diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index cff4831aa8f..26c83e9c11b 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -4,6 +4,8 @@ //! Blanket scalar-function implementation and execution visitors for row functions. use std::marker::PhantomData; +use std::mem::needs_drop; +use std::ops::BitOrAssign; use vortex_error::VortexResult; #[cfg(any(test, feature = "_test-harness"))] @@ -24,8 +26,10 @@ use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ElementTuple; use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::IndexedElementTuple; #[cfg(any(test, feature = "_test-harness"))] use crate::scalar_fn::NullStrategy; +use crate::scalar_fn::OutputElement; use crate::scalar_fn::OutputSink; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; @@ -33,8 +37,10 @@ use crate::scalar_fn::SinkResult; #[cfg(any(test, feature = "_test-harness"))] use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::row::execute::RowExecution; +use crate::scalar_fn::row::execute::execute_row_output_prepared; use crate::scalar_fn::row::execute::execute_row_sink_branch; use crate::scalar_fn::row::execute::execute_row_sink_prepared; +use crate::scalar_fn::row::execute::validate_row_output; use crate::scalar_fn::row::execute::validate_row_sink; use crate::scalar_fn::row::lift::Batch; use crate::scalar_fn::row::lift::BatchPlan; @@ -42,11 +48,8 @@ use crate::scalar_fn::row::lift::KernelArgs; use crate::scalar_fn::row::lift::RowPolicy; use crate::scalar_fn::row::lift::reconcile_return; -/// Compile-time check that a dispatched `(A, S, R)` agrees with `F`'s public metadata. Evaluated by -/// monomorphizing -/// [`visit_prepared_into`](RowVisitor::visit_prepared_into), so even a dispatch arm that never runs -/// is checked. -const fn assert_dispatch_agrees() { +/// Compile-time checks shared by both output capabilities. +const fn assert_input_dispatch_agrees() { assert!( A::ARITY == F::ARG_NAMES.len(), "dispatch visited a tuple whose arity differs from RowFn::ARG_NAMES", @@ -57,6 +60,11 @@ const fn assert_dispatch_agrees() { + assert_input_dispatch_agrees::(); assert!( !R::FALLIBLE || F::FALLIBLE, "dispatch returned an error without declaring RowFn::FALLIBLE", @@ -71,7 +79,30 @@ const fn assert_dispatch_agrees() +where + F: RowFn, + A: IndexedElementTuple, + O: OutputElement, + Failure: Copy + Default + BitOrAssign, +{ + assert_input_dispatch_agrees::(); + assert!( + F::FALLIBLE, + "dispatch deferred an error without declaring RowFn::FALLIBLE", + ); + assert!( + size_of::() <= size_of::(), + "failure evidence must be no wider than the value, or it bounds the vector width", + ); + assert!( + !needs_drop::(), + "owned deferred outputs must not require drop glue", + ); +} + +/// The plan-time visit: validate the dtypes and derive execution from the output capability and row /// closure selected by dispatch. struct PlanRows<'a, F> { args: &'a [DType], @@ -85,16 +116,35 @@ impl private::Sealed for PlanRows<'_, F> {} impl RowVisitor for PlanRows<'_, F> { type Out = BatchPlan; + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(A::ConstElems<'_>) -> P, + _apply: impl Fn(&P, A::Elems<'_>) -> (O, Failure), + _finish_failure: impl FnOnce(Failure) -> VortexResult<()>, + ) -> VortexResult + where + A: IndexedElementTuple, + O: OutputElement, + Failure: 'static + Copy + Default + BitOrAssign, + { + const { assert_deferred_dispatch_agrees::() }; + + Ok(BatchPlan { + output_dtype: validate_row_output::(self.args)?, + policy: RowPolicy::for_deferred_output::(), + }) + } + fn visit_prepared_into( self, _prepare: impl FnOnce(A::ConstElems<'_>) -> P, _apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, ) -> VortexResult { - const { assert_dispatch_agrees::() }; + const { assert_sink_dispatch_agrees::() }; Ok(BatchPlan { - sink_dtype: validate_row_sink::(self.args)?, - policy: RowPolicy::for_dispatch::(), + output_dtype: validate_row_sink::(self.args)?, + policy: RowPolicy::for_sink::(), }) } } @@ -103,8 +153,8 @@ impl RowVisitor for PlanRows<'_, F> { struct ExecuteRows<'a, 'b, F> { args: &'a dyn ExecutionArgs, - /// The sink dtype computed by the planning visit. - sink_dtype: &'a DType, + /// The output dtype computed by the planning visit. + output_dtype: &'a DType, ctx: &'b mut ExecutionCtx, @@ -117,15 +167,36 @@ impl private::Sealed for ExecuteRows<'_, '_, F> {} impl RowVisitor for ExecuteRows<'_, '_, F> { type Out = RowExecution; + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(A::ConstElems<'_>) -> P, + apply: impl Fn(&P, A::Elems<'_>) -> (O, Failure), + finish_failure: impl FnOnce(Failure) -> VortexResult<()>, + ) -> VortexResult + where + A: IndexedElementTuple, + O: OutputElement, + Failure: 'static + Copy + Default + BitOrAssign, + { + const { assert_deferred_dispatch_agrees::() }; + execute_row_output_prepared::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } + fn visit_prepared_into( self, prepare: impl FnOnce(A::ConstElems<'_>) -> P, apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, ) -> VortexResult { - const { assert_dispatch_agrees::() }; + const { assert_sink_dispatch_agrees::() }; execute_row_sink_prepared::( self.args, - self.sink_dtype, + self.output_dtype, self.ctx, prepare, apply, @@ -136,13 +207,13 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { /// The run-time visit for the branch-and-skip null strategy: compute only the conjoined-valid /// rows over unfiltered columns. /// -/// `Ok(None)` means the visit cannot take that strategy because the sink cannot skip rows or an -/// argument has no null-tolerant decode, and the lifting falls back to the filter strategy. +/// `Ok(None)` means the visit requires filtering, a sink cannot skip rows, or an argument has no +/// null-tolerant decode. The lifting then falls back to the filter strategy. struct ExecuteRowsBranch<'a, 'b, F> { args: &'a dyn ExecutionArgs, - /// The sink dtype computed by the planning visit. - sink_dtype: &'a DType, + /// The output dtype computed by the planning visit. + output_dtype: &'a DType, /// The conjoined validity, materialized by the lifting and guaranteed mixed. valid: &'a Mask, @@ -158,15 +229,30 @@ impl private::Sealed for ExecuteRowsBranch<'_, '_, F> {} impl RowVisitor for ExecuteRowsBranch<'_, '_, F> { type Out = Option; + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(A::ConstElems<'_>) -> P, + _apply: impl Fn(&P, A::Elems<'_>) -> (O, Failure), + _finish_failure: impl FnOnce(Failure) -> VortexResult<()>, + ) -> VortexResult> + where + A: IndexedElementTuple, + O: OutputElement, + Failure: 'static + Copy + Default + BitOrAssign, + { + const { assert_deferred_dispatch_agrees::() }; + Ok(None) + } + fn visit_prepared_into( self, prepare: impl FnOnce(A::ConstElems<'_>) -> P, apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R, ) -> VortexResult> { - const { assert_dispatch_agrees::() }; + const { assert_sink_dispatch_agrees::() }; execute_row_sink_branch::( self.args, - self.sink_dtype, + self.output_dtype, self.valid, self.ctx, prepare, @@ -192,7 +278,7 @@ fn execute_rows( args.dtypes, ExecuteRows:: { args: args.execution, - sink_dtype: args.sink_dtype, + output_dtype: args.output_dtype, ctx, row_fn: PhantomData, }, @@ -221,7 +307,7 @@ fn execute_rows_branch( args.dtypes, ExecuteRowsBranch:: { args: args.execution, - sink_dtype: args.sink_dtype, + output_dtype: args.output_dtype, valid, ctx, row_fn: PhantomData, @@ -229,7 +315,7 @@ fn execute_rows_branch( ) } -/// The batch facts for `row_fn` over `args`, derived from its dispatched elements and sink. +/// The batch facts for `row_fn` over `args`, derived from its selected output capability. fn lift_batch<'a, F: RowFn>( row_fn: &F, options: &F::Options, @@ -307,9 +393,9 @@ impl ScalarFnVTable for F { }, )?; - let nullability = - plan.sink_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); - Ok(plan.sink_dtype.with_nullability(nullability)) + let nullability = plan.output_dtype.nullability() + | Nullability::from(args.iter().any(DType::is_nullable)); + Ok(plan.output_dtype.with_nullability(nullability)) } fn execute( @@ -328,7 +414,7 @@ impl ScalarFnVTable for F { execution: args, arrays: &[], dtypes: &[], - sink_dtype: &result_dtype, + output_dtype: &result_dtype, }, ctx, )? @@ -343,9 +429,9 @@ impl ScalarFnVTable for F { ) } - /// Output sinks build an all-valid column, so a row kernel cannot turn a wholly non-null row into - /// a null and the output validity is exactly the conjunction of the inputs'. Letting a sink - /// produce nulls would invalidate this. + /// Row output capabilities build an all-valid column, so a kernel cannot turn a wholly non-null + /// row into a null and the output validity is exactly the conjunction of the inputs'. Letting an + /// output capability produce nulls would invalidate this. fn validity( &self, _options: &Self::Options, diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs index 258913e9fc7..9ede2df69e5 100644 --- a/vortex-compute/src/lane_kernels/map_into.rs +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -5,6 +5,7 @@ //! caller-provided `&mut [MaybeUninit]`. use std::mem::MaybeUninit; +use std::ops::BitOrAssign; use vortex_buffer::BitBuffer; @@ -156,6 +157,50 @@ pub trait IndexedSourceExt: IndexedSource + Sized { } } + /// Write each mapped value and OR-reduce independent failure evidence across the batch. + /// + /// The failure stays local to this method so the optimizer can keep it in a register. The + /// caller receives only whether the batch failed and can attribute errors on a cold retry. + /// **`Failure` must be no wider than `Output`**, or its reduction can limit vector width. + /// + /// # Panics + /// + /// Panics if `out.len() != self.len()`. + #[inline] + fn map_checked_into( + self, + out: &mut [MaybeUninit], + mut apply: Apply, + ) -> Failure + where + Failure: Copy + Default + BitOrAssign, + Apply: FnMut(Self::Item) -> (Output, Failure), + { + const { + assert!( + size_of::() <= size_of::(), + "failure evidence must be no wider than the value, or it bounds the vector width" + ) + }; + + let values = self; + let len = values.len(); + assert_eq!(out.len(), len, "out must have the same length as values"); + + let mut failed = Failure::default(); + for index in 0..len { + // SAFETY: `index < len` by the loop bound. + let value = unsafe { values.get_unchecked(index) }; + let (output, failure) = apply(value); + failed |= failure; + + // SAFETY: `index < len == out.len()`. + unsafe { out.get_unchecked_mut(index).write(output) }; + } + + failed + } + /// Apply the predicate `f(value)` lane-by-lane and bit-pack the results into /// `words`, LSB-first, 64 lanes per `u64`. /// @@ -546,6 +591,25 @@ mod tests { assert!(res.is_ok(), "null lane should bypass the range check"); } + #[test] + fn map_checked_into_writes_all_lanes_and_reduces_failure() { + let mut values: Vec = (0..130).collect(); + let mut output = vec![MaybeUninit::::uninit(); 130]; + let failed = values + .as_slice() + .map_checked_into(&mut output, |value| (value as u32, value > u32::MAX as u64)); + assert!(!failed); + assert_eq!(write_t(output), (0..130u32).collect::>()); + + values[77] = (u32::MAX as u64) + 1; + let mut output = vec![MaybeUninit::::uninit(); 130]; + let failed = values + .as_slice() + .map_checked_into(&mut output, |value| (value as u32, value > u32::MAX as u64)); + assert!(failed); + assert_eq!(write_t(output)[76], 76); + } + #[test] fn map_bits_into_packs_full_and_remainder_words() { let values: Vec = (0..130).collect(); From ecb3826cb43d7f5f3bf010180f1161a8a091fb9d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 19:38:35 -0400 Subject: [PATCH 010/160] Record RowFn x86 vectorization research Signed-off-by: "Connor Tsui" --- NUMERIC_ROWFN_PLAN.md | 5 + SCALAR_FN_HANDOFF.md | 2 + research/rowfn-x86-2026-08-07/README.md | 239 +++++++++++++ .../benchmarks/land-base-1.md | 39 +++ .../benchmarks/land-base-2.md | 39 +++ .../benchmarks/land-candidate-1.md | 39 +++ .../benchmarks/land-candidate-2.md | 39 +++ .../benchmarks/land-final-1.md | 39 +++ .../benchmarks/land-final-2.md | 39 +++ .../benchmarks/stage0-base-1.md | 38 +++ .../benchmarks/stage0-base-2.md | 38 +++ .../benchmarks/stage0-candidate-1.md | 38 +++ .../benchmarks/stage0-candidate-2.md | 38 +++ .../benchmarks/stage1-base-1.md | 38 +++ .../benchmarks/stage1-base-2.md | 38 +++ .../benchmarks/stage1-candidate-1.md | 38 +++ .../benchmarks/stage1-candidate-2.md | 38 +++ .../benchmarks/stage1-owned-1.md | 38 +++ .../benchmarks/stage1-owned-2.md | 38 +++ .../benchmarks/stage2-base-1.md | 39 +++ .../benchmarks/stage2-base-2.md | 39 +++ .../benchmarks/stage2-candidate-1.md | 39 +++ .../benchmarks/stage2-candidate-2.md | 39 +++ .../benchmarks/stage2-indexed-1.md | 39 +++ .../benchmarks/stage2-indexed-2.md | 39 +++ .../codegen/base-codegen-summary.md | 139 ++++++++ .../codegen/candidate-i64-mul-dense-ll.md | 84 +++++ .../codegen/candidate-i64-mul-dense-s.md | 69 ++++ .../codegen/candidate-u64-mul-dense-ll.md | 137 ++++++++ .../codegen/candidate-u64-mul-dense-s.md | 70 ++++ .../codegen/final-i32-mul-constant-ll.md | 86 +++++ .../codegen/final-i32-mul-constant-s.md | 49 +++ .../codegen/final-i64-mul-dense-ll.md | 96 ++++++ .../codegen/final-i64-mul-dense-s.md | 34 ++ .../codegen/final-u64-mul-dense-ll.md | 322 ++++++++++++++++++ .../codegen/final-u64-mul-dense-s.md | 71 ++++ .../codegen/indexed-i64-mul-dense-ll.md | 96 ++++++ .../codegen/indexed-i64-mul-dense-s.md | 35 ++ .../codegen/indexed-u64-mul-dense-ll.md | 322 ++++++++++++++++++ .../codegen/indexed-u64-mul-dense-s.md | 71 ++++ .../codegen/owned-i64-mul-dense-ll.md | 84 +++++ .../codegen/owned-i64-mul-dense-s.md | 46 +++ .../codegen/owned-u64-mul-dense-ll.md | 121 +++++++ .../codegen/owned-u64-mul-dense-s.md | 72 ++++ 44 files changed, 3098 insertions(+) create mode 100644 research/rowfn-x86-2026-08-07/README.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/land-base-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/land-base-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/land-candidate-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/land-candidate-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/land-final-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/land-final-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage0-base-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage0-base-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage0-candidate-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage0-candidate-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage1-base-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage1-base-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage1-candidate-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage1-candidate-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage1-owned-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage1-owned-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage2-base-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage2-base-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage2-candidate-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage2-candidate-2.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage2-indexed-1.md create mode 100644 research/rowfn-x86-2026-08-07/benchmarks/stage2-indexed-2.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/base-codegen-summary.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/candidate-i64-mul-dense-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/candidate-i64-mul-dense-s.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/candidate-u64-mul-dense-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/candidate-u64-mul-dense-s.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/final-i32-mul-constant-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/final-i32-mul-constant-s.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/final-i64-mul-dense-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/final-i64-mul-dense-s.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/final-u64-mul-dense-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/final-u64-mul-dense-s.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/indexed-i64-mul-dense-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/indexed-i64-mul-dense-s.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/indexed-u64-mul-dense-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/indexed-u64-mul-dense-s.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/owned-i64-mul-dense-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/owned-i64-mul-dense-s.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/owned-u64-mul-dense-ll.md create mode 100644 research/rowfn-x86-2026-08-07/codegen/owned-u64-mul-dense-s.md diff --git a/NUMERIC_ROWFN_PLAN.md b/NUMERIC_ROWFN_PLAN.md index 5a304e88102..6fb337afc04 100644 --- a/NUMERIC_ROWFN_PLAN.md +++ b/NUMERIC_ROWFN_PLAN.md @@ -3,6 +3,11 @@ # Plan: fit the numeric binary operators onto `RowFn` +> The later x86 follow-up found that the sink-only API regressed varying `i64`/`u64` multiply and +> added separate owned-output and stateful-sink capabilities. Its complete evidence is in +> [`research/rowfn-x86-2026-08-07/README.md`](research/rowfn-x86-2026-08-07/README.md). Treat that +> record as authoritative where it supersedes the pre-x86 conclusions below. + Working note, branch-only, like `SCALAR_FN_HANDOFF.md`. Written so this survives a conversation compaction: everything needed to start is here, and nothing below depends on chat history. diff --git a/SCALAR_FN_HANDOFF.md b/SCALAR_FN_HANDOFF.md index 8b5d6a6d053..10224704829 100644 --- a/SCALAR_FN_HANDOFF.md +++ b/SCALAR_FN_HANDOFF.md @@ -6,6 +6,8 @@ This is the concise source of truth for the branch. `STRICT_SCALAR_FN_RESEARCH.md` keeps the full design history, rejected alternatives, measurements, and generated-code evidence. `NUMERIC_ROWFN_PLAN.md` records the numeric-binary migration and its narrower performance boundary. +`research/rowfn-x86-2026-08-07/README.md` records the later x86 regression reproduction, the +owned-output and indexed-source experiments, raw benchmark logs, and exact production IR/assembly. All three are branch-only working notes for agents. They are not intended to land with the API. The public design lives in these tracking issues, which now match the implementation: diff --git a/research/rowfn-x86-2026-08-07/README.md b/research/rowfn-x86-2026-08-07/README.md new file mode 100644 index 00000000000..e1df97740ba --- /dev/null +++ b/research/rowfn-x86-2026-08-07/README.md @@ -0,0 +1,239 @@ + + + +# RowFn owned-output and x86 numeric research + +This is the durable record for the investigation that produced the owned-output RowFn path. The +result is not that RowFn is inherently difficult to optimize. The declaration must distinguish an +independent returned value from a stateful output sink, and dense primitive inputs must cross a +validated indexed-source boundary that shared execution can lower directly. + +The selected implementation restores `i64` and `u64` varying multiplication to within about 1% of +the actual merge-base throughput. It does so without a numeric array downcast, `reduce_encoded` +override, numeric-owned allocation, or numeric-specific null and constant policy. + +## Revisions and environment + +- Merge-base baseline: `19f771f2a426103aa7d1bf7153a258bb1bab1e19`. +- Untouched sink-only RowFn: `35098c72118f1b555a24bd2f9b58b0400fa46dc5`. +- Selected implementation: `1a0a055c752b54448c8e1d54af032fe43acf8517`. +- Selected diff fingerprint: + `928e7a0baa2895609d102c98d110c21fb7a12e079b04195b85903277c71537a2`. + +The research branch has older tensor and spatial RowFn users. The result was ported rather than +rebased so that history remains intact. Its port also backports `map_checked_into`, which already +exists at the mergeable branch's base. + +```text +AMD Ryzen 9 7950X +1 socket, 16 physical cores, 32 threads +benchmark logical CPU: 8; SMT sibling: 24 +Linux CTCachyDesktop 7.1.6-1-cachyos, x86_64 +rustc 1.91.0, LLVM 21.1.2 +cargo 1.91.0 +``` + +The CPU reports AVX2 and AVX-512F/DQ/BW/VL. Builds used the default repository target and bench +profile without LTO, `target-cpu=native`, profile changes, or forced inlining. The scaling governor +was `performance`. Timed executions were pinned to CPU 8 and never overlapped compilation. + +```bash +taskset -c 8 "$BENCH" --bench --sample-count 100 --max-time 0.5 --color never \ + mul_i8_nonnull mul_u8_nonnull mul_i16_nonnull mul_u16_nonnull \ + mul_i32_nonnull mul_u32_nonnull mul_i64_nonnull mul_u64_nonnull \ + add_i64_nonnull add_i64_constant sub_i64_constant \ + mul_i32_constant mul_i32_nullable div_i64_nonnull +``` + +Every file in [`benchmarks`](benchmarks) is unedited Divan output wrapped in Markdown. It includes +fastest, slowest, median, mean, samples, and iterations rather than only the selected medians. + +## Stage 0: reproduction + +Order: baseline, candidate, baseline, candidate. Values are median microseconds. + +| Benchmark | Baseline 1 / 2 | Candidate 1 / 2 | Candidate/baseline | +| --- | ---: | ---: | ---: | +| `add_i64_constant` | 8.449 / 8.399 | 9.269 / 9.290 | 1.097 / 1.106 | +| `add_i64_nonnull` | 9.205 / 9.149 | 9.455 / 9.490 | 1.027 / 1.037 | +| `div_i64_nonnull` | 44.850 / 44.800 | 45.090 / 45.160 | 1.005 / 1.008 | +| `mul_i8_nonnull` | 6.184 / 6.199 | 4.694 / 4.699 | 0.759 / 0.758 | +| `mul_i16_nonnull` | 4.099 / 4.119 | 4.269 / 4.269 | 1.041 / 1.036 | +| `mul_i32_constant` | 26.420 / 26.430 | 18.880 / 18.870 | 0.715 / 0.714 | +| `mul_i32_nonnull` | 26.410 / 26.420 | 28.390 / 28.350 | 1.075 / 1.073 | +| `mul_i32_nullable` | 27.350 / 27.400 | 29.180 / 29.150 | 1.067 / 1.064 | +| `mul_i64_nonnull` | 23.220 / 23.200 | 30.020 / 30.080 | **1.293 / 1.297** | +| `mul_u8_nonnull` | 3.319 / 3.329 | 3.539 / 3.529 | 1.066 / 1.060 | +| `mul_u16_nonnull` | 2.599 / 2.599 | 2.429 / 2.429 | 0.935 / 0.935 | +| `mul_u32_nonnull` | 6.939 / 6.949 | 7.069 / 7.059 | 1.019 / 1.016 | +| `mul_u64_nonnull` | 19.210 / 19.190 | 30.430 / 30.490 | **1.584 / 1.589** | +| `sub_i64_constant` | 8.255 / 8.239 | 9.099 / 9.099 | 1.102 / 1.104 | + +The x86 regression reproduced. Raw runs are the four `stage0-*` files. + +## Stage 1: owned output without indexed input + +The closure returned `(output, failure)`, shared execution owned the store, and failure remained a +loop-local OR. This removed the numeric checked sink and materially improved 64-bit cases, but did +not solve the general problem. + +| Benchmark | Baseline 1 / 2 | Owned 1 / 2 | Owned/baseline | +| --- | ---: | ---: | ---: | +| `mul_i64_nonnull` | 23.20 / 23.26 | 25.65 / 25.59 | 1.106 / 1.100 | +| `mul_u64_nonnull` | 19.18 / 19.21 | 19.41 / 19.41 | 1.012 / 1.010 | +| `mul_i32_constant` | 26.43 / 26.44 | 32.36 / 32.38 | 1.224 / 1.225 | +| `mul_i32_nonnull` | 26.42 / 26.41 | 31.24 / 31.23 | 1.182 / 1.183 | +| `mul_i32_nullable` | 27.36 / 27.35 | 32.04 / 32.04 | 1.171 / 1.171 | + +The six `stage1-*` files contain the full matrix. This falsifies output ownership as a complete +explanation: it matters, but does not give LLVM the specialized kernel's input representation. + +## Stage 2: indexed dense input + +`IndexedElementTuple` lets a primitive pair expose `LaneZip<&[Left], &[Right]>` after shared +execution validates both varying lengths once. The generic owned executor calls +`map_checked_into`; numeric code still declares only row types, operation, failure, and error. + +| Benchmark | Baseline 1 / 2 | Indexed 1 / 2 | Candidate 1 / 2 | +| --- | ---: | ---: | ---: | +| `mul_i32_nonnull` | 26.39 / 26.41 | 26.58 / 26.60 | 28.34 / 28.36 | +| `mul_i32_nullable` | 27.37 / 27.38 | 27.41 / 27.43 | 29.20 / 29.17 | +| `mul_i64_nonnull` | 23.22 / 23.24 | 23.43 / 23.44 | 30.02 / 30.10 | +| `mul_u64_nonnull` | 19.22 / 19.21 | 19.41 / 19.42 | 30.41 / 30.43 | +| `div_i64_nonnull` | 44.84 / 44.87 | 45.07 / 45.03 | 45.07 / 45.12 | +| `mul_i32_constant` | 26.42 / 26.43 | 32.38 / 32.39 | 18.88 / 18.88 | + +The indexed source closed the varying and nullable gap. It did not affect mixed constants, which +exposed the next compiler-sensitive detail. + +## Store placement and the `Copy` ablation + +Moving the output store before the failure OR changed `mul_i32_constant` from about 32.38 to +18.68 microseconds. Final assembly records overflow `setb`, output store, then loop-carried OR. This +is compiler scheduling sensitivity, not a semantic difference. + +Adding the descriptive `Output: Copy` bound regressed that case to 29.88/29.86 microseconds. +Replacing it with compile-time `!needs_drop::()` returned it to 18.65/18.67. A generic store +helper did not repair the `Copy` case. The executor needs only the no-drop property for safe panic +cleanup; it never copies an output. The API therefore enforces the actual requirement without the +measured optimizer-visible bound. Re-test this workaround whenever LLVM changes. + +## Final results + +Order: baseline, final, candidate, repeated twice. Values are median microseconds. + +| Benchmark | Baseline 1 / 2 | Candidate 1 / 2 | Final 1 / 2 | Final/baseline | +| --- | ---: | ---: | ---: | ---: | +| `add_i64_constant` | 8.399 / 8.449 | 9.310 / 9.289 | 9.269 / 9.279 | 1.104 / 1.098 | +| `add_i64_nonnull` | 9.159 / 9.239 | 9.374 / 9.449 | 9.379 / 9.389 | 1.024 / 1.016 | +| `div_i64_nonnull` | 44.820 / 44.860 | 45.040 / 45.080 | 45.020 / 45.060 | 1.004 / 1.004 | +| `mul_i8_nonnull` | 6.209 / 6.199 | 4.719 / 4.699 | 6.389 / 6.409 | 1.029 / 1.034 | +| `mul_i16_nonnull` | 4.099 / 4.109 | 4.269 / 4.269 | 4.265 / 4.299 | 1.040 / 1.046 | +| `mul_i32_constant` | 26.440 / 26.440 | 18.890 / 18.840 | 18.690 / 18.700 | **0.707 / 0.707** | +| `mul_i32_nonnull` | 26.410 / 26.420 | 28.350 / 28.390 | 26.590 / 26.640 | 1.007 / 1.008 | +| `mul_i32_nullable` | 27.380 / 27.360 | 29.170 / 29.170 | 27.400 / 27.440 | 1.001 / 1.003 | +| `mul_i64_nonnull` | 23.200 / 23.350 | 30.010 / 30.050 | 23.460 / 23.430 | **1.011 / 1.003** | +| `mul_u8_nonnull` | 3.319 / 3.319 | 3.545 / 3.519 | 3.514 / 3.549 | 1.059 / 1.069 | +| `mul_u16_nonnull` | 2.609 / 2.609 | 2.429 / 2.429 | 2.789 / 2.810 | 1.069 / 1.077 | +| `mul_u32_nonnull` | 6.949 / 6.959 | 7.060 / 7.059 | 7.129 / 7.149 | 1.026 / 1.027 | +| `mul_u64_nonnull` | 19.180 / 19.210 | 30.370 / 30.400 | 19.370 / 19.380 | **1.010 / 1.009** | +| `sub_i64_constant` | 8.239 / 8.259 | 9.114 / 9.079 | 9.149 / 9.159 | 1.110 / 1.109 | + +The six `land-*` logs preserve every final run. Narrow widths avoid the rejected zipped-iterator +experiment's 3x to 9x losses. Constant add/sub retain the untouched candidate's roughly 10% gap; +constant multiplication is faster than merge base. Division stays at parity. + +## Generated code: confirmed evidence + +```bash +CARGO_TARGET_DIR="$TARGET" cargo rustc -p vortex-array --lib --profile bench -- \ + --emit=llvm-ir,asm -C codegen-units=1 -C remark=loop-vectorize +``` + +Full output was about 1.85 GiB IR plus 1.02 GiB assembly and was deleted after extracting exact +production monomorphs into [`codegen`](codegen). These are not fixture or benchmark control loops. + +Baseline, candidate, owned, and final signed `i64` use a scalar one-lane loop: one high/low `imulq`, +one store, `sarq`/`xorq` overflow evidence, register OR, and one backedge. Unsigned `u64` uses two +independent scalar `mulq` groups per backedge plus an odd remainder. Neither final loop has a hot +call, panic edge, bounds check, runtime alias check, or vector body. The second input length check is +an `llvm.assume`; loads and stores carry disjoint alias metadata; failure is a register `phi`. + +Therefore host SIMD did not hide a deficient loop. The default build did not enable optional native +AVX features, and LLVM selected the same essential scalar high-half strategy as merge base. See +[`base summary`](codegen/base-codegen-summary.md), +[`final i64 assembly`](codegen/final-i64-mul-dense-s.md), and +[`final u64 assembly`](codegen/final-u64-mul-dense-s.md). + +`-C remark=loop-vectorize` emitted no remark attributable to the exact dense production loop. The +constant fallback source line had successes for other monomorphs and duplicated cost-model misses, +but diagnostics lacked function identity. Exact IR proves the measured specialization is scalar; +it cannot assign those remarks to it. The merge-base focused remark rebuild was cancelled, so no +merge-base missed-vectorization reason is claimed. + +## Findings + +Confirmed: + +- Bounds checks are not the all-varying blocker; candidate dense multiply had no hot bounds edge. +- `SinkResult` already reduced to a register OR and did not impose a per-row `Result`. +- Output ownership materially helped but was insufficient alone. +- A typed indexed source restored stable parity for varying primitive tuples. +- Store-before-OR and omission of a `Copy` bound materially affect LLVM 21.1.2 constant codegen. +- The default x86 target prefers scalar high-half 64-bit multiply; SIMD is not the recovered speed. + +Still inference: + +- No single alias defect explains the original gap. Baseline and candidate had useful metadata too. +- The nearly identical dense inner loops do not explain all end-to-end timing. Surrounding control + flow, placement, and instruction-cache effects remain candidates. +- Unattributed source-line remarks do not prove a missed-vectorization reason for one monomorph. + +Rejected controls: checked unchecked access only partially helped and regressed some `u8` runs; +direct failure accumulation matched existing IR; safe zipped iterators caused 3x to 9x narrow +losses; a numeric `reduce_encoded` fast path recovered speed by duplicating shared policy; and a +primitive-binary visitor seam moved that specialization into generic execution. Earlier Apple work, +including the non-affine `index & mask` failure, remains in +[`NUMERIC_ROWFN_PLAN.md`](../../NUMERIC_ROWFN_PLAN.md). + +## Why both visitor methods exist + +`visit_prepared_deferred` represents an independent owned value and OR-reducible failure per row. +The executor allocates contiguous output, owns the store, and can use a typed indexed source. It is +intentionally limited to indexed inputs, fixed no-drop output, and a batch-deferred row error. + +`visit_prepared_into` represents stateful construction: shared buffers, runtime-shaped layouts, +multiple coordinated builders, skip-capable output, drop-requiring values, non-indexed tuples, and +ordinary immediate or deferred `SinkResult` forms. Encoding those through the owned method would +either hide a mutable builder reference inside a supposed value, allocate a temporary per row, +forbid legitimate output, or duplicate lifting. Encoding numeric output only through the sink loses +the fact that each value and store are independent. These are distinct capabilities. + +## Indexed source, specialization, and safety + +`InputElement` is open and many elements are not contiguous. Sealed `ElementTuple` is the safe +composition point for unchecked reads after one length validation. Stable Rust cannot overlap a +blanket fallback for every tuple with a more specific associated dense source without +specialization. Runtime erasure would obscure the source type LLVM needs. The indexed capability is +therefore explicit and opt-in; only the proven primitive pair implements it today. + +The executor reserves `row_count` slots and exposes exactly that many `MaybeUninit` values. It +validates varying lengths before `LaneZip`; `map_checked_into` validates output length. Either loop +writes every slot exactly once before `set_len`. On panic the vector length remains zero, and the +compile-time no-drop assertion makes abandoning initialized slots safe. Deferred errors are examined +only after initialization. Nullable lifting retries a deferred error over valid rows, so a failure +shaped value behind null cannot surface. + +## Open improvements + +- Investigate infallible owned output only with a measured caller; avoid a speculative result tree. +- Revisit constant add/sub only with exact production IR and a stable regression. +- Add indexed tuple/element families only for real consumers with a safe source. +- Re-run the store-order and `Copy` ablations after LLVM upgrades. +- Produce an upstream LLVM reproducer for those compiler sensitivities. +- Preserve assembly checks because throughput can hide compensating target-specific instructions. + +The selected branch passed focused checks, 87 numeric tests, 3,385 nextest tests with one skipped, +73 doctests with 13 ignored, nightly formatting, all-target/all-feature clippy, and `diff --check`. +One intermediate 1.85 GiB IR copy hit `ENOSPC`; exact-final codegen later completed. The requested +`ROWFN_FIRST_PR_PROMPT.md` was absent from the repository, fetched refs, home tree, and worktrees. diff --git a/research/rowfn-x86-2026-08-07/benchmarks/land-base-1.md b/research/rowfn-x86-2026-08-07/benchmarks/land-base-1.md new file mode 100644 index 00000000000..2759e0fa9e5 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/land-base-1.md @@ -0,0 +1,39 @@ + + + +# `land-base-1` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.059 µs │ 781.5 µs │ 8.399 µs │ 16.16 µs │ 100 │ 100 +│ 4.065 Gitem/s │ 41.92 Mitem/s │ 3.901 Gitem/s │ 2.027 Gitem/s │ │ +├─ add_i64_nonnull 9.079 µs │ 29.55 µs │ 9.159 µs │ 9.466 µs │ 100 │ 100 +│ 3.608 Gitem/s │ 1.108 Gitem/s │ 3.577 Gitem/s │ 3.461 Gitem/s │ │ +├─ div_i64_nonnull 44.73 µs │ 78.44 µs │ 44.82 µs │ 45.35 µs │ 100 │ 100 +│ 732.4 Mitem/s │ 417.6 Mitem/s │ 731 Mitem/s │ 722.5 Mitem/s │ │ +├─ mul_i8_nonnull 5.819 µs │ 72.73 µs │ 6.209 µs │ 6.939 µs │ 100 │ 100 +│ 5.63 Gitem/s │ 450.5 Mitem/s │ 5.276 Gitem/s │ 4.721 Gitem/s │ │ +├─ mul_i16_nonnull 4.039 µs │ 66.44 µs │ 4.099 µs │ 4.731 µs │ 100 │ 100 +│ 8.111 Gitem/s │ 493.1 Mitem/s │ 7.992 Gitem/s │ 6.926 Gitem/s │ │ +├─ mul_i32_constant 26.37 µs │ 56.27 µs │ 26.44 µs │ 26.97 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 582.2 Mitem/s │ 1.238 Gitem/s │ 1.214 Gitem/s │ │ +├─ mul_i32_nonnull 26.35 µs │ 38.26 µs │ 26.41 µs │ 26.64 µs │ 100 │ 100 +│ 1.243 Gitem/s │ 856.4 Mitem/s │ 1.24 Gitem/s │ 1.229 Gitem/s │ │ +├─ mul_i32_nullable 27.28 µs │ 340.1 µs │ 27.38 µs │ 30.62 µs │ 100 │ 100 +│ 1.2 Gitem/s │ 96.34 Mitem/s │ 1.196 Gitem/s │ 1.069 Gitem/s │ │ +├─ mul_i64_nonnull 23.11 µs │ 45.96 µs │ 23.2 µs │ 23.55 µs │ 100 │ 100 +│ 1.417 Gitem/s │ 712.9 Mitem/s │ 1.411 Gitem/s │ 1.391 Gitem/s │ │ +├─ mul_u8_nonnull 3.259 µs │ 52.07 µs │ 3.319 µs │ 3.838 µs │ 100 │ 100 +│ 10.05 Gitem/s │ 629.1 Mitem/s │ 9.87 Gitem/s │ 8.535 Gitem/s │ │ +├─ mul_u16_nonnull 2.539 µs │ 30.81 µs │ 2.609 µs │ 2.888 µs │ 100 │ 100 +│ 12.9 Gitem/s │ 1.063 Gitem/s │ 12.55 Gitem/s │ 11.34 Gitem/s │ │ +├─ mul_u32_nonnull 6.879 µs │ 26.44 µs │ 6.949 µs │ 7.183 µs │ 100 │ 100 +│ 4.762 Gitem/s │ 1.238 Gitem/s │ 4.714 Gitem/s │ 4.561 Gitem/s │ │ +├─ mul_u64_nonnull 19.11 µs │ 41.4 µs │ 19.18 µs │ 19.53 µs │ 100 │ 100 +│ 1.713 Gitem/s │ 791.4 Mitem/s │ 1.707 Gitem/s │ 1.677 Gitem/s │ │ +╰─ sub_i64_constant 8.109 µs │ 40.38 µs │ 8.239 µs │ 8.604 µs │ 100 │ 100 + 4.04 Gitem/s │ 811.2 Mitem/s │ 3.976 Gitem/s │ 3.808 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/land-base-2.md b/research/rowfn-x86-2026-08-07/benchmarks/land-base-2.md new file mode 100644 index 00000000000..ec40d57a6e2 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/land-base-2.md @@ -0,0 +1,39 @@ + + + +# `land-base-2` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.319 µs │ 39.44 µs │ 8.449 µs │ 8.813 µs │ 100 │ 100 +│ 3.938 Gitem/s │ 830.8 Mitem/s │ 3.877 Gitem/s │ 3.718 Gitem/s │ │ +├─ add_i64_nonnull 9.169 µs │ 12.56 µs │ 9.239 µs │ 9.304 µs │ 100 │ 100 +│ 3.573 Gitem/s │ 2.608 Gitem/s │ 3.546 Gitem/s │ 3.521 Gitem/s │ │ +├─ div_i64_nonnull 44.77 µs │ 50.37 µs │ 44.86 µs │ 45.15 µs │ 100 │ 100 +│ 731.7 Mitem/s │ 650.4 Mitem/s │ 730.2 Mitem/s │ 725.7 Mitem/s │ │ +├─ mul_i8_nonnull 5.829 µs │ 8.179 µs │ 6.199 µs │ 6.265 µs │ 100 │ 100 +│ 5.62 Gitem/s │ 4.005 Gitem/s │ 5.285 Gitem/s │ 5.23 Gitem/s │ │ +├─ mul_i16_nonnull 4.049 µs │ 8.029 µs │ 4.109 µs │ 4.15 µs │ 100 │ 100 +│ 8.091 Gitem/s │ 4.08 Gitem/s │ 7.973 Gitem/s │ 7.895 Gitem/s │ │ +├─ mul_i32_constant 26.37 µs │ 35.43 µs │ 26.44 µs │ 26.77 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 924.6 Mitem/s │ 1.239 Gitem/s │ 1.223 Gitem/s │ │ +├─ mul_i32_nonnull 26.37 µs │ 35.11 µs │ 26.42 µs │ 26.84 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 933 Mitem/s │ 1.239 Gitem/s │ 1.22 Gitem/s │ │ +├─ mul_i32_nullable 27.26 µs │ 40.16 µs │ 27.36 µs │ 27.78 µs │ 100 │ 100 +│ 1.201 Gitem/s │ 815.9 Mitem/s │ 1.197 Gitem/s │ 1.179 Gitem/s │ │ +├─ mul_i64_nonnull 23.24 µs │ 32.01 µs │ 23.35 µs │ 23.63 µs │ 100 │ 100 +│ 1.409 Gitem/s │ 1.023 Gitem/s │ 1.402 Gitem/s │ 1.386 Gitem/s │ │ +├─ mul_u8_nonnull 3.26 µs │ 4.759 µs │ 3.319 µs │ 3.349 µs │ 100 │ 100 +│ 10.04 Gitem/s │ 6.884 Gitem/s │ 9.87 Gitem/s │ 9.783 Gitem/s │ │ +├─ mul_u16_nonnull 2.53 µs │ 7.289 µs │ 2.609 µs │ 2.664 µs │ 100 │ 100 +│ 12.94 Gitem/s │ 4.495 Gitem/s │ 12.55 Gitem/s │ 12.29 Gitem/s │ │ +├─ mul_u32_nonnull 6.889 µs │ 12.89 µs │ 6.959 µs │ 7.027 µs │ 100 │ 100 +│ 4.756 Gitem/s │ 2.54 Gitem/s │ 4.708 Gitem/s │ 4.663 Gitem/s │ │ +├─ mul_u64_nonnull 19.13 µs │ 22.83 µs │ 19.21 µs │ 19.29 µs │ 100 │ 100 +│ 1.712 Gitem/s │ 1.435 Gitem/s │ 1.704 Gitem/s │ 1.698 Gitem/s │ │ +╰─ sub_i64_constant 8.139 µs │ 11.06 µs │ 8.259 µs │ 8.297 µs │ 100 │ 100 + 4.025 Gitem/s │ 2.96 Gitem/s │ 3.967 Gitem/s │ 3.949 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/land-candidate-1.md b/research/rowfn-x86-2026-08-07/benchmarks/land-candidate-1.md new file mode 100644 index 00000000000..8ec200b7df1 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/land-candidate-1.md @@ -0,0 +1,39 @@ + + + +# `land-candidate-1` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.989 µs │ 1.016 ms │ 9.31 µs │ 19.44 µs │ 100 │ 100 +│ 3.645 Gitem/s │ 32.25 Mitem/s │ 3.519 Gitem/s │ 1.684 Gitem/s │ │ +├─ add_i64_nonnull 9.279 µs │ 13.06 µs │ 9.374 µs │ 9.443 µs │ 100 │ 100 +│ 3.531 Gitem/s │ 2.508 Gitem/s │ 3.495 Gitem/s │ 3.469 Gitem/s │ │ +├─ div_i64_nonnull 44.97 µs │ 63.03 µs │ 45.04 µs │ 45.41 µs │ 100 │ 100 +│ 728.5 Mitem/s │ 519.7 Mitem/s │ 727.3 Mitem/s │ 721.5 Mitem/s │ │ +├─ mul_i8_nonnull 4.649 µs │ 60.73 µs │ 4.719 µs │ 5.455 µs │ 100 │ 100 +│ 7.047 Gitem/s │ 539.4 Mitem/s │ 6.942 Gitem/s │ 6.006 Gitem/s │ │ +├─ mul_i16_nonnull 4.219 µs │ 71.24 µs │ 4.269 µs │ 4.959 µs │ 100 │ 100 +│ 7.765 Gitem/s │ 459.9 Mitem/s │ 7.674 Gitem/s │ 6.607 Gitem/s │ │ +├─ mul_i32_constant 18.79 µs │ 72.37 µs │ 18.89 µs │ 19.53 µs │ 100 │ 100 +│ 1.742 Gitem/s │ 452.7 Mitem/s │ 1.734 Gitem/s │ 1.677 Gitem/s │ │ +├─ mul_i32_nonnull 28.24 µs │ 33.43 µs │ 28.35 µs │ 28.48 µs │ 100 │ 100 +│ 1.159 Gitem/s │ 980.1 Mitem/s │ 1.155 Gitem/s │ 1.15 Gitem/s │ │ +├─ mul_i32_nullable 29.01 µs │ 234.9 µs │ 29.17 µs │ 31.37 µs │ 100 │ 100 +│ 1.129 Gitem/s │ 139.4 Mitem/s │ 1.122 Gitem/s │ 1.044 Gitem/s │ │ +├─ mul_i64_nonnull 29.63 µs │ 52.67 µs │ 30.01 µs │ 30.5 µs │ 100 │ 100 +│ 1.105 Gitem/s │ 622 Mitem/s │ 1.091 Gitem/s │ 1.074 Gitem/s │ │ +├─ mul_u8_nonnull 3.459 µs │ 14.69 µs │ 3.545 µs │ 3.659 µs │ 100 │ 100 +│ 9.471 Gitem/s │ 2.229 Gitem/s │ 9.242 Gitem/s │ 8.953 Gitem/s │ │ +├─ mul_u16_nonnull 2.369 µs │ 13.01 µs │ 2.429 µs │ 2.615 µs │ 100 │ 100 +│ 13.82 Gitem/s │ 2.516 Gitem/s │ 13.48 Gitem/s │ 12.52 Gitem/s │ │ +├─ mul_u32_nonnull 6.999 µs │ 18.45 µs │ 7.06 µs │ 7.181 µs │ 100 │ 100 +│ 4.681 Gitem/s │ 1.775 Gitem/s │ 4.641 Gitem/s │ 4.562 Gitem/s │ │ +├─ mul_u64_nonnull 30.27 µs │ 43.99 µs │ 30.37 µs │ 30.65 µs │ 100 │ 100 +│ 1.082 Gitem/s │ 744.8 Mitem/s │ 1.078 Gitem/s │ 1.068 Gitem/s │ │ +╰─ sub_i64_constant 8.959 µs │ 31.53 µs │ 9.114 µs │ 9.385 µs │ 100 │ 100 + 3.657 Gitem/s │ 1.038 Gitem/s │ 3.595 Gitem/s │ 3.491 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/land-candidate-2.md b/research/rowfn-x86-2026-08-07/benchmarks/land-candidate-2.md new file mode 100644 index 00000000000..84f969b8646 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/land-candidate-2.md @@ -0,0 +1,39 @@ + + + +# `land-candidate-2` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.129 µs │ 48.92 µs │ 9.289 µs │ 9.744 µs │ 100 │ 100 +│ 3.589 Gitem/s │ 669.8 Mitem/s │ 3.527 Gitem/s │ 3.362 Gitem/s │ │ +├─ add_i64_nonnull 9.349 µs │ 10.43 µs │ 9.449 µs │ 9.46 µs │ 100 │ 100 +│ 3.504 Gitem/s │ 3.141 Gitem/s │ 3.467 Gitem/s │ 3.463 Gitem/s │ │ +├─ div_i64_nonnull 44.99 µs │ 50.41 µs │ 45.08 µs │ 45.29 µs │ 100 │ 100 +│ 728.1 Mitem/s │ 650 Mitem/s │ 726.8 Mitem/s │ 723.5 Mitem/s │ │ +├─ mul_i8_nonnull 4.639 µs │ 7.969 µs │ 4.699 µs │ 4.761 µs │ 100 │ 100 +│ 7.062 Gitem/s │ 4.111 Gitem/s │ 6.972 Gitem/s │ 6.881 Gitem/s │ │ +├─ mul_i16_nonnull 4.209 µs │ 7.669 µs │ 4.269 µs │ 4.308 µs │ 100 │ 100 +│ 7.783 Gitem/s │ 4.272 Gitem/s │ 7.674 Gitem/s │ 7.605 Gitem/s │ │ +├─ mul_i32_constant 18.75 µs │ 22.77 µs │ 18.84 µs │ 18.95 µs │ 100 │ 100 +│ 1.746 Gitem/s │ 1.439 Gitem/s │ 1.738 Gitem/s │ 1.728 Gitem/s │ │ +├─ mul_i32_nonnull 28.27 µs │ 45.57 µs │ 28.39 µs │ 28.65 µs │ 100 │ 100 +│ 1.158 Gitem/s │ 718.9 Mitem/s │ 1.154 Gitem/s │ 1.143 Gitem/s │ │ +├─ mul_i32_nullable 29.04 µs │ 43.06 µs │ 29.17 µs │ 29.41 µs │ 100 │ 100 +│ 1.128 Gitem/s │ 760.8 Mitem/s │ 1.122 Gitem/s │ 1.113 Gitem/s │ │ +├─ mul_i64_nonnull 29.7 µs │ 35.65 µs │ 30.05 µs │ 30.18 µs │ 100 │ 100 +│ 1.103 Gitem/s │ 918.9 Mitem/s │ 1.09 Gitem/s │ 1.085 Gitem/s │ │ +├─ mul_u8_nonnull 3.459 µs │ 4.579 µs │ 3.519 µs │ 3.532 µs │ 100 │ 100 +│ 9.471 Gitem/s │ 7.154 Gitem/s │ 9.309 Gitem/s │ 9.275 Gitem/s │ │ +├─ mul_u16_nonnull 2.359 µs │ 3.269 µs │ 2.429 µs │ 2.441 µs │ 100 │ 100 +│ 13.88 Gitem/s │ 10.02 Gitem/s │ 13.48 Gitem/s │ 13.42 Gitem/s │ │ +├─ mul_u32_nonnull 6.999 µs │ 11.21 µs │ 7.059 µs │ 7.105 µs │ 100 │ 100 +│ 4.681 Gitem/s │ 2.92 Gitem/s │ 4.641 Gitem/s │ 4.611 Gitem/s │ │ +├─ mul_u64_nonnull 30.33 µs │ 34.65 µs │ 30.4 µs │ 30.53 µs │ 100 │ 100 +│ 1.08 Gitem/s │ 945.4 Mitem/s │ 1.077 Gitem/s │ 1.073 Gitem/s │ │ +╰─ sub_i64_constant 8.949 µs │ 12.59 µs │ 9.079 µs │ 9.155 µs │ 100 │ 100 + 3.661 Gitem/s │ 2.6 Gitem/s │ 3.608 Gitem/s │ 3.578 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/land-final-1.md b/research/rowfn-x86-2026-08-07/benchmarks/land-final-1.md new file mode 100644 index 00000000000..fc941735f9c --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/land-final-1.md @@ -0,0 +1,39 @@ + + + +# `land-final-1` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.979 µs │ 88.65 µs │ 9.269 µs │ 10.12 µs │ 100 │ 100 +│ 3.649 Gitem/s │ 369.6 Mitem/s │ 3.534 Gitem/s │ 3.237 Gitem/s │ │ +├─ add_i64_nonnull 9.299 µs │ 13.4 µs │ 9.379 µs │ 9.444 µs │ 100 │ 100 +│ 3.523 Gitem/s │ 2.443 Gitem/s │ 3.493 Gitem/s │ 3.469 Gitem/s │ │ +├─ div_i64_nonnull 44.96 µs │ 55 µs │ 45.02 µs │ 45.48 µs │ 100 │ 100 +│ 728.8 Mitem/s │ 595.6 Mitem/s │ 727.6 Mitem/s │ 720.4 Mitem/s │ │ +├─ mul_i8_nonnull 5.959 µs │ 44.56 µs │ 6.389 µs │ 6.855 µs │ 100 │ 100 +│ 5.498 Gitem/s │ 735.3 Mitem/s │ 5.128 Gitem/s │ 4.78 Gitem/s │ │ +├─ mul_i16_nonnull 4.199 µs │ 7.599 µs │ 4.265 µs │ 4.321 µs │ 100 │ 100 +│ 7.802 Gitem/s │ 4.311 Gitem/s │ 7.682 Gitem/s │ 7.581 Gitem/s │ │ +├─ mul_i32_constant 18.6 µs │ 22.22 µs │ 18.69 µs │ 18.81 µs │ 100 │ 100 +│ 1.76 Gitem/s │ 1.474 Gitem/s │ 1.753 Gitem/s │ 1.741 Gitem/s │ │ +├─ mul_i32_nonnull 26.52 µs │ 34.76 µs │ 26.59 µs │ 26.77 µs │ 100 │ 100 +│ 1.235 Gitem/s │ 942.6 Mitem/s │ 1.231 Gitem/s │ 1.223 Gitem/s │ │ +├─ mul_i32_nullable 27.3 µs │ 51.11 µs │ 27.4 µs │ 27.77 µs │ 100 │ 100 +│ 1.199 Gitem/s │ 641 Mitem/s │ 1.195 Gitem/s │ 1.179 Gitem/s │ │ +├─ mul_i64_nonnull 23.37 µs │ 27.17 µs │ 23.46 µs │ 23.56 µs │ 100 │ 100 +│ 1.401 Gitem/s │ 1.206 Gitem/s │ 1.396 Gitem/s │ 1.39 Gitem/s │ │ +├─ mul_u8_nonnull 3.459 µs │ 59.52 µs │ 3.514 µs │ 4.079 µs │ 100 │ 100 +│ 9.471 Gitem/s │ 550.5 Mitem/s │ 9.322 Gitem/s │ 8.033 Gitem/s │ │ +├─ mul_u16_nonnull 2.709 µs │ 3.799 µs │ 2.789 µs │ 2.798 µs │ 100 │ 100 +│ 12.09 Gitem/s │ 8.623 Gitem/s │ 11.74 Gitem/s │ 11.71 Gitem/s │ │ +├─ mul_u32_nonnull 7.049 µs │ 10.9 µs │ 7.129 µs │ 7.19 µs │ 100 │ 100 +│ 4.648 Gitem/s │ 3.003 Gitem/s │ 4.595 Gitem/s │ 4.556 Gitem/s │ │ +├─ mul_u64_nonnull 19.26 µs │ 22.46 µs │ 19.37 µs │ 19.45 µs │ 100 │ 100 +│ 1.7 Gitem/s │ 1.458 Gitem/s │ 1.691 Gitem/s │ 1.683 Gitem/s │ │ +╰─ sub_i64_constant 9.009 µs │ 13.94 µs │ 9.149 µs │ 9.215 µs │ 100 │ 100 + 3.636 Gitem/s │ 2.348 Gitem/s │ 3.581 Gitem/s │ 3.555 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/land-final-2.md b/research/rowfn-x86-2026-08-07/benchmarks/land-final-2.md new file mode 100644 index 00000000000..6588d08f1e1 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/land-final-2.md @@ -0,0 +1,39 @@ + + + +# `land-final-2` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.149 µs │ 73.57 µs │ 9.279 µs │ 9.987 µs │ 100 │ 100 +│ 3.581 Gitem/s │ 445.3 Mitem/s │ 3.531 Gitem/s │ 3.28 Gitem/s │ │ +├─ add_i64_nonnull 9.319 µs │ 14.28 µs │ 9.389 µs │ 9.475 µs │ 100 │ 100 +│ 3.515 Gitem/s │ 2.293 Gitem/s │ 3.489 Gitem/s │ 3.458 Gitem/s │ │ +├─ div_i64_nonnull 44.96 µs │ 62.15 µs │ 45.06 µs │ 45.46 µs │ 100 │ 100 +│ 728.6 Mitem/s │ 527.2 Mitem/s │ 727 Mitem/s │ 720.7 Mitem/s │ │ +├─ mul_i8_nonnull 5.979 µs │ 41.01 µs │ 6.409 µs │ 6.853 µs │ 100 │ 100 +│ 5.479 Gitem/s │ 798.8 Mitem/s │ 5.112 Gitem/s │ 4.78 Gitem/s │ │ +├─ mul_i16_nonnull 4.229 µs │ 5.739 µs │ 4.299 µs │ 4.314 µs │ 100 │ 100 +│ 7.746 Gitem/s │ 5.708 Gitem/s │ 7.62 Gitem/s │ 7.595 Gitem/s │ │ +├─ mul_i32_constant 18.6 µs │ 23.91 µs │ 18.7 µs │ 18.79 µs │ 100 │ 100 +│ 1.76 Gitem/s │ 1.369 Gitem/s │ 1.751 Gitem/s │ 1.743 Gitem/s │ │ +├─ mul_i32_nonnull 26.56 µs │ 30.13 µs │ 26.64 µs │ 26.74 µs │ 100 │ 100 +│ 1.233 Gitem/s │ 1.087 Gitem/s │ 1.229 Gitem/s │ 1.225 Gitem/s │ │ +├─ mul_i32_nullable 27.35 µs │ 42.74 µs │ 27.44 µs │ 27.69 µs │ 100 │ 100 +│ 1.197 Gitem/s │ 766.5 Mitem/s │ 1.193 Gitem/s │ 1.183 Gitem/s │ │ +├─ mul_i64_nonnull 23.31 µs │ 27.31 µs │ 23.43 µs │ 23.56 µs │ 100 │ 100 +│ 1.405 Gitem/s │ 1.199 Gitem/s │ 1.398 Gitem/s │ 1.39 Gitem/s │ │ +├─ mul_u8_nonnull 3.479 µs │ 52.76 µs │ 3.549 µs │ 4.103 µs │ 100 │ 100 +│ 9.416 Gitem/s │ 620.9 Mitem/s │ 9.231 Gitem/s │ 7.984 Gitem/s │ │ +├─ mul_u16_nonnull 2.739 µs │ 3.799 µs │ 2.81 µs │ 2.824 µs │ 100 │ 100 +│ 11.96 Gitem/s │ 8.623 Gitem/s │ 11.66 Gitem/s │ 11.6 Gitem/s │ │ +├─ mul_u32_nonnull 7.089 µs │ 10.28 µs │ 7.149 µs │ 7.207 µs │ 100 │ 100 +│ 4.621 Gitem/s │ 3.184 Gitem/s │ 4.583 Gitem/s │ 4.546 Gitem/s │ │ +├─ mul_u64_nonnull 19.32 µs │ 23.55 µs │ 19.38 µs │ 19.47 µs │ 100 │ 100 +│ 1.695 Gitem/s │ 1.391 Gitem/s │ 1.689 Gitem/s │ 1.682 Gitem/s │ │ +╰─ sub_i64_constant 8.939 µs │ 30.07 µs │ 9.159 µs │ 9.386 µs │ 100 │ 100 + 3.665 Gitem/s │ 1.089 Gitem/s │ 3.577 Gitem/s │ 3.49 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage0-base-1.md b/research/rowfn-x86-2026-08-07/benchmarks/stage0-base-1.md new file mode 100644 index 00000000000..b12390b54c2 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage0-base-1.md @@ -0,0 +1,38 @@ + + + +# `stage0-base-1` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.319 µs │ 62.83 µs │ 8.449 µs │ 9.036 µs │ 100 │ 100 +│ 3.938 Gitem/s │ 521.5 Mitem/s │ 3.877 Gitem/s │ 3.626 Gitem/s │ │ +├─ add_i64_nonnull 9.139 µs │ 13.11 µs │ 9.205 µs │ 9.316 µs │ 100 │ 100 +│ 3.585 Gitem/s │ 2.497 Gitem/s │ 3.559 Gitem/s │ 3.517 Gitem/s │ │ +├─ div_i64_nonnull 44.78 µs │ 66.09 µs │ 44.85 µs │ 45.23 µs │ 100 │ 100 +│ 731.5 Mitem/s │ 495.8 Mitem/s │ 730.4 Mitem/s │ 724.4 Mitem/s │ │ +├─ mul_i8_nonnull 5.799 µs │ 17.72 µs │ 6.184 µs │ 6.39 µs │ 100 │ 100 +│ 5.649 Gitem/s │ 1.848 Gitem/s │ 5.298 Gitem/s │ 5.127 Gitem/s │ │ +├─ mul_i16_nonnull 4.009 µs │ 10.13 µs │ 4.099 µs │ 4.214 µs │ 100 │ 100 +│ 8.172 Gitem/s │ 3.234 Gitem/s │ 7.992 Gitem/s │ 7.774 Gitem/s │ │ +├─ mul_i32_constant 26.35 µs │ 30.02 µs │ 26.42 µs │ 26.53 µs │ 100 │ 100 +│ 1.243 Gitem/s │ 1.091 Gitem/s │ 1.239 Gitem/s │ 1.234 Gitem/s │ │ +├─ mul_i32_nonnull 26.36 µs │ 30.26 µs │ 26.41 µs │ 26.51 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 1.082 Gitem/s │ 1.24 Gitem/s │ 1.235 Gitem/s │ │ +├─ mul_i32_nullable 27.26 µs │ 48.92 µs │ 27.35 µs │ 27.7 µs │ 100 │ 100 +│ 1.202 Gitem/s │ 669.8 Mitem/s │ 1.197 Gitem/s │ 1.182 Gitem/s │ │ +├─ mul_i64_nonnull 23.13 µs │ 28.03 µs │ 23.22 µs │ 23.37 µs │ 100 │ 100 +│ 1.416 Gitem/s │ 1.168 Gitem/s │ 1.41 Gitem/s │ 1.402 Gitem/s │ │ +├─ mul_u8_nonnull 3.259 µs │ 6.989 µs │ 3.319 µs │ 3.365 µs │ 100 │ 100 +│ 10.05 Gitem/s │ 4.687 Gitem/s │ 9.87 Gitem/s │ 9.735 Gitem/s │ │ +├─ mul_u16_nonnull 2.539 µs │ 3.489 µs │ 2.599 µs │ 2.613 µs │ 100 │ 100 +│ 12.9 Gitem/s │ 9.389 Gitem/s │ 12.6 Gitem/s │ 12.53 Gitem/s │ │ +├─ mul_u32_nonnull 6.879 µs │ 12.13 µs │ 6.939 µs │ 7.009 µs │ 100 │ 100 +│ 4.762 Gitem/s │ 2.699 Gitem/s │ 4.721 Gitem/s │ 4.674 Gitem/s │ │ +├─ mul_u64_nonnull 19.14 µs │ 23.82 µs │ 19.21 µs │ 19.32 µs │ 100 │ 100 +│ 1.711 Gitem/s │ 1.375 Gitem/s │ 1.704 Gitem/s │ 1.695 Gitem/s │ │ +╰─ sub_i64_constant 8.129 µs │ 12.18 µs │ 8.255 µs │ 8.34 µs │ 100 │ 100 + 4.03 Gitem/s │ 2.688 Gitem/s │ 3.969 Gitem/s │ 3.928 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage0-base-2.md b/research/rowfn-x86-2026-08-07/benchmarks/stage0-base-2.md new file mode 100644 index 00000000000..c2254e0238f --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage0-base-2.md @@ -0,0 +1,38 @@ + + + +# `stage0-base-2` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.089 µs │ 63.25 µs │ 8.399 µs │ 8.966 µs │ 100 │ 100 +│ 4.05 Gitem/s │ 518 Mitem/s │ 3.901 Gitem/s │ 3.654 Gitem/s │ │ +├─ add_i64_nonnull 9.069 µs │ 13.16 µs │ 9.149 µs │ 9.232 µs │ 100 │ 100 +│ 3.612 Gitem/s │ 2.488 Gitem/s │ 3.581 Gitem/s │ 3.549 Gitem/s │ │ +├─ div_i64_nonnull 44.73 µs │ 51.02 µs │ 44.8 µs │ 45.03 µs │ 100 │ 100 +│ 732.4 Mitem/s │ 642.1 Mitem/s │ 731.2 Mitem/s │ 727.6 Mitem/s │ │ +├─ mul_i8_nonnull 5.979 µs │ 11.05 µs │ 6.199 µs │ 6.323 µs │ 100 │ 100 +│ 5.479 Gitem/s │ 2.962 Gitem/s │ 5.285 Gitem/s │ 5.182 Gitem/s │ │ +├─ mul_i16_nonnull 4.049 µs │ 8.709 µs │ 4.119 µs │ 4.205 µs │ 100 │ 100 +│ 8.091 Gitem/s │ 3.762 Gitem/s │ 7.953 Gitem/s │ 7.791 Gitem/s │ │ +├─ mul_i32_constant 26.36 µs │ 29.65 µs │ 26.43 µs │ 26.51 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 1.104 Gitem/s │ 1.239 Gitem/s │ 1.235 Gitem/s │ │ +├─ mul_i32_nonnull 26.37 µs │ 36.95 µs │ 26.42 µs │ 26.61 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 886.8 Mitem/s │ 1.239 Gitem/s │ 1.231 Gitem/s │ │ +├─ mul_i32_nullable 27.3 µs │ 49.11 µs │ 27.4 µs │ 27.76 µs │ 100 │ 100 +│ 1.199 Gitem/s │ 667.1 Mitem/s │ 1.195 Gitem/s │ 1.18 Gitem/s │ │ +├─ mul_i64_nonnull 23.11 µs │ 27.98 µs │ 23.2 µs │ 23.32 µs │ 100 │ 100 +│ 1.417 Gitem/s │ 1.171 Gitem/s │ 1.411 Gitem/s │ 1.404 Gitem/s │ │ +├─ mul_u8_nonnull 3.259 µs │ 4.809 µs │ 3.329 µs │ 3.345 µs │ 100 │ 100 +│ 10.05 Gitem/s │ 6.812 Gitem/s │ 9.84 Gitem/s │ 9.794 Gitem/s │ │ +├─ mul_u16_nonnull 2.549 µs │ 3.649 µs │ 2.599 µs │ 2.611 µs │ 100 │ 100 +│ 12.85 Gitem/s │ 8.978 Gitem/s │ 12.6 Gitem/s │ 12.54 Gitem/s │ │ +├─ mul_u32_nonnull 6.889 µs │ 10.15 µs │ 6.949 µs │ 6.999 µs │ 100 │ 100 +│ 4.756 Gitem/s │ 3.228 Gitem/s │ 4.714 Gitem/s │ 4.681 Gitem/s │ │ +├─ mul_u64_nonnull 19.12 µs │ 24.11 µs │ 19.19 µs │ 19.3 µs │ 100 │ 100 +│ 1.712 Gitem/s │ 1.358 Gitem/s │ 1.706 Gitem/s │ 1.697 Gitem/s │ │ +╰─ sub_i64_constant 8.119 µs │ 12.58 µs │ 8.239 µs │ 8.323 µs │ 100 │ 100 + 4.035 Gitem/s │ 2.602 Gitem/s │ 3.976 Gitem/s │ 3.936 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage0-candidate-1.md b/research/rowfn-x86-2026-08-07/benchmarks/stage0-candidate-1.md new file mode 100644 index 00000000000..c8539b0b1c6 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage0-candidate-1.md @@ -0,0 +1,38 @@ + + + +# `stage0-candidate-1` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.13 µs │ 93.77 µs │ 9.269 µs │ 10.17 µs │ 100 │ 100 +│ 3.588 Gitem/s │ 349.4 Mitem/s │ 3.534 Gitem/s │ 3.22 Gitem/s │ │ +├─ add_i64_nonnull 9.369 µs │ 12.45 µs │ 9.455 µs │ 9.51 µs │ 100 │ 100 +│ 3.497 Gitem/s │ 2.629 Gitem/s │ 3.465 Gitem/s │ 3.445 Gitem/s │ │ +├─ div_i64_nonnull 45.01 µs │ 54.4 µs │ 45.09 µs │ 45.42 µs │ 100 │ 100 +│ 727.8 Mitem/s │ 602.2 Mitem/s │ 726.5 Mitem/s │ 721.4 Mitem/s │ │ +├─ mul_i8_nonnull 4.639 µs │ 12.25 µs │ 4.694 µs │ 4.777 µs │ 100 │ 100 +│ 7.062 Gitem/s │ 2.672 Gitem/s │ 6.979 Gitem/s │ 6.858 Gitem/s │ │ +├─ mul_i16_nonnull 4.219 µs │ 6.999 µs │ 4.269 µs │ 4.33 µs │ 100 │ 100 +│ 7.765 Gitem/s │ 4.681 Gitem/s │ 7.674 Gitem/s │ 7.567 Gitem/s │ │ +├─ mul_i32_constant 18.77 µs │ 21.99 µs │ 18.88 µs │ 19 µs │ 100 │ 100 +│ 1.744 Gitem/s │ 1.489 Gitem/s │ 1.734 Gitem/s │ 1.724 Gitem/s │ │ +├─ mul_i32_nonnull 28.23 µs │ 31.75 µs │ 28.39 µs │ 28.48 µs │ 100 │ 100 +│ 1.16 Gitem/s │ 1.031 Gitem/s │ 1.153 Gitem/s │ 1.15 Gitem/s │ │ +├─ mul_i32_nullable 29.04 µs │ 44.74 µs │ 29.18 µs │ 29.42 µs │ 100 │ 100 +│ 1.128 Gitem/s │ 732.2 Mitem/s │ 1.122 Gitem/s │ 1.113 Gitem/s │ │ +├─ mul_i64_nonnull 29.7 µs │ 34.12 µs │ 30.02 µs │ 30.14 µs │ 100 │ 100 +│ 1.102 Gitem/s │ 960.1 Mitem/s │ 1.091 Gitem/s │ 1.087 Gitem/s │ │ +├─ mul_u8_nonnull 3.489 µs │ 7.519 µs │ 3.539 µs │ 3.602 µs │ 100 │ 100 +│ 9.389 Gitem/s │ 4.357 Gitem/s │ 9.257 Gitem/s │ 9.095 Gitem/s │ │ +├─ mul_u16_nonnull 2.349 µs │ 3.849 µs │ 2.429 µs │ 2.446 µs │ 100 │ 100 +│ 13.94 Gitem/s │ 8.511 Gitem/s │ 13.48 Gitem/s │ 13.39 Gitem/s │ │ +├─ mul_u32_nonnull 6.999 µs │ 9.889 µs │ 7.069 µs │ 7.11 µs │ 100 │ 100 +│ 4.681 Gitem/s │ 3.313 Gitem/s │ 4.634 Gitem/s │ 4.608 Gitem/s │ │ +├─ mul_u64_nonnull 30.36 µs │ 33.89 µs │ 30.43 µs │ 30.55 µs │ 100 │ 100 +│ 1.078 Gitem/s │ 966.6 Mitem/s │ 1.076 Gitem/s │ 1.072 Gitem/s │ │ +╰─ sub_i64_constant 8.979 µs │ 12.15 µs │ 9.099 µs │ 9.159 µs │ 100 │ 100 + 3.649 Gitem/s │ 2.696 Gitem/s │ 3.6 Gitem/s │ 3.577 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage0-candidate-2.md b/research/rowfn-x86-2026-08-07/benchmarks/stage0-candidate-2.md new file mode 100644 index 00000000000..401b1ae9bc8 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage0-candidate-2.md @@ -0,0 +1,38 @@ + + + +# `stage0-candidate-2` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.159 µs │ 88.36 µs │ 9.29 µs │ 10.41 µs │ 100 │ 100 +│ 3.577 Gitem/s │ 370.8 Mitem/s │ 3.527 Gitem/s │ 3.147 Gitem/s │ │ +├─ add_i64_nonnull 9.389 µs │ 12.37 µs │ 9.49 µs │ 9.622 µs │ 100 │ 100 +│ 3.489 Gitem/s │ 2.646 Gitem/s │ 3.452 Gitem/s │ 3.405 Gitem/s │ │ +├─ div_i64_nonnull 45.1 µs │ 48.73 µs │ 45.16 µs │ 45.34 µs │ 100 │ 100 +│ 726.4 Mitem/s │ 672.3 Mitem/s │ 725.5 Mitem/s │ 722.6 Mitem/s │ │ +├─ mul_i8_nonnull 4.639 µs │ 7.529 µs │ 4.699 µs │ 4.751 µs │ 100 │ 100 +│ 7.062 Gitem/s │ 4.351 Gitem/s │ 6.972 Gitem/s │ 6.897 Gitem/s │ │ +├─ mul_i16_nonnull 4.199 µs │ 6.379 µs │ 4.269 µs │ 4.295 µs │ 100 │ 100 +│ 7.802 Gitem/s │ 5.136 Gitem/s │ 7.674 Gitem/s │ 7.629 Gitem/s │ │ +├─ mul_i32_constant 18.77 µs │ 24.05 µs │ 18.87 µs │ 19 µs │ 100 │ 100 +│ 1.744 Gitem/s │ 1.361 Gitem/s │ 1.736 Gitem/s │ 1.724 Gitem/s │ │ +├─ mul_i32_nonnull 28.19 µs │ 31.61 µs │ 28.35 µs │ 28.45 µs │ 100 │ 100 +│ 1.161 Gitem/s │ 1.036 Gitem/s │ 1.155 Gitem/s │ 1.151 Gitem/s │ │ +├─ mul_i32_nullable 29 µs │ 50.06 µs │ 29.15 µs │ 29.47 µs │ 100 │ 100 +│ 1.129 Gitem/s │ 654.5 Mitem/s │ 1.123 Gitem/s │ 1.111 Gitem/s │ │ +├─ mul_i64_nonnull 29.82 µs │ 33.7 µs │ 30.08 µs │ 30.21 µs │ 100 │ 100 +│ 1.098 Gitem/s │ 972 Mitem/s │ 1.089 Gitem/s │ 1.084 Gitem/s │ │ +├─ mul_u8_nonnull 3.469 µs │ 9.249 µs │ 3.529 µs │ 3.592 µs │ 100 │ 100 +│ 9.443 Gitem/s │ 3.542 Gitem/s │ 9.283 Gitem/s │ 9.119 Gitem/s │ │ +├─ mul_u16_nonnull 2.369 µs │ 3.699 µs │ 2.429 µs │ 2.448 µs │ 100 │ 100 +│ 13.82 Gitem/s │ 8.856 Gitem/s │ 13.48 Gitem/s │ 13.38 Gitem/s │ │ +├─ mul_u32_nonnull 6.989 µs │ 9.659 µs │ 7.059 µs │ 7.111 µs │ 100 │ 100 +│ 4.687 Gitem/s │ 3.392 Gitem/s │ 4.641 Gitem/s │ 4.607 Gitem/s │ │ +├─ mul_u64_nonnull 30.41 µs │ 33.95 µs │ 30.49 µs │ 30.63 µs │ 100 │ 100 +│ 1.077 Gitem/s │ 964.9 Mitem/s │ 1.074 Gitem/s │ 1.069 Gitem/s │ │ +╰─ sub_i64_constant 8.989 µs │ 11.76 µs │ 9.099 µs │ 9.169 µs │ 100 │ 100 + 3.645 Gitem/s │ 2.784 Gitem/s │ 3.6 Gitem/s │ 3.573 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage1-base-1.md b/research/rowfn-x86-2026-08-07/benchmarks/stage1-base-1.md new file mode 100644 index 00000000000..4b6e17ff63e --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage1-base-1.md @@ -0,0 +1,38 @@ + + + +# `stage1-base-1` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.089 µs │ 783.5 µs │ 8.419 µs │ 16.21 µs │ 100 │ 100 +│ 4.05 Gitem/s │ 41.81 Mitem/s │ 3.891 Gitem/s │ 2.02 Gitem/s │ │ +├─ add_i64_nonnull 9.089 µs │ 32.99 µs │ 9.189 µs │ 9.53 µs │ 100 │ 100 +│ 3.604 Gitem/s │ 992.9 Mitem/s │ 3.565 Gitem/s │ 3.438 Gitem/s │ │ +├─ div_i64_nonnull 44.76 µs │ 76.23 µs │ 44.84 µs │ 45.51 µs │ 100 │ 100 +│ 731.9 Mitem/s │ 429.8 Mitem/s │ 730.6 Mitem/s │ 719.8 Mitem/s │ │ +├─ mul_i8_nonnull 5.929 µs │ 72.34 µs │ 6.239 µs │ 7.053 µs │ 100 │ 100 +│ 5.526 Gitem/s │ 452.9 Mitem/s │ 5.251 Gitem/s │ 4.645 Gitem/s │ │ +├─ mul_i16_nonnull 4.049 µs │ 61.69 µs │ 4.114 µs │ 4.692 µs │ 100 │ 100 +│ 8.091 Gitem/s │ 531.1 Mitem/s │ 7.963 Gitem/s │ 6.983 Gitem/s │ │ +├─ mul_i32_constant 26.36 µs │ 55.98 µs │ 26.43 µs │ 26.85 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 585.2 Mitem/s │ 1.239 Gitem/s │ 1.219 Gitem/s │ │ +├─ mul_i32_nonnull 26.38 µs │ 37.56 µs │ 26.42 µs │ 26.65 µs │ 100 │ 100 +│ 1.241 Gitem/s │ 872.3 Mitem/s │ 1.239 Gitem/s │ 1.229 Gitem/s │ │ +├─ mul_i32_nullable 27.23 µs │ 340.3 µs │ 27.36 µs │ 30.62 µs │ 100 │ 100 +│ 1.202 Gitem/s │ 96.26 Mitem/s │ 1.197 Gitem/s │ 1.07 Gitem/s │ │ +├─ mul_i64_nonnull 23.11 µs │ 44.46 µs │ 23.2 µs │ 23.5 µs │ 100 │ 100 +│ 1.417 Gitem/s │ 736.8 Mitem/s │ 1.411 Gitem/s │ 1.394 Gitem/s │ │ +├─ mul_u8_nonnull 3.269 µs │ 51.96 µs │ 3.329 µs │ 3.829 µs │ 100 │ 100 +│ 10.02 Gitem/s │ 630.6 Mitem/s │ 9.84 Gitem/s │ 8.556 Gitem/s │ │ +├─ mul_u16_nonnull 2.559 µs │ 30.97 µs │ 2.609 µs │ 2.898 µs │ 100 │ 100 +│ 12.8 Gitem/s │ 1.057 Gitem/s │ 12.55 Gitem/s │ 11.3 Gitem/s │ │ +├─ mul_u32_nonnull 6.88 µs │ 26.3 µs │ 6.959 µs │ 7.202 µs │ 100 │ 100 +│ 4.762 Gitem/s │ 1.245 Gitem/s │ 4.708 Gitem/s │ 4.549 Gitem/s │ │ +├─ mul_u64_nonnull 19.11 µs │ 40.79 µs │ 19.18 µs │ 19.48 µs │ 100 │ 100 +│ 1.713 Gitem/s │ 803.3 Mitem/s │ 1.707 Gitem/s │ 1.681 Gitem/s │ │ +╰─ sub_i64_constant 8.109 µs │ 41.21 µs │ 8.219 µs │ 8.589 µs │ 100 │ 100 + 4.04 Gitem/s │ 794.9 Mitem/s │ 3.986 Gitem/s │ 3.814 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage1-base-2.md b/research/rowfn-x86-2026-08-07/benchmarks/stage1-base-2.md new file mode 100644 index 00000000000..e540fafc12f --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage1-base-2.md @@ -0,0 +1,38 @@ + + + +# `stage1-base-2` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.309 µs │ 51.73 µs │ 8.459 µs │ 8.93 µs │ 100 │ 100 +│ 3.943 Gitem/s │ 633.3 Mitem/s │ 3.873 Gitem/s │ 3.669 Gitem/s │ │ +├─ add_i64_nonnull 9.119 µs │ 19.9 µs │ 9.199 µs │ 9.345 µs │ 100 │ 100 +│ 3.593 Gitem/s │ 1.646 Gitem/s │ 3.561 Gitem/s │ 3.506 Gitem/s │ │ +├─ div_i64_nonnull 44.77 µs │ 49.58 µs │ 44.85 µs │ 45.06 µs │ 100 │ 100 +│ 731.7 Mitem/s │ 660.7 Mitem/s │ 730.5 Mitem/s │ 727.1 Mitem/s │ │ +├─ mul_i8_nonnull 5.779 µs │ 10.19 µs │ 6.15 µs │ 6.267 µs │ 100 │ 100 +│ 5.669 Gitem/s │ 3.212 Gitem/s │ 5.327 Gitem/s │ 5.228 Gitem/s │ │ +├─ mul_i16_nonnull 4.059 µs │ 8.189 µs │ 4.109 µs │ 4.198 µs │ 100 │ 100 +│ 8.071 Gitem/s │ 4.001 Gitem/s │ 7.973 Gitem/s │ 7.804 Gitem/s │ │ +├─ mul_i32_constant 26.37 µs │ 30.04 µs │ 26.44 µs │ 26.54 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 1.09 Gitem/s │ 1.238 Gitem/s │ 1.234 Gitem/s │ │ +├─ mul_i32_nonnull 26.36 µs │ 29.99 µs │ 26.41 µs │ 26.54 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 1.092 Gitem/s │ 1.24 Gitem/s │ 1.234 Gitem/s │ │ +├─ mul_i32_nullable 27.23 µs │ 42.3 µs │ 27.35 µs │ 27.61 µs │ 100 │ 100 +│ 1.202 Gitem/s │ 774.4 Mitem/s │ 1.197 Gitem/s │ 1.186 Gitem/s │ │ +├─ mul_i64_nonnull 23.14 µs │ 27.63 µs │ 23.26 µs │ 23.41 µs │ 100 │ 100 +│ 1.415 Gitem/s │ 1.185 Gitem/s │ 1.408 Gitem/s │ 1.399 Gitem/s │ │ +├─ mul_u8_nonnull 3.269 µs │ 4.809 µs │ 3.319 µs │ 3.339 µs │ 100 │ 100 +│ 10.02 Gitem/s │ 6.812 Gitem/s │ 9.87 Gitem/s │ 9.812 Gitem/s │ │ +├─ mul_u16_nonnull 2.549 µs │ 3.519 µs │ 2.609 µs │ 2.618 µs │ 100 │ 100 +│ 12.85 Gitem/s │ 9.309 Gitem/s │ 12.55 Gitem/s │ 12.51 Gitem/s │ │ +├─ mul_u32_nonnull 6.879 µs │ 11.37 µs │ 6.95 µs │ 7.026 µs │ 100 │ 100 +│ 4.762 Gitem/s │ 2.879 Gitem/s │ 4.714 Gitem/s │ 4.663 Gitem/s │ │ +├─ mul_u64_nonnull 19.16 µs │ 22.33 µs │ 19.21 µs │ 19.3 µs │ 100 │ 100 +│ 1.709 Gitem/s │ 1.466 Gitem/s │ 1.705 Gitem/s │ 1.697 Gitem/s │ │ +╰─ sub_i64_constant 8.139 µs │ 11.6 µs │ 8.259 µs │ 8.319 µs │ 100 │ 100 + 4.025 Gitem/s │ 2.822 Gitem/s │ 3.967 Gitem/s │ 3.938 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage1-candidate-1.md b/research/rowfn-x86-2026-08-07/benchmarks/stage1-candidate-1.md new file mode 100644 index 00000000000..89efbeab66c --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage1-candidate-1.md @@ -0,0 +1,38 @@ + + + +# `stage1-candidate-1` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.149 µs │ 1.038 ms │ 9.279 µs │ 19.64 µs │ 100 │ 100 +│ 3.581 Gitem/s │ 31.54 Mitem/s │ 3.531 Gitem/s │ 1.667 Gitem/s │ │ +├─ add_i64_nonnull 9.399 µs │ 23.57 µs │ 9.459 µs │ 9.66 µs │ 100 │ 100 +│ 3.486 Gitem/s │ 1.389 Gitem/s │ 3.463 Gitem/s │ 3.391 Gitem/s │ │ +├─ div_i64_nonnull 45.03 µs │ 63.07 µs │ 45.1 µs │ 45.48 µs │ 100 │ 100 +│ 727.5 Mitem/s │ 519.4 Mitem/s │ 726.4 Mitem/s │ 720.4 Mitem/s │ │ +├─ mul_i8_nonnull 4.629 µs │ 65.15 µs │ 4.689 µs │ 5.344 µs │ 100 │ 100 +│ 7.077 Gitem/s │ 502.8 Mitem/s │ 6.987 Gitem/s │ 6.131 Gitem/s │ │ +├─ mul_i16_nonnull 4.209 µs │ 71.32 µs │ 4.259 µs │ 4.934 µs │ 100 │ 100 +│ 7.783 Gitem/s │ 459.3 Mitem/s │ 7.692 Gitem/s │ 6.64 Gitem/s │ │ +├─ mul_i32_constant 18.71 µs │ 73.84 µs │ 18.82 µs │ 19.43 µs │ 100 │ 100 +│ 1.75 Gitem/s │ 443.7 Mitem/s │ 1.74 Gitem/s │ 1.685 Gitem/s │ │ +├─ mul_i32_nonnull 28.22 µs │ 32.07 µs │ 28.34 µs │ 28.45 µs │ 100 │ 100 +│ 1.16 Gitem/s │ 1.021 Gitem/s │ 1.155 Gitem/s │ 1.151 Gitem/s │ │ +├─ mul_i32_nullable 29.04 µs │ 237.4 µs │ 29.16 µs │ 31.4 µs │ 100 │ 100 +│ 1.127 Gitem/s │ 138 Mitem/s │ 1.123 Gitem/s │ 1.043 Gitem/s │ │ +├─ mul_i64_nonnull 29.72 µs │ 54.86 µs │ 30.07 µs │ 30.53 µs │ 100 │ 100 +│ 1.102 Gitem/s │ 597.1 Mitem/s │ 1.089 Gitem/s │ 1.073 Gitem/s │ │ +├─ mul_u8_nonnull 3.469 µs │ 15.4 µs │ 3.529 µs │ 3.658 µs │ 100 │ 100 +│ 9.443 Gitem/s │ 2.126 Gitem/s │ 9.283 Gitem/s │ 8.956 Gitem/s │ │ +├─ mul_u16_nonnull 2.339 µs │ 13.45 µs │ 2.419 µs │ 2.574 µs │ 100 │ 100 +│ 14 Gitem/s │ 2.434 Gitem/s │ 13.54 Gitem/s │ 12.72 Gitem/s │ │ +├─ mul_u32_nonnull 6.969 µs │ 19.59 µs │ 7.049 µs │ 7.223 µs │ 100 │ 100 +│ 4.701 Gitem/s │ 1.671 Gitem/s │ 4.648 Gitem/s │ 4.536 Gitem/s │ │ +├─ mul_u64_nonnull 30.36 µs │ 42.47 µs │ 30.45 µs │ 30.69 µs │ 100 │ 100 +│ 1.079 Gitem/s │ 771.3 Mitem/s │ 1.075 Gitem/s │ 1.067 Gitem/s │ │ +╰─ sub_i64_constant 9.009 µs │ 31.68 µs │ 9.119 µs │ 9.431 µs │ 100 │ 100 + 3.636 Gitem/s │ 1.034 Gitem/s │ 3.593 Gitem/s │ 3.474 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage1-candidate-2.md b/research/rowfn-x86-2026-08-07/benchmarks/stage1-candidate-2.md new file mode 100644 index 00000000000..4bb5ec2261e --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage1-candidate-2.md @@ -0,0 +1,38 @@ + + + +# `stage1-candidate-2` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.149 µs │ 65.44 µs │ 9.309 µs │ 9.916 µs │ 100 │ 100 +│ 3.581 Gitem/s │ 500.6 Mitem/s │ 3.519 Gitem/s │ 3.304 Gitem/s │ │ +├─ add_i64_nonnull 9.429 µs │ 18.82 µs │ 9.529 µs │ 9.64 µs │ 100 │ 100 +│ 3.474 Gitem/s │ 1.74 Gitem/s │ 3.438 Gitem/s │ 3.399 Gitem/s │ │ +├─ div_i64_nonnull 45.09 µs │ 51.42 µs │ 45.16 µs │ 45.37 µs │ 100 │ 100 +│ 726.5 Mitem/s │ 637.2 Mitem/s │ 725.4 Mitem/s │ 722.1 Mitem/s │ │ +├─ mul_i8_nonnull 4.669 µs │ 7.159 µs │ 4.729 µs │ 4.781 µs │ 100 │ 100 +│ 7.017 Gitem/s │ 4.576 Gitem/s │ 6.928 Gitem/s │ 6.853 Gitem/s │ │ +├─ mul_i16_nonnull 4.249 µs │ 8.269 µs │ 4.319 µs │ 4.391 µs │ 100 │ 100 +│ 7.71 Gitem/s │ 3.962 Gitem/s │ 7.585 Gitem/s │ 7.461 Gitem/s │ │ +├─ mul_i32_constant 18.75 µs │ 22.35 µs │ 18.85 µs │ 18.93 µs │ 100 │ 100 +│ 1.746 Gitem/s │ 1.465 Gitem/s │ 1.737 Gitem/s │ 1.73 Gitem/s │ │ +├─ mul_i32_nonnull 28.25 µs │ 32.31 µs │ 28.39 µs │ 28.46 µs │ 100 │ 100 +│ 1.159 Gitem/s │ 1.013 Gitem/s │ 1.153 Gitem/s │ 1.151 Gitem/s │ │ +├─ mul_i32_nullable 29.04 µs │ 38.59 µs │ 29.18 µs │ 29.37 µs │ 100 │ 100 +│ 1.127 Gitem/s │ 848.9 Mitem/s │ 1.122 Gitem/s │ 1.115 Gitem/s │ │ +├─ mul_i64_nonnull 29.84 µs │ 33.8 µs │ 30.16 µs │ 30.24 µs │ 100 │ 100 +│ 1.097 Gitem/s │ 969.1 Mitem/s │ 1.086 Gitem/s │ 1.083 Gitem/s │ │ +├─ mul_u8_nonnull 3.509 µs │ 6.339 µs │ 3.579 µs │ 3.604 µs │ 100 │ 100 +│ 9.336 Gitem/s │ 5.168 Gitem/s │ 9.153 Gitem/s │ 9.091 Gitem/s │ │ +├─ mul_u16_nonnull 2.389 µs │ 38.12 µs │ 2.474 µs │ 2.857 µs │ 100 │ 100 +│ 13.71 Gitem/s │ 859.3 Mitem/s │ 13.24 Gitem/s │ 11.46 Gitem/s │ │ +├─ mul_u32_nonnull 7.019 µs │ 8.19 µs │ 7.109 µs │ 7.121 µs │ 100 │ 100 +│ 4.667 Gitem/s │ 4 Gitem/s │ 4.608 Gitem/s │ 4.601 Gitem/s │ │ +├─ mul_u64_nonnull 30.4 µs │ 33.46 µs │ 30.49 µs │ 30.63 µs │ 100 │ 100 +│ 1.077 Gitem/s │ 979 Mitem/s │ 1.074 Gitem/s │ 1.069 Gitem/s │ │ +╰─ sub_i64_constant 9.009 µs │ 10.35 µs │ 9.119 µs │ 9.142 µs │ 100 │ 100 + 3.636 Gitem/s │ 3.163 Gitem/s │ 3.593 Gitem/s │ 3.584 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage1-owned-1.md b/research/rowfn-x86-2026-08-07/benchmarks/stage1-owned-1.md new file mode 100644 index 00000000000..06e474ab7f0 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage1-owned-1.md @@ -0,0 +1,38 @@ + + + +# `stage1-owned-1` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.869 µs │ 97.56 µs │ 9.224 µs │ 10.16 µs │ 100 │ 100 +│ 3.694 Gitem/s │ 335.8 Mitem/s │ 3.552 Gitem/s │ 3.225 Gitem/s │ │ +├─ add_i64_nonnull 9.299 µs │ 13.77 µs │ 9.389 µs │ 9.495 µs │ 100 │ 100 +│ 3.523 Gitem/s │ 2.377 Gitem/s │ 3.489 Gitem/s │ 3.45 Gitem/s │ │ +├─ div_i64_nonnull 44.96 µs │ 54.06 µs │ 45.04 µs │ 45.34 µs │ 100 │ 100 +│ 728.8 Mitem/s │ 606.1 Mitem/s │ 727.5 Mitem/s │ 722.7 Mitem/s │ │ +├─ mul_i8_nonnull 4.559 µs │ 58.61 µs │ 4.619 µs │ 5.202 µs │ 100 │ 100 +│ 7.186 Gitem/s │ 558.9 Mitem/s │ 7.092 Gitem/s │ 6.298 Gitem/s │ │ +├─ mul_i16_nonnull 4.159 µs │ 5.809 µs │ 4.229 µs │ 4.244 µs │ 100 │ 100 +│ 7.877 Gitem/s │ 5.64 Gitem/s │ 7.746 Gitem/s │ 7.72 Gitem/s │ │ +├─ mul_i32_constant 32.23 µs │ 36.15 µs │ 32.36 µs │ 32.49 µs │ 100 │ 100 +│ 1.016 Gitem/s │ 906.2 Mitem/s │ 1.012 Gitem/s │ 1.008 Gitem/s │ │ +├─ mul_i32_nonnull 27.81 µs │ 43.9 µs │ 31.24 µs │ 31.22 µs │ 100 │ 100 +│ 1.177 Gitem/s │ 746.2 Mitem/s │ 1.048 Gitem/s │ 1.049 Gitem/s │ │ +├─ mul_i32_nullable 28.55 µs │ 50.24 µs │ 32.04 µs │ 31.62 µs │ 100 │ 100 +│ 1.147 Gitem/s │ 652.1 Mitem/s │ 1.022 Gitem/s │ 1.036 Gitem/s │ │ +├─ mul_i64_nonnull 25.31 µs │ 29.26 µs │ 25.65 µs │ 25.77 µs │ 100 │ 100 +│ 1.294 Gitem/s │ 1.119 Gitem/s │ 1.277 Gitem/s │ 1.271 Gitem/s │ │ +├─ mul_u8_nonnull 3.399 µs │ 55.89 µs │ 3.469 µs │ 3.999 µs │ 100 │ 100 +│ 9.638 Gitem/s │ 586.1 Mitem/s │ 9.443 Gitem/s │ 8.193 Gitem/s │ │ +├─ mul_u16_nonnull 2.289 µs │ 6.769 µs │ 2.369 µs │ 2.415 µs │ 100 │ 100 +│ 14.31 Gitem/s │ 4.84 Gitem/s │ 13.82 Gitem/s │ 13.56 Gitem/s │ │ +├─ mul_u32_nonnull 6.919 µs │ 9.879 µs │ 7.009 µs │ 7.059 µs │ 100 │ 100 +│ 4.735 Gitem/s │ 3.316 Gitem/s │ 4.674 Gitem/s │ 4.641 Gitem/s │ │ +├─ mul_u64_nonnull 19.34 µs │ 22.38 µs │ 19.41 µs │ 19.5 µs │ 100 │ 100 +│ 1.693 Gitem/s │ 1.463 Gitem/s │ 1.687 Gitem/s │ 1.68 Gitem/s │ │ +╰─ sub_i64_constant 9.419 µs │ 12.37 µs │ 9.554 µs │ 9.607 µs │ 100 │ 100 + 3.478 Gitem/s │ 2.646 Gitem/s │ 3.429 Gitem/s │ 3.41 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage1-owned-2.md b/research/rowfn-x86-2026-08-07/benchmarks/stage1-owned-2.md new file mode 100644 index 00000000000..a653d31d205 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage1-owned-2.md @@ -0,0 +1,38 @@ + + + +# `stage1-owned-2` + +```text +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.919 µs │ 99.93 µs │ 9.229 µs │ 10.19 µs │ 100 │ 100 +│ 3.673 Gitem/s │ 327.8 Mitem/s │ 3.55 Gitem/s │ 3.214 Gitem/s │ │ +├─ add_i64_nonnull 9.279 µs │ 12.53 µs │ 9.339 µs │ 9.417 µs │ 100 │ 100 +│ 3.531 Gitem/s │ 2.613 Gitem/s │ 3.508 Gitem/s │ 3.479 Gitem/s │ │ +├─ div_i64_nonnull 44.91 µs │ 54.36 µs │ 45.01 µs │ 45.32 µs │ 100 │ 100 +│ 729.4 Mitem/s │ 602.6 Mitem/s │ 727.8 Mitem/s │ 722.9 Mitem/s │ │ +├─ mul_i8_nonnull 4.579 µs │ 61.3 µs │ 4.649 µs │ 5.229 µs │ 100 │ 100 +│ 7.154 Gitem/s │ 534.5 Mitem/s │ 7.047 Gitem/s │ 6.265 Gitem/s │ │ +├─ mul_i16_nonnull 4.169 µs │ 7.419 µs │ 4.229 µs │ 4.276 µs │ 100 │ 100 +│ 7.858 Gitem/s │ 4.416 Gitem/s │ 7.746 Gitem/s │ 7.661 Gitem/s │ │ +├─ mul_i32_constant 32.27 µs │ 35.87 µs │ 32.38 µs │ 32.49 µs │ 100 │ 100 +│ 1.015 Gitem/s │ 913.5 Mitem/s │ 1.011 Gitem/s │ 1.008 Gitem/s │ │ +├─ mul_i32_nonnull 27.77 µs │ 32.39 µs │ 31.23 µs │ 30.31 µs │ 100 │ 100 +│ 1.179 Gitem/s │ 1.011 Gitem/s │ 1.048 Gitem/s │ 1.08 Gitem/s │ │ +├─ mul_i32_nullable 28.53 µs │ 49.46 µs │ 32.04 µs │ 31.34 µs │ 100 │ 100 +│ 1.148 Gitem/s │ 662.3 Mitem/s │ 1.022 Gitem/s │ 1.045 Gitem/s │ │ +├─ mul_i64_nonnull 25.25 µs │ 29.06 µs │ 25.59 µs │ 25.68 µs │ 100 │ 100 +│ 1.297 Gitem/s │ 1.127 Gitem/s │ 1.28 Gitem/s │ 1.275 Gitem/s │ │ +├─ mul_u8_nonnull 3.429 µs │ 60.44 µs │ 3.479 µs │ 4.062 µs │ 100 │ 100 +│ 9.553 Gitem/s │ 542 Mitem/s │ 9.416 Gitem/s │ 8.065 Gitem/s │ │ +├─ mul_u16_nonnull 2.319 µs │ 7.879 µs │ 2.389 µs │ 2.446 µs │ 100 │ 100 +│ 14.12 Gitem/s │ 4.158 Gitem/s │ 13.71 Gitem/s │ 13.39 Gitem/s │ │ +├─ mul_u32_nonnull 6.949 µs │ 10.35 µs │ 7.019 µs │ 7.069 µs │ 100 │ 100 +│ 4.714 Gitem/s │ 3.163 Gitem/s │ 4.667 Gitem/s │ 4.635 Gitem/s │ │ +├─ mul_u64_nonnull 19.34 µs │ 22.75 µs │ 19.41 µs │ 19.49 µs │ 100 │ 100 +│ 1.693 Gitem/s │ 1.44 Gitem/s │ 1.687 Gitem/s │ 1.68 Gitem/s │ │ +╰─ sub_i64_constant 9.439 µs │ 12.1 µs │ 9.569 µs │ 9.636 µs │ 100 │ 100 + 3.471 Gitem/s │ 2.705 Gitem/s │ 3.424 Gitem/s │ 3.4 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage2-base-1.md b/research/rowfn-x86-2026-08-07/benchmarks/stage2-base-1.md new file mode 100644 index 00000000000..3053244a92a --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage2-base-1.md @@ -0,0 +1,39 @@ + + + +# `stage2-base-1` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.099 µs │ 766.1 µs │ 8.419 µs │ 16.03 µs │ 100 │ 100 +│ 4.045 Gitem/s │ 42.76 Mitem/s │ 3.891 Gitem/s │ 2.043 Gitem/s │ │ +├─ add_i64_nonnull 9.099 µs │ 30.45 µs │ 9.179 µs │ 9.474 µs │ 100 │ 100 +│ 3.6 Gitem/s │ 1.075 Gitem/s │ 3.569 Gitem/s │ 3.458 Gitem/s │ │ +├─ div_i64_nonnull 44.76 µs │ 75.04 µs │ 44.84 µs │ 45.32 µs │ 100 │ 100 +│ 731.9 Mitem/s │ 436.6 Mitem/s │ 730.6 Mitem/s │ 722.9 Mitem/s │ │ +├─ mul_i8_nonnull 5.809 µs │ 69.77 µs │ 6.209 µs │ 6.952 µs │ 100 │ 100 +│ 5.64 Gitem/s │ 469.5 Mitem/s │ 5.276 Gitem/s │ 4.713 Gitem/s │ │ +├─ mul_i16_nonnull 4.029 µs │ 66.19 µs │ 4.099 µs │ 4.725 µs │ 100 │ 100 +│ 8.131 Gitem/s │ 494.9 Mitem/s │ 7.992 Gitem/s │ 6.934 Gitem/s │ │ +├─ mul_i32_constant 26.36 µs │ 55.21 µs │ 26.42 µs │ 26.88 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 593.4 Mitem/s │ 1.239 Gitem/s │ 1.218 Gitem/s │ │ +├─ mul_i32_nonnull 26.34 µs │ 39.19 µs │ 26.39 µs │ 26.64 µs │ 100 │ 100 +│ 1.243 Gitem/s │ 835.9 Mitem/s │ 1.241 Gitem/s │ 1.229 Gitem/s │ │ +├─ mul_i32_nullable 27.21 µs │ 333.7 µs │ 27.37 µs │ 30.56 µs │ 100 │ 100 +│ 1.203 Gitem/s │ 98.17 Mitem/s │ 1.196 Gitem/s │ 1.072 Gitem/s │ │ +├─ mul_i64_nonnull 23.14 µs │ 43.91 µs │ 23.22 µs │ 23.6 µs │ 100 │ 100 +│ 1.415 Gitem/s │ 746 Mitem/s │ 1.41 Gitem/s │ 1.388 Gitem/s │ │ +├─ mul_u8_nonnull 3.259 µs │ 52.38 µs │ 3.329 µs │ 3.817 µs │ 100 │ 100 +│ 10.05 Gitem/s │ 625.4 Mitem/s │ 9.84 Gitem/s │ 8.582 Gitem/s │ │ +├─ mul_u16_nonnull 2.549 µs │ 30.18 µs │ 2.609 µs │ 2.929 µs │ 100 │ 100 +│ 12.85 Gitem/s │ 1.085 Gitem/s │ 12.55 Gitem/s │ 11.18 Gitem/s │ │ +├─ mul_u32_nonnull 6.869 µs │ 27.32 µs │ 6.939 µs │ 7.147 µs │ 100 │ 100 +│ 4.769 Gitem/s │ 1.198 Gitem/s │ 4.721 Gitem/s │ 4.584 Gitem/s │ │ +├─ mul_u64_nonnull 19.14 µs │ 41.82 µs │ 19.22 µs │ 19.57 µs │ 100 │ 100 +│ 1.711 Gitem/s │ 783.3 Mitem/s │ 1.704 Gitem/s │ 1.674 Gitem/s │ │ +╰─ sub_i64_constant 8.159 µs │ 40.98 µs │ 8.249 µs │ 8.63 µs │ 100 │ 100 + 4.015 Gitem/s │ 799.4 Mitem/s │ 3.971 Gitem/s │ 3.796 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage2-base-2.md b/research/rowfn-x86-2026-08-07/benchmarks/stage2-base-2.md new file mode 100644 index 00000000000..793ece2b6f1 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage2-base-2.md @@ -0,0 +1,39 @@ + + + +# `stage2-base-2` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.299 µs │ 47.08 µs │ 8.429 µs │ 8.916 µs │ 100 │ 100 +│ 3.948 Gitem/s │ 695.8 Mitem/s │ 3.887 Gitem/s │ 3.675 Gitem/s │ │ +├─ add_i64_nonnull 9.139 µs │ 12.35 µs │ 9.209 µs │ 9.286 µs │ 100 │ 100 +│ 3.585 Gitem/s │ 2.651 Gitem/s │ 3.557 Gitem/s │ 3.528 Gitem/s │ │ +├─ div_i64_nonnull 44.8 µs │ 49.64 µs │ 44.87 µs │ 45.09 µs │ 100 │ 100 +│ 731.2 Mitem/s │ 659.9 Mitem/s │ 730.1 Mitem/s │ 726.6 Mitem/s │ │ +├─ mul_i8_nonnull 5.869 µs │ 9.279 µs │ 6.174 µs │ 6.295 µs │ 100 │ 100 +│ 5.582 Gitem/s │ 3.531 Gitem/s │ 5.306 Gitem/s │ 5.204 Gitem/s │ │ +├─ mul_i16_nonnull 4.039 µs │ 7.909 µs │ 4.104 µs │ 4.153 µs │ 100 │ 100 +│ 8.111 Gitem/s │ 4.142 Gitem/s │ 7.982 Gitem/s │ 7.888 Gitem/s │ │ +├─ mul_i32_constant 26.37 µs │ 29.56 µs │ 26.43 µs │ 26.52 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 1.108 Gitem/s │ 1.239 Gitem/s │ 1.235 Gitem/s │ │ +├─ mul_i32_nonnull 26.36 µs │ 30.91 µs │ 26.41 µs │ 26.53 µs │ 100 │ 100 +│ 1.242 Gitem/s │ 1.059 Gitem/s │ 1.24 Gitem/s │ 1.234 Gitem/s │ │ +├─ mul_i32_nullable 27.27 µs │ 42.76 µs │ 27.38 µs │ 27.63 µs │ 100 │ 100 +│ 1.201 Gitem/s │ 766.1 Mitem/s │ 1.196 Gitem/s │ 1.185 Gitem/s │ │ +├─ mul_i64_nonnull 23.14 µs │ 26.52 µs │ 23.24 µs │ 23.33 µs │ 100 │ 100 +│ 1.415 Gitem/s │ 1.235 Gitem/s │ 1.409 Gitem/s │ 1.404 Gitem/s │ │ +├─ mul_u8_nonnull 3.279 µs │ 6.329 µs │ 3.329 µs │ 3.378 µs │ 100 │ 100 +│ 9.99 Gitem/s │ 5.176 Gitem/s │ 9.84 Gitem/s │ 9.698 Gitem/s │ │ +├─ mul_u16_nonnull 2.549 µs │ 3.489 µs │ 2.599 µs │ 2.615 µs │ 100 │ 100 +│ 12.85 Gitem/s │ 9.389 Gitem/s │ 12.6 Gitem/s │ 12.52 Gitem/s │ │ +├─ mul_u32_nonnull 6.879 µs │ 9.609 µs │ 6.939 µs │ 7.003 µs │ 100 │ 100 +│ 4.762 Gitem/s │ 3.409 Gitem/s │ 4.721 Gitem/s │ 4.678 Gitem/s │ │ +├─ mul_u64_nonnull 19.15 µs │ 22.82 µs │ 19.21 µs │ 19.33 µs │ 100 │ 100 +│ 1.71 Gitem/s │ 1.435 Gitem/s │ 1.704 Gitem/s │ 1.694 Gitem/s │ │ +╰─ sub_i64_constant 8.119 µs │ 11.08 µs │ 8.249 µs │ 8.293 µs │ 100 │ 100 + 4.035 Gitem/s │ 2.954 Gitem/s │ 3.971 Gitem/s │ 3.951 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage2-candidate-1.md b/research/rowfn-x86-2026-08-07/benchmarks/stage2-candidate-1.md new file mode 100644 index 00000000000..6364ccde44c --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage2-candidate-1.md @@ -0,0 +1,39 @@ + + + +# `stage2-candidate-1` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.099 µs │ 1.015 ms │ 9.269 µs │ 19.41 µs │ 100 │ 100 +│ 3.6 Gitem/s │ 32.26 Mitem/s │ 3.534 Gitem/s │ 1.687 Gitem/s │ │ +├─ add_i64_nonnull 9.319 µs │ 12.12 µs │ 9.429 µs │ 9.48 µs │ 100 │ 100 +│ 3.515 Gitem/s │ 2.701 Gitem/s │ 3.474 Gitem/s │ 3.456 Gitem/s │ │ +├─ div_i64_nonnull 44.99 µs │ 61.85 µs │ 45.07 µs │ 45.6 µs │ 100 │ 100 +│ 728.1 Mitem/s │ 529.7 Mitem/s │ 726.8 Mitem/s │ 718.4 Mitem/s │ │ +├─ mul_i8_nonnull 4.639 µs │ 61.66 µs │ 4.689 µs │ 5.268 µs │ 100 │ 100 +│ 7.062 Gitem/s │ 531.3 Mitem/s │ 6.987 Gitem/s │ 6.219 Gitem/s │ │ +├─ mul_i16_nonnull 4.209 µs │ 66.06 µs │ 4.279 µs │ 4.931 µs │ 100 │ 100 +│ 7.783 Gitem/s │ 495.9 Mitem/s │ 7.656 Gitem/s │ 6.644 Gitem/s │ │ +├─ mul_i32_constant 18.77 µs │ 71.76 µs │ 18.88 µs │ 19.64 µs │ 100 │ 100 +│ 1.744 Gitem/s │ 456.5 Mitem/s │ 1.734 Gitem/s │ 1.667 Gitem/s │ │ +├─ mul_i32_nonnull 28.21 µs │ 33.28 µs │ 28.34 µs │ 28.48 µs │ 100 │ 100 +│ 1.161 Gitem/s │ 984.3 Mitem/s │ 1.155 Gitem/s │ 1.15 Gitem/s │ │ +├─ mul_i32_nullable 29.02 µs │ 233.9 µs │ 29.2 µs │ 31.37 µs │ 100 │ 100 +│ 1.128 Gitem/s │ 140 Mitem/s │ 1.121 Gitem/s │ 1.044 Gitem/s │ │ +├─ mul_i64_nonnull 29.73 µs │ 52.85 µs │ 30.02 µs │ 30.37 µs │ 100 │ 100 +│ 1.101 Gitem/s │ 619.9 Mitem/s │ 1.091 Gitem/s │ 1.078 Gitem/s │ │ +├─ mul_u8_nonnull 3.449 µs │ 14.5 µs │ 3.529 µs │ 3.679 µs │ 100 │ 100 +│ 9.498 Gitem/s │ 2.258 Gitem/s │ 9.283 Gitem/s │ 8.906 Gitem/s │ │ +├─ mul_u16_nonnull 2.349 µs │ 13.44 µs │ 2.419 µs │ 2.529 µs │ 100 │ 100 +│ 13.94 Gitem/s │ 2.436 Gitem/s │ 13.54 Gitem/s │ 12.95 Gitem/s │ │ +├─ mul_u32_nonnull 6.979 µs │ 18.54 µs │ 7.059 µs │ 7.228 µs │ 100 │ 100 +│ 4.694 Gitem/s │ 1.766 Gitem/s │ 4.641 Gitem/s │ 4.532 Gitem/s │ │ +├─ mul_u64_nonnull 30.31 µs │ 42.57 µs │ 30.41 µs │ 30.68 µs │ 100 │ 100 +│ 1.08 Gitem/s │ 769.5 Mitem/s │ 1.077 Gitem/s │ 1.067 Gitem/s │ │ +╰─ sub_i64_constant 8.979 µs │ 32.68 µs │ 9.064 µs │ 9.358 µs │ 100 │ 100 + 3.649 Gitem/s │ 1.002 Gitem/s │ 3.614 Gitem/s │ 3.501 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage2-candidate-2.md b/research/rowfn-x86-2026-08-07/benchmarks/stage2-candidate-2.md new file mode 100644 index 00000000000..6093fe215ee --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage2-candidate-2.md @@ -0,0 +1,39 @@ + + + +# `stage2-candidate-2` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.159 µs │ 57.71 µs │ 9.269 µs │ 9.826 µs │ 100 │ 100 +│ 3.577 Gitem/s │ 567.7 Mitem/s │ 3.534 Gitem/s │ 3.334 Gitem/s │ │ +├─ add_i64_nonnull 9.359 µs │ 18.58 µs │ 9.454 µs │ 9.765 µs │ 100 │ 100 +│ 3.5 Gitem/s │ 1.762 Gitem/s │ 3.465 Gitem/s │ 3.355 Gitem/s │ │ +├─ div_i64_nonnull 45.03 µs │ 49.11 µs │ 45.12 µs │ 45.32 µs │ 100 │ 100 +│ 727.5 Mitem/s │ 667.1 Mitem/s │ 726 Mitem/s │ 722.9 Mitem/s │ │ +├─ mul_i8_nonnull 4.649 µs │ 8.699 µs │ 4.729 µs │ 4.798 µs │ 100 │ 100 +│ 7.047 Gitem/s │ 3.766 Gitem/s │ 6.928 Gitem/s │ 6.828 Gitem/s │ │ +├─ mul_i16_nonnull 4.219 µs │ 7.469 µs │ 4.309 µs │ 4.374 µs │ 100 │ 100 +│ 7.765 Gitem/s │ 4.386 Gitem/s │ 7.603 Gitem/s │ 7.491 Gitem/s │ │ +├─ mul_i32_constant 18.75 µs │ 22.67 µs │ 18.88 µs │ 18.95 µs │ 100 │ 100 +│ 1.746 Gitem/s │ 1.444 Gitem/s │ 1.734 Gitem/s │ 1.728 Gitem/s │ │ +├─ mul_i32_nonnull 28.2 µs │ 40.32 µs │ 28.36 µs │ 28.69 µs │ 100 │ 100 +│ 1.161 Gitem/s │ 812.5 Mitem/s │ 1.155 Gitem/s │ 1.141 Gitem/s │ │ +├─ mul_i32_nullable 29.02 µs │ 40.92 µs │ 29.17 µs │ 29.4 µs │ 100 │ 100 +│ 1.128 Gitem/s │ 800.5 Mitem/s │ 1.123 Gitem/s │ 1.114 Gitem/s │ │ +├─ mul_i64_nonnull 29.63 µs │ 33.89 µs │ 30.1 µs │ 30.22 µs │ 100 │ 100 +│ 1.105 Gitem/s │ 966.6 Mitem/s │ 1.088 Gitem/s │ 1.084 Gitem/s │ │ +├─ mul_u8_nonnull 3.489 µs │ 4.839 µs │ 3.559 µs │ 3.576 µs │ 100 │ 100 +│ 9.389 Gitem/s │ 6.77 Gitem/s │ 9.205 Gitem/s │ 9.163 Gitem/s │ │ +├─ mul_u16_nonnull 2.379 µs │ 5.529 µs │ 2.439 µs │ 2.489 µs │ 100 │ 100 +│ 13.76 Gitem/s │ 5.925 Gitem/s │ 13.43 Gitem/s │ 13.16 Gitem/s │ │ +├─ mul_u32_nonnull 6.989 µs │ 8.019 µs │ 7.079 µs │ 7.089 µs │ 100 │ 100 +│ 4.687 Gitem/s │ 4.085 Gitem/s │ 4.628 Gitem/s │ 4.621 Gitem/s │ │ +├─ mul_u64_nonnull 30.35 µs │ 34.77 µs │ 30.43 µs │ 30.57 µs │ 100 │ 100 +│ 1.079 Gitem/s │ 942.1 Mitem/s │ 1.076 Gitem/s │ 1.071 Gitem/s │ │ +╰─ sub_i64_constant 8.949 µs │ 10.5 µs │ 9.089 µs │ 9.109 µs │ 100 │ 100 + 3.661 Gitem/s │ 3.117 Gitem/s │ 3.604 Gitem/s │ 3.597 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage2-indexed-1.md b/research/rowfn-x86-2026-08-07/benchmarks/stage2-indexed-1.md new file mode 100644 index 00000000000..440304b08fe --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage2-indexed-1.md @@ -0,0 +1,39 @@ + + + +# `stage2-indexed-1` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 8.889 µs │ 89.9 µs │ 9.244 µs │ 10.11 µs │ 100 │ 100 +│ 3.686 Gitem/s │ 364.4 Mitem/s │ 3.544 Gitem/s │ 3.24 Gitem/s │ │ +├─ add_i64_nonnull 9.279 µs │ 18.28 µs │ 9.399 µs │ 9.605 µs │ 100 │ 100 +│ 3.531 Gitem/s │ 1.791 Gitem/s │ 3.486 Gitem/s │ 3.411 Gitem/s │ │ +├─ div_i64_nonnull 44.92 µs │ 52.66 µs │ 45.07 µs │ 45.39 µs │ 100 │ 100 +│ 729.3 Mitem/s │ 622.1 Mitem/s │ 726.8 Mitem/s │ 721.9 Mitem/s │ │ +├─ mul_i8_nonnull 6.069 µs │ 69.59 µs │ 6.339 µs │ 7.085 µs │ 100 │ 100 +│ 5.398 Gitem/s │ 470.8 Mitem/s │ 5.168 Gitem/s │ 4.624 Gitem/s │ │ +├─ mul_i16_nonnull 4.209 µs │ 5.439 µs │ 4.259 µs │ 4.274 µs │ 100 │ 100 +│ 7.783 Gitem/s │ 6.023 Gitem/s │ 7.692 Gitem/s │ 7.665 Gitem/s │ │ +├─ mul_i32_constant 32.23 µs │ 36.63 µs │ 32.38 µs │ 32.54 µs │ 100 │ 100 +│ 1.016 Gitem/s │ 894.3 Mitem/s │ 1.011 Gitem/s │ 1.006 Gitem/s │ │ +├─ mul_i32_nonnull 26.52 µs │ 31.34 µs │ 26.58 µs │ 26.69 µs │ 100 │ 100 +│ 1.235 Gitem/s │ 1.045 Gitem/s │ 1.232 Gitem/s │ 1.227 Gitem/s │ │ +├─ mul_i32_nullable 27.31 µs │ 47.02 µs │ 27.41 µs │ 27.83 µs │ 100 │ 100 +│ 1.199 Gitem/s │ 696.7 Mitem/s │ 1.195 Gitem/s │ 1.177 Gitem/s │ │ +├─ mul_i64_nonnull 23.33 µs │ 32.13 µs │ 23.43 µs │ 23.74 µs │ 100 │ 100 +│ 1.403 Gitem/s │ 1.019 Gitem/s │ 1.397 Gitem/s │ 1.379 Gitem/s │ │ +├─ mul_u8_nonnull 3.439 µs │ 61.51 µs │ 3.509 µs │ 4.121 µs │ 100 │ 100 +│ 9.526 Gitem/s │ 532.6 Mitem/s │ 9.336 Gitem/s │ 7.951 Gitem/s │ │ +├─ mul_u16_nonnull 2.699 µs │ 6.979 µs │ 2.769 µs │ 2.813 µs │ 100 │ 100 +│ 12.13 Gitem/s │ 4.694 Gitem/s │ 11.83 Gitem/s │ 11.64 Gitem/s │ │ +├─ mul_u32_nonnull 7.029 µs │ 9.939 µs │ 7.109 µs │ 7.16 µs │ 100 │ 100 +│ 4.661 Gitem/s │ 3.296 Gitem/s │ 4.608 Gitem/s │ 4.576 Gitem/s │ │ +├─ mul_u64_nonnull 19.35 µs │ 22.92 µs │ 19.41 µs │ 19.51 µs │ 100 │ 100 +│ 1.692 Gitem/s │ 1.429 Gitem/s │ 1.687 Gitem/s │ 1.679 Gitem/s │ │ +╰─ sub_i64_constant 9.499 µs │ 12.48 µs │ 9.609 µs │ 9.668 µs │ 100 │ 100 + 3.449 Gitem/s │ 2.623 Gitem/s │ 3.409 Gitem/s │ 3.389 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/benchmarks/stage2-indexed-2.md b/research/rowfn-x86-2026-08-07/benchmarks/stage2-indexed-2.md new file mode 100644 index 00000000000..aa7ed2c846a --- /dev/null +++ b/research/rowfn-x86-2026-08-07/benchmarks/stage2-indexed-2.md @@ -0,0 +1,39 @@ + + + +# `stage2-indexed-2` + +```text +Timer precision: 10 ns +binary_ops fastest │ slowest │ median │ mean │ samples │ iters +├─ add_i64_constant 9.009 µs │ 93.1 µs │ 9.259 µs │ 10.14 µs │ 100 │ 100 +│ 3.636 Gitem/s │ 351.9 Mitem/s │ 3.538 Gitem/s │ 3.23 Gitem/s │ │ +├─ add_i64_nonnull 9.309 µs │ 12.56 µs │ 9.379 µs │ 9.426 µs │ 100 │ 100 +│ 3.519 Gitem/s │ 2.606 Gitem/s │ 3.493 Gitem/s │ 3.476 Gitem/s │ │ +├─ div_i64_nonnull 44.96 µs │ 48.24 µs │ 45.03 µs │ 45.22 µs │ 100 │ 100 +│ 728.6 Mitem/s │ 679.1 Mitem/s │ 727.5 Mitem/s │ 724.5 Mitem/s │ │ +├─ mul_i8_nonnull 5.969 µs │ 50.06 µs │ 6.359 µs │ 6.902 µs │ 100 │ 100 +│ 5.488 Gitem/s │ 654.4 Mitem/s │ 5.152 Gitem/s │ 4.747 Gitem/s │ │ +├─ mul_i16_nonnull 4.199 µs │ 15.43 µs │ 4.284 µs │ 4.418 µs │ 100 │ 100 +│ 7.802 Gitem/s │ 2.122 Gitem/s │ 7.647 Gitem/s │ 7.415 Gitem/s │ │ +├─ mul_i32_constant 32.25 µs │ 35.51 µs │ 32.39 µs │ 32.51 µs │ 100 │ 100 +│ 1.015 Gitem/s │ 922.5 Mitem/s │ 1.011 Gitem/s │ 1.007 Gitem/s │ │ +├─ mul_i32_nonnull 26.54 µs │ 29.82 µs │ 26.6 µs │ 26.7 µs │ 100 │ 100 +│ 1.234 Gitem/s │ 1.098 Gitem/s │ 1.231 Gitem/s │ 1.226 Gitem/s │ │ +├─ mul_i32_nullable 27.32 µs │ 40.06 µs │ 27.43 µs │ 27.68 µs │ 100 │ 100 +│ 1.198 Gitem/s │ 817.7 Mitem/s │ 1.194 Gitem/s │ 1.183 Gitem/s │ │ +├─ mul_i64_nonnull 23.33 µs │ 26.7 µs │ 23.44 µs │ 23.53 µs │ 100 │ 100 +│ 1.403 Gitem/s │ 1.226 Gitem/s │ 1.397 Gitem/s │ 1.392 Gitem/s │ │ +├─ mul_u8_nonnull 3.459 µs │ 61.29 µs │ 3.519 µs │ 4.131 µs │ 100 │ 100 +│ 9.471 Gitem/s │ 534.5 Mitem/s │ 9.309 Gitem/s │ 7.931 Gitem/s │ │ +├─ mul_u16_nonnull 2.709 µs │ 6.939 µs │ 2.789 µs │ 2.828 µs │ 100 │ 100 +│ 12.09 Gitem/s │ 4.721 Gitem/s │ 11.74 Gitem/s │ 11.58 Gitem/s │ │ +├─ mul_u32_nonnull 7.039 µs │ 10.56 µs │ 7.119 µs │ 7.18 µs │ 100 │ 100 +│ 4.654 Gitem/s │ 3.1 Gitem/s │ 4.602 Gitem/s │ 4.563 Gitem/s │ │ +├─ mul_u64_nonnull 19.34 µs │ 23.2 µs │ 19.42 µs │ 19.48 µs │ 100 │ 100 +│ 1.693 Gitem/s │ 1.411 Gitem/s │ 1.686 Gitem/s │ 1.681 Gitem/s │ │ +╰─ sub_i64_constant 9.509 µs │ 12.69 µs │ 9.619 µs │ 9.668 µs │ 100 │ 100 + 3.445 Gitem/s │ 2.58 Gitem/s │ 3.406 Gitem/s │ 3.389 Gitem/s │ │ + + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/base-codegen-summary.md b/research/rowfn-x86-2026-08-07/codegen/base-codegen-summary.md new file mode 100644 index 00000000000..1296a6b3bb1 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/base-codegen-summary.md @@ -0,0 +1,139 @@ + + + +# Merge-base production numeric multiply code generation + +Revision: `19f771f2a426103aa7d1bf7153a258bb1bab1e19` + +Command: + +```text +CARGO_TARGET_DIR=/tmp/rowfn-x86.ccCdz5/target-base-codegen \ + cargo rustc -p vortex-array --lib --profile bench -- \ + --emit=llvm-ir,asm -C codegen-units=1 +``` + +Artifacts: + +```text +/tmp/rowfn-x86.ccCdz5/target-base-codegen/release/deps/vortex_array-4e5fe3dd7af89793.ll +/tmp/rowfn-x86.ccCdz5/target-base-codegen/release/deps/vortex_array-4e5fe3dd7af89793.s +``` + +## Production symbols + +```text +i64 execute_checked_typed: 648ec4b22808a2d4 +i64 checked_op_lanes (varying x varying): df33f84e66e75a91 +u64 execute_checked_typed: 8da1eac40a9b0934 +u64 checked_op_lanes (varying x varying): 84edf83ddcd3fe05 +``` + +## i64 hot loop + +Assembly source begins at line 6,698,561 in the `.s` artifact. The loop is +`.LBB6227_8`: + +```asm +movq (%rdi,%rsi,8), %rax +imulq (%r15,%rsi,8) +movq %rax, (%r13,%rsi,8) +incq %rsi +sarq $63, %rax +xorq %rdx, %rax +orq %rax, %rcx +cmpq %rsi, %rbx +jne .LBB6227_8 +``` + +This is one lane per backedge. The one-operand `imulq` produces the signed +128-bit product in `RDX:RAX`; the low half is stored and the high half is +compared with the low-half sign extension through `sarq`/`xorq`. Failure stays +in register `%rcx`. There is no `vector.body`, unroll, or separate remainder. + +## u64 hot loop + +Assembly source begins at line 6,646,461 in the `.s` artifact. The loop is +`.LBB6175_10`: + +```asm +movq (%rbx,%rdi,8), %rax +mulq (%r11,%rdi,8) +movq %rdx, %rsi +movq %rax, -8(%r9,%rdi,8) +movq 8(%rbx,%rdi,8), %rax +mulq 8(%r11,%rdi,8) +orq %rcx, %rsi +movq %rax, (%r9,%rdi,8) +addq $2, %rdi +movq %rdx, %rcx +orq %rsi, %rcx +cmpq %r10, %rdi +jne .LBB6175_10 +``` + +This is scalar unsigned high-half multiplication unrolled by two, followed by +a one-lane remainder when the row count is odd. The two loads, multiplies, and +stores are independent except for the register OR reduction. There is no +`vector.body` in this fast value loop. + +## IR facts + +The all-varying functions are internal and take the source structure through a +`noalias readonly` pointer and return storage through a `noalias writeonly` +pointer. The allocated output stores carry a distinct `!alias.scope` and +`!noalias`; both input loads carry input-side `!noalias`. The second input +length check has become `llvm.assume`, so no panic branch remains in either hot +loop. A slice/assert failure edge exists before the loop at the output-length +validation boundary. + +The u64 IR loop is unrolled by two and reduces two i128 high halves through +scalar `or i64`; it has a one-lane epilogue. The i64 IR loop is scalar and uses +an i128 signed multiply, truncation, arithmetic sign extraction, XOR, and a +loop-carried register OR. Neither fast loop contains a call. + +Both monomorphs have this parameter-level ownership shape (metadata IDs differ +between them): + +```llvm +define internal fastcc void @checked_op_lanes( + ptr noalias writable writeonly %output, + ptr noalias readonly %source, + i64 %valid_rows_tag, + ptr readonly %valid_rows_data) +``` + +The relevant u64 body is structurally: + +```llvm +%failed = phi i64 [ 0, %preheader ], [ %failed_2, %loop ] +%lhs_0 = load i64, ptr %lhs_ptr_0, !noalias !input_scope +%rhs_0 = load i64, ptr %rhs_ptr_0, !noalias !input_scope +%low_0 = mul i64 %rhs_0, %lhs_0 +%wide_0 = mul nuw i128 (zext i64 %rhs_0), (zext i64 %lhs_0) +%high_0 = trunc i128 (lshr i128 %wide_0, 64) to i64 +%failed_1 = or i64 %failed, %high_0 +store i64 %low_0, ptr %output_0, !alias.scope !output_scope, !noalias !output_noalias + +%lhs_1 = load i64, ptr %lhs_ptr_1, !noalias !input_scope +%rhs_1 = load i64, ptr %rhs_ptr_1, !noalias !input_scope +%low_1 = mul i64 %rhs_1, %lhs_1 +%wide_1 = mul nuw i128 (zext i64 %rhs_1), (zext i64 %lhs_1) +%high_1 = trunc i128 (lshr i128 %wide_1, 64) to i64 +%failed_2 = or i64 %failed_1, %high_1 +store i64 %low_1, ptr %output_1, !alias.scope !output_scope, !noalias !output_noalias +``` + +The relevant i64 body is structurally: + +```llvm +%failed = phi i64 [ 0, %preheader ], [ %failed_next, %loop ] +%lhs = load i64, ptr %lhs_ptr, !noalias !input_scope +%rhs = load i64, ptr %rhs_ptr, !noalias !input_scope +%wide = mul nsw i128 (sext i64 %rhs), (sext i64 %lhs) +%low = trunc i128 %wide to i64 +%high = trunc i128 (lshr i128 %wide, 64) to i64 +%discarded_mismatch = xor i64 (ashr i64 %low, 63), %high +%failed_next = or i64 %discarded_mismatch, %failed +store i64 %low, ptr %output, !alias.scope !output_scope, !noalias !output_noalias +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/candidate-i64-mul-dense-ll.md b/research/rowfn-x86-2026-08-07/codegen/candidate-i64-mul-dense-ll.md new file mode 100644 index 00000000000..b93c15fd04e --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/candidate-i64-mul-dense-ll.md @@ -0,0 +1,84 @@ + + + +# `candidate-i64-mul-dense.ll` + +```ll + %_16456.i = phi i64 [ 1, %bb28.lr.ph.i ], [ %_164.i, %bb28.i ] + %iter.sroa.0.055.i = phi i64 [ 0, %bb28.lr.ph.i ], [ %_16456.i, %bb28.i ] + %accumulated.sroa.0.054.i = phi i64 [ 0, %bb28.lr.ph.i ], [ %77, %bb28.i ] + #dbg_value(i64 %iter.sroa.0.055.i, !561376, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !561961) + #dbg_value(i64 %accumulated.sroa.0.054.i, !561355, !DIExpression(), !561631) + #dbg_value(i64 %iter.sroa.0.055.i, !561378, !DIExpression(), !562028) + #dbg_value(ptr undef, !558643, !DIExpression(), !561458) + #dbg_value(i64 %iter.sroa.0.055.i, !558649, !DIExpression(), !561458) + #dbg_value(ptr poison, !559346, !DIExpression(), !562029) + #dbg_value(i64 %iter.sroa.0.055.i, !559351, !DIExpression(), !562029) + #dbg_value(ptr poison, !559346, !DIExpression(), !562031) + #dbg_value(i64 %iter.sroa.0.055.i, !559351, !DIExpression(), !562031) + #dbg_value(ptr poison, !561841, !DIExpression(), !562033) + #dbg_value(i64 %iter.sroa.0.055.i, !561851, !DIExpression(), !562033) + %75 = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.055.i, !dbg !562035 + %_0.i5.i.i = load i64, ptr %75, align 8, !dbg !562035, !noalias !562036, !noundef !23 + %76 = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.055.i, !dbg !562039 + %_0.i.i123.i = load i64, ptr %76, align 8, !dbg !562039, !noalias !562036, !noundef !23 + %_3.i126.i = getelementptr inbounds nuw i64, ptr %ptr.i.i, i64 %iter.sroa.0.055.i, !dbg !562040 + #dbg_value(ptr poison, !561857, !DIExpression(), !562041) + #dbg_value(ptr poison, !561867, !DIExpression(), !562041) + #dbg_value(i64 %_0.i.i123.i, !561868, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !562041) + #dbg_value(i64 %_0.i5.i.i, !561868, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !562041) + #dbg_value(ptr %_3.i126.i, !561863, !DIExpression(), !562041) + #dbg_value(i64 %_0.i.i123.i, !561864, !DIExpression(), !562043) + #dbg_value(i64 %_0.i.i123.i, !561873, !DIExpression(), !562044) + #dbg_value(i64 %_0.i5.i.i, !561866, !DIExpression(), !562043) + #dbg_value(i64 %_0.i5.i.i, !561880, !DIExpression(), !562044) + #dbg_value(ptr %_3.i126.i, !561879, !DIExpression(), !562044) + #dbg_value(ptr %_3.i126.i, !561886, !DIExpression(), !562046) + #dbg_value(i64 %_0.i.i123.i, !561892, !DIExpression(), !562048) + #dbg_value(i64 %_0.i5.i.i, !561901, !DIExpression(), !562048) + #dbg_value(i64 %_0.i.i123.i, !561904, !DIExpression(), !562050) + #dbg_value(i64 %_0.i.i123.i, !561910, !DIExpression(), !562052) + #dbg_value(i64 %_0.i5.i.i, !561907, !DIExpression(), !562050) + #dbg_value(i64 %_0.i5.i.i, !561913, !DIExpression(), !562052) + %_0.i.i128.i = mul i64 %_0.i.i123.i, %_0.i5.i.i, !dbg !562054 + #dbg_value(i64 %_0.i.i123.i, !561917, !DIExpression(), !562055) + #dbg_value(i64 %_0.i.i123.i, !561923, !DIExpression(), !562057) + #dbg_value(i64 %_0.i5.i.i, !561922, !DIExpression(), !562055) + #dbg_value(i64 %_0.i5.i.i, !561925, !DIExpression(), !562057) + %_4.i1.i.i = sext i64 %_0.i.i123.i to i128, !dbg !562058 + %_5.i.i.i = sext i64 %_0.i5.i.i to i128, !dbg !562059 + %wide.i.i.i = mul nsw i128 %_4.i1.i.i, %_5.i.i.i, !dbg !562058 + #dbg_value(i128 %wide.i.i.i, !561926, !DIExpression(), !562060) + %kept.i.i.i = trunc i128 %wide.i.i.i to i64, !dbg !562061 + #dbg_value(i64 %kept.i.i.i, !561928, !DIExpression(), !562062) + %_8.i.i.i = lshr i128 %wide.i.i.i, 64, !dbg !562063 + %discarded.i.i.i = trunc nuw i128 %_8.i.i.i to i64, !dbg !562064 + #dbg_value(i64 %discarded.i.i.i, !561930, !DIExpression(), !562065) + %_10.i.i.i = ashr i64 %kept.i.i.i, 63, !dbg !562066 + %_9.i.i.i = xor i64 %_10.i.i.i, %discarded.i.i.i, !dbg !562067 + #dbg_value(i64 %_0.i.i128.i, !561881, !DIExpression(), !562068) + #dbg_value(i64 %_0.i.i128.i, !561889, !DIExpression(), !562046) + #dbg_value(i64 %_9.i.i.i, !561883, !DIExpression(), !562068) + store i64 %_0.i.i128.i, ptr %_3.i126.i, align 8, !dbg !562069, !alias.scope !562070, !noalias !561484 + #dbg_value(i64 %_9.i.i.i, !561469, !DIExpression(), !561472) + #dbg_value(ptr undef, !561463, !DIExpression(), !561472) + %77 = or i64 %_9.i.i.i, %accumulated.sroa.0.054.i, !dbg !562073 + #dbg_value(i64 %77, !561355, !DIExpression(), !561631) + #dbg_value(i64 %_16456.i, !561376, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !561961) + #dbg_value(ptr undef, !561426, !DIExpression(), !561451) + #dbg_value(ptr undef, !561414, !DIExpression(), !561447) + #dbg_value(ptr undef, !561430, !DIExpression(), !561452) + #dbg_value(ptr poison, !561433, !DIExpression(), !561452) + %_164.i = add i64 %_16456.i, 1, !dbg !562074 + #dbg_value(i64 poison, !561376, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !561961) + %exitcond.not.i = icmp eq i64 %_16456.i, %4, !dbg !561962 + br i1 %exitcond.not.i, label %bb54.i, label %bb28.i, !dbg !561963 + +bb59.i: ; preds = %bb24.i + call void @llvm.memcpy.p0.p0.i64(ptr noundef nonnull align 8 dereferenceable(48) %71, ptr noundef nonnull align 8 dereferenceable(48) %_59.i, i64 48, i1 false), !dbg !562075, !noalias !561484 + call void @llvm.lifetime.end.p0(i64 48, ptr nonnull %_59.i), !dbg !561556, !noalias !561565 + %_49.sroa.4.0..sroa_idx.i = getelementptr inbounds nuw i8, ptr %_0, i64 16, !dbg !561964 + call void @llvm.memcpy.p0.p0.i64(ptr noundef nonnull align 8 dereferenceable(24) %_49.sroa.4.0..sroa_idx.i, ptr noundef nonnull align 8 dereferenceable(24) %_57.i, i64 24, i1 false), !dbg !561556, !noalias !561605 + call void @llvm.lifetime.end.p0(i64 24, ptr nonnull %_57.i), !dbg !561556, !noalias !561565 + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/candidate-i64-mul-dense-s.md b/research/rowfn-x86-2026-08-07/codegen/candidate-i64-mul-dense-s.md new file mode 100644 index 00000000000..9d2c024bac0 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/candidate-i64-mul-dense-s.md @@ -0,0 +1,69 @@ + + + +# `candidate-i64-mul-dense.s` + +```s + movq -376(%rbp), %r15 +.Ltmp108913: + .loc 524 86 32 + testq %r15, %r15 + .loc 524 86 16 is_stmt 0 + je .LBB1673_26 +.Ltmp108914: + .loc 563 318 19 is_stmt 1 + xorq %r14, %rdi +.Ltmp108915: + .loc 563 0 19 is_stmt 0 + xorq %r14, %r10 +.Ltmp108916: + .loc 524 88 17 is_stmt 1 + orq %rdi, %r10 +.Ltmp108917: + jne .LBB1673_42 +.Ltmp108918: + .loc 182 1904 50 + testq %r14, %r14 +.Ltmp108919: + .loc 524 92 26 + je .LBB1673_41 +.Ltmp108920: + .loc 524 0 26 is_stmt 0 + movq -320(%rbp), %rsi +.Ltmp108921: + xorl %edi, %edi + xorl %ecx, %ecx +.Ltmp108922: + .p2align 4 +.LBB1673_25: + .loc 564 62 9 is_stmt 1 + movq (%r15,%rdi,8), %rax +.Ltmp108923: + .loc 565 193 24 + imulq (%rsi,%rdi,8) +.Ltmp108924: + .loc 207 475 9 + movq %rax, (%r9,%rdi,8) +.Ltmp108925: + .loc 565 197 26 + sarq $63, %rax +.Ltmp108926: + .loc 565 197 13 is_stmt 0 + xorq %rdx, %rax +.Ltmp108927: + .loc 566 109 21 is_stmt 1 + orq %rax, %rcx +.Ltmp108928: + .loc 182 1904 50 + incq %rdi +.Ltmp108929: + cmpq %rdi, %r14 + jne .LBB1673_25 + jmp .LBB1673_62 +.Ltmp108930: +.LBB1673_26: + .loc 563 90 47 + cmpq %r14, %rdi + sete %cl + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/candidate-u64-mul-dense-ll.md b/research/rowfn-x86-2026-08-07/codegen/candidate-u64-mul-dense-ll.md new file mode 100644 index 00000000000..05db6cccb49 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/candidate-u64-mul-dense-ll.md @@ -0,0 +1,137 @@ + + + +# `candidate-u64-mul-dense.ll` + +```ll +terminate.i81.i: ; preds = %cleanup.i80.i + %114 = landingpad { ptr, i32 } + filter [0 x ptr] zeroinitializer +; call core::panicking::panic_in_cleanup + call void @_ZN4core9panicking16panic_in_cleanup17h8f68387bb6cbbf54E() #88, !dbg !588172, !noalias !587626 + unreachable, !dbg !588172 + +bb28.i: ; preds = %bb28.i, %bb28.lr.ph.i.new + %_16456.i = phi i64 [ 1, %bb28.lr.ph.i.new ], [ %_164.i.1, %bb28.i ] + %iter.sroa.0.055.i = phi i64 [ 0, %bb28.lr.ph.i.new ], [ %_164.i, %bb28.i ] + %accumulated.sroa.0.054.i = phi i64 [ 0, %bb28.lr.ph.i.new ], [ %120, %bb28.i ] + %niter = phi i64 [ 0, %bb28.lr.ph.i.new ], [ %niter.next.1, %bb28.i ] + #dbg_value(i64 %iter.sroa.0.055.i, !587525, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !588106) + #dbg_value(i64 %accumulated.sroa.0.054.i, !587504, !DIExpression(), !587773) + #dbg_value(i64 %iter.sroa.0.055.i, !587527, !DIExpression(), !588173) + #dbg_value(ptr undef, !579662, !DIExpression(), !587607) + #dbg_value(i64 %iter.sroa.0.055.i, !579668, !DIExpression(), !587607) + #dbg_value(ptr poison, !580362, !DIExpression(), !588174) + #dbg_value(i64 %iter.sroa.0.055.i, !580367, !DIExpression(), !588174) + #dbg_value(ptr poison, !580362, !DIExpression(), !588176) + #dbg_value(i64 %iter.sroa.0.055.i, !580367, !DIExpression(), !588176) + #dbg_value(ptr poison, !587985, !DIExpression(), !588178) + #dbg_value(i64 %iter.sroa.0.055.i, !587986, !DIExpression(), !588178) + %115 = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.055.i, !dbg !588180 + %_0.i5.i.i = load i64, ptr %115, align 8, !dbg !588180, !noalias !588181, !noundef !23 + %116 = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.055.i, !dbg !588184 + %_0.i.i123.i = load i64, ptr %116, align 8, !dbg !588184, !noalias !588181, !noundef !23 + %_3.i126.i = getelementptr inbounds nuw i64, ptr %ptr.i.i, i64 %iter.sroa.0.055.i, !dbg !588185 + #dbg_value(ptr poison, !588025, !DIExpression(), !588186) + #dbg_value(ptr poison, !588026, !DIExpression(), !588186) + #dbg_value(i64 %_0.i.i123.i, !588027, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !588186) + #dbg_value(i64 %_0.i5.i.i, !588027, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !588186) + #dbg_value(ptr %_3.i126.i, !588022, !DIExpression(), !588186) + #dbg_value(i64 %_0.i.i123.i, !588023, !DIExpression(), !588188) + #dbg_value(i64 %_0.i.i123.i, !588010, !DIExpression(), !588189) + #dbg_value(i64 %_0.i5.i.i, !588024, !DIExpression(), !588188) + #dbg_value(i64 %_0.i5.i.i, !588011, !DIExpression(), !588189) + #dbg_value(ptr %_3.i126.i, !588009, !DIExpression(), !588189) + #dbg_value(ptr %_3.i126.i, !588036, !DIExpression(), !588191) + #dbg_value(i64 %_0.i.i123.i, !588001, !DIExpression(), !588193) + #dbg_value(i64 %_0.i5.i.i, !588002, !DIExpression(), !588193) + #dbg_value(i64 %_0.i.i123.i, !588067, !DIExpression(), !588195) + #dbg_value(i64 %_0.i.i123.i, !588073, !DIExpression(), !588197) + #dbg_value(i64 %_0.i5.i.i, !588070, !DIExpression(), !588195) + #dbg_value(i64 %_0.i5.i.i, !588076, !DIExpression(), !588197) + %_0.i3.i.i = mul i64 %_0.i.i123.i, %_0.i5.i.i, !dbg !588199 + #dbg_value(i64 %_0.i.i123.i, !587992, !DIExpression(), !588200) + #dbg_value(i64 %_0.i.i123.i, !587994, !DIExpression(), !588202) + #dbg_value(i64 %_0.i5.i.i, !587993, !DIExpression(), !588200) + #dbg_value(i64 %_0.i5.i.i, !587995, !DIExpression(), !588202) + %_5.i.i.i = zext i64 %_0.i.i123.i to i128, !dbg !588203 + %_6.i.i.i = zext i64 %_0.i5.i.i to i128, !dbg !588204 + %_4.i1.i.i = mul nuw i128 %_5.i.i.i, %_6.i.i.i, !dbg !588205 + %_3.i2.i.i = lshr i128 %_4.i1.i.i, 64, !dbg !588206 + %_0.i.i128.i = trunc nuw i128 %_3.i2.i.i to i64, !dbg !588207 + #dbg_value(i64 %_0.i3.i.i, !588012, !DIExpression(), !588208) + #dbg_value(i64 %_0.i3.i.i, !588037, !DIExpression(), !588191) + #dbg_value(i64 %_0.i.i128.i, !588014, !DIExpression(), !588208) + store i64 %_0.i3.i.i, ptr %_3.i126.i, align 8, !dbg !588209, !alias.scope !588210, !noalias !587626 + #dbg_value(i64 %_0.i.i128.i, !561469, !DIExpression(), !587614) + #dbg_value(ptr undef, !561463, !DIExpression(), !587614) + %117 = or i64 %accumulated.sroa.0.054.i, %_0.i.i128.i, !dbg !588213 + #dbg_value(i64 %117, !587504, !DIExpression(), !587773) + #dbg_value(i64 %_16456.i, !587525, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !588106) + #dbg_value(ptr undef, !587575, !DIExpression(), !587600) + #dbg_value(ptr undef, !587563, !DIExpression(), !587596) + #dbg_value(ptr undef, !587579, !DIExpression(), !587601) + #dbg_value(ptr poison, !587582, !DIExpression(), !587601) + %_164.i = add i64 %_16456.i, 1, !dbg !588214 + #dbg_value(i64 poison, !587525, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !588106) + #dbg_value(i64 %_16456.i, !587525, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !588106) + #dbg_value(i64 %_16456.i, !587527, !DIExpression(), !588173) + #dbg_value(i64 %_16456.i, !579668, !DIExpression(), !587607) + #dbg_value(i64 %_16456.i, !580367, !DIExpression(), !588174) + #dbg_value(i64 %_16456.i, !580367, !DIExpression(), !588176) + #dbg_value(i64 %_16456.i, !587986, !DIExpression(), !588178) + %118 = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %_16456.i, !dbg !588180 + %_0.i5.i.i.1 = load i64, ptr %118, align 8, !dbg !588180, !noalias !588181, !noundef !23 + %119 = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %_16456.i, !dbg !588184 + %_0.i.i123.i.1 = load i64, ptr %119, align 8, !dbg !588184, !noalias !588181, !noundef !23 + %_3.i126.i.1 = getelementptr inbounds nuw i64, ptr %ptr.i.i, i64 %_16456.i, !dbg !588185 + #dbg_value(i64 %_0.i.i123.i.1, !588027, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !588186) + #dbg_value(i64 %_0.i5.i.i.1, !588027, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !588186) + #dbg_value(ptr %_3.i126.i.1, !588022, !DIExpression(), !588186) + #dbg_value(i64 %_0.i.i123.i.1, !588023, !DIExpression(), !588188) + #dbg_value(i64 %_0.i.i123.i.1, !588010, !DIExpression(), !588189) + #dbg_value(i64 %_0.i5.i.i.1, !588024, !DIExpression(), !588188) + #dbg_value(i64 %_0.i5.i.i.1, !588011, !DIExpression(), !588189) + #dbg_value(ptr %_3.i126.i.1, !588009, !DIExpression(), !588189) + #dbg_value(ptr %_3.i126.i.1, !588036, !DIExpression(), !588191) + #dbg_value(i64 %_0.i.i123.i.1, !588001, !DIExpression(), !588193) + #dbg_value(i64 %_0.i5.i.i.1, !588002, !DIExpression(), !588193) + #dbg_value(i64 %_0.i.i123.i.1, !588067, !DIExpression(), !588195) + #dbg_value(i64 %_0.i.i123.i.1, !588073, !DIExpression(), !588197) + #dbg_value(i64 %_0.i5.i.i.1, !588070, !DIExpression(), !588195) + #dbg_value(i64 %_0.i5.i.i.1, !588076, !DIExpression(), !588197) + %_0.i3.i.i.1 = mul i64 %_0.i.i123.i.1, %_0.i5.i.i.1, !dbg !588199 + #dbg_value(i64 %_0.i.i123.i.1, !587992, !DIExpression(), !588200) + #dbg_value(i64 %_0.i.i123.i.1, !587994, !DIExpression(), !588202) + #dbg_value(i64 %_0.i5.i.i.1, !587993, !DIExpression(), !588200) + #dbg_value(i64 %_0.i5.i.i.1, !587995, !DIExpression(), !588202) + %_5.i.i.i.1 = zext i64 %_0.i.i123.i.1 to i128, !dbg !588203 + %_6.i.i.i.1 = zext i64 %_0.i5.i.i.1 to i128, !dbg !588204 + %_4.i1.i.i.1 = mul nuw i128 %_5.i.i.i.1, %_6.i.i.i.1, !dbg !588205 + %_3.i2.i.i.1 = lshr i128 %_4.i1.i.i.1, 64, !dbg !588206 + %_0.i.i128.i.1 = trunc nuw i128 %_3.i2.i.i.1 to i64, !dbg !588207 + #dbg_value(i64 %_0.i3.i.i.1, !588012, !DIExpression(), !588208) + #dbg_value(i64 %_0.i3.i.i.1, !588037, !DIExpression(), !588191) + #dbg_value(i64 %_0.i.i128.i.1, !588014, !DIExpression(), !588208) + store i64 %_0.i3.i.i.1, ptr %_3.i126.i.1, align 8, !dbg !588209, !alias.scope !588210, !noalias !587626 + #dbg_value(i64 %_0.i.i128.i.1, !561469, !DIExpression(), !587614) + %120 = or i64 %117, %_0.i.i128.i.1, !dbg !588213 + #dbg_value(i64 %120, !587504, !DIExpression(), !587773) + #dbg_value(i64 %_164.i, !587525, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !588106) + %_164.i.1 = add i64 %_16456.i, 2, !dbg !588214 + #dbg_value(i64 poison, !587525, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !588106) + %niter.next.1 = add i64 %niter, 2, !dbg !588108 + %niter.ncmp.1 = icmp eq i64 %niter.next.1, %unroll_iter, !dbg !588108 + br i1 %niter.ncmp.1, label %bb54.i.loopexit135.unr-lcssa, label %bb28.i, !dbg !588108 + +bb59.i: ; preds = %bb24.i + call void @llvm.memcpy.p0.p0.i64(ptr noundef nonnull align 8 dereferenceable(48) %111, ptr noundef nonnull align 8 dereferenceable(48) %_59.i, i64 48, i1 false), !dbg !588215, !noalias !587626 + call void @llvm.lifetime.end.p0(i64 48, ptr nonnull %_59.i), !dbg !587698, !noalias !587707 + %_49.sroa.4.0..sroa_idx.i = getelementptr inbounds nuw i8, ptr %_0, i64 16, !dbg !588109 + call void @llvm.memcpy.p0.p0.i64(ptr noundef nonnull align 8 dereferenceable(24) %_49.sroa.4.0..sroa_idx.i, ptr noundef nonnull align 8 dereferenceable(24) %_57.i, i64 24, i1 false), !dbg !587698, !noalias !587747 + call void @llvm.lifetime.end.p0(i64 24, ptr nonnull %_57.i), !dbg !587698, !noalias !587707 + br label %bb61.i, !dbg !587940 + +bb39.i: ; preds = %bb2.i109.i, %"_ZN12vortex_array9scalar_fn3row7element5tuple18ArgColumn$LT$T$GT$14addresses_rows17h8cb4442712b37c9aE.exit.i.i" + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/candidate-u64-mul-dense-s.md b/research/rowfn-x86-2026-08-07/codegen/candidate-u64-mul-dense-s.md new file mode 100644 index 00000000000..0676296c7a7 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/candidate-u64-mul-dense-s.md @@ -0,0 +1,70 @@ + + + +# `candidate-u64-mul-dense.s` + +```s + .loc 524 92 26 + andq $-2, %r15 + leaq (%r12,%rdx), %r11 + addq $8, %r11 + xorl %ecx, %ecx + xorl %r10d, %r10d +.Ltmp117915: +.LBB1693_70: + .loc 564 62 9 + movq (%r14,%r10,8), %rax +.Ltmp117916: + .loc 565 175 44 + mulq (%r8,%r10,8) +.Ltmp117917: + movq %rdx, %rdi +.Ltmp117918: + .loc 207 475 9 + movq %rax, -8(%r11,%r10,8) +.Ltmp117919: + .loc 564 62 9 + movq 8(%r14,%r10,8), %rax +.Ltmp117920: + .loc 565 175 44 + mulq 8(%r8,%r10,8) +.Ltmp117921: + .loc 566 109 21 + orq %rcx, %rdi +.Ltmp117922: + .loc 207 475 9 + movq %rax, (%r11,%r10,8) +.Ltmp117923: + .loc 565 175 44 + movq %rdx, %rcx +.Ltmp117924: + .loc 566 109 21 + orq %rdi, %rcx +.Ltmp117925: + .loc 524 92 26 + addq $2, %r10 + cmpq %r10, %r15 + jne .LBB1693_70 +.Ltmp117926: +.LBB1693_71: + testb $1, %sil + je .LBB1693_87 +.Ltmp117927: + .loc 564 62 9 + movq (%r14,%r10,8), %rax +.Ltmp117928: + .loc 565 175 44 + mulq (%r8,%r10,8) +.Ltmp117929: +.LBB1693_73: + .loc 207 475 9 + movq %rax, (%r9,%r10,8) +.Ltmp117930: + .loc 566 109 21 + orq %rdx, %rcx +.Ltmp117931: + .loc 566 0 21 is_stmt 0 + jmp .LBB1693_87 +.Ltmp117932: + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/final-i32-mul-constant-ll.md b/research/rowfn-x86-2026-08-07/codegen/final-i32-mul-constant-ll.md new file mode 100644 index 00000000000..7513df3b465 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/final-i32-mul-constant-ll.md @@ -0,0 +1,86 @@ + + + +# `final-i32-mul-constant.ll` + +```ll +bb27.preheader.i.split.us: ; preds = %bb27.preheader.i + br i1 %_3.i5.not.i.i, label %panic.i5.i5.i.invoke.i, label %bb27.i.us.preheader + +bb27.i.us.preheader: ; preds = %bb27.preheader.i.split.us + %57 = add nuw nsw i64 %len3.i.i.i, 1, !dbg !566996 + br label %bb27.i.us, !dbg !566996 + +bb27.i.us: ; preds = %bb27.i.us.preheader, %"_ZN12vortex_array9scalar_fn3row7element9primitive83_$LT$impl$u20$vortex_array..scalar_fn..row..element..InputElement$u20$for$u20$T$GT$3get17h2f14351c5788fe60E.exit8.i.i.i.us" + %_15854.i.us = phi i64 [ %_158.i.us, %"_ZN12vortex_array9scalar_fn3row7element9primitive83_$LT$impl$u20$vortex_array..scalar_fn..row..element..InputElement$u20$for$u20$T$GT$3get17h2f14351c5788fe60E.exit8.i.i.i.us" ], [ 1, %bb27.i.us.preheader ] + %iter.sroa.0.053.i.us = phi i64 [ %_15854.i.us, %"_ZN12vortex_array9scalar_fn3row7element9primitive83_$LT$impl$u20$vortex_array..scalar_fn..row..element..InputElement$u20$for$u20$T$GT$3get17h2f14351c5788fe60E.exit8.i.i.i.us" ], [ 0, %bb27.i.us.preheader ] + %accumulated.sroa.0.052.i.us = phi i1 [ %60, %"_ZN12vortex_array9scalar_fn3row7element9primitive83_$LT$impl$u20$vortex_array..scalar_fn..row..element..InputElement$u20$for$u20$T$GT$3get17h2f14351c5788fe60E.exit8.i.i.i.us" ], [ false, %bb27.i.us.preheader ] + #dbg_value(i64 %iter.sroa.0.053.i.us, !566730, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !566993) + #dbg_value(i64 %iter.sroa.0.053.i.us, !566732, !DIExpression(), !567057) + #dbg_value(ptr %columns.i, !545259, !DIExpression(), !567058) + #dbg_value(i64 %iter.sroa.0.053.i.us, !545260, !DIExpression(), !567058) + #dbg_value(ptr %columns.i, !545249, !DIExpression(), !567059) + #dbg_value(i64 %iter.sroa.0.053.i.us, !545250, !DIExpression(), !567059) + #dbg_value(ptr %columns.i, !545120, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !567060) + #dbg_value(ptr %columns.i, !545120, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !567062) + #dbg_value(ptr %columns.i, !545251, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !567063) + #dbg_value(i64 %iter.sroa.0.053.i.us, !545125, !DIExpression(), !567062) + %exitcond35.not = icmp eq i64 %_15854.i.us, %57, !dbg !566996 + br i1 %exitcond35.not, label %panic.i5.i5.i.invoke.i, label %"_ZN12vortex_array9scalar_fn3row7element9primitive83_$LT$impl$u20$vortex_array..scalar_fn..row..element..InputElement$u20$for$u20$T$GT$3get17h2f14351c5788fe60E.exit8.i.i.i.us", !dbg !566996 + +"_ZN12vortex_array9scalar_fn3row7element9primitive83_$LT$impl$u20$vortex_array..scalar_fn..row..element..InputElement$u20$for$u20$T$GT$3get17h2f14351c5788fe60E.exit8.i.i.i.us": ; preds = %bb27.i.us + #dbg_value(ptr %columns.i, !545251, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !567063) + #dbg_value(ptr %columns.i, !545120, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !567062) + %58 = getelementptr inbounds nuw i32, ptr %data.i6.i.i.i, i64 %iter.sroa.0.053.i.us, !dbg !566996 + %_0.sroa.0.0.i.i.i.us = load i32, ptr %58, align 4, !dbg !567000, !noalias !566784, !noundef !23 + #dbg_value(ptr %14, !545249, !DIExpression(), !567064) + #dbg_value(i64 %iter.sroa.0.053.i.us, !545250, !DIExpression(), !567064) + #dbg_value(ptr %14, !545120, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !567065) + #dbg_value(ptr %14, !545120, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !567067) + #dbg_value(ptr %14, !545252, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !567068) + #dbg_value(i64 0, !545125, !DIExpression(), !567065) + %_0.sroa.0.0.i9.i.i.us = load i32, ptr %data.i6.i7.i.i, align 4, !dbg !567005, !noalias !566784, !noundef !23 + #dbg_value(ptr poison, !567031, !DIExpression(), !567069) + #dbg_value(ptr poison, !567032, !DIExpression(), !567069) + #dbg_value(i32 %_0.sroa.0.0.i.i.i.us, !567033, !DIExpression(DW_OP_LLVM_fragment, 0, 32), !567069) + #dbg_value(i32 %_0.sroa.0.0.i9.i.i.us, !567033, !DIExpression(DW_OP_LLVM_fragment, 32, 32), !567069) + #dbg_value(i32 %_0.sroa.0.0.i.i.i.us, !567029, !DIExpression(), !567070) + #dbg_value(i32 %_0.sroa.0.0.i9.i.i.us, !567030, !DIExpression(), !567070) + #dbg_value(i32 %_0.sroa.0.0.i.i.i.us, !567020, !DIExpression(), !567071) + #dbg_value(i32 %_0.sroa.0.0.i9.i.i.us, !567021, !DIExpression(), !567071) + #dbg_value(i32 %_0.sroa.0.0.i.i.i.us, !567015, !DIExpression(), !567072) + #dbg_value(i32 %_0.sroa.0.0.i.i.i.us, !567010, !DIExpression(), !567073) + #dbg_value(i32 %_0.sroa.0.0.i9.i.i.us, !567016, !DIExpression(), !567072) + #dbg_value(i32 %_0.sroa.0.0.i9.i.i.us, !567011, !DIExpression(), !567073) + %_0.i.i160.i.us = mul i32 %_0.sroa.0.0.i9.i.i.us, %_0.sroa.0.0.i.i.i.us, !dbg !567007 + #dbg_value(i32 %_0.sroa.0.0.i.i.i.us, !567039, !DIExpression(), !567074) + #dbg_value(i32 %_0.sroa.0.0.i.i.i.us, !567041, !DIExpression(), !567075) + #dbg_value(i32 %_0.sroa.0.0.i9.i.i.us, !567040, !DIExpression(), !567074) + #dbg_value(i32 %_0.sroa.0.0.i9.i.i.us, !567042, !DIExpression(), !567075) + %_4.i1.i.i.us = sext i32 %_0.sroa.0.0.i.i.i.us to i64, !dbg !567035 + %_5.i.i.i.us = sext i32 %_0.sroa.0.0.i9.i.i.us to i64, !dbg !567046 + %product.i.i.i.us = mul nsw i64 %_5.i.i.i.us, %_4.i1.i.i.us, !dbg !567035 + #dbg_value(i64 %product.i.i.i.us, !567043, !DIExpression(), !567076) + %59 = add nsw i64 %product.i.i.i.us, -2147483648, !dbg !567047 + %_0.sroa.0.0.i.i161.i.us = icmp ult i64 %59, -4294967296, !dbg !567047 + #dbg_value(i1 %_0.sroa.0.0.i.i161.i.us, !566736, !DIExpression(DW_OP_LLVM_convert, 1, DW_ATE_unsigned, DW_OP_LLVM_convert, 8, DW_ATE_unsigned, DW_OP_stack_value), !567077) + #dbg_value(i32 %_0.i.i160.i.us, !566734, !DIExpression(), !567077) + %self34.i.us = getelementptr inbounds nuw i32, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.053.i.us, !dbg !567048 + #dbg_value(ptr %self34.i.us, !567052, !DIExpression(), !567078) + #dbg_value(i32 %_0.i.i160.i.us, !567053, !DIExpression(), !567078) + store i32 %_0.i.i160.i.us, ptr %self34.i.us, align 4, !dbg !567049, !noalias !566784 + #dbg_value(ptr undef, !541560, !DIExpression(), !566763) + #dbg_value(i1 %_0.sroa.0.0.i.i161.i.us, !541568, !DIExpression(DW_OP_LLVM_convert, 1, DW_ATE_unsigned, DW_OP_LLVM_convert, 8, DW_ATE_unsigned, DW_OP_stack_value), !566763) + %60 = or i1 %accumulated.sroa.0.052.i.us, %_0.sroa.0.0.i.i161.i.us, !dbg !567055 + #dbg_value(i8 poison, !566728, !DIExpression(), !566992) + #dbg_value(i64 %_15854.i.us, !566730, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !566993) + #dbg_value(ptr undef, !566753, !DIExpression(), !566756) + #dbg_value(ptr undef, !566744, !DIExpression(), !566749) + #dbg_value(ptr undef, !566757, !DIExpression(), !566761) + #dbg_value(ptr poison, !566760, !DIExpression(), !566761) + %_158.i.us = add i64 %_15854.i.us, 1, !dbg !567079 + #dbg_value(i64 poison, !566730, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !566993) + %exitcond.not.i.us = icmp eq i64 %_15854.i.us, %len3.i.i.i, !dbg !566994 + br i1 %exitcond.not.i.us, label %bb33.i, label %bb27.i.us, !dbg !566995 + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/final-i32-mul-constant-s.md b/research/rowfn-x86-2026-08-07/codegen/final-i32-mul-constant-s.md new file mode 100644 index 00000000000..f82f250cdd1 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/final-i32-mul-constant-s.md @@ -0,0 +1,49 @@ + + + +# `final-i32-mul-constant.s` + +```s +.LBB1677_31: + .loc 564 47 9 is_stmt 1 + cmpq %rdi, %rdx + je .LBB1677_89 +.Ltmp110238: + .loc 564 47 9 is_stmt 0 + movslq (%r13,%rdi,4), %rcx +.Ltmp110239: + .loc 564 47 9 + movslq (%r12), %r8 +.Ltmp110240: + .loc 462 2133 13 is_stmt 1 + movl %r8d, %r10d + imull %ecx, %r10d +.Ltmp110241: + .loc 565 185 27 + imulq %rcx, %r8 +.Ltmp110242: + .loc 565 185 35 is_stmt 0 + addq $-2147483648, %r8 +.Ltmp110243: + cmpq %rax, %r8 + setb %cl +.Ltmp110244: + .loc 565 0 35 + movq -48(%rbp), %r8 +.Ltmp110245: + .loc 207 475 9 is_stmt 1 + movl %r10d, (%r8,%rdi,4) +.Ltmp110246: + .loc 566 821 53 + orb %cl, %r9b +.Ltmp110247: + .loc 182 1904 50 + incq %rdi +.Ltmp110248: + cmpq %rdi, %rdx +.Ltmp110249: + .loc 562 124 26 + jne .LBB1677_31 + jmp .LBB1677_64 + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/final-i64-mul-dense-ll.md b/research/rowfn-x86-2026-08-07/codegen/final-i64-mul-dense-ll.md new file mode 100644 index 00000000000..6b28482775c --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/final-i64-mul-dense-ll.md @@ -0,0 +1,96 @@ + + + +# `final-i64-mul-dense.ll` + +```ll +bb15.i.i: ; preds = %bb11.i, %bb15.i.i + %iter.sroa.0.012.i.i = phi i64 [ %_36.i.i, %bb15.i.i ], [ 0, %bb11.i ] + %failed.sroa.0.011.i.i = phi i64 [ %79, %bb15.i.i ], [ 0, %bb11.i ] + #dbg_value(i64 %iter.sroa.0.012.i.i, !564425, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564479) + #dbg_value(i64 %failed.sroa.0.011.i.i, !564423, !DIExpression(), !564478) + #dbg_value(i64 %iter.sroa.0.012.i.i, !564463, !DIExpression(), !564694) + #dbg_value(i64 %iter.sroa.0.012.i.i, !564456, !DIExpression(), !564457) + #dbg_value(i64 %iter.sroa.0.012.i.i, !564473, !DIExpression(), !564474) + %_36.i.i = add nuw i64 %iter.sroa.0.012.i.i, 1, !dbg !564695 + #dbg_value(i64 %_36.i.i, !564425, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564479) + #dbg_value(i64 %iter.sroa.0.012.i.i, !564427, !DIExpression(), !564696) + #dbg_value(i64 %iter.sroa.0.012.i.i, !564450, !DIExpression(), !564451) + #dbg_value(i64 %iter.sroa.0.012.i.i, !564697, !DIExpression(), !564701) + #dbg_value(ptr undef, !547076, !DIExpression(), !564445) + #dbg_value(i64 %iter.sroa.0.012.i.i, !547082, !DIExpression(), !564445) + #dbg_value(ptr poison, !547154, !DIExpression(), !564703) + #dbg_value(i64 %iter.sroa.0.012.i.i, !547155, !DIExpression(), !564703) + #dbg_value(i64 %iter.sroa.0.012.i.i, !547147, !DIExpression(), !564705) + #dbg_value(i64 %iter.sroa.0.012.i.i, !547139, !DIExpression(), !564707) + #dbg_value(ptr %column.val.i.i, !547146, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564705) + #dbg_value(ptr %column.val.i.i, !547140, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564707) + #dbg_value(i64 %len3.i.i.i, !547146, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !564705) + #dbg_value(i64 %len3.i.i.i, !547140, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !564707) + %_4.i.i.i.i = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.012.i.i, !dbg !564709 + %_0.i.i.i.i = load i64, ptr %_4.i.i.i.i, align 8, !dbg !564710, !noalias !564711, !noundef !23 + #dbg_value(ptr poison, !547154, !DIExpression(), !564715) + #dbg_value(i64 %iter.sroa.0.012.i.i, !547155, !DIExpression(), !564715) + #dbg_value(i64 %iter.sroa.0.012.i.i, !547147, !DIExpression(), !564717) + #dbg_value(i64 %iter.sroa.0.012.i.i, !547139, !DIExpression(), !564719) + #dbg_value(ptr %column5.val.i.i, !547146, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564717) + #dbg_value(ptr %column5.val.i.i, !547140, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564719) + #dbg_value(i64 %len3.i.i.i, !547146, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !564717) + #dbg_value(i64 %len3.i.i.i, !547140, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !564719) + %_5.i3.i.i.i = icmp ult i64 %iter.sroa.0.012.i.i, %len3.i.i.i, !dbg !564721 + tail call void @llvm.assume(i1 %_5.i3.i.i.i), !dbg !564722 + %_4.i4.i.i.i = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.012.i.i, !dbg !564723 + %_0.i5.i.i.i = load i64, ptr %_4.i4.i.i.i, align 8, !dbg !564724, !noalias !564711, !noundef !23 + #dbg_value(i64 %_0.i.i.i.i, !564429, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564725) + #dbg_value(i64 %_0.i5.i.i.i, !564429, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !564725) + #dbg_value(i64 %_0.i.i.i.i, !564726, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564734) + #dbg_value(i64 %_0.i5.i.i.i, !564726, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !564734) + #dbg_value(ptr poison, !564315, !DIExpression(), !564736) + #dbg_value(ptr poison, !564316, !DIExpression(), !564736) + #dbg_value(i64 %_0.i.i.i.i, !564317, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564736) + #dbg_value(i64 %_0.i5.i.i.i, !564317, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !564736) + #dbg_value(i64 %_0.i.i.i.i, !564313, !DIExpression(), !564738) + #dbg_value(i64 %_0.i5.i.i.i, !564314, !DIExpression(), !564738) + #dbg_value(i64 %_0.i.i.i.i, !564304, !DIExpression(), !564739) + #dbg_value(i64 %_0.i5.i.i.i, !564305, !DIExpression(), !564739) + #dbg_value(i64 %_0.i.i.i.i, !564293, !DIExpression(), !564741) + #dbg_value(i64 %_0.i.i.i.i, !564288, !DIExpression(), !564743) + #dbg_value(i64 %_0.i5.i.i.i, !564294, !DIExpression(), !564741) + #dbg_value(i64 %_0.i5.i.i.i, !564289, !DIExpression(), !564743) + %_0.i.i.i.i.i = mul i64 %_0.i5.i.i.i, %_0.i.i.i.i, !dbg !564745 + #dbg_value(i64 %_0.i.i.i.i, !564325, !DIExpression(), !564746) + #dbg_value(i64 %_0.i.i.i.i, !564327, !DIExpression(), !564748) + #dbg_value(i64 %_0.i5.i.i.i, !564326, !DIExpression(), !564746) + #dbg_value(i64 %_0.i5.i.i.i, !564328, !DIExpression(), !564748) + %_4.i1.i.i.i.i = sext i64 %_0.i.i.i.i to i128, !dbg !564749 + %_5.i.i.i.i.i = sext i64 %_0.i5.i.i.i to i128, !dbg !564750 + %wide.i.i.i.i.i = mul nsw i128 %_5.i.i.i.i.i, %_4.i1.i.i.i.i, !dbg !564749 + #dbg_value(i128 %wide.i.i.i.i.i, !564329, !DIExpression(), !564751) + %kept.i.i.i.i.i = trunc i128 %wide.i.i.i.i.i to i64, !dbg !564752 + #dbg_value(i64 %kept.i.i.i.i.i, !564331, !DIExpression(), !564753) + %_8.i.i.i.i.i = lshr i128 %wide.i.i.i.i.i, 64, !dbg !564754 + %discarded.i.i.i.i.i = trunc nuw i128 %_8.i.i.i.i.i to i64, !dbg !564755 + #dbg_value(i64 %discarded.i.i.i.i.i, !564333, !DIExpression(), !564756) + %_10.i.i.i.i.i = ashr i64 %kept.i.i.i.i.i, 63, !dbg !564757 + %_9.i.i.i.i.i = xor i64 %_10.i.i.i.i.i, %discarded.i.i.i.i.i, !dbg !564758 + #dbg_value(i64 poison, !564431, !DIExpression(), !564759) + #dbg_value(i64 %_9.i.i.i.i.i, !564433, !DIExpression(), !564759) + #dbg_value(ptr undef, !564034, !DIExpression(), !564443) + #dbg_value(i64 %_9.i.i.i.i.i, !564040, !DIExpression(), !564443) + %79 = or i64 %_9.i.i.i.i.i, %failed.sroa.0.011.i.i, !dbg !564760 + #dbg_value(i64 %79, !564423, !DIExpression(), !564478) + #dbg_value(i64 %_0.i.i.i.i.i, !564431, !DIExpression(), !564759) + #dbg_value(ptr %_4.sroa.10.0.i.i, !564700, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !564701) + #dbg_value(i64 %index.i, !564700, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !564701) + %self4.i.i = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.012.i.i, !dbg !564761 + #dbg_value(ptr %self4.i.i, !564762, !DIExpression(), !564766) + #dbg_value(i64 %_0.i.i.i.i.i, !564765, !DIExpression(), !564766) + store i64 %_0.i.i.i.i.i, ptr %self4.i.i, align 8, !dbg !564768, !alias.scope !564439, !noalias !564769 + #dbg_value(ptr undef, !564467, !DIExpression(), !564480) + #dbg_value(ptr undef, !564462, !DIExpression(), !564481) + #dbg_value(ptr undef, !564482, !DIExpression(), !564486) + #dbg_value(ptr poison, !564485, !DIExpression(), !564486) + %exitcond.not.i.i = icmp eq i64 %_36.i.i, %len3.i4.i.i.fr, !dbg !564770 + br i1 %exitcond.not.i.i, label %bb33.i, label %bb15.i.i, !dbg !564488 + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/final-i64-mul-dense-s.md b/research/rowfn-x86-2026-08-07/codegen/final-i64-mul-dense-s.md new file mode 100644 index 00000000000..56206a6ec40 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/final-i64-mul-dense-s.md @@ -0,0 +1,34 @@ + + + +# `final-i64-mul-dense.s` + +```s +.LBB1675_25: + .loc 567 39 18 is_stmt 1 + movq (%r13,%rsi,8), %rax +.Ltmp109477: + .loc 565 194 24 + imulq (%r12,%rsi,8) +.Ltmp109478: + .loc 207 475 9 + movq %rax, (%rdi,%rsi,8) +.Ltmp109479: + .loc 156 717 17 + incq %rsi +.Ltmp109480: + .loc 565 198 26 + sarq $63, %rax +.Ltmp109481: + .loc 565 198 13 is_stmt 0 + xorq %rdx, %rax +.Ltmp109482: + .loc 566 821 53 is_stmt 1 + orq %rax, %rcx +.Ltmp109483: + .loc 182 1904 50 + cmpq %rsi, %r9 + jne .LBB1675_25 + jmp .LBB1675_60 + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/final-u64-mul-dense-ll.md b/research/rowfn-x86-2026-08-07/codegen/final-u64-mul-dense-ll.md new file mode 100644 index 00000000000..4e40cfb8f2c --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/final-u64-mul-dense-ll.md @@ -0,0 +1,322 @@ + + + +# `final-u64-mul-dense.ll` + +```ll +bb15.i.i: ; preds = %bb15.i.i, %bb15.i.i.preheader.new + %iter.sroa.0.012.i.i = phi i64 [ 0, %bb15.i.i.preheader.new ], [ %_36.i.i.1, %bb15.i.i ] + %failed.sroa.0.011.i.i = phi i64 [ 0, %bb15.i.i.preheader.new ], [ %116, %bb15.i.i ] + %niter = phi i64 [ 0, %bb15.i.i.preheader.new ], [ %niter.next.1, %bb15.i.i ] + #dbg_value(i64 %iter.sroa.0.012.i.i, !569875, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !569928) + #dbg_value(i64 %failed.sroa.0.011.i.i, !569873, !DIExpression(), !569927) + #dbg_value(i64 %iter.sroa.0.012.i.i, !569912, !DIExpression(), !570143) + #dbg_value(i64 %iter.sroa.0.012.i.i, !569905, !DIExpression(), !569906) + #dbg_value(i64 %iter.sroa.0.012.i.i, !569922, !DIExpression(), !569923) + %_36.i.i = or disjoint i64 %iter.sroa.0.012.i.i, 1, !dbg !570144 + #dbg_value(i64 %_36.i.i, !569875, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !569928) + #dbg_value(i64 %iter.sroa.0.012.i.i, !569877, !DIExpression(), !570145) + #dbg_value(i64 %iter.sroa.0.012.i.i, !569899, !DIExpression(), !569900) + #dbg_value(i64 %iter.sroa.0.012.i.i, !570146, !DIExpression(), !570150) + #dbg_value(ptr undef, !551471, !DIExpression(), !569894) + #dbg_value(i64 %iter.sroa.0.012.i.i, !551477, !DIExpression(), !569894) + #dbg_value(ptr poison, !551549, !DIExpression(), !570152) + #dbg_value(i64 %iter.sroa.0.012.i.i, !551550, !DIExpression(), !570152) + #dbg_value(i64 %iter.sroa.0.012.i.i, !551542, !DIExpression(), !570154) + #dbg_value(i64 %iter.sroa.0.012.i.i, !551534, !DIExpression(), !570156) + #dbg_value(ptr %column.val.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570154) + #dbg_value(ptr %column.val.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570156) + #dbg_value(i64 %len3.i.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570154) + #dbg_value(i64 %len3.i.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570156) + %_4.i.i.i.i = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.012.i.i, !dbg !570158 + %_0.i.i.i.i = load i64, ptr %_4.i.i.i.i, align 8, !dbg !570159, !noalias !570160, !noundef !23 + #dbg_value(ptr poison, !551549, !DIExpression(), !570164) + #dbg_value(i64 %iter.sroa.0.012.i.i, !551550, !DIExpression(), !570164) + #dbg_value(i64 %iter.sroa.0.012.i.i, !551542, !DIExpression(), !570166) + #dbg_value(i64 %iter.sroa.0.012.i.i, !551534, !DIExpression(), !570168) + #dbg_value(ptr %column5.val.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570166) + #dbg_value(ptr %column5.val.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570168) + #dbg_value(i64 %len3.i.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570166) + #dbg_value(i64 %len3.i.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570168) + %_5.i3.i.i.i = icmp ult i64 %iter.sroa.0.012.i.i, %len3.i.i.i, !dbg !570170 + tail call void @llvm.assume(i1 %_5.i3.i.i.i), !dbg !570171 + %_4.i4.i.i.i = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.012.i.i, !dbg !570172 + %_0.i5.i.i.i = load i64, ptr %_4.i4.i.i.i, align 8, !dbg !570173, !noalias !570160, !noundef !23 + #dbg_value(i64 %_0.i.i.i.i, !569879, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570174) + #dbg_value(i64 %_0.i5.i.i.i, !569879, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570174) + #dbg_value(i64 %_0.i.i.i.i, !570175, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570183) + #dbg_value(i64 %_0.i5.i.i.i, !570175, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570183) + #dbg_value(ptr poison, !569758, !DIExpression(), !570185) + #dbg_value(ptr poison, !569759, !DIExpression(), !570185) + #dbg_value(i64 %_0.i.i.i.i, !569760, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570185) + #dbg_value(i64 %_0.i5.i.i.i, !569760, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570185) + #dbg_value(i64 %_0.i.i.i.i, !569756, !DIExpression(), !570187) + #dbg_value(i64 %_0.i5.i.i.i, !569757, !DIExpression(), !570187) + #dbg_value(i64 %_0.i.i.i.i, !569747, !DIExpression(), !570188) + #dbg_value(i64 %_0.i5.i.i.i, !569748, !DIExpression(), !570188) + #dbg_value(i64 %_0.i.i.i.i, !569740, !DIExpression(), !570190) + #dbg_value(i64 %_0.i.i.i.i, !569735, !DIExpression(), !570192) + #dbg_value(i64 %_0.i5.i.i.i, !569741, !DIExpression(), !570190) + #dbg_value(i64 %_0.i5.i.i.i, !569736, !DIExpression(), !570192) + %_0.i3.i.i.i.i = mul i64 %_0.i5.i.i.i, %_0.i.i.i.i, !dbg !570194 + #dbg_value(i64 %_0.i.i.i.i, !569766, !DIExpression(), !570195) + #dbg_value(i64 %_0.i.i.i.i, !569768, !DIExpression(), !570197) + #dbg_value(i64 %_0.i5.i.i.i, !569767, !DIExpression(), !570195) + #dbg_value(i64 %_0.i5.i.i.i, !569769, !DIExpression(), !570197) + %_5.i.i.i.i.i = zext i64 %_0.i.i.i.i to i128, !dbg !570198 + %_6.i.i.i.i.i = zext i64 %_0.i5.i.i.i to i128, !dbg !570199 + %_4.i1.i.i.i.i = mul nuw i128 %_6.i.i.i.i.i, %_5.i.i.i.i.i, !dbg !570200 + %_3.i2.i.i.i.i = lshr i128 %_4.i1.i.i.i.i, 64, !dbg !570201 + %_0.i.i.i.i.i = trunc nuw i128 %_3.i2.i.i.i.i to i64, !dbg !570202 + #dbg_value(i64 poison, !569881, !DIExpression(), !570203) + #dbg_value(i64 %_0.i.i.i.i.i, !569883, !DIExpression(), !570203) + #dbg_value(ptr undef, !564034, !DIExpression(), !569892) + #dbg_value(i64 %_0.i.i.i.i.i, !564040, !DIExpression(), !569892) + %115 = or i64 %failed.sroa.0.011.i.i, %_0.i.i.i.i.i, !dbg !570204 + #dbg_value(i64 %115, !569873, !DIExpression(), !569927) + #dbg_value(i64 %_0.i3.i.i.i.i, !569881, !DIExpression(), !570203) + #dbg_value(ptr %_4.sroa.10.0.i.i, !570149, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570150) + #dbg_value(i64 %index.i, !570149, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570150) + %self4.i.i = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.012.i.i, !dbg !570205 + #dbg_value(ptr %self4.i.i, !570206, !DIExpression(), !570210) + #dbg_value(i64 %_0.i3.i.i.i.i, !570209, !DIExpression(), !570210) + store i64 %_0.i3.i.i.i.i, ptr %self4.i.i, align 8, !dbg !570212, !alias.scope !569888, !noalias !570213 + #dbg_value(ptr undef, !569916, !DIExpression(), !569929) + #dbg_value(ptr undef, !569911, !DIExpression(), !569930) + #dbg_value(ptr undef, !569931, !DIExpression(), !569935) + #dbg_value(ptr poison, !569934, !DIExpression(), !569935) + #dbg_value(i64 %_36.i.i, !569912, !DIExpression(), !570143) + #dbg_value(i64 %_36.i.i, !569905, !DIExpression(), !569906) + #dbg_value(i64 %_36.i.i, !569922, !DIExpression(), !569923) + %_36.i.i.1 = add nuw i64 %iter.sroa.0.012.i.i, 2, !dbg !570144 + #dbg_value(i64 %_36.i.i.1, !569875, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !569928) + #dbg_value(i64 %_36.i.i, !569877, !DIExpression(), !570145) + #dbg_value(i64 %_36.i.i, !569899, !DIExpression(), !569900) + #dbg_value(i64 %_36.i.i, !570146, !DIExpression(), !570150) + #dbg_value(i64 %_36.i.i, !551477, !DIExpression(), !569894) + #dbg_value(i64 %_36.i.i, !551550, !DIExpression(), !570152) + #dbg_value(i64 %_36.i.i, !551542, !DIExpression(), !570154) + #dbg_value(i64 %_36.i.i, !551534, !DIExpression(), !570156) + #dbg_value(ptr %column.val.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570154) + #dbg_value(ptr %column.val.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570156) + #dbg_value(i64 %len3.i.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570154) + #dbg_value(i64 %len3.i.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570156) + %_4.i.i.i.i.1 = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %_36.i.i, !dbg !570158 + %_0.i.i.i.i.1 = load i64, ptr %_4.i.i.i.i.1, align 8, !dbg !570159, !noalias !570160, !noundef !23 + #dbg_value(i64 %_36.i.i, !551550, !DIExpression(), !570164) + #dbg_value(i64 %_36.i.i, !551542, !DIExpression(), !570166) + #dbg_value(i64 %_36.i.i, !551534, !DIExpression(), !570168) + #dbg_value(ptr %column5.val.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570166) + #dbg_value(ptr %column5.val.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570168) + #dbg_value(i64 %len3.i.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570166) + #dbg_value(i64 %len3.i.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570168) + %_5.i3.i.i.i.1 = icmp ult i64 %_36.i.i, %len3.i.i.i, !dbg !570170 + tail call void @llvm.assume(i1 %_5.i3.i.i.i.1), !dbg !570171 + %_4.i4.i.i.i.1 = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %_36.i.i, !dbg !570172 + %_0.i5.i.i.i.1 = load i64, ptr %_4.i4.i.i.i.1, align 8, !dbg !570173, !noalias !570160, !noundef !23 + #dbg_value(i64 %_0.i.i.i.i.1, !569879, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570174) + #dbg_value(i64 %_0.i5.i.i.i.1, !569879, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570174) + #dbg_value(i64 %_0.i.i.i.i.1, !570175, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570183) + #dbg_value(i64 %_0.i5.i.i.i.1, !570175, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570183) + #dbg_value(i64 %_0.i.i.i.i.1, !569760, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570185) + #dbg_value(i64 %_0.i5.i.i.i.1, !569760, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570185) + #dbg_value(i64 %_0.i.i.i.i.1, !569756, !DIExpression(), !570187) + #dbg_value(i64 %_0.i5.i.i.i.1, !569757, !DIExpression(), !570187) + #dbg_value(i64 %_0.i.i.i.i.1, !569747, !DIExpression(), !570188) + #dbg_value(i64 %_0.i5.i.i.i.1, !569748, !DIExpression(), !570188) + #dbg_value(i64 %_0.i.i.i.i.1, !569740, !DIExpression(), !570190) + #dbg_value(i64 %_0.i.i.i.i.1, !569735, !DIExpression(), !570192) + #dbg_value(i64 %_0.i5.i.i.i.1, !569741, !DIExpression(), !570190) + #dbg_value(i64 %_0.i5.i.i.i.1, !569736, !DIExpression(), !570192) + %_0.i3.i.i.i.i.1 = mul i64 %_0.i5.i.i.i.1, %_0.i.i.i.i.1, !dbg !570194 + #dbg_value(i64 %_0.i.i.i.i.1, !569766, !DIExpression(), !570195) + #dbg_value(i64 %_0.i.i.i.i.1, !569768, !DIExpression(), !570197) + #dbg_value(i64 %_0.i5.i.i.i.1, !569767, !DIExpression(), !570195) + #dbg_value(i64 %_0.i5.i.i.i.1, !569769, !DIExpression(), !570197) + %_5.i.i.i.i.i.1 = zext i64 %_0.i.i.i.i.1 to i128, !dbg !570198 + %_6.i.i.i.i.i.1 = zext i64 %_0.i5.i.i.i.1 to i128, !dbg !570199 + %_4.i1.i.i.i.i.1 = mul nuw i128 %_6.i.i.i.i.i.1, %_5.i.i.i.i.i.1, !dbg !570200 + %_3.i2.i.i.i.i.1 = lshr i128 %_4.i1.i.i.i.i.1, 64, !dbg !570201 + %_0.i.i.i.i.i.1 = trunc nuw i128 %_3.i2.i.i.i.i.1 to i64, !dbg !570202 + #dbg_value(i64 poison, !569881, !DIExpression(), !570203) + #dbg_value(i64 %_0.i.i.i.i.i.1, !569883, !DIExpression(), !570203) + #dbg_value(i64 %_0.i.i.i.i.i.1, !564040, !DIExpression(), !569892) + %116 = or i64 %115, %_0.i.i.i.i.i.1, !dbg !570204 + #dbg_value(i64 %116, !569873, !DIExpression(), !569927) + #dbg_value(i64 %_0.i3.i.i.i.i.1, !569881, !DIExpression(), !570203) + #dbg_value(ptr %_4.sroa.10.0.i.i, !570149, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570150) + #dbg_value(i64 %index.i, !570149, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570150) + %self4.i.i.1 = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %_36.i.i, !dbg !570205 + #dbg_value(ptr %self4.i.i.1, !570206, !DIExpression(), !570210) + #dbg_value(i64 %_0.i3.i.i.i.i.1, !570209, !DIExpression(), !570210) + store i64 %_0.i3.i.i.i.i.1, ptr %self4.i.i.1, align 8, !dbg !570212, !alias.scope !569888, !noalias !570213 + %niter.next.1 = add i64 %niter, 2, !dbg !569937 + %niter.ncmp.1 = icmp eq i64 %niter.next.1, %unroll_iter, !dbg !569937 + br i1 %niter.ncmp.1, label %bb33.i.loopexit144.unr-lcssa, label %bb15.i.i, !dbg !569937 + +bb33.thread.i: ; preds = %bb26.preheader.i.thread, %bb11.i, %bb26.preheader.i + #dbg_value(i64 0, !569430, !DIExpression(), !570214) + #dbg_value(i64 %index.i, !569423, !DIExpression(DW_OP_LLVM_fragment, 128, 64), !569651) + call void @llvm.lifetime.start.p0(i64 24, ptr nonnull %value.i.i), !dbg !570215, !noalias !569589 + #dbg_value(i64 0, !569401, !DIExpression(), !570217) + #dbg_declare(ptr poison, !569405, !DIExpression(), !570218) + #dbg_declare(ptr %value.i.i, !570219, !DIExpression(), !570222) + #dbg_value(ptr undef, !564775, !DIExpression(), !570225) + #dbg_value(ptr undef, !564776, !DIExpression(), !570225) + br label %bb36.i, !dbg !570226 + +bb33.i.loopexit.unr-lcssa: ; preds = %bb27.us.i.us, %bb27.us.i.us.preheader + %.lcssa.ph = phi i64 [ poison, %bb27.us.i.us.preheader ], [ %88, %bb27.us.i.us ] + %iter.sroa.0.046.us.i.us.unr = phi i64 [ 0, %bb27.us.i.us.preheader ], [ %_158.us.i.us, %bb27.us.i.us ] + %accumulated.sroa.0.045.us.i.us.unr = phi i64 [ 0, %bb27.us.i.us.preheader ], [ %88, %bb27.us.i.us ] + %lcmp.mod148.not = icmp eq i64 %xtraiter147, 0, !dbg !569720 + br i1 %lcmp.mod148.not, label %bb33.i, label %bb27.us.i.us.epil, !dbg !569720 + +bb27.us.i.us.epil: ; preds = %bb33.i.loopexit.unr-lcssa + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !569449, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !569718) + #dbg_value(i64 %accumulated.sroa.0.045.us.i.us.unr, !569447, !DIExpression(), !569717) + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !569451, !DIExpression(), !569793) + #dbg_value(ptr %columns.i, !551276, !DIExpression(), !569794) + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !551277, !DIExpression(), !569794) + #dbg_value(ptr %columns.i, !551266, !DIExpression(), !569795) + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !551267, !DIExpression(), !569795) + #dbg_value(ptr %columns.i, !551137, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !569798) + #dbg_value(i64 0, !551142, !DIExpression(), !569796) + #dbg_value(ptr %columns.i, !551269, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !569827) + #dbg_value(ptr %columns.i, !551137, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !569796) + %_0.sroa.0.0.i.i.us.i.us.epil = load i64, ptr %data.i.i.i.us.i, align 8, !dbg !569725, !noalias !569509, !noundef !23 + #dbg_value(ptr %14, !551266, !DIExpression(), !569800) + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !551267, !DIExpression(), !569800) + #dbg_value(ptr %14, !551137, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !569801) + #dbg_value(ptr %14, !551137, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !569803) + #dbg_value(ptr %14, !551269, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !569804) + #dbg_value(i64 0, !551142, !DIExpression(), !569801) + %_0.sroa.0.0.i9.i.us.i.us.epil = load i64, ptr %data.i6.i7.i.us.i, align 8, !dbg !569730, !noalias !569509, !noundef !23 + #dbg_value(ptr poison, !569758, !DIExpression(), !569805) + #dbg_value(ptr poison, !569759, !DIExpression(), !569805) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !569760, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !569805) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !569760, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !569805) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !569756, !DIExpression(), !569806) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !569757, !DIExpression(), !569806) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !569747, !DIExpression(), !569807) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !569748, !DIExpression(), !569807) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !569740, !DIExpression(), !569808) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !569735, !DIExpression(), !569809) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !569741, !DIExpression(), !569808) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !569736, !DIExpression(), !569809) + %_0.i3.i.us.i.us.epil = mul i64 %_0.sroa.0.0.i9.i.us.i.us.epil, %_0.sroa.0.0.i.i.us.i.us.epil, !dbg !569732 + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !569766, !DIExpression(), !569810) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !569768, !DIExpression(), !569811) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !569767, !DIExpression(), !569810) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !569769, !DIExpression(), !569811) + %_5.i.i.us.i.us.epil = zext i64 %_0.sroa.0.0.i.i.us.i.us.epil to i128, !dbg !569762 + %_6.i.i.us.i.us.epil = zext i64 %_0.sroa.0.0.i9.i.us.i.us.epil to i128, !dbg !569771 + %_4.i1.i.us.i.us.epil = mul nuw i128 %_6.i.i.us.i.us.epil, %_5.i.i.us.i.us.epil, !dbg !569772 + %_3.i2.i.us.i.us.epil = lshr i128 %_4.i1.i.us.i.us.epil, 64, !dbg !569773 + %_0.i.i159.us.i.us.epil = trunc nuw i128 %_3.i2.i.us.i.us.epil to i64, !dbg !569774 + #dbg_value(i64 %_0.i.i159.us.i.us.epil, !569455, !DIExpression(), !569812) + #dbg_value(i64 %_0.i3.i.us.i.us.epil, !569453, !DIExpression(), !569812) + %self34.us.i.us.epil = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.046.us.i.us.unr, !dbg !569775 + #dbg_value(ptr %self34.us.i.us.epil, !569779, !DIExpression(), !569813) + #dbg_value(i64 %_0.i3.i.us.i.us.epil, !569780, !DIExpression(), !569813) + store i64 %_0.i3.i.us.i.us.epil, ptr %self34.us.i.us.epil, align 8, !dbg !569776, !noalias !569509 + #dbg_value(ptr undef, !564034, !DIExpression(), !569482) + #dbg_value(i64 %_0.i.i159.us.i.us.epil, !564040, !DIExpression(), !569482) + %117 = or i64 %accumulated.sroa.0.045.us.i.us.unr, %_0.i.i159.us.i.us.epil, !dbg !569782 + #dbg_value(i64 %117, !569447, !DIExpression(), !569717) + #dbg_value(i64 poison, !569449, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !569718) + #dbg_value(ptr undef, !569472, !DIExpression(), !569475) + #dbg_value(ptr undef, !569463, !DIExpression(), !569468) + #dbg_value(ptr undef, !569476, !DIExpression(), !569480) + #dbg_value(ptr poison, !569479, !DIExpression(), !569480) + #dbg_value(i64 poison, !569449, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !569718) + br label %bb33.i, !dbg !570227 + +bb33.i.loopexit144.unr-lcssa: ; preds = %bb15.i.i, %bb15.i.i.preheader + %.lcssa145.ph = phi i64 [ poison, %bb15.i.i.preheader ], [ %116, %bb15.i.i ] + %iter.sroa.0.012.i.i.unr = phi i64 [ 0, %bb15.i.i.preheader ], [ %_36.i.i.1, %bb15.i.i ] + %failed.sroa.0.011.i.i.unr = phi i64 [ 0, %bb15.i.i.preheader ], [ %116, %bb15.i.i ] + %lcmp.mod.not = icmp eq i64 %xtraiter, 0, !dbg !569937 + br i1 %lcmp.mod.not, label %bb33.i, label %bb15.i.i.epil, !dbg !569937 + +bb15.i.i.epil: ; preds = %bb33.i.loopexit144.unr-lcssa + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !569875, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !569928) + #dbg_value(i64 %failed.sroa.0.011.i.i.unr, !569873, !DIExpression(), !569927) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !569912, !DIExpression(), !570143) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !569905, !DIExpression(), !569906) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !569922, !DIExpression(), !569923) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !569875, !DIExpression(DW_OP_plus_uconst, 1, DW_OP_stack_value, DW_OP_LLVM_fragment, 0, 64), !569928) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !569877, !DIExpression(), !570145) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !569899, !DIExpression(), !569900) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !570146, !DIExpression(), !570150) + #dbg_value(ptr undef, !551471, !DIExpression(), !569894) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !551477, !DIExpression(), !569894) + #dbg_value(ptr poison, !551549, !DIExpression(), !570152) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !551550, !DIExpression(), !570152) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !551542, !DIExpression(), !570154) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !551534, !DIExpression(), !570156) + #dbg_value(ptr %column.val.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570154) + #dbg_value(ptr %column.val.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570156) + #dbg_value(i64 %len3.i.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570154) + #dbg_value(i64 %len3.i.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570156) + %_4.i.i.i.i.epil = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.012.i.i.unr, !dbg !570158 + %_0.i.i.i.i.epil = load i64, ptr %_4.i.i.i.i.epil, align 8, !dbg !570159, !noalias !570160, !noundef !23 + #dbg_value(ptr poison, !551549, !DIExpression(), !570164) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !551550, !DIExpression(), !570164) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !551542, !DIExpression(), !570166) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !551534, !DIExpression(), !570168) + #dbg_value(ptr %column5.val.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570166) + #dbg_value(ptr %column5.val.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570168) + #dbg_value(i64 %len3.i.i.i, !551541, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570166) + #dbg_value(i64 %len3.i.i.i, !551535, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570168) + %_5.i3.i.i.i.epil = icmp ult i64 %iter.sroa.0.012.i.i.unr, %len3.i.i.i, !dbg !570170 + tail call void @llvm.assume(i1 %_5.i3.i.i.i.epil), !dbg !570171 + %_4.i4.i.i.i.epil = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.012.i.i.unr, !dbg !570172 + %_0.i5.i.i.i.epil = load i64, ptr %_4.i4.i.i.i.epil, align 8, !dbg !570173, !noalias !570160, !noundef !23 + #dbg_value(i64 %_0.i.i.i.i.epil, !569879, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570174) + #dbg_value(i64 %_0.i5.i.i.i.epil, !569879, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570174) + #dbg_value(i64 %_0.i.i.i.i.epil, !570175, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570183) + #dbg_value(i64 %_0.i5.i.i.i.epil, !570175, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570183) + #dbg_value(ptr poison, !569758, !DIExpression(), !570185) + #dbg_value(ptr poison, !569759, !DIExpression(), !570185) + #dbg_value(i64 %_0.i.i.i.i.epil, !569760, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570185) + #dbg_value(i64 %_0.i5.i.i.i.epil, !569760, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570185) + #dbg_value(i64 %_0.i.i.i.i.epil, !569756, !DIExpression(), !570187) + #dbg_value(i64 %_0.i5.i.i.i.epil, !569757, !DIExpression(), !570187) + #dbg_value(i64 %_0.i.i.i.i.epil, !569747, !DIExpression(), !570188) + #dbg_value(i64 %_0.i5.i.i.i.epil, !569748, !DIExpression(), !570188) + #dbg_value(i64 %_0.i.i.i.i.epil, !569740, !DIExpression(), !570190) + #dbg_value(i64 %_0.i.i.i.i.epil, !569735, !DIExpression(), !570192) + #dbg_value(i64 %_0.i5.i.i.i.epil, !569741, !DIExpression(), !570190) + #dbg_value(i64 %_0.i5.i.i.i.epil, !569736, !DIExpression(), !570192) + %_0.i3.i.i.i.i.epil = mul i64 %_0.i5.i.i.i.epil, %_0.i.i.i.i.epil, !dbg !570194 + #dbg_value(i64 %_0.i.i.i.i.epil, !569766, !DIExpression(), !570195) + #dbg_value(i64 %_0.i.i.i.i.epil, !569768, !DIExpression(), !570197) + #dbg_value(i64 %_0.i5.i.i.i.epil, !569767, !DIExpression(), !570195) + #dbg_value(i64 %_0.i5.i.i.i.epil, !569769, !DIExpression(), !570197) + %_5.i.i.i.i.i.epil = zext i64 %_0.i.i.i.i.epil to i128, !dbg !570198 + %_6.i.i.i.i.i.epil = zext i64 %_0.i5.i.i.i.epil to i128, !dbg !570199 + %_4.i1.i.i.i.i.epil = mul nuw i128 %_6.i.i.i.i.i.epil, %_5.i.i.i.i.i.epil, !dbg !570200 + %_3.i2.i.i.i.i.epil = lshr i128 %_4.i1.i.i.i.i.epil, 64, !dbg !570201 + %_0.i.i.i.i.i.epil = trunc nuw i128 %_3.i2.i.i.i.i.epil to i64, !dbg !570202 + #dbg_value(i64 poison, !569881, !DIExpression(), !570203) + #dbg_value(i64 %_0.i.i.i.i.i.epil, !569883, !DIExpression(), !570203) + #dbg_value(ptr undef, !564034, !DIExpression(), !569892) + #dbg_value(i64 %_0.i.i.i.i.i.epil, !564040, !DIExpression(), !569892) + %118 = or i64 %failed.sroa.0.011.i.i.unr, %_0.i.i.i.i.i.epil, !dbg !570204 + #dbg_value(i64 %118, !569873, !DIExpression(), !569927) + #dbg_value(i64 %_0.i3.i.i.i.i.epil, !569881, !DIExpression(), !570203) + #dbg_value(ptr %_4.sroa.10.0.i.i, !570149, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !570150) + #dbg_value(i64 %index.i, !570149, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !570150) + %self4.i.i.epil = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.012.i.i.unr, !dbg !570205 + #dbg_value(ptr %self4.i.i.epil, !570206, !DIExpression(), !570210) + #dbg_value(i64 %_0.i3.i.i.i.i.epil, !570209, !DIExpression(), !570210) + store i64 %_0.i3.i.i.i.i.epil, ptr %self4.i.i.epil, align 8, !dbg !570212, !alias.scope !569888, !noalias !570213 + #dbg_value(ptr undef, !569916, !DIExpression(), !569929) + #dbg_value(ptr undef, !569911, !DIExpression(), !569930) + #dbg_value(ptr undef, !569931, !DIExpression(), !569935) + #dbg_value(ptr poison, !569934, !DIExpression(), !569935) + br label %bb33.i, !dbg !570227 + + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/final-u64-mul-dense-s.md b/research/rowfn-x86-2026-08-07/codegen/final-u64-mul-dense-s.md new file mode 100644 index 00000000000..e30db12f65a --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/final-u64-mul-dense-s.md @@ -0,0 +1,71 @@ + + + +# `final-u64-mul-dense.s` + +```s +.LBB1679_67: + .loc 181 765 12 + andq $-2, %r10 + xorl %edi, %edi + xorl %ecx, %ecx + movq -48(%rbp), %r8 +.Ltmp111183: +.LBB1679_68: + .loc 567 39 18 + movq (%r13,%rdi,8), %rax +.Ltmp111184: + .loc 565 176 44 + mulq (%r12,%rdi,8) +.Ltmp111185: + movq %rdx, %rsi +.Ltmp111186: + .loc 207 475 9 + movq %rax, (%r8,%rdi,8) +.Ltmp111187: + .loc 567 39 18 + movq 8(%r13,%rdi,8), %rax +.Ltmp111188: + .loc 565 176 44 + mulq 8(%r12,%rdi,8) +.Ltmp111189: + .loc 566 821 53 + orq %rcx, %rsi +.Ltmp111190: + .loc 207 475 9 + movq %rax, 8(%r8,%rdi,8) +.Ltmp111191: + .loc 156 717 17 + addq $2, %rdi +.Ltmp111192: + .loc 565 176 44 + movq %rdx, %rcx +.Ltmp111193: + .loc 566 821 53 + orq %rsi, %rcx +.Ltmp111194: + .loc 181 765 12 + cmpq %r10, %rdi + jne .LBB1679_68 +.Ltmp111195: +.LBB1679_69: + testb $1, %r9b + je .LBB1679_84 +.Ltmp111196: + .loc 567 39 18 + movq (%r13,%rdi,8), %rax +.Ltmp111197: + .loc 565 176 44 + mulq (%r12,%rdi,8) +.Ltmp111198: + .loc 566 821 53 + orq %rdx, %rcx +.Ltmp111199: + .loc 566 0 53 is_stmt 0 + movq -48(%rbp), %rdx +.Ltmp111200: + .loc 207 475 9 is_stmt 1 + movq %rax, (%rdx,%rdi,8) +.Ltmp111201: + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/indexed-i64-mul-dense-ll.md b/research/rowfn-x86-2026-08-07/codegen/indexed-i64-mul-dense-ll.md new file mode 100644 index 00000000000..4cd73bddfad --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/indexed-i64-mul-dense-ll.md @@ -0,0 +1,96 @@ + + + +# `indexed-i64-mul-dense.ll` + +```ll +bb15.i.i: ; preds = %bb11.i, %bb15.i.i + %iter.sroa.0.012.i.i = phi i64 [ %_36.i.i, %bb15.i.i ], [ 0, %bb11.i ] + %failed.sroa.0.011.i.i = phi i64 [ %79, %bb15.i.i ], [ 0, %bb11.i ] + #dbg_value(i64 %iter.sroa.0.012.i.i, !576781, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !576835) + #dbg_value(i64 %failed.sroa.0.011.i.i, !576779, !DIExpression(), !576834) + #dbg_value(i64 %iter.sroa.0.012.i.i, !576819, !DIExpression(), !577050) + #dbg_value(i64 %iter.sroa.0.012.i.i, !576812, !DIExpression(), !576813) + #dbg_value(i64 %iter.sroa.0.012.i.i, !576829, !DIExpression(), !576830) + %_36.i.i = add nuw i64 %iter.sroa.0.012.i.i, 1, !dbg !577051 + #dbg_value(i64 %_36.i.i, !576781, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !576835) + #dbg_value(i64 %iter.sroa.0.012.i.i, !576783, !DIExpression(), !577052) + #dbg_value(i64 %iter.sroa.0.012.i.i, !576806, !DIExpression(), !576807) + #dbg_value(i64 %iter.sroa.0.012.i.i, !577053, !DIExpression(), !577057) + #dbg_value(ptr undef, !553722, !DIExpression(), !576801) + #dbg_value(i64 %iter.sroa.0.012.i.i, !553728, !DIExpression(), !576801) + #dbg_value(ptr poison, !553800, !DIExpression(), !577059) + #dbg_value(i64 %iter.sroa.0.012.i.i, !553801, !DIExpression(), !577059) + #dbg_value(i64 %iter.sroa.0.012.i.i, !553793, !DIExpression(), !577061) + #dbg_value(i64 %iter.sroa.0.012.i.i, !553785, !DIExpression(), !577063) + #dbg_value(ptr %column.val.i.i, !553792, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577061) + #dbg_value(ptr %column.val.i.i, !553786, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577063) + #dbg_value(i64 %len3.i.i.i, !553792, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577061) + #dbg_value(i64 %len3.i.i.i, !553786, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577063) + %_4.i.i.i.i = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.012.i.i, !dbg !577065 + %_0.i.i.i.i = load i64, ptr %_4.i.i.i.i, align 8, !dbg !577066, !noalias !577067, !noundef !23 + #dbg_value(ptr poison, !553800, !DIExpression(), !577071) + #dbg_value(i64 %iter.sroa.0.012.i.i, !553801, !DIExpression(), !577071) + #dbg_value(i64 %iter.sroa.0.012.i.i, !553793, !DIExpression(), !577073) + #dbg_value(i64 %iter.sroa.0.012.i.i, !553785, !DIExpression(), !577075) + #dbg_value(ptr %column5.val.i.i, !553792, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577073) + #dbg_value(ptr %column5.val.i.i, !553786, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577075) + #dbg_value(i64 %len3.i.i.i, !553792, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577073) + #dbg_value(i64 %len3.i.i.i, !553786, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577075) + %_5.i3.i.i.i = icmp ult i64 %iter.sroa.0.012.i.i, %len3.i.i.i, !dbg !577077 + tail call void @llvm.assume(i1 %_5.i3.i.i.i), !dbg !577078 + %_4.i4.i.i.i = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.012.i.i, !dbg !577079 + %_0.i5.i.i.i = load i64, ptr %_4.i4.i.i.i, align 8, !dbg !577080, !noalias !577067, !noundef !23 + #dbg_value(i64 %_0.i.i.i.i, !576785, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577081) + #dbg_value(i64 %_0.i5.i.i.i, !576785, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577081) + #dbg_value(i64 %_0.i.i.i.i, !577082, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577090) + #dbg_value(i64 %_0.i5.i.i.i, !577082, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577090) + #dbg_value(ptr poison, !576671, !DIExpression(), !577092) + #dbg_value(ptr poison, !576672, !DIExpression(), !577092) + #dbg_value(i64 %_0.i.i.i.i, !576673, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577092) + #dbg_value(i64 %_0.i5.i.i.i, !576673, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577092) + #dbg_value(i64 %_0.i.i.i.i, !576669, !DIExpression(), !577094) + #dbg_value(i64 %_0.i5.i.i.i, !576670, !DIExpression(), !577094) + #dbg_value(i64 %_0.i.i.i.i, !576660, !DIExpression(), !577095) + #dbg_value(i64 %_0.i5.i.i.i, !576661, !DIExpression(), !577095) + #dbg_value(i64 %_0.i.i.i.i, !576649, !DIExpression(), !577097) + #dbg_value(i64 %_0.i.i.i.i, !576644, !DIExpression(), !577099) + #dbg_value(i64 %_0.i5.i.i.i, !576650, !DIExpression(), !577097) + #dbg_value(i64 %_0.i5.i.i.i, !576645, !DIExpression(), !577099) + %_0.i.i.i.i.i = mul i64 %_0.i5.i.i.i, %_0.i.i.i.i, !dbg !577101 + #dbg_value(i64 %_0.i.i.i.i, !576681, !DIExpression(), !577102) + #dbg_value(i64 %_0.i.i.i.i, !576683, !DIExpression(), !577104) + #dbg_value(i64 %_0.i5.i.i.i, !576682, !DIExpression(), !577102) + #dbg_value(i64 %_0.i5.i.i.i, !576684, !DIExpression(), !577104) + %_4.i1.i.i.i.i = sext i64 %_0.i.i.i.i to i128, !dbg !577105 + %_5.i.i.i.i.i = sext i64 %_0.i5.i.i.i to i128, !dbg !577106 + %wide.i.i.i.i.i = mul nsw i128 %_5.i.i.i.i.i, %_4.i1.i.i.i.i, !dbg !577105 + #dbg_value(i128 %wide.i.i.i.i.i, !576685, !DIExpression(), !577107) + %kept.i.i.i.i.i = trunc i128 %wide.i.i.i.i.i to i64, !dbg !577108 + #dbg_value(i64 %kept.i.i.i.i.i, !576687, !DIExpression(), !577109) + %_8.i.i.i.i.i = lshr i128 %wide.i.i.i.i.i, 64, !dbg !577110 + %discarded.i.i.i.i.i = trunc nuw i128 %_8.i.i.i.i.i to i64, !dbg !577111 + #dbg_value(i64 %discarded.i.i.i.i.i, !576689, !DIExpression(), !577112) + %_10.i.i.i.i.i = ashr i64 %kept.i.i.i.i.i, 63, !dbg !577113 + %_9.i.i.i.i.i = xor i64 %_10.i.i.i.i.i, %discarded.i.i.i.i.i, !dbg !577114 + #dbg_value(i64 poison, !576787, !DIExpression(), !577115) + #dbg_value(i64 %_9.i.i.i.i.i, !576789, !DIExpression(), !577115) + #dbg_value(ptr undef, !576390, !DIExpression(), !576799) + #dbg_value(i64 %_9.i.i.i.i.i, !576396, !DIExpression(), !576799) + %79 = or i64 %_9.i.i.i.i.i, %failed.sroa.0.011.i.i, !dbg !577116 + #dbg_value(i64 %79, !576779, !DIExpression(), !576834) + #dbg_value(i64 %_0.i.i.i.i.i, !576787, !DIExpression(), !577115) + #dbg_value(ptr %_4.sroa.10.0.i.i, !577056, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577057) + #dbg_value(i64 %index.i, !577056, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577057) + %self4.i.i = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.012.i.i, !dbg !577117 + #dbg_value(ptr %self4.i.i, !577118, !DIExpression(), !577122) + #dbg_value(i64 %_0.i.i.i.i.i, !577121, !DIExpression(), !577122) + store i64 %_0.i.i.i.i.i, ptr %self4.i.i, align 8, !dbg !577124, !alias.scope !576795, !noalias !577125 + #dbg_value(ptr undef, !576823, !DIExpression(), !576836) + #dbg_value(ptr undef, !576818, !DIExpression(), !576837) + #dbg_value(ptr undef, !576838, !DIExpression(), !576842) + #dbg_value(ptr poison, !576841, !DIExpression(), !576842) + %exitcond.not.i.i = icmp eq i64 %_36.i.i, %len3.i4.i.i.fr, !dbg !577126 + br i1 %exitcond.not.i.i, label %bb33.i, label %bb15.i.i, !dbg !576844 + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/indexed-i64-mul-dense-s.md b/research/rowfn-x86-2026-08-07/codegen/indexed-i64-mul-dense-s.md new file mode 100644 index 00000000000..4f95e7d0668 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/indexed-i64-mul-dense-s.md @@ -0,0 +1,35 @@ + + + +# `indexed-i64-mul-dense.s` + +```s + .p2align 4 +.LBB1685_25: + .loc 566 39 18 is_stmt 1 + movq (%r13,%rsi,8), %rax +.Ltmp113564: + .loc 565 194 24 + imulq (%r12,%rsi,8) +.Ltmp113565: + .loc 207 475 9 + movq %rax, (%rdi,%rsi,8) +.Ltmp113566: + .loc 156 717 17 + incq %rsi +.Ltmp113567: + .loc 565 198 26 + sarq $63, %rax +.Ltmp113568: + .loc 565 198 13 is_stmt 0 + xorq %rdx, %rax +.Ltmp113569: + .loc 568 821 53 is_stmt 1 + orq %rax, %rcx +.Ltmp113570: + .loc 182 1904 50 + cmpq %rsi, %r9 + jne .LBB1685_25 + jmp .LBB1685_60 + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/indexed-u64-mul-dense-ll.md b/research/rowfn-x86-2026-08-07/codegen/indexed-u64-mul-dense-ll.md new file mode 100644 index 00000000000..4b254e7d98a --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/indexed-u64-mul-dense-ll.md @@ -0,0 +1,322 @@ + + + +# `indexed-u64-mul-dense.ll` + +```ll +bb15.i.i: ; preds = %bb15.i.i, %bb15.i.i.preheader.new + %iter.sroa.0.012.i.i = phi i64 [ 0, %bb15.i.i.preheader.new ], [ %_36.i.i.1, %bb15.i.i ] + %failed.sroa.0.011.i.i = phi i64 [ 0, %bb15.i.i.preheader.new ], [ %116, %bb15.i.i ] + %niter = phi i64 [ 0, %bb15.i.i.preheader.new ], [ %niter.next.1, %bb15.i.i ] + #dbg_value(i64 %iter.sroa.0.012.i.i, !580573, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580626) + #dbg_value(i64 %failed.sroa.0.011.i.i, !580571, !DIExpression(), !580625) + #dbg_value(i64 %iter.sroa.0.012.i.i, !580610, !DIExpression(), !580841) + #dbg_value(i64 %iter.sroa.0.012.i.i, !580603, !DIExpression(), !580604) + #dbg_value(i64 %iter.sroa.0.012.i.i, !580620, !DIExpression(), !580621) + %_36.i.i = or disjoint i64 %iter.sroa.0.012.i.i, 1, !dbg !580842 + #dbg_value(i64 %_36.i.i, !580573, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580626) + #dbg_value(i64 %iter.sroa.0.012.i.i, !580575, !DIExpression(), !580843) + #dbg_value(i64 %iter.sroa.0.012.i.i, !580597, !DIExpression(), !580598) + #dbg_value(i64 %iter.sroa.0.012.i.i, !580844, !DIExpression(), !580848) + #dbg_value(ptr undef, !563670, !DIExpression(), !580592) + #dbg_value(i64 %iter.sroa.0.012.i.i, !563676, !DIExpression(), !580592) + #dbg_value(ptr poison, !563748, !DIExpression(), !580850) + #dbg_value(i64 %iter.sroa.0.012.i.i, !563749, !DIExpression(), !580850) + #dbg_value(i64 %iter.sroa.0.012.i.i, !563741, !DIExpression(), !580852) + #dbg_value(i64 %iter.sroa.0.012.i.i, !563733, !DIExpression(), !580854) + #dbg_value(ptr %column.val.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580852) + #dbg_value(ptr %column.val.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580854) + #dbg_value(i64 %len3.i.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580852) + #dbg_value(i64 %len3.i.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580854) + %_4.i.i.i.i = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.012.i.i, !dbg !580856 + %_0.i.i.i.i = load i64, ptr %_4.i.i.i.i, align 8, !dbg !580857, !noalias !580858, !noundef !23 + #dbg_value(ptr poison, !563748, !DIExpression(), !580862) + #dbg_value(i64 %iter.sroa.0.012.i.i, !563749, !DIExpression(), !580862) + #dbg_value(i64 %iter.sroa.0.012.i.i, !563741, !DIExpression(), !580864) + #dbg_value(i64 %iter.sroa.0.012.i.i, !563733, !DIExpression(), !580866) + #dbg_value(ptr %column5.val.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580864) + #dbg_value(ptr %column5.val.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580866) + #dbg_value(i64 %len3.i.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580864) + #dbg_value(i64 %len3.i.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580866) + %_5.i3.i.i.i = icmp ult i64 %iter.sroa.0.012.i.i, %len3.i.i.i, !dbg !580868 + tail call void @llvm.assume(i1 %_5.i3.i.i.i), !dbg !580869 + %_4.i4.i.i.i = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.012.i.i, !dbg !580870 + %_0.i5.i.i.i = load i64, ptr %_4.i4.i.i.i, align 8, !dbg !580871, !noalias !580858, !noundef !23 + #dbg_value(i64 %_0.i.i.i.i, !580577, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580872) + #dbg_value(i64 %_0.i5.i.i.i, !580577, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580872) + #dbg_value(i64 %_0.i.i.i.i, !580873, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580881) + #dbg_value(i64 %_0.i5.i.i.i, !580873, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580881) + #dbg_value(ptr poison, !580456, !DIExpression(), !580883) + #dbg_value(ptr poison, !580457, !DIExpression(), !580883) + #dbg_value(i64 %_0.i.i.i.i, !580458, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580883) + #dbg_value(i64 %_0.i5.i.i.i, !580458, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580883) + #dbg_value(i64 %_0.i.i.i.i, !580454, !DIExpression(), !580885) + #dbg_value(i64 %_0.i5.i.i.i, !580455, !DIExpression(), !580885) + #dbg_value(i64 %_0.i.i.i.i, !580445, !DIExpression(), !580886) + #dbg_value(i64 %_0.i5.i.i.i, !580446, !DIExpression(), !580886) + #dbg_value(i64 %_0.i.i.i.i, !580438, !DIExpression(), !580888) + #dbg_value(i64 %_0.i.i.i.i, !580433, !DIExpression(), !580890) + #dbg_value(i64 %_0.i5.i.i.i, !580439, !DIExpression(), !580888) + #dbg_value(i64 %_0.i5.i.i.i, !580434, !DIExpression(), !580890) + %_0.i3.i.i.i.i = mul i64 %_0.i5.i.i.i, %_0.i.i.i.i, !dbg !580892 + #dbg_value(i64 %_0.i.i.i.i, !580464, !DIExpression(), !580893) + #dbg_value(i64 %_0.i.i.i.i, !580466, !DIExpression(), !580895) + #dbg_value(i64 %_0.i5.i.i.i, !580465, !DIExpression(), !580893) + #dbg_value(i64 %_0.i5.i.i.i, !580467, !DIExpression(), !580895) + %_5.i.i.i.i.i = zext i64 %_0.i.i.i.i to i128, !dbg !580896 + %_6.i.i.i.i.i = zext i64 %_0.i5.i.i.i to i128, !dbg !580897 + %_4.i1.i.i.i.i = mul nuw i128 %_6.i.i.i.i.i, %_5.i.i.i.i.i, !dbg !580898 + %_3.i2.i.i.i.i = lshr i128 %_4.i1.i.i.i.i, 64, !dbg !580899 + %_0.i.i.i.i.i = trunc nuw i128 %_3.i2.i.i.i.i to i64, !dbg !580900 + #dbg_value(i64 poison, !580579, !DIExpression(), !580901) + #dbg_value(i64 %_0.i.i.i.i.i, !580581, !DIExpression(), !580901) + #dbg_value(ptr undef, !576390, !DIExpression(), !580590) + #dbg_value(i64 %_0.i.i.i.i.i, !576396, !DIExpression(), !580590) + %115 = or i64 %failed.sroa.0.011.i.i, %_0.i.i.i.i.i, !dbg !580902 + #dbg_value(i64 %115, !580571, !DIExpression(), !580625) + #dbg_value(i64 %_0.i3.i.i.i.i, !580579, !DIExpression(), !580901) + #dbg_value(ptr %_4.sroa.10.0.i.i, !580847, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580848) + #dbg_value(i64 %index.i, !580847, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580848) + %self4.i.i = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.012.i.i, !dbg !580903 + #dbg_value(ptr %self4.i.i, !580904, !DIExpression(), !580908) + #dbg_value(i64 %_0.i3.i.i.i.i, !580907, !DIExpression(), !580908) + store i64 %_0.i3.i.i.i.i, ptr %self4.i.i, align 8, !dbg !580910, !alias.scope !580586, !noalias !580911 + #dbg_value(ptr undef, !580614, !DIExpression(), !580627) + #dbg_value(ptr undef, !580609, !DIExpression(), !580628) + #dbg_value(ptr undef, !580629, !DIExpression(), !580633) + #dbg_value(ptr poison, !580632, !DIExpression(), !580633) + #dbg_value(i64 %_36.i.i, !580610, !DIExpression(), !580841) + #dbg_value(i64 %_36.i.i, !580603, !DIExpression(), !580604) + #dbg_value(i64 %_36.i.i, !580620, !DIExpression(), !580621) + %_36.i.i.1 = add nuw i64 %iter.sroa.0.012.i.i, 2, !dbg !580842 + #dbg_value(i64 %_36.i.i.1, !580573, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580626) + #dbg_value(i64 %_36.i.i, !580575, !DIExpression(), !580843) + #dbg_value(i64 %_36.i.i, !580597, !DIExpression(), !580598) + #dbg_value(i64 %_36.i.i, !580844, !DIExpression(), !580848) + #dbg_value(i64 %_36.i.i, !563676, !DIExpression(), !580592) + #dbg_value(i64 %_36.i.i, !563749, !DIExpression(), !580850) + #dbg_value(i64 %_36.i.i, !563741, !DIExpression(), !580852) + #dbg_value(i64 %_36.i.i, !563733, !DIExpression(), !580854) + #dbg_value(ptr %column.val.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580852) + #dbg_value(ptr %column.val.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580854) + #dbg_value(i64 %len3.i.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580852) + #dbg_value(i64 %len3.i.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580854) + %_4.i.i.i.i.1 = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %_36.i.i, !dbg !580856 + %_0.i.i.i.i.1 = load i64, ptr %_4.i.i.i.i.1, align 8, !dbg !580857, !noalias !580858, !noundef !23 + #dbg_value(i64 %_36.i.i, !563749, !DIExpression(), !580862) + #dbg_value(i64 %_36.i.i, !563741, !DIExpression(), !580864) + #dbg_value(i64 %_36.i.i, !563733, !DIExpression(), !580866) + #dbg_value(ptr %column5.val.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580864) + #dbg_value(ptr %column5.val.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580866) + #dbg_value(i64 %len3.i.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580864) + #dbg_value(i64 %len3.i.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580866) + %_5.i3.i.i.i.1 = icmp ult i64 %_36.i.i, %len3.i.i.i, !dbg !580868 + tail call void @llvm.assume(i1 %_5.i3.i.i.i.1), !dbg !580869 + %_4.i4.i.i.i.1 = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %_36.i.i, !dbg !580870 + %_0.i5.i.i.i.1 = load i64, ptr %_4.i4.i.i.i.1, align 8, !dbg !580871, !noalias !580858, !noundef !23 + #dbg_value(i64 %_0.i.i.i.i.1, !580577, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580872) + #dbg_value(i64 %_0.i5.i.i.i.1, !580577, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580872) + #dbg_value(i64 %_0.i.i.i.i.1, !580873, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580881) + #dbg_value(i64 %_0.i5.i.i.i.1, !580873, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580881) + #dbg_value(i64 %_0.i.i.i.i.1, !580458, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580883) + #dbg_value(i64 %_0.i5.i.i.i.1, !580458, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580883) + #dbg_value(i64 %_0.i.i.i.i.1, !580454, !DIExpression(), !580885) + #dbg_value(i64 %_0.i5.i.i.i.1, !580455, !DIExpression(), !580885) + #dbg_value(i64 %_0.i.i.i.i.1, !580445, !DIExpression(), !580886) + #dbg_value(i64 %_0.i5.i.i.i.1, !580446, !DIExpression(), !580886) + #dbg_value(i64 %_0.i.i.i.i.1, !580438, !DIExpression(), !580888) + #dbg_value(i64 %_0.i.i.i.i.1, !580433, !DIExpression(), !580890) + #dbg_value(i64 %_0.i5.i.i.i.1, !580439, !DIExpression(), !580888) + #dbg_value(i64 %_0.i5.i.i.i.1, !580434, !DIExpression(), !580890) + %_0.i3.i.i.i.i.1 = mul i64 %_0.i5.i.i.i.1, %_0.i.i.i.i.1, !dbg !580892 + #dbg_value(i64 %_0.i.i.i.i.1, !580464, !DIExpression(), !580893) + #dbg_value(i64 %_0.i.i.i.i.1, !580466, !DIExpression(), !580895) + #dbg_value(i64 %_0.i5.i.i.i.1, !580465, !DIExpression(), !580893) + #dbg_value(i64 %_0.i5.i.i.i.1, !580467, !DIExpression(), !580895) + %_5.i.i.i.i.i.1 = zext i64 %_0.i.i.i.i.1 to i128, !dbg !580896 + %_6.i.i.i.i.i.1 = zext i64 %_0.i5.i.i.i.1 to i128, !dbg !580897 + %_4.i1.i.i.i.i.1 = mul nuw i128 %_6.i.i.i.i.i.1, %_5.i.i.i.i.i.1, !dbg !580898 + %_3.i2.i.i.i.i.1 = lshr i128 %_4.i1.i.i.i.i.1, 64, !dbg !580899 + %_0.i.i.i.i.i.1 = trunc nuw i128 %_3.i2.i.i.i.i.1 to i64, !dbg !580900 + #dbg_value(i64 poison, !580579, !DIExpression(), !580901) + #dbg_value(i64 %_0.i.i.i.i.i.1, !580581, !DIExpression(), !580901) + #dbg_value(i64 %_0.i.i.i.i.i.1, !576396, !DIExpression(), !580590) + %116 = or i64 %115, %_0.i.i.i.i.i.1, !dbg !580902 + #dbg_value(i64 %116, !580571, !DIExpression(), !580625) + #dbg_value(i64 %_0.i3.i.i.i.i.1, !580579, !DIExpression(), !580901) + #dbg_value(ptr %_4.sroa.10.0.i.i, !580847, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580848) + #dbg_value(i64 %index.i, !580847, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580848) + %self4.i.i.1 = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %_36.i.i, !dbg !580903 + #dbg_value(ptr %self4.i.i.1, !580904, !DIExpression(), !580908) + #dbg_value(i64 %_0.i3.i.i.i.i.1, !580907, !DIExpression(), !580908) + store i64 %_0.i3.i.i.i.i.1, ptr %self4.i.i.1, align 8, !dbg !580910, !alias.scope !580586, !noalias !580911 + %niter.next.1 = add i64 %niter, 2, !dbg !580635 + %niter.ncmp.1 = icmp eq i64 %niter.next.1, %unroll_iter, !dbg !580635 + br i1 %niter.ncmp.1, label %bb33.i.loopexit144.unr-lcssa, label %bb15.i.i, !dbg !580635 + +bb33.thread.i: ; preds = %bb26.preheader.i.thread, %bb11.i, %bb26.preheader.i + #dbg_value(i64 0, !580128, !DIExpression(), !580912) + #dbg_value(i64 %index.i, !580121, !DIExpression(DW_OP_LLVM_fragment, 128, 64), !580349) + call void @llvm.lifetime.start.p0(i64 24, ptr nonnull %value.i.i), !dbg !580913, !noalias !580287 + #dbg_value(i64 0, !580099, !DIExpression(), !580915) + #dbg_declare(ptr poison, !580103, !DIExpression(), !580916) + #dbg_declare(ptr %value.i.i, !580917, !DIExpression(), !580920) + #dbg_value(ptr undef, !577131, !DIExpression(), !580923) + #dbg_value(ptr undef, !577132, !DIExpression(), !580923) + br label %bb36.i, !dbg !580924 + +bb33.i.loopexit.unr-lcssa: ; preds = %bb27.us.i.us, %bb27.us.i.us.preheader + %.lcssa.ph = phi i64 [ poison, %bb27.us.i.us.preheader ], [ %88, %bb27.us.i.us ] + %iter.sroa.0.046.us.i.us.unr = phi i64 [ 0, %bb27.us.i.us.preheader ], [ %_157.us.i.us, %bb27.us.i.us ] + %accumulated.sroa.0.045.us.i.us.unr = phi i64 [ 0, %bb27.us.i.us.preheader ], [ %88, %bb27.us.i.us ] + %lcmp.mod148.not = icmp eq i64 %xtraiter147, 0, !dbg !580418 + br i1 %lcmp.mod148.not, label %bb33.i, label %bb27.us.i.us.epil, !dbg !580418 + +bb27.us.i.us.epil: ; preds = %bb33.i.loopexit.unr-lcssa + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !580147, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580416) + #dbg_value(i64 %accumulated.sroa.0.045.us.i.us.unr, !580145, !DIExpression(), !580415) + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !580149, !DIExpression(), !580491) + #dbg_value(ptr %columns.i, !563475, !DIExpression(), !580492) + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !563476, !DIExpression(), !580492) + #dbg_value(ptr %columns.i, !563465, !DIExpression(), !580493) + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !563466, !DIExpression(), !580493) + #dbg_value(ptr %columns.i, !563336, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !580496) + #dbg_value(i64 0, !563341, !DIExpression(), !580494) + #dbg_value(ptr %columns.i, !563468, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !580525) + #dbg_value(ptr %columns.i, !563336, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !580494) + %_0.sroa.0.0.i.i.us.i.us.epil = load i64, ptr %data.i.i.i.us.i, align 8, !dbg !580423, !noalias !580207, !noundef !23 + #dbg_value(ptr %14, !563465, !DIExpression(), !580498) + #dbg_value(i64 %iter.sroa.0.046.us.i.us.unr, !563466, !DIExpression(), !580498) + #dbg_value(ptr %14, !563336, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !580499) + #dbg_value(ptr %14, !563336, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !580501) + #dbg_value(ptr %14, !563468, !DIExpression(DW_OP_plus_uconst, 8, DW_OP_stack_value), !580502) + #dbg_value(i64 0, !563341, !DIExpression(), !580499) + %_0.sroa.0.0.i9.i.us.i.us.epil = load i64, ptr %data.i6.i7.i.us.i, align 8, !dbg !580428, !noalias !580207, !noundef !23 + #dbg_value(ptr poison, !580456, !DIExpression(), !580503) + #dbg_value(ptr poison, !580457, !DIExpression(), !580503) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !580458, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580503) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !580458, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580503) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !580454, !DIExpression(), !580504) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !580455, !DIExpression(), !580504) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !580445, !DIExpression(), !580505) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !580446, !DIExpression(), !580505) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !580438, !DIExpression(), !580506) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !580433, !DIExpression(), !580507) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !580439, !DIExpression(), !580506) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !580434, !DIExpression(), !580507) + %_0.i3.i.us.i.us.epil = mul i64 %_0.sroa.0.0.i9.i.us.i.us.epil, %_0.sroa.0.0.i.i.us.i.us.epil, !dbg !580430 + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !580464, !DIExpression(), !580508) + #dbg_value(i64 %_0.sroa.0.0.i.i.us.i.us.epil, !580466, !DIExpression(), !580509) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !580465, !DIExpression(), !580508) + #dbg_value(i64 %_0.sroa.0.0.i9.i.us.i.us.epil, !580467, !DIExpression(), !580509) + %_5.i.i.us.i.us.epil = zext i64 %_0.sroa.0.0.i.i.us.i.us.epil to i128, !dbg !580460 + %_6.i.i.us.i.us.epil = zext i64 %_0.sroa.0.0.i9.i.us.i.us.epil to i128, !dbg !580469 + %_4.i1.i.us.i.us.epil = mul nuw i128 %_6.i.i.us.i.us.epil, %_5.i.i.us.i.us.epil, !dbg !580470 + %_3.i2.i.us.i.us.epil = lshr i128 %_4.i1.i.us.i.us.epil, 64, !dbg !580471 + %_0.i.i159.us.i.us.epil = trunc nuw i128 %_3.i2.i.us.i.us.epil to i64, !dbg !580472 + #dbg_value(i64 %_0.i3.i.us.i.us.epil, !580151, !DIExpression(), !580510) + #dbg_value(i64 %_0.i.i159.us.i.us.epil, !580153, !DIExpression(), !580510) + #dbg_value(ptr undef, !576390, !DIExpression(), !580180) + #dbg_value(i64 %_0.i.i159.us.i.us.epil, !576396, !DIExpression(), !580180) + %117 = or i64 %accumulated.sroa.0.045.us.i.us.unr, %_0.i.i159.us.i.us.epil, !dbg !580473 + #dbg_value(i64 %117, !580145, !DIExpression(), !580415) + %self34.us.i.us.epil = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.046.us.i.us.unr, !dbg !580474 + #dbg_value(ptr %self34.us.i.us.epil, !580478, !DIExpression(), !580511) + #dbg_value(i64 %_0.i3.i.us.i.us.epil, !580479, !DIExpression(), !580511) + store i64 %_0.i3.i.us.i.us.epil, ptr %self34.us.i.us.epil, align 8, !dbg !580475, !noalias !580207 + #dbg_value(i64 poison, !580147, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580416) + #dbg_value(ptr undef, !580170, !DIExpression(), !580173) + #dbg_value(ptr undef, !580161, !DIExpression(), !580166) + #dbg_value(ptr undef, !580174, !DIExpression(), !580178) + #dbg_value(ptr poison, !580177, !DIExpression(), !580178) + #dbg_value(i64 poison, !580147, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580416) + br label %bb33.i, !dbg !580925 + +bb33.i.loopexit144.unr-lcssa: ; preds = %bb15.i.i, %bb15.i.i.preheader + %.lcssa145.ph = phi i64 [ poison, %bb15.i.i.preheader ], [ %116, %bb15.i.i ] + %iter.sroa.0.012.i.i.unr = phi i64 [ 0, %bb15.i.i.preheader ], [ %_36.i.i.1, %bb15.i.i ] + %failed.sroa.0.011.i.i.unr = phi i64 [ 0, %bb15.i.i.preheader ], [ %116, %bb15.i.i ] + %lcmp.mod.not = icmp eq i64 %xtraiter, 0, !dbg !580635 + br i1 %lcmp.mod.not, label %bb33.i, label %bb15.i.i.epil, !dbg !580635 + +bb15.i.i.epil: ; preds = %bb33.i.loopexit144.unr-lcssa + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !580573, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580626) + #dbg_value(i64 %failed.sroa.0.011.i.i.unr, !580571, !DIExpression(), !580625) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !580610, !DIExpression(), !580841) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !580603, !DIExpression(), !580604) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !580620, !DIExpression(), !580621) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !580573, !DIExpression(DW_OP_plus_uconst, 1, DW_OP_stack_value, DW_OP_LLVM_fragment, 0, 64), !580626) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !580575, !DIExpression(), !580843) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !580597, !DIExpression(), !580598) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !580844, !DIExpression(), !580848) + #dbg_value(ptr undef, !563670, !DIExpression(), !580592) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !563676, !DIExpression(), !580592) + #dbg_value(ptr poison, !563748, !DIExpression(), !580850) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !563749, !DIExpression(), !580850) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !563741, !DIExpression(), !580852) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !563733, !DIExpression(), !580854) + #dbg_value(ptr %column.val.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580852) + #dbg_value(ptr %column.val.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580854) + #dbg_value(i64 %len3.i.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580852) + #dbg_value(i64 %len3.i.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580854) + %_4.i.i.i.i.epil = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.012.i.i.unr, !dbg !580856 + %_0.i.i.i.i.epil = load i64, ptr %_4.i.i.i.i.epil, align 8, !dbg !580857, !noalias !580858, !noundef !23 + #dbg_value(ptr poison, !563748, !DIExpression(), !580862) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !563749, !DIExpression(), !580862) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !563741, !DIExpression(), !580864) + #dbg_value(i64 %iter.sroa.0.012.i.i.unr, !563733, !DIExpression(), !580866) + #dbg_value(ptr %column5.val.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580864) + #dbg_value(ptr %column5.val.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580866) + #dbg_value(i64 %len3.i.i.i, !563740, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580864) + #dbg_value(i64 %len3.i.i.i, !563734, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580866) + %_5.i3.i.i.i.epil = icmp ult i64 %iter.sroa.0.012.i.i.unr, %len3.i.i.i, !dbg !580868 + tail call void @llvm.assume(i1 %_5.i3.i.i.i.epil), !dbg !580869 + %_4.i4.i.i.i.epil = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.012.i.i.unr, !dbg !580870 + %_0.i5.i.i.i.epil = load i64, ptr %_4.i4.i.i.i.epil, align 8, !dbg !580871, !noalias !580858, !noundef !23 + #dbg_value(i64 %_0.i.i.i.i.epil, !580577, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580872) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580577, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580872) + #dbg_value(i64 %_0.i.i.i.i.epil, !580873, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580881) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580873, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580881) + #dbg_value(ptr poison, !580456, !DIExpression(), !580883) + #dbg_value(ptr poison, !580457, !DIExpression(), !580883) + #dbg_value(i64 %_0.i.i.i.i.epil, !580458, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580883) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580458, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580883) + #dbg_value(i64 %_0.i.i.i.i.epil, !580454, !DIExpression(), !580885) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580455, !DIExpression(), !580885) + #dbg_value(i64 %_0.i.i.i.i.epil, !580445, !DIExpression(), !580886) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580446, !DIExpression(), !580886) + #dbg_value(i64 %_0.i.i.i.i.epil, !580438, !DIExpression(), !580888) + #dbg_value(i64 %_0.i.i.i.i.epil, !580433, !DIExpression(), !580890) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580439, !DIExpression(), !580888) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580434, !DIExpression(), !580890) + %_0.i3.i.i.i.i.epil = mul i64 %_0.i5.i.i.i.epil, %_0.i.i.i.i.epil, !dbg !580892 + #dbg_value(i64 %_0.i.i.i.i.epil, !580464, !DIExpression(), !580893) + #dbg_value(i64 %_0.i.i.i.i.epil, !580466, !DIExpression(), !580895) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580465, !DIExpression(), !580893) + #dbg_value(i64 %_0.i5.i.i.i.epil, !580467, !DIExpression(), !580895) + %_5.i.i.i.i.i.epil = zext i64 %_0.i.i.i.i.epil to i128, !dbg !580896 + %_6.i.i.i.i.i.epil = zext i64 %_0.i5.i.i.i.epil to i128, !dbg !580897 + %_4.i1.i.i.i.i.epil = mul nuw i128 %_6.i.i.i.i.i.epil, %_5.i.i.i.i.i.epil, !dbg !580898 + %_3.i2.i.i.i.i.epil = lshr i128 %_4.i1.i.i.i.i.epil, 64, !dbg !580899 + %_0.i.i.i.i.i.epil = trunc nuw i128 %_3.i2.i.i.i.i.epil to i64, !dbg !580900 + #dbg_value(i64 poison, !580579, !DIExpression(), !580901) + #dbg_value(i64 %_0.i.i.i.i.i.epil, !580581, !DIExpression(), !580901) + #dbg_value(ptr undef, !576390, !DIExpression(), !580590) + #dbg_value(i64 %_0.i.i.i.i.i.epil, !576396, !DIExpression(), !580590) + %118 = or i64 %failed.sroa.0.011.i.i.unr, %_0.i.i.i.i.i.epil, !dbg !580902 + #dbg_value(i64 %118, !580571, !DIExpression(), !580625) + #dbg_value(i64 %_0.i3.i.i.i.i.epil, !580579, !DIExpression(), !580901) + #dbg_value(ptr %_4.sroa.10.0.i.i, !580847, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !580848) + #dbg_value(i64 %index.i, !580847, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !580848) + %self4.i.i.epil = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.012.i.i.unr, !dbg !580903 + #dbg_value(ptr %self4.i.i.epil, !580904, !DIExpression(), !580908) + #dbg_value(i64 %_0.i3.i.i.i.i.epil, !580907, !DIExpression(), !580908) + store i64 %_0.i3.i.i.i.i.epil, ptr %self4.i.i.epil, align 8, !dbg !580910, !alias.scope !580586, !noalias !580911 + #dbg_value(ptr undef, !580614, !DIExpression(), !580627) + #dbg_value(ptr undef, !580609, !DIExpression(), !580628) + #dbg_value(ptr undef, !580629, !DIExpression(), !580633) + #dbg_value(ptr poison, !580632, !DIExpression(), !580633) + br label %bb33.i, !dbg !580925 + + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/indexed-u64-mul-dense-s.md b/research/rowfn-x86-2026-08-07/codegen/indexed-u64-mul-dense-s.md new file mode 100644 index 00000000000..f0981969a0e --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/indexed-u64-mul-dense-s.md @@ -0,0 +1,71 @@ + + + +# `indexed-u64-mul-dense.s` + +```s +.LBB1688_67: + .loc 181 765 12 + andq $-2, %r10 + xorl %edi, %edi + xorl %ecx, %ecx + movq -48(%rbp), %r8 +.Ltmp114798: +.LBB1688_68: + .loc 566 39 18 + movq (%r13,%rdi,8), %rax +.Ltmp114799: + .loc 565 176 44 + mulq (%r12,%rdi,8) +.Ltmp114800: + movq %rdx, %rsi +.Ltmp114801: + .loc 207 475 9 + movq %rax, (%r8,%rdi,8) +.Ltmp114802: + .loc 566 39 18 + movq 8(%r13,%rdi,8), %rax +.Ltmp114803: + .loc 565 176 44 + mulq 8(%r12,%rdi,8) +.Ltmp114804: + .loc 568 821 53 + orq %rcx, %rsi +.Ltmp114805: + .loc 207 475 9 + movq %rax, 8(%r8,%rdi,8) +.Ltmp114806: + .loc 156 717 17 + addq $2, %rdi +.Ltmp114807: + .loc 565 176 44 + movq %rdx, %rcx +.Ltmp114808: + .loc 568 821 53 + orq %rsi, %rcx +.Ltmp114809: + .loc 181 765 12 + cmpq %r10, %rdi + jne .LBB1688_68 +.Ltmp114810: +.LBB1688_69: + testb $1, %r9b + je .LBB1688_84 +.Ltmp114811: + .loc 566 39 18 + movq (%r13,%rdi,8), %rax +.Ltmp114812: + .loc 565 176 44 + mulq (%r12,%rdi,8) +.Ltmp114813: + .loc 568 821 53 + orq %rdx, %rcx +.Ltmp114814: + .loc 568 0 53 is_stmt 0 + movq -48(%rbp), %rdx +.Ltmp114815: + .loc 207 475 9 is_stmt 1 + movq %rax, (%rdx,%rdi,8) +.Ltmp114816: + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/owned-i64-mul-dense-ll.md b/research/rowfn-x86-2026-08-07/codegen/owned-i64-mul-dense-ll.md new file mode 100644 index 00000000000..c4a917571df --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/owned-i64-mul-dense-ll.md @@ -0,0 +1,84 @@ + + + +# `owned-i64-mul-dense.ll` + +```ll +terminate.i: ; preds = %bb57.i + %78 = landingpad { ptr, i32 } + filter [0 x ptr] zeroinitializer +; call core::panicking::panic_in_cleanup + call void @_ZN4core9panicking16panic_in_cleanup17h8f68387bb6cbbf54E() #88, !dbg !573712, !noalias !573567 + unreachable, !dbg !573712 + +bb18.i: ; preds = %bb18.i, %bb18.lr.ph.i + %_15552.i = phi i64 [ 1, %bb18.lr.ph.i ], [ %_155.i, %bb18.i ] + %iter.sroa.0.051.i = phi i64 [ 0, %bb18.lr.ph.i ], [ %_15552.i, %bb18.i ] + %failed.sroa.0.050.i = phi i64 [ 0, %bb18.lr.ph.i ], [ %81, %bb18.i ] + #dbg_value(i64 %iter.sroa.0.051.i, !573477, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !573890) + #dbg_value(i64 %failed.sroa.0.050.i, !573466, !DIExpression(), !573746) + #dbg_value(i64 %iter.sroa.0.051.i, !573479, !DIExpression(), !574113) + #dbg_value(ptr undef, !552190, !DIExpression(), !573546) + #dbg_value(i64 %iter.sroa.0.051.i, !552196, !DIExpression(), !573546) + #dbg_value(ptr poison, !552716, !DIExpression(), !574114) + #dbg_value(i64 %iter.sroa.0.051.i, !552717, !DIExpression(), !574114) + #dbg_value(ptr poison, !552716, !DIExpression(), !574116) + #dbg_value(i64 %iter.sroa.0.051.i, !552717, !DIExpression(), !574116) + %79 = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.051.i, !dbg !574118 + %_0.i.i97.i = load i64, ptr %79, align 8, !dbg !574118, !noalias !574119, !noundef !23 + %80 = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.051.i, !dbg !574122 + %_0.i5.i.i = load i64, ptr %80, align 8, !dbg !574122, !noalias !574119, !noundef !23 + #dbg_value(ptr poison, !573820, !DIExpression(), !574123) + #dbg_value(ptr poison, !573821, !DIExpression(), !574123) + #dbg_value(i64 %_0.i.i97.i, !573822, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !574123) + #dbg_value(i64 %_0.i5.i.i, !573822, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !574123) + #dbg_value(i64 %_0.i.i97.i, !573818, !DIExpression(), !574125) + #dbg_value(i64 %_0.i5.i.i, !573819, !DIExpression(), !574125) + #dbg_value(i64 %_0.i.i97.i, !573809, !DIExpression(), !574126) + #dbg_value(i64 %_0.i5.i.i, !573810, !DIExpression(), !574126) + #dbg_value(i64 %_0.i.i97.i, !573798, !DIExpression(), !574128) + #dbg_value(i64 %_0.i.i97.i, !573793, !DIExpression(), !574130) + #dbg_value(i64 %_0.i5.i.i, !573799, !DIExpression(), !574128) + #dbg_value(i64 %_0.i5.i.i, !573794, !DIExpression(), !574130) + %_0.i.i111.i = mul i64 %_0.i5.i.i, %_0.i.i97.i, !dbg !574132 + #dbg_value(i64 %_0.i.i97.i, !573830, !DIExpression(), !574133) + #dbg_value(i64 %_0.i.i97.i, !573832, !DIExpression(), !574135) + #dbg_value(i64 %_0.i5.i.i, !573831, !DIExpression(), !574133) + #dbg_value(i64 %_0.i5.i.i, !573833, !DIExpression(), !574135) + %_4.i1.i.i = sext i64 %_0.i.i97.i to i128, !dbg !574136 + %_5.i.i.i = sext i64 %_0.i5.i.i to i128, !dbg !574137 + %wide.i.i.i = mul nsw i128 %_5.i.i.i, %_4.i1.i.i, !dbg !574136 + #dbg_value(i128 %wide.i.i.i, !573834, !DIExpression(), !574138) + %kept.i.i.i = trunc i128 %wide.i.i.i to i64, !dbg !574139 + #dbg_value(i64 %kept.i.i.i, !573836, !DIExpression(), !574140) + %_8.i.i.i = lshr i128 %wide.i.i.i, 64, !dbg !574141 + %discarded.i.i.i = trunc nuw i128 %_8.i.i.i to i64, !dbg !574142 + #dbg_value(i64 %discarded.i.i.i, !573838, !DIExpression(), !574143) + %_10.i.i.i = ashr i64 %kept.i.i.i, 63, !dbg !574144 + %_9.i.i.i = xor i64 %_10.i.i.i, %discarded.i.i.i, !dbg !574145 + #dbg_value(i64 %_0.i.i111.i, !573481, !DIExpression(), !574146) + #dbg_value(i64 %_9.i.i.i, !573483, !DIExpression(), !574146) + #dbg_value(ptr undef, !573548, !DIExpression(), !573557) + #dbg_value(i64 %_9.i.i.i, !573554, !DIExpression(), !573557) + %81 = or i64 %_9.i.i.i, %failed.sroa.0.050.i, !dbg !574147 + #dbg_value(i64 %81, !573466, !DIExpression(), !573746) + %self32.i = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.051.i, !dbg !574148 + #dbg_value(ptr %self32.i, !573852, !DIExpression(), !574149) + #dbg_value(i64 %_0.i.i111.i, !573853, !DIExpression(), !574149) + store i64 %_0.i.i111.i, ptr %self32.i, align 8, !dbg !574151, !noalias !573567 + #dbg_value(i64 %_15552.i, !573477, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !573890) + #dbg_value(ptr undef, !573517, !DIExpression(), !573539) + #dbg_value(ptr undef, !573505, !DIExpression(), !573535) + #dbg_value(ptr undef, !573521, !DIExpression(), !573540) + #dbg_value(ptr poison, !573524, !DIExpression(), !573540) + %_155.i = add i64 %_15552.i, 1, !dbg !574152 + #dbg_value(i64 poison, !573477, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !573890) + %exitcond.not.i = icmp eq i64 %_15552.i, %len3.i4.i.i.fr, !dbg !574153 + br i1 %exitcond.not.i, label %bb38.i, label %bb18.i, !dbg !573891 + +bb38.thread.i: ; preds = %bb31.preheader.i.thread, %bb17.preheader.split.i, %bb31.preheader.i + #dbg_value(i64 0, !573466, !DIExpression(), !573746) + #dbg_value(i64 %index.i, !573459, !DIExpression(DW_OP_LLVM_fragment, 128, 64), !573709) + call void @llvm.lifetime.start.p0(i64 24, ptr nonnull %value.i.i), !dbg !574154, !noalias !573647 + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/owned-i64-mul-dense-s.md b/research/rowfn-x86-2026-08-07/codegen/owned-i64-mul-dense-s.md new file mode 100644 index 00000000000..db24d697b4e --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/owned-i64-mul-dense-s.md @@ -0,0 +1,46 @@ + + + +# `owned-i64-mul-dense.s` + +```s + .loc 562 112 26 is_stmt 1 + je .LBB1685_70 +.Ltmp114190: + .loc 562 0 26 is_stmt 0 + movq -128(%rbp), %r13 +.Ltmp114191: + xorl %esi, %esi +.Ltmp114192: + xorl %ecx, %ecx + movq -48(%rbp), %rdi +.Ltmp114193: + .p2align 4 +.LBB1685_25: + .loc 564 62 9 is_stmt 1 + movq (%r13,%rsi,8), %rax +.Ltmp114194: + .loc 565 194 24 + imulq (%r12,%rsi,8) +.Ltmp114195: + .loc 207 475 9 + movq %rax, (%rdi,%rsi,8) +.Ltmp114196: + .loc 565 198 26 + sarq $63, %rax +.Ltmp114197: + .loc 565 198 13 is_stmt 0 + xorq %rdx, %rax +.Ltmp114198: + .loc 566 821 53 is_stmt 1 + orq %rax, %rcx +.Ltmp114199: + .loc 182 1904 50 + incq %rsi +.Ltmp114200: + cmpq %rsi, %r9 + jne .LBB1685_25 + jmp .LBB1685_60 +.Ltmp114201: + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/owned-u64-mul-dense-ll.md b/research/rowfn-x86-2026-08-07/codegen/owned-u64-mul-dense-ll.md new file mode 100644 index 00000000000..0659a92c119 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/owned-u64-mul-dense-ll.md @@ -0,0 +1,121 @@ + + + +# `owned-u64-mul-dense.ll` + +```ll +bb18.i: ; preds = %bb18.i, %bb18.lr.ph.i.new + %_15552.i = phi i64 [ 1, %bb18.lr.ph.i.new ], [ %_155.i.1, %bb18.i ] + %iter.sroa.0.051.i = phi i64 [ 0, %bb18.lr.ph.i.new ], [ %_155.i, %bb18.i ] + %failed.sroa.0.050.i = phi i64 [ 0, %bb18.lr.ph.i.new ], [ %120, %bb18.i ] + %niter = phi i64 [ 0, %bb18.lr.ph.i.new ], [ %niter.next.1, %bb18.i ] + #dbg_value(i64 %iter.sroa.0.051.i, !576964, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577378) + #dbg_value(i64 %failed.sroa.0.050.i, !576953, !DIExpression(), !577231) + #dbg_value(i64 %iter.sroa.0.051.i, !576966, !DIExpression(), !577601) + #dbg_value(ptr undef, !561311, !DIExpression(), !577032) + #dbg_value(i64 %iter.sroa.0.051.i, !561317, !DIExpression(), !577032) + #dbg_value(ptr poison, !561836, !DIExpression(), !577602) + #dbg_value(i64 %iter.sroa.0.051.i, !561837, !DIExpression(), !577602) + #dbg_value(ptr poison, !561836, !DIExpression(), !577604) + #dbg_value(i64 %iter.sroa.0.051.i, !561837, !DIExpression(), !577604) + %115 = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %iter.sroa.0.051.i, !dbg !577606 + %_0.i.i92.i = load i64, ptr %115, align 8, !dbg !577606, !noalias !577607, !noundef !23 + %116 = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %iter.sroa.0.051.i, !dbg !577610 + %_0.i5.i.i = load i64, ptr %116, align 8, !dbg !577610, !noalias !577607, !noundef !23 + #dbg_value(ptr poison, !577301, !DIExpression(), !577611) + #dbg_value(ptr poison, !577302, !DIExpression(), !577611) + #dbg_value(i64 %_0.i.i92.i, !577303, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577611) + #dbg_value(i64 %_0.i5.i.i, !577303, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577611) + #dbg_value(i64 %_0.i.i92.i, !577299, !DIExpression(), !577613) + #dbg_value(i64 %_0.i5.i.i, !577300, !DIExpression(), !577613) + #dbg_value(i64 %_0.i.i92.i, !577290, !DIExpression(), !577614) + #dbg_value(i64 %_0.i5.i.i, !577291, !DIExpression(), !577614) + #dbg_value(i64 %_0.i.i92.i, !577283, !DIExpression(), !577616) + #dbg_value(i64 %_0.i.i92.i, !577278, !DIExpression(), !577618) + #dbg_value(i64 %_0.i5.i.i, !577284, !DIExpression(), !577616) + #dbg_value(i64 %_0.i5.i.i, !577279, !DIExpression(), !577618) + %_0.i3.i.i = mul i64 %_0.i5.i.i, %_0.i.i92.i, !dbg !577620 + #dbg_value(i64 %_0.i.i92.i, !577309, !DIExpression(), !577621) + #dbg_value(i64 %_0.i.i92.i, !577311, !DIExpression(), !577623) + #dbg_value(i64 %_0.i5.i.i, !577310, !DIExpression(), !577621) + #dbg_value(i64 %_0.i5.i.i, !577312, !DIExpression(), !577623) + %_5.i.i.i = zext i64 %_0.i.i92.i to i128, !dbg !577624 + %_6.i.i.i = zext i64 %_0.i5.i.i to i128, !dbg !577625 + %_4.i1.i.i = mul nuw i128 %_6.i.i.i, %_5.i.i.i, !dbg !577626 + %_3.i2.i.i = lshr i128 %_4.i1.i.i, 64, !dbg !577627 + %_0.i.i106.i = trunc nuw i128 %_3.i2.i.i to i64, !dbg !577628 + #dbg_value(i64 %_0.i3.i.i, !576968, !DIExpression(), !577629) + #dbg_value(i64 %_0.i.i106.i, !576970, !DIExpression(), !577629) + #dbg_value(ptr undef, !573548, !DIExpression(), !577036) + #dbg_value(i64 %_0.i.i106.i, !573554, !DIExpression(), !577036) + %117 = or i64 %failed.sroa.0.050.i, %_0.i.i106.i, !dbg !577630 + #dbg_value(i64 %117, !576953, !DIExpression(), !577231) + %self32.i = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %iter.sroa.0.051.i, !dbg !577631 + #dbg_value(ptr %self32.i, !577323, !DIExpression(), !577632) + #dbg_value(i64 %_0.i3.i.i, !577324, !DIExpression(), !577632) + store i64 %_0.i3.i.i, ptr %self32.i, align 8, !dbg !577634, !noalias !577052 + #dbg_value(i64 %_15552.i, !576964, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577378) + #dbg_value(ptr undef, !577003, !DIExpression(), !577025) + #dbg_value(ptr undef, !576991, !DIExpression(), !577021) + #dbg_value(ptr undef, !577007, !DIExpression(), !577026) + #dbg_value(ptr poison, !577010, !DIExpression(), !577026) + %_155.i = add i64 %_15552.i, 1, !dbg !577635 + #dbg_value(i64 poison, !576964, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577378) + #dbg_value(i64 %_15552.i, !576964, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577378) + #dbg_value(i64 %_15552.i, !576966, !DIExpression(), !577601) + #dbg_value(i64 %_15552.i, !561317, !DIExpression(), !577032) + #dbg_value(i64 %_15552.i, !561837, !DIExpression(), !577602) + #dbg_value(i64 %_15552.i, !561837, !DIExpression(), !577604) + %118 = getelementptr inbounds nuw i64, ptr %column.val.i.i, i64 %_15552.i, !dbg !577606 + %_0.i.i92.i.1 = load i64, ptr %118, align 8, !dbg !577606, !noalias !577607, !noundef !23 + %119 = getelementptr inbounds nuw i64, ptr %column5.val.i.i, i64 %_15552.i, !dbg !577610 + %_0.i5.i.i.1 = load i64, ptr %119, align 8, !dbg !577610, !noalias !577607, !noundef !23 + #dbg_value(i64 %_0.i.i92.i.1, !577303, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577611) + #dbg_value(i64 %_0.i5.i.i.1, !577303, !DIExpression(DW_OP_LLVM_fragment, 64, 64), !577611) + #dbg_value(i64 %_0.i.i92.i.1, !577299, !DIExpression(), !577613) + #dbg_value(i64 %_0.i5.i.i.1, !577300, !DIExpression(), !577613) + #dbg_value(i64 %_0.i.i92.i.1, !577290, !DIExpression(), !577614) + #dbg_value(i64 %_0.i5.i.i.1, !577291, !DIExpression(), !577614) + #dbg_value(i64 %_0.i.i92.i.1, !577283, !DIExpression(), !577616) + #dbg_value(i64 %_0.i.i92.i.1, !577278, !DIExpression(), !577618) + #dbg_value(i64 %_0.i5.i.i.1, !577284, !DIExpression(), !577616) + #dbg_value(i64 %_0.i5.i.i.1, !577279, !DIExpression(), !577618) + %_0.i3.i.i.1 = mul i64 %_0.i5.i.i.1, %_0.i.i92.i.1, !dbg !577620 + #dbg_value(i64 %_0.i.i92.i.1, !577309, !DIExpression(), !577621) + #dbg_value(i64 %_0.i.i92.i.1, !577311, !DIExpression(), !577623) + #dbg_value(i64 %_0.i5.i.i.1, !577310, !DIExpression(), !577621) + #dbg_value(i64 %_0.i5.i.i.1, !577312, !DIExpression(), !577623) + %_5.i.i.i.1 = zext i64 %_0.i.i92.i.1 to i128, !dbg !577624 + %_6.i.i.i.1 = zext i64 %_0.i5.i.i.1 to i128, !dbg !577625 + %_4.i1.i.i.1 = mul nuw i128 %_6.i.i.i.1, %_5.i.i.i.1, !dbg !577626 + %_3.i2.i.i.1 = lshr i128 %_4.i1.i.i.1, 64, !dbg !577627 + %_0.i.i106.i.1 = trunc nuw i128 %_3.i2.i.i.1 to i64, !dbg !577628 + #dbg_value(i64 %_0.i3.i.i.1, !576968, !DIExpression(), !577629) + #dbg_value(i64 %_0.i.i106.i.1, !576970, !DIExpression(), !577629) + #dbg_value(i64 %_0.i.i106.i.1, !573554, !DIExpression(), !577036) + %120 = or i64 %117, %_0.i.i106.i.1, !dbg !577630 + #dbg_value(i64 %120, !576953, !DIExpression(), !577231) + %self32.i.1 = getelementptr inbounds nuw i64, ptr %_4.sroa.10.0.i.i, i64 %_15552.i, !dbg !577631 + #dbg_value(ptr %self32.i.1, !577323, !DIExpression(), !577632) + #dbg_value(i64 %_0.i3.i.i.1, !577324, !DIExpression(), !577632) + store i64 %_0.i3.i.i.1, ptr %self32.i.1, align 8, !dbg !577634, !noalias !577052 + #dbg_value(i64 %_155.i, !576964, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577378) + %_155.i.1 = add i64 %_15552.i, 2, !dbg !577635 + #dbg_value(i64 poison, !576964, !DIExpression(DW_OP_LLVM_fragment, 0, 64), !577378) + %niter.next.1 = add i64 %niter, 2, !dbg !577379 + %niter.ncmp.1 = icmp eq i64 %niter.next.1, %unroll_iter, !dbg !577379 + br i1 %niter.ncmp.1, label %bb38.i.loopexit144.unr-lcssa, label %bb18.i, !dbg !577379 + +bb38.thread.i: ; preds = %bb31.preheader.i.thread, %bb17.preheader.split.i, %bb31.preheader.i + #dbg_value(i64 0, !576953, !DIExpression(), !577231) + #dbg_value(i64 %index.i, !576946, !DIExpression(DW_OP_LLVM_fragment, 128, 64), !577194) + call void @llvm.lifetime.start.p0(i64 24, ptr nonnull %value.i.i), !dbg !577636, !noalias !577132 + #dbg_value(i64 0, !576924, !DIExpression(), !577638) + #dbg_declare(ptr poison, !576928, !DIExpression(), !577639) + #dbg_declare(ptr %value.i.i, !577640, !DIExpression(), !577643) + #dbg_value(ptr undef, !574157, !DIExpression(), !577646) + #dbg_value(ptr undef, !574158, !DIExpression(), !577646) + br label %bb41.i, !dbg !577647 + + +``` diff --git a/research/rowfn-x86-2026-08-07/codegen/owned-u64-mul-dense-s.md b/research/rowfn-x86-2026-08-07/codegen/owned-u64-mul-dense-s.md new file mode 100644 index 00000000000..47957a445c2 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/owned-u64-mul-dense-s.md @@ -0,0 +1,72 @@ + + + +# `owned-u64-mul-dense.s` + +```s +.LBB1688_67: + .loc 562 112 26 + andq $-2, %r10 + xorl %ecx, %ecx + xorl %edi, %edi + movq -48(%rbp), %r8 +.Ltmp115437: +.LBB1688_68: + .loc 564 62 9 + movq (%r13,%rdi,8), %rax +.Ltmp115438: + .loc 565 176 44 + mulq (%r12,%rdi,8) +.Ltmp115439: + movq %rdx, %rsi +.Ltmp115440: + .loc 207 475 9 + movq %rax, (%r8,%rdi,8) +.Ltmp115441: + .loc 564 62 9 + movq 8(%r13,%rdi,8), %rax +.Ltmp115442: + .loc 565 176 44 + mulq 8(%r12,%rdi,8) +.Ltmp115443: + .loc 566 821 53 + orq %rcx, %rsi +.Ltmp115444: + .loc 565 176 44 + movq %rdx, %rcx +.Ltmp115445: + .loc 566 821 53 + orq %rsi, %rcx +.Ltmp115446: + .loc 207 475 9 + movq %rax, 8(%r8,%rdi,8) +.Ltmp115447: + .loc 562 112 26 + addq $2, %rdi + cmpq %rdi, %r10 + jne .LBB1688_68 +.Ltmp115448: +.LBB1688_69: + testb $1, %r9b + je .LBB1688_84 +.Ltmp115449: + .loc 564 62 9 + movq (%r13,%rdi,8), %rax +.Ltmp115450: + .loc 565 176 44 + mulq (%r12,%rdi,8) +.Ltmp115451: + .loc 566 821 53 + orq %rdx, %rcx +.Ltmp115452: + .loc 566 0 53 is_stmt 0 + movq -48(%rbp), %rdx +.Ltmp115453: + .loc 207 475 9 is_stmt 1 + movq %rax, (%rdx,%rdi,8) +.Ltmp115454: +.LBB1688_84: + .loc 182 1868 54 + testq %rcx, %rcx + +``` From d0a3186502ee230ac1d90c8e421c4faaff08c11b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 19:39:06 -0400 Subject: [PATCH 011/160] Clarify RowFn compiler ablation evidence Signed-off-by: "Connor Tsui" --- research/rowfn-x86-2026-08-07/README.md | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/research/rowfn-x86-2026-08-07/README.md b/research/rowfn-x86-2026-08-07/README.md index e1df97740ba..fefa6f973d0 100644 --- a/research/rowfn-x86-2026-08-07/README.md +++ b/research/rowfn-x86-2026-08-07/README.md @@ -109,8 +109,9 @@ exposed the next compiler-sensitive detail. ## Store placement and the `Copy` ablation Moving the output store before the failure OR changed `mul_i32_constant` from about 32.38 to -18.68 microseconds. Final assembly records overflow `setb`, output store, then loop-carried OR. This -is compiler scheduling sensitivity, not a semantic difference. +18.68 microseconds. Final assembly records overflow `setb`, output store, then loop-carried OR. The +source-order ablation therefore changes whole-function code generation, but the final instruction +order is **not** a confirmed explanation for the throughput change. Adding the descriptive `Output: Copy` bound regressed that case to 29.88/29.86 microseconds. Replacing it with compile-time `!needs_drop::()` returned it to 18.65/18.67. A generic store @@ -118,6 +119,13 @@ helper did not repair the `Copy` case. The executor needs only the no-drop prope cleanup; it never copies an output. The API therefore enforces the actual requirement without the measured optimizer-visible bound. Re-test this workaround whenever LLVM changes. +An isolated exact-loop ablation on the same CPU contradicted the simple scheduling story. LLVM-MCA +estimated both final instruction orders at 2.7 cycles per iteration. Direct CPU 8 measurements made +OR-before-store slightly faster at 0.75-0.77 nanoseconds per row than store-before-OR at +0.823-0.825 nanoseconds per row. The production source-order effect is therefore an unresolved +whole-function compiler interaction, such as layout, surrounding control flow, or another +optimization decision. Do not justify the chosen source order as intrinsically better scheduling. + ## Final results Order: baseline, final, candidate, repeated twice. Values are median microseconds. @@ -165,6 +173,12 @@ AVX features, and LLVM selected the same essential scalar high-half strategy as [`final i64 assembly`](codegen/final-i64-mul-dense-s.md), and [`final u64 assembly`](codegen/final-u64-mul-dense-s.md). +A separate minimal `target-cpu=native` experiment did form `<8 x i128>` operations in LLVM IR for +the widened `u64` product. The x86 backend still scalarized them into eight `mulq`/`imulq` +instructions, then used ZMM registers only to pack and reduce the scalar results. x86 has no true +wide 64-by-64-to-128 integer multiply here. Seeing a vector IR type or ZMM instruction is therefore +not evidence that the expensive multiply itself executed as SIMD. + `-C remark=loop-vectorize` emitted no remark attributable to the exact dense production loop. The constant fallback source line had successes for other monomorphs and duplicated cost-model misses, but diagnostics lacked function identity. Exact IR proves the measured specialization is scalar; @@ -179,7 +193,8 @@ Confirmed: - `SinkResult` already reduced to a register OR and did not impose a per-row `Result`. - Output ownership materially helped but was insufficient alone. - A typed indexed source restored stable parity for varying primitive tuples. -- Store-before-OR and omission of a `Copy` bound materially affect LLVM 21.1.2 constant codegen. +- Source store/OR order and omission of a `Copy` bound materially affect LLVM 21.1.2 production + codegen; the mechanism behind the source-order effect remains unresolved. - The default x86 target prefers scalar high-half 64-bit multiply; SIMD is not the recovered speed. Still inference: From 53fff3325046ac24d1d6a5e5d595ebf1f6502f5b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 7 Aug 2026 19:58:32 -0400 Subject: [PATCH 012/160] Record completed RowFn Copy ablation Signed-off-by: "Connor Tsui" --- research/rowfn-x86-2026-08-07/README.md | 64 +++++++++++++------ .../codegen/copy-ablation.md | 52 +++++++++++++++ 2 files changed, 95 insertions(+), 21 deletions(-) create mode 100644 research/rowfn-x86-2026-08-07/codegen/copy-ablation.md diff --git a/research/rowfn-x86-2026-08-07/README.md b/research/rowfn-x86-2026-08-07/README.md index fefa6f973d0..5a657f47fa0 100644 --- a/research/rowfn-x86-2026-08-07/README.md +++ b/research/rowfn-x86-2026-08-07/README.md @@ -106,25 +106,46 @@ execution validates both varying lengths once. The generic owned executor calls The indexed source closed the varying and nullable gap. It did not affect mixed constants, which exposed the next compiler-sensitive detail. -## Store placement and the `Copy` ablation - -Moving the output store before the failure OR changed `mul_i32_constant` from about 32.38 to -18.68 microseconds. Final assembly records overflow `setb`, output store, then loop-carried OR. The -source-order ablation therefore changes whole-function code generation, but the final instruction -order is **not** a confirmed explanation for the throughput change. - -Adding the descriptive `Output: Copy` bound regressed that case to 29.88/29.86 microseconds. -Replacing it with compile-time `!needs_drop::()` returned it to 18.65/18.67. A generic store -helper did not repair the `Copy` case. The executor needs only the no-drop property for safe panic -cleanup; it never copies an output. The API therefore enforces the actual requirement without the -measured optimizer-visible bound. Re-test this workaround whenever LLVM changes. - -An isolated exact-loop ablation on the same CPU contradicted the simple scheduling story. LLVM-MCA -estimated both final instruction orders at 2.7 cycles per iteration. Direct CPU 8 measurements made -OR-before-store slightly faster at 0.75-0.77 nanoseconds per row than store-before-OR at -0.823-0.825 nanoseconds per row. The production source-order effect is therefore an unresolved -whole-function compiler interaction, such as layout, surrounding control flow, or another -optimization decision. Do not justify the chosen source order as intrinsically better scheduling. +## Compiler ablations: `Copy`, source order, and whole-function sensitivity + +The completed ablation matrix isolates the public `Output: Copy` bound as a reliable trigger, while +falsifying the simpler explanations considered during the initial investigation: + +| Variant | `mul_i32_constant` run 1 / 2 | +| --- | ---: | +| No `Copy` bound | 18.77 / 18.72 us | +| Inert private marker bound | 18.77 / 18.72 us | +| `Output: Copy` | 29.94 / 29.93 us | +| `Output: Copy`, `codegen-units=1` | 29.87 / 29.89 us | + +The `i64` and `u64` controls did not move. The inert private marker is important: an arbitrary +where-clause or source perturbation is insufficient to trigger the loss. The result is specific to +the optimizer-visible `Copy` constraint, though the mechanism is not yet known. + +The default-CGU DWARF ranges show large whole-function differences for the exact `i32 CheckedMul` +monomorph. The `Copy` function spans `0xe58c90..0xe59adc` (`0xe4c` bytes); no-Copy spans +`0xe7a1c0..0xe7b6d0` (`0x1510` bytes). The `Copy` hot loop at `0xe58f90` is only 16-byte aligned and +computes the low multiply before the widened chain. The no-Copy loop at `0xe7b260` is 32-byte +aligned, computes the widened chain first, and delays the low multiply. LLVM-MCA nevertheless +predicts the smaller `Copy` loop slightly better, 2.5 versus 2.7 cycles. Alignment and final loop +scheduling therefore do not explain the measured direction. + +A fresh `Copy` plus `codegen-units=1` build makes this conclusion stronger: its optimized IR already +has store-before-OR, yet the linked benchmark remains at about 29.9 microseconds. Store-before-OR is +neither sufficient nor established as causal. The earlier source-order edit changed production +performance, but it must be described only as another trigger for a whole-function compiler +interaction. An isolated exact-loop hardware ablation also found OR-before-store slightly faster +(0.75-0.77 ns/row) than store-before-OR (0.823-0.825 ns/row), while LLVM-MCA rated both at 2.7 +cycles. The loop's local instruction order cannot explain the production result. + +A standalone generic `MaybeUninit` loop emits identical optimized IR and assembly with and without +`Copy`. The sensitivity therefore needs the real trait, closure, `Vec`, and monomorphization context. +This is currently evidence of compiler phase-order or code-quality sensitivity, not enough to claim +a rustc correctness bug or a specific LLVM bug. The next upstream step is to reduce the real +monomorph while retaining both the timing and whole-function delta, then bisect MIR/LLVM passes and +compiler versions. The executor needs only the no-drop property, so the selected API continues to +enforce `!needs_drop::()` without exposing the harmful, unnecessary `Copy` bound. See the +compact [Copy-ablation evidence](codegen/copy-ablation.md). ## Final results @@ -193,8 +214,9 @@ Confirmed: - `SinkResult` already reduced to a register OR and did not impose a per-row `Result`. - Output ownership materially helped but was insufficient alone. - A typed indexed source restored stable parity for varying primitive tuples. -- Source store/OR order and omission of a `Copy` bound materially affect LLVM 21.1.2 production - codegen; the mechanism behind the source-order effect remains unresolved. +- An `Output: Copy` bound reliably triggers slower LLVM 21.1.2 production codegen; an inert marker + does not. Source store/OR order is neither sufficient nor established as causal. The mechanism is + an unresolved whole-function compiler interaction. - The default x86 target prefers scalar high-half 64-bit multiply; SIMD is not the recovered speed. Still inference: diff --git a/research/rowfn-x86-2026-08-07/codegen/copy-ablation.md b/research/rowfn-x86-2026-08-07/codegen/copy-ablation.md new file mode 100644 index 00000000000..185018ab6a9 --- /dev/null +++ b/research/rowfn-x86-2026-08-07/codegen/copy-ablation.md @@ -0,0 +1,52 @@ + + + +# `Copy`-bound compiler ablation + +This note records the compact evidence behind the no-drop assertion in owned RowFn execution. + +## Timings + +All public `binary_ops` runs used CPU 8 and the same default repository flags. + +```text +no Copy bound mul_i32_constant 18.77 / 18.72 us +inert private marker bound mul_i32_constant 18.77 / 18.72 us +Output: Copy mul_i32_constant 29.94 / 29.93 us +Output: Copy, CGU=1 mul_i32_constant 29.87 / 29.89 us +i64/u64 controls unchanged +``` + +The inert marker rules out a generic “any where-clause/source change perturbs codegen” explanation. +`codegen-units=1` rules out the default partitioning choice as a repair. + +## Exact production functions + +Default-CGU DWARF identified the measured `i32 CheckedMul` monomorphs: + +```text +Copy: 0xe58c90..0xe59adc, size 0xe4c +no-Copy: 0xe7a1c0..0xe7b6d0, size 0x1510 + +Copy hot loop: 0xe58f90, 16-byte but not 32-byte aligned +no-Copy hot loop: 0xe7b260, 32-byte aligned +``` + +The Copy loop schedules the low `imul` before the widened-product chain. No-Copy schedules the +widened chain first and delays the low multiply. LLVM-MCA predicts Copy slightly better at 2.5 +cycles versus 2.7, contradicting scheduling as the cause of its 1.6x wall-time loss. + +Fresh Copy-plus-CGU1 optimized IR contains store-before-OR and still runs at 29.9 microseconds. +Therefore store-before-OR is not sufficient. Exact isolated loops also contradict causality: + +```text +LLVM-MCA: both orders 2.7 cycles +OR before store: 0.75-0.77 ns/row +store before OR: 0.823-0.825 ns/row +``` + +The standalone generic `MaybeUninit` loop produces identical Copy/no-Copy IR and assembly. The +remaining hypothesis is phase-order or code-quality sensitivity requiring the real trait, closure, +`Vec`, and monomorphization context. Do not label this a correctness bug or assign it to a specific +rustc/LLVM pass without a reduced reproducer. Reduce the real monomorph while retaining timing and +whole-function changes, then bisect MIR/LLVM passes and compiler versions. From 0a0ad0db146c1b761b3727f40f1c78ef02f62baa Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 10:42:32 -0400 Subject: [PATCH 013/160] Add the RowFn scalar function framework Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/mod.rs | 8 + vortex-array/src/scalar_fn/row/batch/args.rs | 66 +++ .../src/scalar_fn/row/batch/execution.rs | 440 ++++++++++++++++++ vortex-array/src/scalar_fn/row/batch/mod.rs | 22 + .../src/scalar_fn/row/batch/policy.rs | 102 ++++ vortex-array/src/scalar_fn/row/execute/mod.rs | 52 +++ .../src/scalar_fn/row/execute/owned.rs | 112 +++++ .../src/scalar_fn/row/execute/sink.rs | 192 ++++++++ vortex-array/src/scalar_fn/row/mod.rs | 36 ++ vortex-array/src/scalar_fn/row/row_fn.rs | 90 ++++ .../src/scalar_fn/row/types/element/bool.rs | 68 +++ .../src/scalar_fn/row/types/element/mod.rs | 129 +++++ .../scalar_fn/row/types/element/primitive.rs | 74 +++ .../src/scalar_fn/row/types/element/tuple.rs | 414 ++++++++++++++++ vortex-array/src/scalar_fn/row/types/mod.rs | 23 + .../src/scalar_fn/row/types/result.rs | 124 +++++ vortex-array/src/scalar_fn/row/types/sink.rs | 139 ++++++ .../src/scalar_fn/row/visitor/check.rs | 126 +++++ .../src/scalar_fn/row/visitor/execute.rs | 226 +++++++++ vortex-array/src/scalar_fn/row/visitor/mod.rs | 157 +++++++ .../src/scalar_fn/row/visitor/plan.rs | 104 +++++ vortex-array/src/scalar_fn/row/vtable.rs | 171 +++++++ 22 files changed, 2875 insertions(+) create mode 100644 vortex-array/src/scalar_fn/row/batch/args.rs create mode 100644 vortex-array/src/scalar_fn/row/batch/execution.rs create mode 100644 vortex-array/src/scalar_fn/row/batch/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/batch/policy.rs create mode 100644 vortex-array/src/scalar_fn/row/execute/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/execute/owned.rs create mode 100644 vortex-array/src/scalar_fn/row/execute/sink.rs create mode 100644 vortex-array/src/scalar_fn/row/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/row_fn.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/bool.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/primitive.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/tuple.rs create mode 100644 vortex-array/src/scalar_fn/row/types/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/types/result.rs create mode 100644 vortex-array/src/scalar_fn/row/types/sink.rs create mode 100644 vortex-array/src/scalar_fn/row/visitor/check.rs create mode 100644 vortex-array/src/scalar_fn/row/visitor/execute.rs create mode 100644 vortex-array/src/scalar_fn/row/visitor/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/visitor/plan.rs create mode 100644 vortex-array/src/scalar_fn/row/vtable.rs diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 003dbc2ca48..5e73caefdfa 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -6,6 +6,11 @@ //! This module contains the [`ScalarFnVTable`] trait and all built-in scalar function //! implementations. Expressions ([`crate::expr::Expression`]) reference scalar functions //! at each node. +//! +//! Use [`RowFn`] for strict functions whose natural kernel computes one row at a time. It derives +//! decoding, constant handling, null propagation, output construction, and validity. Implement +//! [`ScalarFnVTable`] directly when the natural kernel is columnar, aliases an input, or may +//! produce null from otherwise valid inputs. use vortex_session::registry::Id; @@ -35,6 +40,9 @@ pub use options::*; mod signature; pub use signature::*; +mod row; +pub use row::*; + pub mod fns; pub mod internal; pub mod session; diff --git a/vortex-array/src/scalar_fn/row/batch/args.rs b/vortex-array/src/scalar_fn/row/batch/args.rs new file mode 100644 index 00000000000..262b8252928 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/args.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Input views and planning metadata passed to a row kernel. + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; + +/// The arguments handed to one kernel invocation. +/// +/// `arrays` may be filtered or sliced, while `dtypes` and `output_dtype` always describe the +/// original planned batch. Keeping them together prevents an execution path from accidentally +/// pairing an input view with unrelated planning metadata. +#[derive(Clone, Copy)] +pub struct KernelArgs<'a> { + /// The executor-facing view, including the row count for this invocation. + pub execution: &'a dyn ExecutionArgs, + + /// The same inputs as concrete arrays for encoding-aware rewrites. + pub arrays: &'a [ArrayRef], + + /// The original input dtypes used to select the row implementation. + pub dtypes: &'a [DType], + + /// The non-nullable dtype built by the selected output capability. + pub output_dtype: &'a DType, +} + +/// An [`ExecutionArgs`] view over borrowed arrays with an explicit row count. +pub(super) struct BorrowedExecutionArgs<'a> { + /// The arrays exposed through this execution view. + inputs: &'a [ArrayRef], + + /// The row count reported for this execution view. + row_count: usize, +} + +impl<'a> BorrowedExecutionArgs<'a> { + pub(super) fn new(inputs: &'a [ArrayRef], row_count: usize) -> Self { + Self { inputs, row_count } + } +} + +impl ExecutionArgs for BorrowedExecutionArgs<'_> { + fn get(&self, index: usize) -> VortexResult { + self.inputs.get(index).cloned().ok_or_else(|| { + vortex_err!( + "Input index {} out of bounds (num_inputs={})", + index, + self.inputs.len() + ) + }) + } + + fn num_inputs(&self) -> usize { + self.inputs.len() + } + + fn row_count(&self) -> usize { + self.row_count + } +} diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs new file mode 100644 index 00000000000..30670f968a1 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -0,0 +1,440 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Null propagation, constant folding, and strategy execution for one columnar batch. + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::args::BorrowedExecutionArgs; +use super::args::KernelArgs; +use super::policy::BatchPlan; +use super::policy::RowPolicy; +use super::policy::skipping_beats_filtering; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::row::execute::RowExecution; +use crate::scalar_fn::row::types::batch_constant; +use crate::validity::Validity; + +/// How far [`Batch::resolve_validity`] got before a mixed-mask strategy became necessary. +enum ResolvedMask { + /// The all-valid or all-null batch was answered without a mixed-mask strategy. + Decided(ArrayRef), + + /// A mask with both set and unset bits, which a strategy must now execute. + Mixed(Mask), +} + +/// One batch of inputs and the metadata needed before its row kernel runs. +pub struct Batch<'a> { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The arguments as the execution layer handed them over. Every path but the filter strategy + /// gives the kernel these untouched, so it sees the original encodings. + args: &'a dyn ExecutionArgs, + + /// The input columns, collected once: constant folding inspects them and the filter strategy + /// filters them. + inputs: SmallVec<[ArrayRef; 4]>, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: SmallVec<[DType; 4]>, + + /// The conjoined input validity, so a row of the output is valid iff it is valid in every + /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. + validity: Validity, + + /// The dtype the function declares for these inputs, which the kernel's output is reconciled + /// against. Already widened to nullable if any input is nullable. + result_dtype: DType, + + /// The non-nullable dtype the dispatched output capability builds, computed while planning. + output_dtype: DType, + + /// How the concrete dispatch executes nullable rows. + policy: RowPolicy, +} + +impl<'a> Batch<'a> { + /// Collect the inputs and derive their dtype, validity, and execution policy. + /// + /// **Not** for a nullary function: with no inputs there is no validity to propagate and no + /// per-row work to fold, and the all-constant check below would vacuously pass. + pub fn new( + id: ScalarFnId, + args: &'a dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>()?; + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let nullability = plan.output_dtype.nullability() + | Nullability::from(arg_dtypes.iter().any(DType::is_nullable)); + let result_dtype = plan.output_dtype.with_nullability(nullability); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + args, + inputs, + arg_dtypes, + validity, + result_dtype, + output_dtype: plan.output_dtype, + policy: plan.policy, + }) + } + + /// Add null propagation, constant folding, and strategy selection around `kernel`. + /// + /// The kernel may ignore input validity. It receives valid-only rows when required, and its + /// output **must** match the planned dtype up to nullability. `try_unfiltered` receives the + /// original inputs plus a mixed validity mask. `Ok(None)` selects filter-and-scatter. + pub fn execute( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: any null-constant input forces an all-null result without evaluating the + // kernel. + if self + .inputs + .iter() + .any(|input| input.as_constant().is_some_and(|scalar| scalar.is_null())) + { + return Ok(self.all_null()); + } + + // All inputs constant, and their conjoined validity proves every row non-null. This sees + // through extension and masked wrappers just like argument decoding does. + if self.args.row_count() > 0 + && self.validity.definitely_no_nulls() + && self + .inputs + .iter() + .all(|input| batch_constant(input).is_some()) + { + return self.broadcast_one_row(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, false, ctx), + RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), + RowPolicy::ValidOnly { + filtered_decode_cost, + } => self.execute_valid_only(kernel, try_unfiltered, filtered_decode_cost, ctx), + } + } + + /// Evaluate a single row of all-constant inputs and broadcast its value. + /// + /// Reconciling the row's dtype before reading the scalar keeps this path on the same + /// kernel/declaration agreement check as the dense and filter paths, rather than letting `cast` + /// paper over a disagreement. + fn broadcast_one_row( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one_row: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.slice(0..1)) + .collect::>()?; + + let args = BorrowedExecutionArgs::new(&one_row, 1); + let result = VortexResult::from(kernel(self.kernel_args(&args, &one_row), ctx)?)?; + let scalar = self.finalize_output(result, 1)?.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.args.row_count()).into_array()) + } + + /// Run the kernel over every row, including the rows behind nulls, then mask its result. + /// + /// The arguments reach the kernel untouched, so the inputs keep their original encoding, and + /// the conjoined validity is handed to `mask` as an array rather than materialized into a + /// [`Mask`] first. + fn execute_dense( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + retry_deferred_error: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Every row is null, so the kernel has nothing to contribute. + if matches!(self.validity, Validity::AllInvalid) { + return Ok(self.all_null()); + } + + let values = match kernel(self.kernel_args(self.args, &self.inputs), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self + .validity + .clone() + .execute_mask(self.args.row_count(), ctx)?; + + // Unlike `resolve_validity`, all-true preserves the deferred error and all-false + // suppresses evidence that came entirely from null rows. An empty loop cannot + // produce deferred evidence, so the ambiguous empty mask cannot reach this arm. + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + // Deferred retry receives only the dense kernel. Filter first so the retry cannot + // evaluate null rows again. + return self.filter_and_scatter(kernel, &valid, ctx); + } + RowExecution::DeferredError(error) => return Err(error), + }; + + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.args.row_count()) + } + Validity::Array(valid) => { + self.finalize_output(values.mask(valid)?, self.args.row_count()) + } + // Handled by the guard above, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } + + /// Materialize validity and answer all-valid or all-null batches before selecting a mixed-mask + /// strategy. + fn resolve_validity( + &self, + kernel: &impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self + .validity + .clone() + .execute_mask(self.args.row_count(), ctx)?; + + // Check all-true before all-false: an empty mask is both, and must not be treated as + // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). + if valid.all_true() { + return self + .finalize_output( + VortexResult::from(kernel(self.kernel_args(self.args, &self.inputs), ctx)?)?, + self.args.row_count(), + ) + .map(ResolvedMask::Decided); + } + + if valid.all_false() { + return Ok(ResolvedMask::Decided(self.all_null())); + } + + Ok(ResolvedMask::Mixed(valid)) + } + + /// Resolve validity, try unfiltered execution when worthwhile, then fall back to filtering. + fn execute_valid_only( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + filtered_decode_cost: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedMask::Decided(result) => return Ok(result), + ResolvedMask::Mixed(valid) => valid, + }; + + if skipping_beats_filtering(filtered_decode_cost, &valid) + && let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? + { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Try execution against the original inputs, then mask a returned full-length result. + fn try_execute_unfiltered( + &self, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(execution) = + try_unfiltered(self.kernel_args(self.args, &self.inputs), valid, ctx)? + else { + return Ok(None); + }; + let values = VortexResult::from(execution)?; + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + self.finalize_output(values.mask(mask)?, valid.len()) + .map(Some) + } + + /// The filter strategy for a mixed mask: filter every input down to the rows set in `valid`, + /// run the kernel over those, and scatter its results back into a null-padded output. + fn filter_and_scatter( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(valid.clone())) + .collect::>()?; + + let args = BorrowedExecutionArgs::new(&filtered, valid.true_count()); + let values = VortexResult::from(kernel(self.kernel_args(&args, &filtered), ctx)?)?; + + self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) + } + + /// An all-null result of the function's declared return dtype. + fn all_null(&self) -> ArrayRef { + ConstantArray::new( + Scalar::null(self.result_dtype.clone()), + self.args.row_count(), + ) + .into_array() + } + + /// Pair an input view with this batch's planning metadata. + fn kernel_args<'b>( + &'b self, + execution: &'b dyn ExecutionArgs, + arrays: &'b [ArrayRef], + ) -> KernelArgs<'b> { + KernelArgs { + execution, + arrays, + dtypes: &self.arg_dtypes, + output_dtype: &self.output_dtype, + } + } + + /// Finalize an output against this batch's expected length and declared return dtype. + fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { + finalize_kernel_output(self.id, &self.result_dtype, expected_len, values) + } + + /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set + /// bits, producing an array of length `valid.len()` that is null at every unset position. + fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult { + vortex_ensure_eq!( + values.len(), + valid.true_count(), + "the {} kernel produced {} rows for {} filtered rows", + self.id, + values.len(), + valid.true_count(), + ); + + let AllOr::Some(slices) = valid.slices() else { + // The caller handles the all-true and all-false masks. + vortex_bail!("scatter_valid requires a mixed mask"); + }; + + // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index + // 0, and any in-bounds index would do since they are masked out below (values is non-empty + // here). + let mut indices = vec![0u64; valid.len()]; + let mut rank = 0u64; + for &(start, end) in slices { + for index in &mut indices[start..end] { + *index = rank; + rank += 1; + } + } + let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array(); + + let scattered = values.take(indices)?; + + // A kernel that produced nulls of its own (only `reduce_encoded` may) cannot be wrapped, + // since a `Masked` child must be all valid. Those nulls have to be unioned with the + // batch validity, which is what the general masking pass does. + if scattered.dtype().is_nullable() { + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + return scattered.mask(mask); + } + + // The gathered values are all valid, so attaching validity is sufficient. + Ok(MaskedArray::try_new( + scattered, + Validity::from_mask(valid.clone(), Nullability::Nullable), + )? + .into_array()) + } +} + +/// Validate a kernel output, then cast it to the row function's declared nullability. +/// +/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring +/// nullability. The kernel may omit nullability because batch execution owns strict null +/// propagation, so a nullability-only difference is cast to `result_dtype`. +pub fn finalize_kernel_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, +) -> VortexResult { + vortex_ensure_eq!( + values.len(), + expected_len, + "the {id} kernel produced {} rows for {expected_len} input rows", + values.len(), + ); + vortex_ensure!( + values.dtype().eq_ignore_nullability(result_dtype), + "the {id} kernel produced {} but the function declares {result_dtype}", + values.dtype(), + ); + + if values.dtype() == result_dtype { + Ok(values) + } else { + values.cast(result_dtype.clone()) + } +} diff --git a/vortex-array/src/scalar_fn/row/batch/mod.rs b/vortex-array/src/scalar_fn/row/batch/mod.rs new file mode 100644 index 00000000000..1b492f1ff6f --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/mod.rs @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batch execution around a non-null row kernel. +//! +//! A row kernel handles typed values for one row. This module adds the columnar concerns around it: +//! planning the output and null strategy, preserving batch constants and encodings, propagating +//! strict validity, selecting an execution strategy, and validating the finished output. +//! +//! [`policy`] derives the nullable execution strategy from a concrete dispatch. [`execution`] +//! applies that strategy, and [`args`] pairs each kernel invocation with its planning metadata. + +mod args; +pub(super) use args::KernelArgs; + +mod execution; +pub(super) use execution::Batch; +pub(super) use execution::finalize_kernel_output; + +mod policy; +pub(super) use policy::BatchPlan; +pub(super) use policy::RowPolicy; diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs new file mode 100644 index 00000000000..1b6f3f1f6cc --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Nullable execution strategies derived from a concrete row dispatch. + +use vortex_mask::Mask; + +use crate::dtype::DType; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::SinkResult; + +/// The execution policy and output dtype selected by a planning visit. +pub struct BatchPlan { + /// The non-nullable dtype built by the selected output capability. + pub output_dtype: DType, + + /// How this concrete dispatch executes nullable rows. + pub policy: RowPolicy, +} + +/// The nullable execution policy derived from one concrete dispatch. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RowPolicy { + /// Evaluate all rows and mask the result. + Dense, + + /// Evaluate all rows, retrying only valid rows if a deferred error is raised. + DenseWithRetry, + + /// Execute only valid rows, choosing between skip-invalid execution and filtering based on the + /// mask and decode cost. + ValidOnly { + /// Relative per-row decode work that filtering would avoid. + filtered_decode_cost: usize, + }, +} + +impl RowPolicy { + /// The policy for an infallible owned output. + pub const fn for_owned_output() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + Self::Dense + } else { + Self::ValidOnly { + filtered_decode_cost: Args::FILTERED_DECODE_COST, + } + } + } + + /// The policy for an owned output carrying batch-deferred failure evidence. + pub const fn for_deferred_output() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + Self::DenseWithRetry + } else { + Self::ValidOnly { + filtered_decode_cost: Args::FILTERED_DECODE_COST, + } + } + } + + /// The policy one concrete dispatch executes nullable rows under. + /// + /// This deliberately ignores [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. Batch execution tries + /// [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) against the original arrays + /// before it tries the sink or filters the inputs. Skipping that probe can change the result of + /// an encoding-aware function. + /// + /// [`OutputSink::SUPPORTS_SKIPPED_ROWS`]: crate::scalar_fn::OutputSink::SUPPORTS_SKIPPED_ROWS + pub const fn for_sink() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { + if ApplyResult::DEFERRED { + Self::DenseWithRetry + } else { + Self::Dense + } + } else { + Self::ValidOnly { + filtered_decode_cost: Args::FILTERED_DECODE_COST, + } + } + } +} + +/// Minimum surviving-row fractions for skipping when filtering avoids per-row decode work. +/// The thresholds distinguish one costly decode from multiple costly decodes. +const ONE_DECODE_SKIP_MIN_SURVIVING_FRACTION: f64 = 0.50; +const MULTI_DECODE_SKIP_MIN_SURVIVING_FRACTION: f64 = 0.85; + +/// Whether skipping invalid rows should be preferred over filtering for a mixed mask. +pub(super) fn skipping_beats_filtering(filtered_decode_cost: usize, valid: &Mask) -> bool { + if filtered_decode_cost == 0 { + return true; + } + + let minimum = if filtered_decode_cost == 1 { + ONE_DECODE_SKIP_MIN_SURVIVING_FRACTION + } else { + MULTI_DECODE_SKIP_MIN_SURVIVING_FRACTION + }; + + valid.true_count() as f64 >= valid.len() as f64 * minimum +} diff --git a/vortex-array/src/scalar_fn/row/execute/mod.rs b/vortex-array/src/scalar_fn/row/execute/mod.rs new file mode 100644 index 00000000000..cdfad84ee7a --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/mod.rs @@ -0,0 +1,52 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Row-loop execution for owned outputs and output sinks. +//! +//! [`owned`] stores one independent value per row and can reduce failure evidence. [`sink`] +//! drives output builders whose row handles may refer to shared batch state. + +use vortex_error::VortexError; +use vortex_error::VortexResult; + +use crate::ArrayRef; + +mod owned; +pub(super) use owned::execute_owned; +pub(super) use owned::execute_owned_infallible; + +mod sink; +pub(super) use sink::execute_sink; +pub(super) use sink::execute_sink_valid_rows; + +/// The outcome of a row loop before batch execution decides whether an error is observable. +/// +/// Together with the surrounding [`VortexResult`], this represents three outcomes: +/// +/// - `Err(error)` is a non-retryable execution or immediate row error. +/// - [`Output`](Self::Output) is a successful row loop. +/// - [`DeferredError`](Self::DeferredError) is failure evidence from a completed row loop. +/// +/// A dense loop may evaluate values behind nulls. Its deferred error is therefore not necessarily +/// observable: batch execution can retry over only valid rows, suppressing an error that came from +/// a null row while preserving one from a valid row. A plain `VortexResult` would lose +/// the distinction between that retryable error and an error for which retrying cannot help. +/// +/// Once execution is known to contain only valid rows, converting this outcome into a +/// `VortexResult` turns [`DeferredError`](Self::DeferredError) into an ordinary error. +pub enum RowExecution { + /// The successfully built, full-length output column. + Output(ArrayRef), + + /// An error constructed from failure evidence reduced across a completed row loop. + DeferredError(VortexError), +} + +impl From for VortexResult { + fn from(execution: RowExecution) -> Self { + match execution { + RowExecution::Output(output) => Ok(output), + RowExecution::DeferredError(error) => Err(error), + } + } +} diff --git a/vortex-array/src/scalar_fn/row/execute/owned.rs b/vortex-array/src/scalar_fn/row/execute/owned.rs new file mode 100644 index 00000000000..f2246875143 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/owned.rs @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution that stores one owned output value per row. + +use std::ops::BitOrAssign; + +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::row::visitor::assert_owned_output_needs_no_drop; + +/// Zero-sized evidence used to erase failure reduction from infallible owned visits. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + +/// Decode every input column once, then store one infallible owned output per row. +pub fn execute_owned_infallible( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, +{ + execute_owned::( + args, + ctx, + prepare, + move |prepared, args| (apply(prepared, args), NoFailure), + |_| Ok(()), + ) +} + +/// Decode every input column once, then store owned row outputs and reduce deferred failures. +pub fn execute_owned( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, +{ + const { assert_owned_output_needs_no_drop::() }; + + // Keep the vector length at zero until every row succeeds. An unwind then abandons partially + // initialized spare capacity without treating it as initialized output. The no-drop assertion + // above proves that no initialized value requires its destructor to run. + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let failure; + + { + let output = &mut values.spare_capacity_mut()[..row_count]; + + // When every input varies, the indexed source removes argument-shape dispatch from the hot + // loop and lets the lane kernel optimize the traversal as one operation. + if let Some(varying) = Args::varying(&columns) { + vortex_ensure!( + Args::varying_len_matches(&varying, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + failure = Args::indexed_source(&varying) + .map_checked_into(output, |elements| apply(&prepared, elements)); + } else { + // A batch-constant input was collapsed to one row during decoding. This path reads that + // row repeatedly while indexing only the inputs that vary. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut accumulated = Fail::default(); + for index in 0..row_count { + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + output[index].write(value); + accumulated |= row_failure; + } + failure = accumulated; + } + } + + // SAFETY: normal completion of either loop initializes every slot in `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + // Failure evidence is reduced inside the loop so its richer error construction stays cold. + // Preserve that provenance so batch execution may retry over only valid rows. + match finish_failure(failure) { + Ok(()) => Ok(RowExecution::Output(Out::build(values))), + Err(error) => Ok(RowExecution::DeferredError(error)), + } +} diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs new file mode 100644 index 00000000000..abb58f95b02 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution that writes through an output sink. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::DeferredError; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::SinkResult; + +/// Decode every input column once, allocate the sink once, then write one row at a time. +/// +/// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state +/// does not need to be captured by the closure. +pub fn execute_sink( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, +{ + let row_count = args.row_count(); + let mut sink = Sink::with_capacity(row_count, sink_dtype)?; + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let mut accumulated = ApplyResult::Accumulated::default(); + + { + // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This + // scope releases the borrow before `finish_sink` consumes the sink. + let mut rows = sink.rows(); + vortex_ensure!( + Sink::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + // The all-varying representation removes argument-shape dispatch from the hot loop. The + // mixed path instead reads collapsed batch constants at row zero. + if let Some(varying) = Args::varying(&columns) { + vortex_ensure!( + Args::varying_len_matches(&varying, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + for index in 0..row_count { + apply( + &prepared, + Args::get_varying(&varying, index), + Sink::row(&mut rows, index), + ) + .accumulate(&mut accumulated)?; + } + } else { + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + for index in 0..row_count { + apply( + &prepared, + Args::get(&columns, index), + Sink::row(&mut rows, index), + ) + .accumulate(&mut accumulated)?; + } + } + } + + // Immediate failures returned above. Only reduced, deferred failure evidence reaches finish. + finish_sink(sink, DeferredError::new(ApplyResult::occurred(accumulated))) +} + +/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +pub fn execute_sink_valid_rows( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + valid: &Mask, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, +) -> VortexResult> +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, +{ + // Batch execution needs a full-length result before applying the validity mask. Decline when + // the sink cannot leave legal placeholders in positions this loop skips. + if !Sink::SUPPORTS_SKIPPED_ROWS { + return Ok(None); + } + + // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An + // element representation may decline when it cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + let prepared = prepare(Args::constants(&columns)); + let row_count = args.row_count(); + let mut sink = Sink::with_capacity(row_count, sink_dtype)?; + let mut accumulated = ApplyResult::Accumulated::default(); + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid) = valid.bit_buffer() else { + vortex_bail!("execute_sink_valid_rows requires a mixed mask"); + }; + + { + let mut rows = sink.rows(); + vortex_ensure!( + Sink::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + let varying = Args::varying(&columns); + let lens_match = match &varying { + Some(varying) => Args::varying_len_matches(varying, row_count), + None => Args::decoded_lens_match(&columns, row_count), + }; + vortex_ensure!( + lens_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + // The loop writes only valid indices, but the sink still finishes a full-length output. + // Initialize placeholders now; batch execution masks them before the result escapes. + Sink::initialize_skipped_rows(&mut rows); + + // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first + // immediate error, turn later callbacks into no-ops, and return before finishing the sink. + let mut error = None; + valid.for_each_set_index(|index| { + if error.is_some() { + return; + } + + let result = match &varying { + Some(varying) => apply( + &prepared, + Args::get_varying(varying, index), + Sink::row(&mut rows, index), + ), + None => apply( + &prepared, + Args::get(&columns, index), + Sink::row(&mut rows, index), + ), + }; + if let Err(err) = result.accumulate(&mut accumulated) { + error = Some(err); + } + }); + + if let Some(error) = error { + return Err(error); + } + } + + // Immediate failures returned above. Only reduced, deferred failure evidence reaches finish. + finish_sink(sink, DeferredError::new(ApplyResult::occurred(accumulated))).map(Some) +} + +/// Classify a sink error as retryable only when row accumulation recorded a deferred failure. +/// +/// The sink contract requires [`OutputSink::finish`] to surface recorded failure evidence. Without +/// that evidence, its error is structural and retrying over a different set of rows cannot help. +fn finish_sink( + sink: S, + deferred_error: DeferredError, +) -> VortexResult { + match sink.finish(deferred_error) { + Ok(output) => Ok(RowExecution::Output(output)), + Err(error) if deferred_error.occurred() => Ok(RowExecution::DeferredError(error)), + Err(error) => Err(error), + } +} diff --git a/vortex-array/src/scalar_fn/row/mod.rs b/vortex-array/src/scalar_fn/row/mod.rs new file mode 100644 index 00000000000..3351c24d100 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions computed one row at a time. +//! +//! Start with [`RowFn`] to define an operation and [`RowVisitor`] to select one of its output +//! capabilities. [`InputElement`] and [`ElementTuple`] describe how columns become row values. +//! [`OutputElement`] covers independent owned values, while [`OutputSink`] supports outputs that +//! need row handles or shared batch state. [`SinkResult`] and [`DeferredError`] describe how a +//! sink-writing closure reports errors. +//! +//! The internal executor owns decoding, batch constants, null propagation, allocation, and +//! validity. A visitor's prepare closure may derive shared state from constant operands once per +//! batch. + +mod execute; + +mod batch; + +mod row_fn; +pub use row_fn::RowFn; + +mod types; +pub use types::DeferredError; +pub use types::ElementTuple; +pub use types::IndexedElementTuple; +pub use types::InputElement; +pub use types::OutputElement; +pub use types::OutputSink; +pub use types::SinkResult; +pub use types::UninitElementSink; + +mod visitor; +pub use visitor::RowVisitor; + +mod vtable; diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs new file mode 100644 index 00000000000..68d92c192ad --- /dev/null +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions computed one row at a time. + +use std::fmt::Debug; +use std::fmt::Display; +use std::hash::Hash; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use super::visitor::RowVisitor; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ScalarFnId; + +/// A scalar function computed one row at a time. +/// +/// Declare the argument names and use [`dispatch`](Self::dispatch) to choose concrete element and +/// sink types for each accepted dtype combination. Implement +/// [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) directly for columnar kernels. +pub trait RowFn: 'static + Sized + Clone + Send + Sync { + /// Options for this function, if any. Use [`EmptyOptions`](crate::scalar_fn::EmptyOptions) + /// for none. + type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash; + + /// The arguments in display order. Its length is the function's exact arity. + const ARG_NAMES: &'static [&'static str]; + + /// Whether any legal dispatch can raise a semantic error as defined by + /// [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible). + /// + /// The framework checks this at compile time for every fallible dispatched element or result. + /// A conservative `true` is allowed when only some dtype choices are fallible. + const FALLIBLE: bool = false; + + /// Returns the ID of the scalar function. + fn id(&self) -> ScalarFnId; + + /// Serialize this function's options, or return `None` when the function is not serializable. + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + _ = options; + Ok(None) + } + + /// Restore options written by [`serialize`](Self::serialize). + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + vortex_bail!("Expression {} is not deserializable", self.id()) + } + + /// Choose element types for these input dtypes and visit the framework with them. + /// + /// Plan time and run time both call this method, so the choice **must** be a pure function of + /// `options` and `args`. Cross-argument dtype validation belongs here. + fn dispatch( + &self, + options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult; + + /// Try an encoding-aware implementation before decoding the inputs into row elements. + /// + /// `None` continues to the dispatched row loop. `Some(output)` skips that loop. The output can + /// remain encoded or lazy. Filter-and-scatter execution can pass compacted inputs. + /// + /// # Requirements + /// + /// - `output.len()` **must** equal `args[0].len()`. + /// - The output dtype **must** match the planned dtype when ignoring nullability. + /// - The output **must not** introduce a null where every input is valid. + /// + /// The framework skips this hook for nullary functions. + fn reduce_encoded( + &self, + options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + _ = (options, args, ctx); + Ok(None) + } +} diff --git a/vortex-array/src/scalar_fn/row/types/element/bool.rs b/vortex-array/src/scalar_fn/row/types/element/bool.rs new file mode 100644 index 00000000000..e5c29f756b5 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/bool.rs @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::InputElement; +use crate::scalar_fn::OutputElement; +use crate::validity::Validity; + +impl InputElement for bool { + type Column = BitBuffer; + type Varying<'a> = &'a BitBuffer; + type Elem<'a> = bool; + + // Every bit of the buffer is readable, valid or not. + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + matches!(dtype, DType::Bool(_)), + "expected a Bool column, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_bit_buffer()) + } + + fn get(column: &Self::Column, index: usize) -> bool { + column.value(index) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> bool + where + Self: 'a, + { + column.value(index) + } +} + +impl OutputElement for bool { + fn element_dtype() -> DType { + DType::Bool(Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + // `From>` uses the bulk bit-packing path. + BoolArray::new(BitBuffer::from(values), Validity::NonNullable).into_array() + } +} diff --git a/vortex-array/src/scalar_fn/row/types/element/mod.rs b/vortex-array/src/scalar_fn/row/types/element/mod.rs new file mode 100644 index 00000000000..3aa92df9121 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -0,0 +1,129 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The element types a row function can read and produce. +//! +//! [`InputElement::Elem`] may borrow from its decoded column. [`OutputElement`] is returned by an +//! owned row computation; runtime-shaped output uses an +//! [`OutputSink`](crate::scalar_fn::OutputSink). + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; + +mod bool; + +mod primitive; + +mod tuple; +pub use tuple::ElementTuple; +pub use tuple::IndexedElementTuple; +pub use tuple::batch_constant; + +/// An element type that can be read row-wise out of an input column. +pub trait InputElement: 'static { + /// The decoded column representation supporting `O(1)` row access. + type Column; + + /// The view of a varying decoded column read by the hot row loop. + /// + /// This may borrow a cheaper representation than [`Column`](Self::Column). Primitive elements, + /// for example, expose a slice so its pointer and length are loop invariants rather than + /// re-reading a [`Buffer`](vortex_buffer::Buffer) descriptor for every row. + type Varying<'a>; + + /// The borrowed element value handed to a row closure. + type Elem<'a>; + + /// Whether every dense decode and access path tolerates rows that are null in the input. + /// + /// Arrays only guarantee payloads for valid rows. This is `false` when a null row's stored + /// offset or pointer may not address anything, and `true` only when [`decode`](Self::decode), + /// [`get`](Self::get), [`varying`](Self::varying), [`varying_len`](Self::varying_len), and + /// [`get_varying`](Self::get_varying) remain safe and correct for null rows. + /// + /// Dense execution requires this of every argument; otherwise the row layer executes only + /// valid rows. + const DENSE_SAFE: bool = false; + + /// Whether [`decode`](Self::decode) can fail on _legal_ input data. + /// + /// This excludes infrastructural failures such as IO or allocation. Set it when legal input may + /// contain a value that the decoder rejects. + const DECODE_FALLIBLE: bool = true; + + /// A relative unit count for per-row decode work avoided by filtering this argument first. + /// + /// Leave this at zero for bulk canonicalization. Use a positive value when filtering first + /// avoids meaningful per-row decode work. The executor adds this cost across arguments when it + /// chooses between skipping invalid rows and filtering. + const FILTERED_DECODE_COST: usize = 0; + + /// Validate that `dtype` is an acceptable input column dtype for this element type. + fn validate(dtype: &DType) -> VortexResult<()>; + + /// Decode `array` into its column representation. Called once per batch. + /// + /// Hoist dtype checks, downcasts, and other batch-invariant work into this method. + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Decode `array` _without_ assuming every row is valid, or `Ok(None)` when this element + /// cannot for this particular array. + /// + /// An element with [`DENSE_SAFE`](Self::DENSE_SAFE) set **should not** override this: its + /// ordinary decode already tolerates null payloads, so the default is already correct and an + /// override just restates it. Overriding is for an element that is _not_ dense-safe but can + /// still write an arbitrary placeholder into null slots; the caller guarantees + /// [`get`](Self::get) is never called for such a row. The skip-invalid strategy uses this + /// representation to avoid filtering the input. + /// + /// Return `Ok(None)` rather than an error when an array has no null-tolerant decode; the + /// batch execution falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if Self::DENSE_SAFE { + Self::decode(array, ctx).map(Some) + } else { + Ok(None) + } + } + + /// Read the element at `index`, the one function called once per row. + /// + /// This must not repeat work that is constant across the batch; do that work in + /// [`decode`](Self::decode). + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>; + + /// Borrow the representation used when this argument varies within the batch. + /// + /// Called once before the hot loop. Constants do not use this view because the tuple adapter + /// keeps their one-row decoded representation separate. + fn varying(column: &Self::Column) -> Self::Varying<'_>; + + /// Number of rows addressable through a [`Varying`](Self::Varying) view. + fn varying_len(column: &Self::Varying<'_>) -> usize; + + /// Read one row from a [`Varying`](Self::Varying) view. + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a; +} + +/// An owned row value that can be built into an all-valid column. +pub trait OutputElement: 'static + Sized { + /// The dtype of columns built from this element type. **Must** be non-nullable: nullability is + /// derived from the inputs by batch execution. + /// + /// Taking no arguments confines an element's dtype to a property of its Rust type, so an output + /// whose dtype depends on runtime data (a tensor, whose dtype carries its shape) cannot be an + /// element. Such an output uses an [`OutputSink`](crate::scalar_fn::OutputSink), whose + /// [`sink_dtype`](crate::scalar_fn::OutputSink::sink_dtype) does see the input dtypes. + fn element_dtype() -> DType; + + /// Build a column from one value per row. Called once per batch. + fn build(values: Vec) -> ArrayRef; +} diff --git a/vortex-array/src/scalar_fn/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/row/types/element/primitive.rs new file mode 100644 index 00000000000..05fddbd25e4 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/primitive.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::InputElement; +use crate::scalar_fn::OutputElement; +use crate::validity::Validity; + +impl InputElement for T { + type Column = Buffer; + type Varying<'a> = &'a [T]; + type Elem<'a> = T; + + // Every lane of the buffer holds a `T`, valid or not. + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + let expected = T::PTYPE; + let DType::Primitive(ptype, _) = dtype else { + vortex_bail!("expected a {expected} column, got {dtype}"); + }; + vortex_ensure_eq!( + *ptype, + expected, + "expected a {expected} column, got {dtype}" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_buffer::()) + } + + fn get(column: &Self::Column, index: usize) -> T { + column[index] + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column.as_slice() + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> T + where + Self: 'a, + { + column[index] + } +} + +impl OutputElement for T { + fn element_dtype() -> DType { + DType::Primitive(T::PTYPE, Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + PrimitiveArray::new(values, Validity::NonNullable).into_array() + } +} diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple.rs b/vortex-array/src/scalar_fn/row/types/element/tuple.rs new file mode 100644 index 00000000000..becde073a0b --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/tuple.rs @@ -0,0 +1,414 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Argument lists built from [`InputElement`]s, and the per-argument decode behind them. + +use vortex_compute::lane_kernels::IndexedSource; +use vortex_compute::lane_kernels::LaneZip; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Extension; +use crate::arrays::Masked; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::InputElement; + +mod private { + pub trait Sealed {} +} + +/// One decoded input, collapsed to a single row when it is constant for the batch. +pub struct ArgColumn( + /// The decoded column, classified by whether it varies within the batch. + ArgColumnKind, +); + +enum ArgColumnKind { + Varying(T::Column), + Constant(T::Column), +} + +impl ArgColumn { + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // An empty input has no row 0 to slice, and its row loop runs zero times either way. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?))); + } + + Ok(Self(ArgColumnKind::Varying(T::decode(array, ctx)?))) + } + + fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + // Batch execution short-circuits null constants before selecting a strategy, so a + // constant reaching this path is non-null and can use the ordinary decode. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Some(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?)))); + } + + Ok(T::decode_null_tolerant(array, ctx)? + .map(ArgColumnKind::Varying) + .map(Self)) + } + + fn get(&self, index: usize) -> T::Elem<'_> { + match &self.0 { + ArgColumnKind::Varying(column) => T::get(column, index), + ArgColumnKind::Constant(column) => T::get(column, 0), + } + } + + fn varying(&self) -> Option<&T::Column> { + match &self.0 { + ArgColumnKind::Varying(column) => Some(column), + ArgColumnKind::Constant(_) => None, + } + } + + fn addresses_rows(&self, row_count: usize) -> bool { + // A constant is always read at index zero, so it addresses any batch length. + match &self.0 { + ArgColumnKind::Varying(column) => T::varying_len(&T::varying(column)) == row_count, + ArgColumnKind::Constant(_) => true, + } + } + + fn constant(&self) -> Option> { + match &self.0 { + ArgColumnKind::Varying(_) => None, + ArgColumnKind::Constant(column) => Some(T::get(column, 0)), + } + } +} + +/// Return the batch-constant array, looking through masked and extension wrappers. +/// +/// Batch execution owns mask validity, so a masked constant may expose its constant child here. An +/// extension over constant storage remains wrapped to preserve its extension dtype. +pub fn batch_constant(array: &ArrayRef) -> Option { + if array.as_constant().is_some() { + return Some(array.clone()); + } + + if let Some(masked) = array.as_opt::() { + return Some(masked.child().clone()).filter(|child| child.as_constant().is_some()); + } + + array + .as_opt::() + .is_some_and(|ext| ext.storage_array().as_constant().is_some()) + .then(|| array.clone()) +} + +/// Typed argument tuples for arities zero through twelve. +/// +/// This trait is sealed; add a new row representation by implementing [`InputElement`] and placing +/// it in one of the supplied tuples. +pub trait ElementTuple: 'static + private::Sealed { + /// The decoded column representations. + type Columns; + + /// Direct references to decoded columns when every argument varies within the batch. + type VaryingColumns<'a>; + + /// The borrowed row of element values. + type Elems<'a>; + + /// The batch-constant element values: [`Elems`](Self::Elems) with every argument wrapped in + /// `Option`. + /// + /// `Some` marks an argument whose operand is constant for the batch and carries the element + /// every row reads; `None` marks one that varies by row. This is what + /// [`RowVisitor`](crate::scalar_fn::RowVisitor) hands to a visit's prepare closure, so a kernel + /// can hoist work that depends only on a constant argument out of the row loop. + type ConstElems<'a>; + + /// The number of arguments. + const ARITY: usize; + + /// Whether every argument is [`InputElement::DENSE_SAFE`]. + const DENSE_SAFE: bool; + + /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. + const DECODE_FALLIBLE: bool; + + /// The additive cost of per-row decode work avoided by filtering the arguments first. + const FILTERED_DECODE_COST: usize; + + /// Validate the input dtypes, including that `dtypes` has exactly `ARITY` entries. + /// + /// The expression layer checks the count against [`Arity`](crate::scalar_fn::Arity) before it + /// builds a call, but this is also the entry point of the public + /// [`return_dtype`](crate::scalar_fn::ScalarFnVTable::return_dtype), so the count is enforced + /// here rather than assumed. + fn validate(dtypes: &[DType]) -> VortexResult<()>; + + /// Decode every input column once. Called once per batch. + fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Decode every input column once while tolerating null rows. + /// + /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid + /// strategy calls this once per batch. + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; + + /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; + + /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// + /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple + /// gives the optimizer ordinary contiguous column access without a per-row constant check. + fn varying(columns: &Self::Columns) -> Option>; + + /// Whether every varying column contains exactly `row_count` rows. + fn varying_len_matches(columns: &Self::VaryingColumns<'_>, row_count: usize) -> bool; + + /// Whether every argument that varies within the batch contains exactly `row_count` rows. + /// + /// The same guarantee as [`varying_len_matches`](Self::varying_len_matches), for the mixed case + /// [`varying`](Self::varying) declines: a batch-constant argument is exempt because it was + /// collapsed to one row, while every argument beside it still has to address the whole batch. + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; + + /// Read one row from columns already known to vary within the batch. + fn get_varying<'a>(columns: &Self::VaryingColumns<'a>, index: usize) -> Self::Elems<'a>; + + /// Read the batch-constant elements out of the decoded columns. Called once per batch. + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; +} + +/// An argument tuple that supports a validated dense indexed traversal. +/// +/// This is separate from [`ElementTuple`] because many row elements have no contiguous source, and +/// stable Rust cannot provide a blanket fallback plus a more specific primitive implementation. +/// The trait is sealed so shared execution can rely on its unchecked-read contract. A tuple only +/// implements it when the source can be validated once and every lane can then be read +/// independently. +pub trait IndexedElementTuple: ElementTuple { + /// The source shared execution uses for a dense all-varying loop. + /// + /// Its length must be the common varying-column length. For every valid index it must preserve + /// row order, return the same value as [`ElementTuple::get_varying`], and uphold the unchecked + /// read contract of [`IndexedSource`]. + type Source<'a>: IndexedSource>; + + /// Borrow a source from columns already validated to vary within the batch. + fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a>; +} + +/// An indexed native slice yielding the one-tuples expected by a unary row closure. +#[derive(Clone, Copy)] +pub struct UnaryTupleSource<'a, T>( + /// The native values read by the row loop. + &'a [T], +); + +impl IndexedSource for UnaryTupleSource<'_, T> { + type Item = (T,); + + fn len(&self) -> usize { + self.0.len() + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the caller guarantees that `index` is in bounds. + (unsafe { *self.0.get_unchecked(index) },) + } +} + +impl private::Sealed for () {} + +impl ElementTuple for () { + type Columns = (); + type VaryingColumns<'a> = (); + type Elems<'a> = (); + type ConstElems<'a> = (); + + const ARITY: usize = 0; + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + const FILTERED_DECODE_COST: usize = 0; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + 0, + "expected 0 argument dtypes, got {}", + dtypes.len(), + ); + Ok(()) + } + + fn decode(_args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(()) + } + + fn decode_null_tolerant( + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(())) + } + + fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} + + fn varying(_columns: &Self::Columns) -> Option> { + Some(()) + } + + fn varying_len_matches(_columns: &Self::VaryingColumns<'_>, _row_count: usize) -> bool { + true + } + + fn decoded_lens_match(_columns: &Self::Columns, _row_count: usize) -> bool { + true + } + + fn get_varying<'a>(_columns: &Self::VaryingColumns<'a>, _index: usize) -> Self::Elems<'a> {} + + fn constants(_columns: &Self::Columns) -> Self::ConstElems<'_> {} +} + +macro_rules! element_tuple { + ($arity:literal; $($t:ident : $idx:tt),+) => { + impl<$($t: InputElement),+> private::Sealed for ($($t,)+) {} + + impl<$($t: InputElement),+> ElementTuple for ($($t,)+) { + type Columns = ($(ArgColumn<$t>,)+); + type VaryingColumns<'a> = ($($t::Varying<'a>,)+); + type Elems<'a> = ($($t::Elem<'a>,)+); + type ConstElems<'a> = ($(Option<$t::Elem<'a>>,)+); + + const ARITY: usize = $arity; + const DENSE_SAFE: bool = $($t::DENSE_SAFE &&)+ true; + const DECODE_FALLIBLE: bool = $($t::DECODE_FALLIBLE ||)+ false; + const FILTERED_DECODE_COST: usize = $($t::FILTERED_DECODE_COST +)+ 0; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + $arity, + "expected {} argument dtypes, got {}", + $arity, + dtypes.len(), + ); + + $($t::validate(&dtypes[$idx])?;)+ + Ok(()) + } + + fn decode( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(($(ArgColumn::<$t>::decode(args.get($idx)?, ctx)?,)+)) + } + + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(( + $(match ArgColumn::<$t>::decode_null_tolerant(args.get($idx)?, ctx)? { + Some(column) => column, + None => return Ok(None), + },)+ + ))) + } + + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_> { + ($(columns.$idx.get(index),)+) + } + + fn varying(columns: &Self::Columns) -> Option> { + Some(($($t::varying(columns.$idx.varying()?),)+)) + } + + fn varying_len_matches( + columns: &Self::VaryingColumns<'_>, + row_count: usize, + ) -> bool { + $($t::varying_len(&columns.$idx) == row_count &&)+ true + } + + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { + $(columns.$idx.addresses_rows(row_count) &&)+ true + } + + fn get_varying<'a>( + columns: &Self::VaryingColumns<'a>, + index: usize, + ) -> Self::Elems<'a> { + ($($t::get_varying(&columns.$idx, index),)+) + } + + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { + ($(columns.$idx.constant(),)+) + } + } + }; +} + +element_tuple!(1; A:0); +element_tuple!(2; A:0, B:1); +element_tuple!(3; A:0, B:1, C:2); +element_tuple!(4; A:0, B:1, C:2, D:3); +element_tuple!(5; A:0, B:1, C:2, D:3, E:4); +element_tuple!(6; A:0, B:1, C:2, D:3, E:4, F:5); +element_tuple!(7; A:0, B:1, C:2, D:3, E:4, F:5, G:6); +element_tuple!(8; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7); +element_tuple!(9; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8); +element_tuple!(10; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9); +element_tuple!(11; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10); +element_tuple!(12; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11); + +impl IndexedElementTuple for (T,) { + type Source<'a> = UnaryTupleSource<'a, T>; + + fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a> { + UnaryTupleSource(columns.0) + } +} + +impl IndexedElementTuple for (Left, Right) { + type Source<'a> = LaneZip<&'a [Left], &'a [Right]>; + + fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a> { + LaneZip::new(columns.0, columns.1) + } +} + +#[cfg(test)] +mod tests { + use vortex_compute::lane_kernels::IndexedSource; + + use super::UnaryTupleSource; + + #[test] + fn unary_tuple_source_reads_one_tuple_per_row() { + let source = UnaryTupleSource(&[10, 20, 30]); + assert_eq!(source.len(), 3); + + // SAFETY: index one is within the three-element source. + assert_eq!(unsafe { source.get_unchecked(1) }, (20,)); + } +} diff --git a/vortex-array/src/scalar_fn/row/types/mod.rs b/vortex-array/src/scalar_fn/row/types/mod.rs new file mode 100644 index 00000000000..2998e19ccbd --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/mod.rs @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Input decoding and output construction for row functions. +//! +//! [`element`] defines the Rust values decoded from input columns and built into simple output +//! columns. [`sink`] handles outputs that need row handles or batch-wide state. [`result`] defines +//! the immediate and deferred outcomes returned by sink-writing row closures. + +mod element; +pub use element::ElementTuple; +pub use element::IndexedElementTuple; +pub use element::InputElement; +pub use element::OutputElement; +pub(super) use element::batch_constant; + +mod result; +pub use result::DeferredError; +pub use result::SinkResult; + +mod sink; +pub use sink::OutputSink; +pub use sink::UninitElementSink; diff --git a/vortex-array/src/scalar_fn/row/types/result.rs b/vortex-array/src/scalar_fn/row/types/result.rs new file mode 100644 index 00000000000..efd841cc969 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/result.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What a sink-writing row closure may return. + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +mod private { + pub trait Sealed {} +} + +/// A value-dependent failure bit reduced across the row loop and handed to the output sink. +/// +/// Unlike [`VortexResult`], this never exits the loop. Use it when every row can write a safe +/// provisional value and report failure once at the end. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct DeferredError( + /// Whether this value records a deferred error. + bool, +); + +impl DeferredError { + /// Record whether this row encountered an error. + pub const fn new(failed: bool) -> Self { + Self(failed) + } + + /// Whether any row accumulated into this value failed. + pub const fn occurred(self) -> bool { + self.0 + } +} + +impl BitOrAssign for DeferredError { + fn bitor_assign(&mut self, rhs: Self) { + self.0 |= rhs.0; + } +} + +/// The result of writing one row: success, an immediate error, or deferred error evidence. +/// +/// The executor OR-reduces [`Accumulated`](Self::Accumulated) in a loop-local. The accumulated word +/// should be no wider than the computed element so error tracking does not constrain vector width. +/// This trait is sealed; row functions choose one of its supplied implementations. +pub trait SinkResult: 'static + private::Sealed { + /// The word this result reduces into, kept in a loop-local by the executor. + type Accumulated: 'static + Copy + Default; + + /// Whether this return type can carry an error. + const FALLIBLE: bool; + + /// Whether this result defers failure reporting until the sink finishes. + const DEFERRED: bool; + + /// Merge this row's outcome into the batch-wide reduction. + fn accumulate(self, accumulated: &mut Self::Accumulated) -> VortexResult<()>; + + /// Whether the finished reduction means some row failed. + fn occurred(accumulated: Self::Accumulated) -> bool; +} + +impl private::Sealed for () {} + +impl SinkResult for () { + type Accumulated = (); + + const FALLIBLE: bool = false; + const DEFERRED: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + Ok(()) + } + + fn occurred(_accumulated: ()) -> bool { + false + } +} + +impl private::Sealed for VortexResult<()> {} + +impl SinkResult for VortexResult<()> { + type Accumulated = (); + + const FALLIBLE: bool = true; + const DEFERRED: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + self + } + + fn occurred(_accumulated: ()) -> bool { + false + } +} + +/// The evidence widths a row closure may reduce. `bool` is the ordinary answer; the unsigned +/// integers exist for a kernel whose per-row comparison would cost it its vectorization. +macro_rules! impl_sink_result_word { + ($($word:ty),+ $(,)?) => { + $( + impl private::Sealed for $word {} + + impl SinkResult for $word { + type Accumulated = $word; + + const FALLIBLE: bool = false; + const DEFERRED: bool = true; + + fn accumulate(self, accumulated: &mut $word) -> VortexResult<()> { + *accumulated |= self; + Ok(()) + } + + fn occurred(accumulated: $word) -> bool { + accumulated != <$word>::default() + } + } + )+ + }; +} + +impl_sink_result_word!(bool, u8, u16, u32, u64); diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs new file mode 100644 index 00000000000..712d209f7e9 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The column builders a row function can write its output into. + +use std::mem::MaybeUninit; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::DeferredError; +use crate::scalar_fn::OutputElement; + +/// A column allocated once per batch that a row closure writes into, one row at a time. +/// +/// A sink may use the input dtypes to build a runtime-shaped output or own shared batch state. The +/// executor passes each row slot into an [`Fn`] closure, keeping mutable state out of its capture. +/// +/// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; +/// skip-invalid execution can omit invalid rows when +/// [`SUPPORTS_SKIPPED_ROWS`](Self::SUPPORTS_SKIPPED_ROWS) is `true`. +pub trait OutputSink: 'static + Sized { + /// Whether this sink accepts [`DeferredError`] from its row closure instead of requiring a + /// per-row [`VortexResult`]. + /// + /// A supporting sink must return an error from [`finish`](Self::finish) when its deferred error + /// argument occurred. + const ERRORS_ARE_DEFERRED: bool = false; + + /// Whether this sink can finish a full-length output when some rows were never visited. + /// + /// A supporting sink must use [`initialize_skipped_rows`](Self::initialize_skipped_rows) to + /// leave a legal arbitrary value at every skipped row. Batch execution masks those values. + const SUPPORTS_SKIPPED_ROWS: bool = false; + + /// A loop-local view of all output rows. + /// + /// Borrowed once before execution so the sink's buffer descriptor and shape become loop + /// invariants rather than being re-read through `&mut Self` for every row. + type Rows<'a> + where + Self: 'a; + + /// The place a row closure writes one row through, borrowed from the sink. + type Row<'a> + where + Self: 'a; + + /// The dtype of the column this sink builds, given the function's input dtypes. + /// + /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the + /// result, and masks the null rows. + fn sink_dtype(args: &[DType]) -> VortexResult; + + /// Allocate a sink for `rows` rows of `dtype`, which is this sink's own + /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; + + /// Borrow all output rows for the hot loop. + fn rows(&mut self) -> Self::Rows<'_>; + + /// Whether every index in `0..row_count` is addressable through [`row`](Self::row). + /// + /// Called once before the hot loop. Besides validating the sink contract, this gives the + /// optimizer the output bounds it needs to remove the bounds check hidden in each row accessor. + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool; + + /// Initialize output positions that skip-invalid execution can omit. + /// + /// Called only when [`SUPPORTS_SKIPPED_ROWS`](Self::SUPPORTS_SKIPPED_ROWS) is `true`. The + /// default is for sinks whose allocation already contains legal values. + fn initialize_skipped_rows(_rows: &mut Self::Rows<'_>) {} + + /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; + + /// Finish into the built column, whose dtype **must** be this sink's + /// [`sink_dtype`](Self::sink_dtype). Called once per batch with whether any deferred row error + /// occurred. + fn finish(self, error: DeferredError) -> VortexResult; +} + +/// An element sink that leaves dense output uninitialized before the row loop. +/// +/// Skip-invalid execution initializes placeholders before omitting rows. Immediate failures are +/// safe because [`OutputSink::finish`] is not called after one. +pub struct UninitElementSink { + /// Spare storage written in increasing row order. + values: Vec, + + /// The number of slots exposed to the row loop and initialized before finishing. + row_count: usize, +} + +impl OutputSink for UninitElementSink { + const SUPPORTS_SKIPPED_ROWS: bool = true; + + type Rows<'a> = &'a mut [MaybeUninit]; + type Row<'a> = &'a mut MaybeUninit; + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(T::element_dtype()) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self { + values: Vec::with_capacity(rows), + row_count: rows, + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.values.spare_capacity_mut()[..self.row_count] + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn initialize_skipped_rows(rows: &mut Self::Rows<'_>) { + for row in rows.iter_mut() { + row.write(T::default()); + } + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + fn finish(mut self, _error: DeferredError) -> VortexResult { + // SAFETY: dense execution writes every row, while skip-invalid execution initializes every + // row before overwriting valid ones. The executor calls `finish` only after successful + // execution, and the allocation reserved every slot in `0..row_count`. + unsafe { self.values.set_len(self.row_count) }; + + Ok(T::build(self.values)) + } +} diff --git a/vortex-array/src/scalar_fn/row/visitor/check.rs b/vortex-array/src/scalar_fn/row/visitor/check.rs new file mode 100644 index 00000000000..a000ce38369 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/check.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Contract checks shared by planning and execution visits. +//! +//! Const assertions reject invalid generic visits during compilation. The validators compare a +//! selected visit with the input dtypes during planning and return its output dtype. + +use std::mem::needs_drop; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::dtype::DType; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::SinkResult; + +/// Assert the no-drop contract that makes partially initialized output safe to abandon on unwind. +pub(in crate::scalar_fn::row) const fn assert_owned_output_needs_no_drop() { + assert!( + !needs_drop::(), + "owned row outputs must not require drop glue" + ); +} + +/// Assert that the input arity and decode fallibility match the function-wide declarations. +const fn assert_input_visit_contract() { + assert!( + Args::ARITY == F::ARG_NAMES.len(), + "the visited argument tuple must have the arity declared by RowFn::ARG_NAMES", + ); + // Dictionary pushdown treats an infallible function as safe to evaluate over values no code + // references, so every dispatch must fit the function-wide declaration. + assert!( + !Args::DECODE_FALLIBLE || F::FALLIBLE, + "RowFn::FALLIBLE must be true when input decoding can fail", + ); +} + +/// Assert the input contract and that owned output values do not require drop glue. +pub(super) const fn assert_owned_visit_contract() +where + Function: RowFn, + Args: IndexedElementTuple, + Out: OutputElement, +{ + assert_input_visit_contract::(); + assert_owned_output_needs_no_drop::(); +} + +/// Assert that a sink visit obeys the input, fallibility, and deferred-error contracts. +pub(super) const fn assert_sink_visit_contract() +where + Function: RowFn, + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, +{ + assert_input_visit_contract::(); + assert!( + !ApplyResult::FALLIBLE || Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result can fail", + ); + assert!( + !ApplyResult::DEFERRED || Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result defers failure evidence", + ); + assert!( + Sink::ERRORS_ARE_DEFERRED == ApplyResult::DEFERRED, + "OutputSink::ERRORS_ARE_DEFERRED must match SinkResult::DEFERRED", + ); +} + +/// Assert the owned-output contract, fallibility declaration, and failure-evidence width bound. +pub(super) const fn assert_deferred_visit_contract() +where + Function: RowFn, + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, +{ + assert_owned_visit_contract::(); + assert!( + Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result defers failure evidence", + ); + assert!( + size_of::() <= size_of::(), + "failure evidence must be no wider than the value, or it bounds the vector width", + ); +} + +/// Validate the input dtypes and return the non-nullable dtype built by `Out`. +pub(super) fn validate_owned_visit( + dtypes: &[DType], +) -> VortexResult { + Args::validate(dtypes)?; + + let dtype = Out::element_dtype(); + vortex_ensure!( + !dtype.is_nullable(), + "row output elements must declare a non-nullable dtype, got {dtype}", + ); + + Ok(dtype) +} + +/// Validate the input dtypes and return the non-nullable dtype built by `Sink`. +pub(super) fn validate_sink_visit( + dtypes: &[DType], +) -> VortexResult { + Args::validate(dtypes)?; + + let dtype = Sink::sink_dtype(dtypes)?; + vortex_ensure!( + !dtype.is_nullable(), + "row output sinks must declare a non-nullable dtype, got {dtype}", + ); + + Ok(dtype) +} diff --git a/vortex-array/src/scalar_fn/row/visitor/execute.rs b/vortex-array/src/scalar_fn/row/visitor/execute.rs new file mode 100644 index 00000000000..36fd03af5d2 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/execute.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visitors that execute dense and skip-invalid row loops. +//! +//! Each method verifies that execution selected the same visit shape as planning before handing +//! its typed closures to the matching loop. Valid-row execution can decline without running a loop; +//! batch execution then filters the inputs and retries the dense loop. + +use std::marker::PhantomData; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::private; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::SinkResult; +use crate::scalar_fn::row::execute::RowExecution; +use crate::scalar_fn::row::execute::execute_owned; +use crate::scalar_fn::row::execute::execute_owned_infallible; +use crate::scalar_fn::row::execute::execute_sink; +use crate::scalar_fn::row::execute::execute_sink_valid_rows; + +/// The run-time visit that decodes every column once and runs the selected row loop. +pub struct ExecuteRows<'args, 'ctx, F> { + /// The inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'args, 'ctx, F> ExecuteRows<'args, 'ctx, F> { + pub fn new( + args: &'args dyn ExecutionArgs, + output_dtype: &'args DType, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + output_dtype, + ctx, + function: PhantomData, + } + } +} + +impl private::Sealed for ExecuteRows<'_, '_, F> {} + +impl RowVisitor for ExecuteRows<'_, '_, F> { + type VisitResult = RowExecution; + + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + + execute_owned_infallible::(self.args, self.ctx, prepare, apply) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + + execute_sink::( + self.args, + self.output_dtype, + self.ctx, + prepare, + apply, + ) + } + + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + + execute_owned::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} + +/// The run-time visit that tries skip-invalid execution over the original input columns. +/// +/// Only output sinks have a contract for skipped output positions. Owned visits therefore decline +/// so batch execution can use its filter-and-scatter fallback. +pub struct ExecuteValidRows<'args, 'ctx, F> { + /// The original inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The conjoined validity, materialized by batch execution and guaranteed mixed. + valid: &'args Mask, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'args, 'ctx, F> ExecuteValidRows<'args, 'ctx, F> { + pub fn new( + args: &'args dyn ExecutionArgs, + output_dtype: &'args DType, + valid: &'args Mask, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + output_dtype, + valid, + ctx, + function: PhantomData, + } + } +} + +impl private::Sealed for ExecuteValidRows<'_, '_, F> {} + +impl RowVisitor for ExecuteValidRows<'_, '_, F> { + type VisitResult = Option; + + fn visit_prepared( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + + // Owned execution has no sink that can initialize skipped output positions. Decline so + // batch execution filters the inputs and retries with the dense visitor. + Ok(None) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + + execute_sink_valid_rows::( + self.args, + self.output_dtype, + self.valid, + self.ctx, + prepare, + apply, + ) + } + + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + _finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + + // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. + Ok(None) + } +} diff --git a/vortex-array/src/scalar_fn/row/visitor/mod.rs b/vortex-array/src/scalar_fn/row/visitor/mod.rs new file mode 100644 index 00000000000..bf172103b5d --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/mod.rs @@ -0,0 +1,157 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visits that plan or execute the concrete row signature selected by [`RowFn::dispatch`]. +//! +//! [`RowFn::dispatch`]: crate::scalar_fn::RowFn::dispatch + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::SinkResult; + +mod check; +pub(super) use check::assert_owned_output_needs_no_drop; + +mod execute; +pub(super) use execute::ExecuteRows; +pub(super) use execute::ExecuteValidRows; + +mod plan; +pub(super) use plan::PlanRows; + +/// A planning or execution visit at concrete input and output types. +/// +/// Only the framework implements this trait. The `visit_prepared*` methods derive shared state +/// from constant arguments before visiting any rows. +pub trait RowVisitor: private::Sealed + Sized { + /// The framework result of visiting one concrete row signature. + /// + /// This is a batch plan or execution result, not the per-row `Out` returned by [`visit`] and + /// [`visit_deferred`](Self::visit_deferred). + type VisitResult; + + /// Visit an infallible row computation that returns one independent output value. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of + /// [`RowFn::ARG_NAMES`](crate::scalar_fn::RowFn::ARG_NAMES). + /// - [`RowFn::FALLIBLE`](crate::scalar_fn::RowFn::FALLIBLE) **must** be `true` when decoding + /// `Args` can fail. + /// - `Out` **must not** require drop glue. + fn visit( + self, + apply: impl Fn(Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + self.visit_prepared::(|_| (), move |&(), args| apply(args)) + } + + /// The prepared form of [`visit`](Self::visit), with the same prerequisites. + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement; + + /// Visit a row computation that writes through a sink-provided row handle. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of + /// [`RowFn::ARG_NAMES`](crate::scalar_fn::RowFn::ARG_NAMES). + /// - [`RowFn::FALLIBLE`](crate::scalar_fn::RowFn::FALLIBLE) **must** be `true` when decoding + /// `Args` or computing the result can fail. + /// - [`OutputSink::ERRORS_ARE_DEFERRED`] **must** match [`SinkResult::DEFERRED`] for the + /// selected `Sink` and `ApplyResult`. + fn visit_into( + self, + apply: impl Fn(Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + self.visit_prepared_into::( + |_| (), + move |&(), args, row| apply(args, row), + ) + } + + /// The prepared form of [`visit_into`](Self::visit_into), with the same prerequisites. + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult; + + /// Visit a row computation that returns an owned output and deferred failure evidence. + /// + /// `Fail` is OR-reduced across rows and handed to `finish_failure`. The value from + /// [`Default::default`] **must** mean success, including for an empty batch. The compiler + /// cannot check this semantic requirement. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of + /// [`RowFn::ARG_NAMES`](crate::scalar_fn::RowFn::ARG_NAMES). + /// - [`RowFn::FALLIBLE`](crate::scalar_fn::RowFn::FALLIBLE) **must** be `true`. + /// - `Out` **must not** require drop glue. + /// - `Out` **must** be at least as wide as `Fail` so failure tracking does not reduce the + /// vector width. + fn visit_deferred( + self, + apply: impl Fn(Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + self.visit_prepared_deferred::( + |_| (), + move |&(), args| apply(args), + finish_failure, + ) + } + + /// The prepared form of [`visit_deferred`](Self::visit_deferred), with the same prerequisites. + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign; +} + +mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/row/visitor/plan.rs b/vortex-array/src/scalar_fn/row/visitor/plan.rs new file mode 100644 index 00000000000..caa17ad651b --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/plan.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The visitor that validates a concrete dispatch and plans its nullable execution. + +use std::marker::PhantomData; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; +use super::private; +use crate::dtype::DType; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::SinkResult; +use crate::scalar_fn::row::batch::BatchPlan; +use crate::scalar_fn::row::batch::RowPolicy; + +/// The plan-time visit that validates dtypes and derives the nullable execution policy. +pub struct PlanRows<'a, F> { + /// The input dtypes for this plan. + dtypes: &'a [DType], + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'a, F> PlanRows<'a, F> { + pub fn new(dtypes: &'a [DType]) -> Self { + Self { + dtypes, + function: PhantomData, + } + } +} + +impl private::Sealed for PlanRows<'_, F> {} + +impl RowVisitor for PlanRows<'_, F> { + type VisitResult = BatchPlan; + + fn visit_prepared( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + + Ok(BatchPlan { + output_dtype: validate_owned_visit::(self.dtypes)?, + policy: RowPolicy::for_owned_output::(), + }) + } + + fn visit_prepared_into( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + + Ok(BatchPlan { + output_dtype: validate_sink_visit::(self.dtypes)?, + policy: RowPolicy::for_sink::(), + }) + } + + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + _finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + + Ok(BatchPlan { + output_dtype: validate_owned_visit::(self.dtypes)?, + policy: RowPolicy::for_deferred_output::(), + }) + } +} diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs new file mode 100644 index 00000000000..98c2cdc9671 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`ScalarFnVTable`] adapter shared by every [`RowFn`]. +//! +//! The [`visitor`](super::visitor) module validates and executes the concrete row signature +//! selected by dispatch. This module connects those visits to batch execution and exposes the +//! resulting scalar function behavior to the rest of the compute stack. + +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use super::row_fn::RowFn; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::expr::Expression; +use crate::expr::union_child_validities; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::row::batch::Batch; +use crate::scalar_fn::row::batch::KernelArgs; +use crate::scalar_fn::row::batch::finalize_kernel_output; +use crate::scalar_fn::row::execute::RowExecution; +use crate::scalar_fn::row::visitor::ExecuteRows; +use crate::scalar_fn::row::visitor::ExecuteValidRows; +use crate::scalar_fn::row::visitor::PlanRows; + +/// Implement [`ScalarFnVTable`] for every [`RowFn`]. +impl ScalarFnVTable for F { + type Options = F::Options; + + fn id(&self) -> ScalarFnId { + RowFn::id(self) + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + RowFn::serialize(self, options) + } + + fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { + RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(F::ARG_NAMES.len()) + } + + fn child_name(&self, _options: &Self::Options, child_index: usize) -> ChildName { + ChildName::from(F::ARG_NAMES[child_index]) + } + + fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { + let plan = self.dispatch(options, args, PlanRows::::new(args))?; + + // Union the output nullability with the nullability of the inputs. This is required for + // strict scalar function semantics. + let nullability = plan.output_dtype.nullability() + | Nullability::from(args.iter().any(DType::is_nullable)); + + Ok(plan.output_dtype.with_nullability(nullability)) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Nullary functions have no input validity to propagate, so they skip batch execution. + if args.num_inputs() == 0 { + let result_dtype = ScalarFnVTable::return_dtype(self, options, &[])?; + let nullary_args = KernelArgs { + execution: args, + arrays: &[], + dtypes: &[], + output_dtype: &result_dtype, + }; + + let execution = execute_rows(self, options, nullary_args, ctx)?; + let values = VortexResult::from(execution)?; + + return finalize_kernel_output( + RowFn::id(self), + &result_dtype, + args.row_count(), + values, + ); + } + + let batch = prepare_batch(self, options, args)?; + batch.execute( + |args, ctx| execute_rows(self, options, args, ctx), + |args, valid, ctx| try_execute_rows_unfiltered(self, options, args, valid, ctx), + ctx, + ) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + F::FALLIBLE + } +} + +/// Run the encoding-aware rewrite when available, or execute the selected row loop. +fn execute_rows( + function: &F, + options: &F::Options, + args: KernelArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !args.arrays.is_empty() + && let Some(reduced) = function.reduce_encoded(options, args.arrays, ctx)? + { + return Ok(RowExecution::Output(reduced)); + } + + function.dispatch( + options, + args.dtypes, + ExecuteRows::::new(args.execution, args.output_dtype, ctx), + ) +} + +/// Try execution against the original inputs, returning `None` when batch execution must filter. +fn try_execute_rows_unfiltered( + function: &F, + options: &F::Options, + args: KernelArgs<'_>, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + // Try the encoding-aware path before filtering changes the inputs. The caller masks its + // full-length result with `valid` before returning it. + if let Some(reduced) = function.reduce_encoded(options, args.arrays, ctx)? { + return Ok(Some(RowExecution::Output(reduced))); + } + + function.dispatch( + options, + args.dtypes, + ExecuteValidRows::::new(args.execution, args.output_dtype, valid, ctx), + ) +} + +/// Prepare the batch inputs and execution plan for `function`. +fn prepare_batch<'args, F: RowFn>( + function: &F, + options: &F::Options, + args: &'args dyn ExecutionArgs, +) -> VortexResult> { + Batch::new(RowFn::id(function), args, |arg_dtypes| { + function.dispatch(options, arg_dtypes, PlanRows::::new(arg_dtypes)) + }) +} From 89fd28bc137a39acfa2bdc939497b21baa6e9002 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 10:42:44 -0400 Subject: [PATCH 014/160] Execute primitive numeric operators with RowFn Signed-off-by: Connor Tsui --- vortex-array/benches/binary_ops.rs | 8 + .../typed_view/primitive/numeric_operator.rs | 2 +- .../scalar_fn/fns/binary/numeric/checked.rs | 88 +---- .../src/scalar_fn/fns/binary/numeric/mod.rs | 12 +- .../scalar_fn/fns/binary/numeric/primitive.rs | 355 ++++-------------- .../src/scalar_fn/fns/binary/numeric/row.rs | 131 +++++++ .../src/scalar_fn/fns/binary/numeric/tests.rs | 9 +- 7 files changed, 236 insertions(+), 369 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/numeric/row.rs diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 6a07d03f50b..3bd466da0b1 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -170,6 +170,14 @@ fn div_i64_nonnull(bencher: Bencher) { bench_primitive(bencher, lhs, rhs, Operator::Div); } +#[divan::bench] +fn div_i64_nullable(bencher: Bencher) { + let lhs = primitive_nullable(1_000_000, 7).into_array(); + let rhs = primitive_nullable(17, 5).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Div); +} + #[divan::bench] fn sub_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index 6b389a0554b..d7bcc8d0b4c 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_err; use crate::scalar_fn::fns::operators::Operator; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. pub enum NumericOperator { /// Binary element-wise addition of two arrays or of two scalars. diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index afb709abe1c..054846b7ef7 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Checked-lane execution for numeric kernels, driven by the shared +//! Checked-lane execution for the decimal kernels, driven by the shared //! `vortex-compute` lane kernels. - -use std::ops::BitOrAssign; +//! +//! The primitive widths do not come through here: they are computed one row at a time by +//! [`row`](super::row), which writes a value for every row and reduces failure evidence without +//! scanning the finished output. use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -13,33 +15,18 @@ use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_mask::AllOr; use vortex_mask::Mask; -/// Evidence that a lane failed, anything other than [`Default`] meaning failure. -/// -/// `bool` is the ordinary choice; [`map_checked_into`] explains why the wider members exist and -/// asserts the width bound that membership here does **not** imply. -/// -/// [`map_checked_into`]: IndexedSourceExt::map_checked_into -pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - -impl Failure for bool {} -impl Failure for u8 {} -impl Failure for u16 {} -impl Failure for u32 {} -impl Failure for u64 {} - /// Apply the fallible `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it returns `None` on a _valid_ lane. +/// `Err(first_failing_valid_lane)` only when it returns `None` on a valid lane. /// /// `apply` also runs on invalid lanes, whose failures are masked out and whose values are /// unspecified, so it must be total: no panics or side effects on any stored lane value. /// -/// This drives the one-pass early-exit kernels, which abort at the end of the enclosing 64-lane -/// chunk. Use it when per-lane failure handling is cheap relative to the operation, as in integer -/// division. Prefer [`checked_apply_lanes`] when the failure check itself vectorizes. +/// This drives the one-pass early-exit kernels: failures abort at the end of the enclosing +/// 64-lane chunk. It suits an operation whose per-lane failure handling is cheap relative to the +/// operation itself, which is what the decimal kernels and their per-lane casts are. /// -/// `#[inline]`: this must inline into the caller that builds the closure, so that a captured -/// constant operand flattens into a register rather than living behind a pointer the loop reloads -/// on every lane, which blocks vectorization. +/// Keep this wrapper inlineable so captured constants can become loop invariants in the caller. +/// The lane kernels retain their own inlining decisions. #[inline] pub(super) fn checked_lanes( source: S, @@ -61,7 +48,6 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, @@ -73,55 +59,3 @@ where Ok(values.freeze()) } - -/// Apply the split value/failure `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it flags a _valid_ lane. -/// -/// The hot pass writes every value unconditionally and OR-reduces one piece of evidence, leaving -/// the loop free of per-lane selects and per-chunk exit branches. Only if a lane flagged does a -/// cold second pass re-run `apply` through the early-exit kernels, which drop null-lane failures -/// and attribute the first valid one. -/// -/// Like [`checked_lanes`], `apply` runs on invalid lanes and must be total. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -pub(super) fn checked_apply_lanes( - source: S, - valid_rows: &Mask, - mut apply: Apply, -) -> Result, usize> -where - S: IndexedSource + Copy, - T: Copy + Default, - Fail: Failure, - Apply: FnMut(S::Item) -> (T, Fail), -{ - let len = source.len(); - debug_assert_eq!(len, valid_rows.len()); - - let valid_bits = match valid_rows.bit_buffer() { - AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), - AllOr::Some(valid_bits) => Some(valid_bits), - }; - - let mut values = BufferMut::::with_capacity(len); - - let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &mut apply) != Fail::default() { - let mut checked = |item: S::Item| { - let (value, failure) = apply(item); - (failure == Fail::default()).then_some(value) - }; - match valid_bits { - None => source.try_map_into(out, &mut checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?, - } - } - - // SAFETY: the kernels initialize every lane in `out`. - unsafe { values.set_len(len) }; - - Ok(values.freeze()) -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 6dc0de0fbea..c7ae86b93c9 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,16 +4,19 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! +//! The primitive widths are computed by a [`RowFn`](crate::scalar_fn::RowFn), which owns null +//! handling, constants, and validity for them; see [`row`]. Decimal keeps its own columnar +//! implementation in [`decimal`]. +//! //! [`Binary`]: super::Binary mod checked; mod decimal; mod primitive; -#[cfg(test)] -mod tests; +mod row; use decimal::execute_numeric_decimal; -use primitive::execute_numeric_primitive; +use row::execute_numeric_primitive; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -81,3 +84,6 @@ fn build_empty_result( Ok(Canonical::empty(&result_dtype).into_array()) } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 8fd53d15216..42fe3fd3e03 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -1,73 +1,48 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::checked::Failure; -use super::checked::checked_apply_lanes; -use super::checked::checked_lanes; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; +//! Checked arithmetic for one primitive row. + +use std::ops::BitOrAssign; + use crate::dtype::NativePType; -use crate::dtype::PType; use crate::dtype::half::f16; -use crate::match_each_native_ptype; -use crate::scalar::NumericOperator; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; -use crate::validity::Validity; -struct CheckedAdd; +/// Checked addition, failing on integer overflow. +pub(super) struct CheckedAdd; -struct CheckedSub; +/// Checked subtraction, failing on integer overflow. +pub(super) struct CheckedSub; -struct CheckedMul; +/// Checked multiplication, failing on integer overflow. +pub(super) struct CheckedMul; -struct CheckedDiv; +/// Checked division, failing on integer division by zero and on `MIN / -1`. +pub(super) struct CheckedDiv; -trait CheckedPrimitiveOp: Sized { - /// Message for the error raised when this operation fails on a valid lane. - const ERROR: &'static str; +/// OR-reducible evidence that a row failed, with [`Default`] meaning success. +pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - /// Whether to check inside the value loop and exit early, rather than reducing evidence over - /// the whole batch and re-running only once some lane has flagged. - const CHECKED_VALUE_LOOP: bool = false; +impl Failure for T {} - /// How this operation reports a failing lane. See [`Failure`]. - type Failure: Failure; - - /// Compute a lane's value and its failure evidence together. - /// - /// Overflowing-style rather than `Option`, so the vectorizable kernels can write every - /// lane's value unconditionally and reduce the evidence separately. [`Self::checked`] is the - /// shape for the scalar and early-exit paths. - fn apply(lhs: T, rhs: T) -> (T, Self::Failure); +/// One arithmetic operator at one width, split into its value and failure evidence. +pub(super) trait CheckedPrimitiveOp: 'static + Sized { + /// The error reported for a batch in which some valid row failed. + const ERROR: &'static str; - /// [`Self::apply`] folded into the `Option` that the early-exit kernels take. - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - let (value, failed) = Self::apply(lhs, rhs); + /// How this operation reports a failing row. See [`Failure`]. + type Fail: Failure; - (failed == Self::Failure::default()).then_some(value) - } + /// The result of this operation, paired with evidence of whether the row failed. + fn apply(lhs: T, rhs: T) -> (T, Self::Fail); } impl CheckedPrimitiveOp for CheckedAdd { const ERROR: &'static str = "integer overflow in checked add"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.add_value(rhs), lhs.add_error(rhs)) } @@ -76,9 +51,9 @@ impl CheckedPrimitiveOp for CheckedAdd { impl CheckedPrimitiveOp for CheckedSub { const ERROR: &'static str = "integer overflow in checked sub"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.sub_value(rhs), lhs.sub_error(rhs)) } @@ -87,9 +62,9 @@ impl CheckedPrimitiveOp for CheckedSub { impl CheckedPrimitiveOp for CheckedMul { const ERROR: &'static str = "integer overflow in checked mul"; - type Failure = T::MulFailure; + type Fail = T::MulFailure; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, T::MulFailure) { (lhs.mul_value(rhs), lhs.mul_failure(rhs)) } @@ -97,16 +72,10 @@ impl CheckedPrimitiveOp for CheckedMul { impl CheckedPrimitiveOp for CheckedDiv { const ERROR: &'static str = "integer division by zero or overflow in checked div"; - // Integer division still lowers to scalar divides, so the split - // value/error-scan loop used to auto-vectorize add/sub/mul only adds a - // second full scan. Use the one-pass early-exit checked kernel for integer - // division, matching Arrow/Velox. Float division has no checked errors and - // stays on the split/vectorizable default path. - const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { let failed = lhs.div_error(rhs); let value = if failed { @@ -116,151 +85,13 @@ impl CheckedPrimitiveOp for CheckedDiv { }; (value, failed) } - - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - lhs.div_checked(rhs) - } -} - -/// Execute a numeric operation between two primitive-typed arrays. -pub(super) fn execute_numeric_primitive( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - - match_each_native_ptype!(ptype, |T| { - match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), - } - }) -} - -fn execute_checked_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType, - Op: CheckedPrimitiveOp, - Scalar: From, - Scalar: From>, -{ - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - let len = lhs.len(); - debug_assert_eq!(len, rhs.len()); - - let validity = lhs.validity().and(rhs.validity())?; - let valid_rows = validity.execute_mask(len, ctx)?; - - let values = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_op_lanes::<_, T, Op>( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - &valid_rows, - |(lhs, rhs)| (lhs, rhs), - ), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Capture the constant by value so it stays hoisted out of the lane loop. - let rhs = *rhs; - checked_op_lanes::<_, T, Op>(lhs.as_slice(), &valid_rows, move |lhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let lhs = *lhs; - checked_op_lanes::<_, T, Op>(rhs.as_slice(), &valid_rows, move |rhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let value = Op::checked(*lhs, *rhs) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - return Ok(constant_result_array(value, len, &result_dtype)); - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => Ok(Buffer::zeroed(len)), - } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - - primitive_result_array::(values, validity, &result_dtype) } -/// Run `Op` over the lanes of `source` in the loop shape it declares: the one-pass early-exit -/// kernel when `Op::CHECKED_VALUE_LOOP` is set, and the split value/failure kernel otherwise. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -fn checked_op_lanes( - source: S, - valid_rows: &Mask, - mut to_operands: impl FnMut(S::Item) -> (T, T), -) -> Result, usize> -where - S: IndexedSource + Copy, - T: NativePType, - Op: CheckedPrimitiveOp, -{ - if Op::CHECKED_VALUE_LOOP { - checked_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::checked(lhs, rhs) - }) - } else { - checked_apply_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::apply(lhs, rhs) - }) - } -} - -fn primitive_result_array( - values: Buffer, - validity: Validity, - dtype: &DType, -) -> VortexResult { - let array = PrimitiveArray::new(values, validity).into_array(); - if array.dtype() == dtype { - return Ok(array); - } - array.cast(dtype.clone()) -} - -fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef -where - T: NativePType, - Scalar: From + From>, -{ - if dtype.is_nullable() { - ConstantArray::new(Some(value), len).into_array() - } else { - ConstantArray::new(value, len).into_array() - } -} - -trait CheckedArithmetic: NativePType { - const DIV_CHECKS_IN_VALUE_LOOP: bool; - - /// How multiplication reports a failing lane, which is width-dependent in a way the other - /// operations are not. `impl_checked_unsigned!` and `impl_checked_signed!` say why each family - /// reports what it reports. +/// Per-width checked arithmetic. Every value method **must** be total over stored lane values. +pub(super) trait CheckedArithmetic: NativePType { + /// How multiplication reports a failing row. + /// + /// This may be a word rather than `bool` when narrowing evidence would block vectorization. type MulFailure: Failure; fn add_value(self, rhs: Self) -> Self; @@ -271,16 +102,9 @@ trait CheckedArithmetic: NativePType { fn mul_failure(self, rhs: Self) -> Self::MulFailure; fn div_value(self, rhs: Self) -> Self; fn div_error(self, rhs: Self) -> bool; - fn div_checked(self, rhs: Self) -> Option; } -/// The integer arithmetic every width shares, given the two things that actually differ between -/// them: how multiplication reports a failing lane, and how add/sub/div detect one. -/// -/// Only `mul_failure` genuinely varies per width, so it is the macro's parameter and everything -/// else is written once. The `$mul_failure_ty` and body are supplied by the caller because the -/// choice between a widened high half and a `bool` is exactly the vectorization decision described -/// on [`Failure`]. +/// Generate the shared integer operations from their failure predicates. macro_rules! impl_checked_integer { ( $ty:ty, @@ -291,67 +115,57 @@ macro_rules! impl_checked_integer { = |$mf_lhs:ident, $mf_rhs:ident| $mul_failure:expr, ) => { impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = true; - type MulFailure = $mul_failure_ty; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self.wrapping_add(rhs) } - #[inline(always)] + #[inline] fn add_error(self, rhs: Self) -> bool { let ($add_lhs, $add_rhs) = (self, rhs); $add_error } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self.wrapping_sub(rhs) } - #[inline(always)] + #[inline] fn sub_error(self, rhs: Self) -> bool { let ($sub_lhs, $sub_rhs) = (self, rhs); $sub_error } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self.wrapping_mul(rhs) } - #[inline(always)] + #[inline] $(#[$mul_failure_attr])* fn mul_failure(self, rhs: Self) -> $mul_failure_ty { let ($mf_lhs, $mf_rhs) = (self, rhs); $mul_failure } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, rhs: Self) -> bool { let ($div_lhs, $div_rhs) = (self, rhs); $div_error } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - self.checked_div(rhs) - } } }; } -/// The unsigned widths. `add`, `sub` and `div` are written once here, and `widening_mul` covers -/// every width including the 64-bit one, which widens into `u128`: the discarded high half of the -/// widened product is the failure evidence, and costs none of the comparison LLVM folds into -/// `umul.with.overflow`. +/// Unsigned multiplication reports its discarded high half as failure evidence. macro_rules! impl_checked_unsigned { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_integer!( @@ -364,12 +178,7 @@ macro_rules! impl_checked_unsigned { }; } -/// The signed widths, on the same principle. `widening_mul` is the shorthand for the narrow widths, -/// whose two-sided range check over a wider product is not the shape LLVM folds into an overflow -/// intrinsic, so they vectorize while reporting a plain `bool`. The 64-bit width cannot: deriving a -/// `bool` there costs the comparison that scalarizes the loop, so `high_half_mul` hands back the -/// discarded high half as a word. Both arms derive every shift and bound from `$ty`, so -/// instantiating one at a new width cannot silently keep another width's constants. +/// Signed widths use a range check or discarded high-half evidence. macro_rules! impl_checked_signed { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_signed!($ty, mul_failure: bool = |lhs, rhs| { @@ -377,9 +186,6 @@ macro_rules! impl_checked_signed { product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide }); }; - // Zero exactly when the product fits: a signed multiply overflows iff the high half of the true - // product differs from the sign extension of the half that was kept, so the two XOR to zero on - // the lanes that fit. `tests::test_multiply_overflow_boundaries` pins the boundaries. ($ty:ty, high_half_mul: $wide:ty => $failure:ty) => { impl_checked_signed!($ty, mul_failure: #[expect( clippy::cast_possible_truncation, @@ -395,7 +201,7 @@ macro_rules! impl_checked_signed { ( $ty:ty, mul_failure: $(#[$mul_failure_attr:meta])* $mul_failure_ty:ty - = |$l:ident, $r:ident| $mul_failure:expr + = |$lhs:ident, $rhs:ident| $mul_failure:expr ) => { impl_checked_integer!( $ty, @@ -408,7 +214,7 @@ macro_rules! impl_checked_signed { ((lhs ^ rhs) & (lhs ^ value)) < 0 }, div_error: |lhs, rhs| rhs == 0 || (lhs == <$ty>::MIN && rhs == -1), - mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$l, $r| $mul_failure, + mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$lhs, $rhs| $mul_failure, ); }; } @@ -417,54 +223,47 @@ macro_rules! impl_checked_float { ($($ty:ty),+ $(,)?) => { $( impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = false; - type MulFailure = bool; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self + rhs } - #[inline(always)] + #[inline] fn add_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self - rhs } - #[inline(always)] + #[inline] fn sub_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self * rhs } - #[inline(always)] + #[inline] fn mul_failure(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, _rhs: Self) -> bool { false } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - Some(self / rhs) - } } )+ }; @@ -484,30 +283,25 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; - /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change - /// pivots, and a spread of magnitudes that straddles the 64-bit split. const PROBES: &[i64] = &[ - 0, // - 1, // - -1, // - 2, // - -2, // - 3, // - i64::MIN, // - i64::MIN + 1, // - i64::MAX, // - i64::MAX - 1, // - 1 << 31, // - 1 << 32, // - 1 << 62, // - -(1 << 62), // - 0x7FFF_FFFF, // - -0x8000_0000, // + 0, + 1, + -1, + 2, + -2, + 3, + i64::MIN, + i64::MIN + 1, + i64::MAX, + i64::MAX - 1, + 1 << 31, + 1 << 32, + 1 << 62, + -(1 << 62), + 0x7FFF_FFFF, + -0x8000_0000, ]; - /// Every `mul_failure` impl is either a bit trick or a two-sided range check, so hold each - /// against the obvious reference: `reference` is the width's own `checked_mul`, whose `None` - /// _is_ the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { let failed = lhs.mul_failure(rhs) != ::default(); @@ -522,14 +316,11 @@ mod tests { assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); let (lhs, rhs) = (lhs as u64, rhs as u64); - assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); } } } - /// The 8-bit widths are cheap enough to check exhaustively, which pins the shift of the - /// unsigned formula and the range check of the signed one against every product that exists. #[test] fn mul_failure_agrees_with_checked_mul_exhaustively_at_8_bits() { for lhs in u8::MIN..=u8::MAX { diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs new file mode 100644 index 00000000000..a8e78bdcea2 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive arithmetic execution through [`RowFn`]. +//! +//! `Binary` keeps its registered contract; [`NumericBinary`] is only an execution helper. Decimal +//! arithmetic remains on its existing columnar path. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::primitive::CheckedAdd; +use super::primitive::CheckedArithmetic; +use super::primitive::CheckedDiv; +use super::primitive::CheckedMul; +use super::primitive::CheckedPrimitiveOp; +use super::primitive::CheckedSub; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::NumericOperator; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::RowVisitor; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::row::UninitElementSink; + +pub(super) fn execute_numeric_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + ScalarFnVTable::execute(&NumericBinary, &op, &args, ctx) +} + +/// Internal row execution for the primitive arithmetic operators. +#[derive(Clone)] +struct NumericBinary; + +impl RowFn for NumericBinary { + type Options = NumericOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + // Fallibility is queried without input dtypes, so this conservatively covers integer widths. + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.numeric_binary"); + *ID + } + + fn dispatch( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = PType::try_from( + args.first() + .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?, + )?; + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_div::(visitor), + } + }) + } +} + +fn visit_checked(visitor: V) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + V: RowVisitor, +{ + visitor.visit_deferred::<(T, T), T, Op::Fail>( + |(lhs, rhs)| Op::apply(lhs, rhs), + |failure| { + if failure != ::default() { + return Err(numeric_error(Op::ERROR)); + } + + Ok(()) + }, + ) +} + +fn visit_div(visitor: V) -> VortexResult +where + T: CheckedArithmetic, + V: RowVisitor, +{ + if T::PTYPE.is_float() { + return visit_checked::(visitor); + } + + // Integer division is scalar and expensive, so deferring its cheap failure check preserves no + // vectorization. Check each divide immediately and stop at the first failure. + // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when + // they need to skip invalid rows. + visitor.visit_into::<(T, T), UninitElementSink, VortexResult<()>>(|(lhs, rhs), output| { + let (value, failed) = CheckedDiv::apply(lhs, rhs); + if failed { + return Err(numeric_error(>::ERROR)); + } + + output.write(value); + Ok(()) + }) +} + +/// Keep rich error construction out of row closures so the closures remain inlineable. +#[cold] +#[inline(never)] +fn numeric_error(message: &'static str) -> VortexError { + vortex_err!(InvalidArgument: "{message}") +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 7ee98ae717e..3813c8612b3 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -201,8 +201,7 @@ fn test_integer_array_array_errors_on_valid_lanes() { assert!(result.is_err()); } -/// Multiply two non-nullable lanes of `lhs` by two of `rhs`, expecting `Some(product)` where the -/// product fits and `None` where the checked kernel must report overflow. +/// Assert one checked multiplication through the complete array execution path. #[track_caller] fn assert_multiply(lhs: T, rhs: T, expected: Option) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -297,13 +296,11 @@ fn test_multiply_overflow_on_null_lane_ignored( Ok(()) } -/// The hot pass OR-reduces evidence across whole 64-lane chunks before anything looks at it, so an -/// overflow in a late chunk must still be caught, and must still be suppressed when its lane is -/// null. Every other test here fits in a single chunk and cannot show either. +/// An overflow late in the batch must still be reported, unless its row is null. #[rstest] #[case::reported(true)] #[case::suppressed_by_null(false)] -fn test_multiply_overflow_survives_chunk_reduction( +fn test_multiply_overflow_survives_batch_reduction( #[case] lane_is_valid: bool, ) -> VortexResult<()> { const LEN: u32 = 1000; From 59c4578ef4b44d9fa8d5b78c92e0e9b2a9d5f86e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 10:42:53 -0400 Subject: [PATCH 015/160] Add focused RowFn executor benchmarks Signed-off-by: Connor Tsui --- vortex-array/Cargo.toml | 8 + vortex-array/benches/row_fn_executor.rs | 280 ++++++++++++++++++++++++ vortex-array/benches/strict_validity.rs | 214 ++++++++++++++++++ 3 files changed, 502 insertions(+) create mode 100644 vortex-array/benches/row_fn_executor.rs create mode 100644 vortex-array/benches/strict_validity.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index d00b811a387..6bb251bd808 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -129,6 +129,10 @@ harness = false name = "binary_ops" harness = false +[[bench]] +name = "row_fn_executor" +harness = false + [[bench]] name = "kleene_bool" harness = false @@ -203,6 +207,10 @@ harness = false name = "validity_is_valid" harness = false +[[bench]] +name = "strict_validity" +harness = false + [[bench]] name = "dict_unreferenced_mask" harness = false diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs new file mode 100644 index 00000000000..e9aa9e8d522 --- /dev/null +++ b/vortex-array/benches/row_fn_executor.rs @@ -0,0 +1,280 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares owned-output, sink-writing, and hand-written primitive row loops. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::DeferredError; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::OutputSink; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const ROWS: usize = 65_536; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +#[derive(Clone)] +struct RowWrappingAdd; + +impl RowFn for RowWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_wrapping_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64), i64>(|(lhs, rhs)| lhs.wrapping_add(rhs)) + } +} + +#[derive(Clone)] +struct RowCheckedAdd; + +impl RowFn for RowCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_checked_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(i64, i64), i64, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(checked_add_error()); + } + Ok(()) + }, + ) + } +} + +/// Keep error construction out of the benchmarked success path. +#[cold] +#[inline(never)] +fn checked_add_error() -> VortexError { + vortex_err!("integer overflow in row checked add") +} + +/// A benchmark sink that writes one `i64` per row. +struct I64Sink( + /// The output values written by the row loop. + BufferMut, +); + +impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + fn finish(self, _error: DeferredError) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +#[derive(Clone)] +struct RowSinkWrappingAdd; + +impl RowFn for RowSinkWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_wrapping_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64, i64), I64Sink, _>(|(lhs, rhs), out| { + *out = lhs.wrapping_add(rhs); + }) + } +} + +fn inputs() -> (ArrayRef, ArrayRef) { + let lhs = (0..ROWS) + .map(|index| index as i64) + .collect::>() + .into_array(); + let rhs = (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>() + .into_array(); + (lhs, rhs) +} + +fn constant_inputs() -> (ArrayRef, ArrayRef) { + let (lhs, _) = inputs(); + let rhs = ConstantArray::new(Scalar::from(7i64), ROWS).into_array(); + (lhs, rhs) +} + +fn nullable_inputs() -> (ArrayRef, ArrayRef) { + let lhs = PrimitiveArray::new( + (0..ROWS).map(|index| index as i64).collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(5))), + ) + .into_array(); + let rhs = PrimitiveArray::new( + (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(7))), + ) + .into_array(); + (lhs, rhs) +} + +fn bench_row_fn(bencher: Bencher, function: F, make_inputs: fn() -> (ArrayRef, ArrayRef)) +where + F: RowFn, +{ + bencher + .with_inputs(make_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + function + .clone() + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, inputs); +} + +#[divan::bench] +fn handrolled_sink_wrapping_add(bencher: Bencher) { + bencher + .with_inputs(inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = lhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let rhs = rhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let mut output = BufferMut::zeroed(ROWS); + for ((out, lhs), rhs) in output + .as_mut_slice() + .iter_mut() + .zip(lhs.as_slice()) + .zip(rhs.as_slice()) + { + *out = lhs.wrapping_add(*rhs); + } + PrimitiveArray::new(output.freeze(), Validity::NonNullable).into_array() + }); +} + +#[divan::bench] +fn row_checked_add(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, inputs); +} + +#[divan::bench] +fn row_wrapping_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, nullable_inputs); +} + +#[divan::bench] +fn row_wrapping_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, nullable_inputs); +} diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs new file mode 100644 index 00000000000..bd2d7bbef10 --- /dev/null +++ b/vortex-array/benches/strict_validity.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks how the row lifting applies validity to a dense kernel. +//! +//! Both arms run the same kernel over the same nullable input and differ only in how the output's +//! nulls are restored: +//! +//! - `lazy` is what the lifting behind [`RowFn`] does: conjoin the input validities (itself lazy) +//! and hand the resulting boolean array straight to `mask`. +//! - `eager` is the strategy it replaced: materialize the conjunction into a `Mask`, copy it into a +//! `BitBuffer`, wrap that in a `BoolArray`, and mask with it. +//! +//! `chain` composes three of the same function, which is where a per-call materialization +//! compounds. + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const SIZES: &[usize] = &[65_536, 1 << 20]; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +/// The shared kernel: double every lane, ignoring validity. +fn doubled(input: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let input = input.clone().execute::(ctx)?; + let values: Buffer = input + .as_slice::() + .iter() + .map(|value| value.wrapping_mul(2)) + .collect(); + Ok(PrimitiveArray::new(values, Validity::NonNullable).into_array()) +} + +/// Dense row function: the lifting decides how validity is applied. +/// +/// Its [`reduce_encoded`](RowFn::reduce_encoded) answers with [`doubled`] before the row loop is +/// reached, so this arm runs exactly the kernel [`EagerDouble`] runs and the two differ only in +/// validity. The row closure below is what makes it a [`RowFn`] at all, and never executes. +#[derive(Clone)] +struct LazyDouble; + +impl RowFn for LazyDouble { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.lazy_double"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i32,), i32>(|(value,)| value.wrapping_mul(2)) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + doubled(&args[0], ctx).map(Some) + } +} + +/// The same function, applying validity the way the adapter used to: materialize a mask first. +#[derive(Clone)] +struct EagerDouble; + +impl ScalarFnVTable for EagerDouble { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.eager_double"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { + ChildName::from("input") + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + Ok(DType::Primitive(PType::I32, args[0].nullability())) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let valid = input.validity()?.execute_mask(args.row_count(), ctx)?; + let values = doubled(&input, ctx)?; + + if valid.all_true() { + return values.cast(DType::Primitive(PType::I32, input.dtype().nullability())); + } + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + values.mask(mask) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// A nullable i32 column with array-backed validity (~10% nulls). +fn nullable_input(len: usize) -> ArrayRef { + PrimitiveArray::new( + (0..len as i32).collect::>(), + Validity::from_iter((0..len).map(|index| !index.is_multiple_of(10))), + ) + .into_array() +} + +fn bench_depth(bencher: Bencher, function: F, len: usize, depth: usize) +where + F: ScalarFnVTable + Clone, +{ + bencher + .with_inputs(|| nullable_input(len)) + .bench_local_values(|input| { + let mut ctx = SESSION.create_execution_ctx(); + let mut array = input; + for _ in 0..depth { + array = function + .clone() + .try_new_array(len, EmptyOptions, [array]) + .unwrap() + .into_array(); + } + array.execute::(&mut ctx).unwrap() + }); +} + +#[divan::bench(args = SIZES)] +fn lazy(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn eager(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn lazy_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 3); +} + +#[divan::bench(args = SIZES)] +fn eager_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 3); +} From 5c02036a234b7c0d39ce801655b3ea40be2cbc9b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 16:54:19 -0400 Subject: [PATCH 016/160] Refine RowFn execution contracts Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/row/batch/args.rs | 38 ------------- .../src/scalar_fn/row/batch/execution.rs | 16 ++---- .../src/scalar_fn/row/batch/policy.rs | 56 ++++++------------- vortex-array/src/scalar_fn/row/execute/mod.rs | 20 +++++++ .../src/scalar_fn/row/execute/owned.rs | 16 ++---- .../src/scalar_fn/row/execute/sink.rs | 24 ++------ .../src/scalar_fn/row/types/element/mod.rs | 10 +--- .../src/scalar_fn/row/types/element/tuple.rs | 7 +-- vortex-array/src/scalar_fn/row/visitor/mod.rs | 9 +++ vortex-array/src/scalar_fn/row/vtable.rs | 8 +-- vortex-array/src/scalar_fn/vtable.rs | 43 +++++++++++--- 11 files changed, 100 insertions(+), 147 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/batch/args.rs b/vortex-array/src/scalar_fn/row/batch/args.rs index 262b8252928..520d364df47 100644 --- a/vortex-array/src/scalar_fn/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/row/batch/args.rs @@ -3,9 +3,6 @@ //! Input views and planning metadata passed to a row kernel. -use vortex_error::VortexResult; -use vortex_error::vortex_err; - use crate::ArrayRef; use crate::dtype::DType; use crate::scalar_fn::ExecutionArgs; @@ -29,38 +26,3 @@ pub struct KernelArgs<'a> { /// The non-nullable dtype built by the selected output capability. pub output_dtype: &'a DType, } - -/// An [`ExecutionArgs`] view over borrowed arrays with an explicit row count. -pub(super) struct BorrowedExecutionArgs<'a> { - /// The arrays exposed through this execution view. - inputs: &'a [ArrayRef], - - /// The row count reported for this execution view. - row_count: usize, -} - -impl<'a> BorrowedExecutionArgs<'a> { - pub(super) fn new(inputs: &'a [ArrayRef], row_count: usize) -> Self { - Self { inputs, row_count } - } -} - -impl ExecutionArgs for BorrowedExecutionArgs<'_> { - fn get(&self, index: usize) -> VortexResult { - self.inputs.get(index).cloned().ok_or_else(|| { - vortex_err!( - "Input index {} out of bounds (num_inputs={})", - index, - self.inputs.len() - ) - }) - } - - fn num_inputs(&self) -> usize { - self.inputs.len() - } - - fn row_count(&self) -> usize { - self.row_count - } -} diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index 30670f968a1..1bdbaa4be44 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -11,11 +11,9 @@ use vortex_error::vortex_ensure_eq; use vortex_mask::AllOr; use vortex_mask::Mask; -use super::args::BorrowedExecutionArgs; use super::args::KernelArgs; use super::policy::BatchPlan; use super::policy::RowPolicy; -use super::policy::skipping_beats_filtering; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; @@ -27,6 +25,7 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; +use crate::scalar_fn::BorrowedExecutionArgs; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::row::execute::RowExecution; @@ -90,9 +89,7 @@ impl<'a> Batch<'a> { let arg_dtypes: SmallVec<[DType; 4]> = inputs.iter().map(|input| input.dtype().clone()).collect(); let plan = plan(&arg_dtypes)?; - let nullability = plan.output_dtype.nullability() - | Nullability::from(arg_dtypes.iter().any(DType::is_nullable)); - let result_dtype = plan.output_dtype.with_nullability(nullability); + let result_dtype = plan.result_dtype(&arg_dtypes); let mut validity = Validity::NonNullable; for input in &inputs { @@ -151,9 +148,7 @@ impl<'a> Batch<'a> { match self.policy { RowPolicy::Dense => self.execute_dense(kernel, false, ctx), RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), - RowPolicy::ValidOnly { - filtered_decode_cost, - } => self.execute_valid_only(kernel, try_unfiltered, filtered_decode_cost, ctx), + RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx), } } @@ -272,7 +267,6 @@ impl<'a> Batch<'a> { &Mask, &mut ExecutionCtx, ) -> VortexResult>, - filtered_decode_cost: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { let valid = match self.resolve_validity(&kernel, ctx)? { @@ -280,9 +274,7 @@ impl<'a> Batch<'a> { ResolvedMask::Mixed(valid) => valid, }; - if skipping_beats_filtering(filtered_decode_cost, &valid) - && let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? - { + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { return Ok(result); } diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs index 1b6f3f1f6cc..be3589a2335 100644 --- a/vortex-array/src/scalar_fn/row/batch/policy.rs +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -3,8 +3,6 @@ //! Nullable execution strategies derived from a concrete row dispatch. -use vortex_mask::Mask; - use crate::dtype::DType; use crate::scalar_fn::ElementTuple; use crate::scalar_fn::SinkResult; @@ -18,6 +16,16 @@ pub struct BatchPlan { pub policy: RowPolicy, } +impl BatchPlan { + /// Return the output dtype widened with strict input nullability. + pub fn result_dtype(&self, args: &[DType]) -> DType { + let nullability = self.output_dtype.nullability() + | crate::dtype::Nullability::from(args.iter().any(DType::is_nullable)); + + self.output_dtype.with_nullability(nullability) + } +} + /// The nullable execution policy derived from one concrete dispatch. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum RowPolicy { @@ -27,12 +35,8 @@ pub enum RowPolicy { /// Evaluate all rows, retrying only valid rows if a deferred error is raised. DenseWithRetry, - /// Execute only valid rows, choosing between skip-invalid execution and filtering based on the - /// mask and decode cost. - ValidOnly { - /// Relative per-row decode work that filtering would avoid. - filtered_decode_cost: usize, - }, + /// Execute only valid rows, trying skip-invalid execution before filtering. + ValidOnly, } impl RowPolicy { @@ -41,9 +45,7 @@ impl RowPolicy { if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { Self::Dense } else { - Self::ValidOnly { - filtered_decode_cost: Args::FILTERED_DECODE_COST, - } + Self::ValidOnly } } @@ -52,16 +54,14 @@ impl RowPolicy { if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { Self::DenseWithRetry } else { - Self::ValidOnly { - filtered_decode_cost: Args::FILTERED_DECODE_COST, - } + Self::ValidOnly } } /// The policy one concrete dispatch executes nullable rows under. /// - /// This deliberately ignores [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. Batch execution tries - /// [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) against the original arrays + /// This deliberately ignores [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. Batch execution always + /// tries [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) against the original arrays /// before it tries the sink or filters the inputs. Skipping that probe can change the result of /// an encoding-aware function. /// @@ -74,29 +74,7 @@ impl RowPolicy { Self::Dense } } else { - Self::ValidOnly { - filtered_decode_cost: Args::FILTERED_DECODE_COST, - } + Self::ValidOnly } } } - -/// Minimum surviving-row fractions for skipping when filtering avoids per-row decode work. -/// The thresholds distinguish one costly decode from multiple costly decodes. -const ONE_DECODE_SKIP_MIN_SURVIVING_FRACTION: f64 = 0.50; -const MULTI_DECODE_SKIP_MIN_SURVIVING_FRACTION: f64 = 0.85; - -/// Whether skipping invalid rows should be preferred over filtering for a mixed mask. -pub(super) fn skipping_beats_filtering(filtered_decode_cost: usize, valid: &Mask) -> bool { - if filtered_decode_cost == 0 { - return true; - } - - let minimum = if filtered_decode_cost == 1 { - ONE_DECODE_SKIP_MIN_SURVIVING_FRACTION - } else { - MULTI_DECODE_SKIP_MIN_SURVIVING_FRACTION - }; - - valid.true_count() as f64 >= valid.len() as f64 * minimum -} diff --git a/vortex-array/src/scalar_fn/row/execute/mod.rs b/vortex-array/src/scalar_fn/row/execute/mod.rs index cdfad84ee7a..66f967be4b7 100644 --- a/vortex-array/src/scalar_fn/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/row/execute/mod.rs @@ -8,8 +8,10 @@ use vortex_error::VortexError; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::ArrayRef; +use crate::scalar_fn::ElementTuple; mod owned; pub(super) use owned::execute_owned; @@ -50,3 +52,21 @@ impl From for VortexResult { } } } + +/// Ensure that every decoded varying column addresses the complete row loop. +pub(super) fn ensure_decoded_lengths( + columns: &Args::Columns, + varying: Option<&Args::VaryingColumns<'_>>, + row_count: usize, +) -> VortexResult<()> { + let lengths_match = match varying { + Some(varying) => Args::varying_len_matches(varying, row_count), + None => Args::decoded_lens_match(columns, row_count), + }; + vortex_ensure!( + lengths_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/execute/owned.rs b/vortex-array/src/scalar_fn/row/execute/owned.rs index f2246875143..aa79e9c1076 100644 --- a/vortex-array/src/scalar_fn/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/row/execute/owned.rs @@ -7,9 +7,9 @@ use std::ops::BitOrAssign; use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use super::RowExecution; +use super::ensure_decoded_lengths; use crate::ExecutionCtx; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::IndexedElementTuple; @@ -66,6 +66,8 @@ where let mut values = Vec::::with_capacity(row_count); let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); + let varying = Args::varying(&columns); + ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; let failure; { @@ -73,22 +75,12 @@ where // When every input varies, the indexed source removes argument-shape dispatch from the hot // loop and lets the lane kernel optimize the traversal as one operation. - if let Some(varying) = Args::varying(&columns) { - vortex_ensure!( - Args::varying_len_matches(&varying, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - + if let Some(varying) = varying { failure = Args::indexed_source(&varying) .map_checked_into(output, |elements| apply(&prepared, elements)); } else { // A batch-constant input was collapsed to one row during decoding. This path reads that // row repeatedly while indexing only the inputs that vary. - vortex_ensure!( - Args::decoded_lens_match(&columns, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - let mut accumulated = Fail::default(); for index in 0..row_count { let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index abb58f95b02..272d48c464e 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -10,6 +10,7 @@ use vortex_mask::AllOr; use vortex_mask::Mask; use super::RowExecution; +use super::ensure_decoded_lengths; use crate::ExecutionCtx; use crate::dtype::DType; use crate::scalar_fn::DeferredError; @@ -38,6 +39,8 @@ where let mut sink = Sink::with_capacity(row_count, sink_dtype)?; let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); + let varying = Args::varying(&columns); + ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; let mut accumulated = ApplyResult::Accumulated::default(); { @@ -51,12 +54,7 @@ where // The all-varying representation removes argument-shape dispatch from the hot loop. The // mixed path instead reads collapsed batch constants at row zero. - if let Some(varying) = Args::varying(&columns) { - vortex_ensure!( - Args::varying_len_matches(&varying, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - + if let Some(varying) = varying { for index in 0..row_count { apply( &prepared, @@ -66,11 +64,6 @@ where .accumulate(&mut accumulated)?; } } else { - vortex_ensure!( - Args::decoded_lens_match(&columns, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - for index in 0..row_count { apply( &prepared, @@ -129,14 +122,7 @@ where ); let varying = Args::varying(&columns); - let lens_match = match &varying { - Some(varying) => Args::varying_len_matches(varying, row_count), - None => Args::decoded_lens_match(&columns, row_count), - }; - vortex_ensure!( - lens_match, - "a decoded row input does not address exactly {row_count} rows", - ); + ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; // The loop writes only valid indices, but the sink still finishes a full-length output. // Initialize placeholders now; batch execution masks them before the result escapes. diff --git a/vortex-array/src/scalar_fn/row/types/element/mod.rs b/vortex-array/src/scalar_fn/row/types/element/mod.rs index 3aa92df9121..82e177de606 100644 --- a/vortex-array/src/scalar_fn/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -46,6 +46,9 @@ pub trait InputElement: 'static { /// /// Dense execution requires this of every argument; otherwise the row layer executes only /// valid rows. + /// + /// A dense row closure can receive unspecified values from null rows. The closure must be + /// total over every stored value: it must not panic or have side effects. const DENSE_SAFE: bool = false; /// Whether [`decode`](Self::decode) can fail on _legal_ input data. @@ -54,13 +57,6 @@ pub trait InputElement: 'static { /// contain a value that the decoder rejects. const DECODE_FALLIBLE: bool = true; - /// A relative unit count for per-row decode work avoided by filtering this argument first. - /// - /// Leave this at zero for bulk canonicalization. Use a positive value when filtering first - /// avoids meaningful per-row decode work. The executor adds this cost across arguments when it - /// chooses between skipping invalid rows and filtering. - const FILTERED_DECODE_COST: usize = 0; - /// Validate that `dtype` is an acceptable input column dtype for this element type. fn validate(dtype: &DType) -> VortexResult<()>; diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple.rs b/vortex-array/src/scalar_fn/row/types/element/tuple.rs index becde073a0b..2a9b5f666a5 100644 --- a/vortex-array/src/scalar_fn/row/types/element/tuple.rs +++ b/vortex-array/src/scalar_fn/row/types/element/tuple.rs @@ -147,9 +147,6 @@ pub trait ElementTuple: 'static + private::Sealed { /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. const DECODE_FALLIBLE: bool; - /// The additive cost of per-row decode work avoided by filtering the arguments first. - const FILTERED_DECODE_COST: usize; - /// Validate the input dtypes, including that `dtypes` has exactly `ARITY` entries. /// /// The expression layer checks the count against [`Arity`](crate::scalar_fn::Arity) before it @@ -246,7 +243,6 @@ impl ElementTuple for () { const ARITY: usize = 0; const DENSE_SAFE: bool = true; const DECODE_FALLIBLE: bool = false; - const FILTERED_DECODE_COST: usize = 0; fn validate(dtypes: &[DType]) -> VortexResult<()> { vortex_ensure_eq!( @@ -301,7 +297,6 @@ macro_rules! element_tuple { const ARITY: usize = $arity; const DENSE_SAFE: bool = $($t::DENSE_SAFE &&)+ true; const DECODE_FALLIBLE: bool = $($t::DECODE_FALLIBLE ||)+ false; - const FILTERED_DECODE_COST: usize = $($t::FILTERED_DECODE_COST +)+ 0; fn validate(dtypes: &[DType]) -> VortexResult<()> { vortex_ensure_eq!( @@ -404,7 +399,7 @@ mod tests { use super::UnaryTupleSource; #[test] - fn unary_tuple_source_reads_one_tuple_per_row() { + fn test_unary_tuple_source_reads_one_tuple_per_row() { let source = UnaryTupleSource(&[10, 20, 30]); assert_eq!(source.len(), 3); diff --git a/vortex-array/src/scalar_fn/row/visitor/mod.rs b/vortex-array/src/scalar_fn/row/visitor/mod.rs index bf172103b5d..7d7e615c8aa 100644 --- a/vortex-array/src/scalar_fn/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/row/visitor/mod.rs @@ -38,6 +38,9 @@ pub trait RowVisitor: private::Sealed + Sized { /// Visit an infallible row computation that returns one independent output value. /// + /// `apply` must be total over every stored element value: it must not panic or have side + /// effects. Dense execution can pass unspecified values from null rows. + /// /// # Prerequisites /// /// The framework checks these at compile time: @@ -70,6 +73,9 @@ pub trait RowVisitor: private::Sealed + Sized { /// Visit a row computation that writes through a sink-provided row handle. /// + /// `apply` must be total over every stored element value: it must not panic or have side + /// effects. Dense execution can pass unspecified values from null rows. + /// /// # Prerequisites /// /// The framework checks these at compile time: @@ -108,6 +114,9 @@ pub trait RowVisitor: private::Sealed + Sized { /// Visit a row computation that returns an owned output and deferred failure evidence. /// + /// `apply` must be total over every stored element value: it must not panic or have side + /// effects. Dense execution can pass unspecified values from null rows. + /// /// `Fail` is OR-reduced across rows and handed to `finish_failure`. The value from /// [`Default::default`] **must** mean success, including for an empty batch. The compiler /// cannot check this semantic requirement. diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index 98c2cdc9671..15d8242dfec 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -15,7 +15,6 @@ use super::row_fn::RowFn; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; -use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::union_child_validities; use crate::scalar_fn::Arity; @@ -58,12 +57,7 @@ impl ScalarFnVTable for F { fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { let plan = self.dispatch(options, args, PlanRows::::new(args))?; - // Union the output nullability with the nullability of the inputs. This is required for - // strict scalar function semantics. - let nullability = plan.output_dtype.nullability() - | Nullability::from(args.iter().any(DType::is_nullable)); - - Ok(plan.output_dtype.with_nullability(nullability)) + Ok(plan.result_dtype(args)) } fn execute( diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 5d3561ff039..d8395bd35b8 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -328,20 +328,22 @@ pub trait ExecutionArgs { fn row_count(&self) -> usize; } -/// A concrete [`ExecutionArgs`] backed by a `Vec`. -pub struct VecExecutionArgs { - inputs: Vec, +/// An [`ExecutionArgs`] view over borrowed arrays with an explicit row count. +pub(crate) struct BorrowedExecutionArgs<'a> { + /// The arrays exposed through this execution view. + inputs: &'a [ArrayRef], + + /// The row count reported for this execution view. row_count: usize, } -impl VecExecutionArgs { - /// Create a new `VecExecutionArgs`. - pub fn new(inputs: Vec, row_count: usize) -> Self { +impl<'a> BorrowedExecutionArgs<'a> { + pub(crate) fn new(inputs: &'a [ArrayRef], row_count: usize) -> Self { Self { inputs, row_count } } } -impl ExecutionArgs for VecExecutionArgs { +impl ExecutionArgs for BorrowedExecutionArgs<'_> { fn get(&self, index: usize) -> VortexResult { self.inputs.get(index).cloned().ok_or_else(|| { vortex_err!( @@ -361,6 +363,33 @@ impl ExecutionArgs for VecExecutionArgs { } } +/// A concrete [`ExecutionArgs`] backed by a `Vec`. +pub struct VecExecutionArgs { + inputs: Vec, + row_count: usize, +} + +impl VecExecutionArgs { + /// Create a new `VecExecutionArgs`. + pub fn new(inputs: Vec, row_count: usize) -> Self { + Self { inputs, row_count } + } +} + +impl ExecutionArgs for VecExecutionArgs { + fn get(&self, index: usize) -> VortexResult { + BorrowedExecutionArgs::new(&self.inputs, self.row_count).get(index) + } + + fn num_inputs(&self) -> usize { + BorrowedExecutionArgs::new(&self.inputs, self.row_count).num_inputs() + } + + fn row_count(&self) -> usize { + BorrowedExecutionArgs::new(&self.inputs, self.row_count).row_count() + } +} + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct EmptyOptions; impl Display for EmptyOptions { From a236e0b9d52668ab0eb108be3fd33e289e18715f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 16:56:17 -0400 Subject: [PATCH 017/160] Make RowFn kernel arguments self-contained Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 1 + .../src/scalar_fn/fns/binary/numeric/row.rs | 6 +- vortex-array/src/scalar_fn/row/batch/args.rs | 13 ++-- .../src/scalar_fn/row/batch/execution.rs | 68 ++++++++----------- .../src/scalar_fn/row/batch/policy.rs | 6 +- vortex-array/src/scalar_fn/row/execute/mod.rs | 2 +- .../src/scalar_fn/row/execute/sink.rs | 4 +- vortex-array/src/scalar_fn/row/mod.rs | 1 + .../src/scalar_fn/row/types/element/mod.rs | 4 +- vortex-array/src/scalar_fn/row/types/mod.rs | 1 + .../src/scalar_fn/row/types/result.rs | 26 +++++++ vortex-array/src/scalar_fn/row/types/sink.rs | 33 ++++++++- .../src/scalar_fn/row/visitor/execute.rs | 4 +- vortex-array/src/scalar_fn/row/visitor/mod.rs | 9 +-- .../src/scalar_fn/row/visitor/plan.rs | 2 +- vortex-array/src/scalar_fn/row/vtable.rs | 17 +++-- vortex-array/src/scalar_fn/vtable.rs | 8 ++- 17 files changed, 128 insertions(+), 77 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index e9aa9e8d522..4c42aaa37f4 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -114,6 +114,7 @@ struct I64Sink( impl OutputSink for I64Sink { type Rows<'a> = &'a mut [i64]; type Row<'a> = &'a mut i64; + type WriteToken = (); fn sink_dtype(_args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index a8e78bdcea2..e0b7a658d01 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -29,6 +29,7 @@ use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::row::InitializedElement; use crate::scalar_fn::row::UninitElementSink; pub(super) fn execute_numeric_primitive( @@ -112,14 +113,13 @@ where // vectorization. Check each divide immediately and stop at the first failure. // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when // they need to skip invalid rows. - visitor.visit_into::<(T, T), UninitElementSink, VortexResult<()>>(|(lhs, rhs), output| { + visitor.visit_into::<(T, T), UninitElementSink, _>(|(lhs, rhs), output| { let (value, failed) = CheckedDiv::apply(lhs, rhs); if failed { return Err(numeric_error(>::ERROR)); } - output.write(value); - Ok(()) + Ok(InitializedElement::write(output, value)) }) } diff --git a/vortex-array/src/scalar_fn/row/batch/args.rs b/vortex-array/src/scalar_fn/row/batch/args.rs index 520d364df47..c908b91e00e 100644 --- a/vortex-array/src/scalar_fn/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/row/batch/args.rs @@ -5,21 +5,20 @@ use crate::ArrayRef; use crate::dtype::DType; -use crate::scalar_fn::ExecutionArgs; /// The arguments handed to one kernel invocation. /// /// `arrays` may be filtered or sliced, while `dtypes` and `output_dtype` always describe the -/// original planned batch. Keeping them together prevents an execution path from accidentally -/// pairing an input view with unrelated planning metadata. +/// original planned batch. Keeping them together prevents an execution path from pairing an input +/// view with unrelated planning metadata. #[derive(Clone, Copy)] pub struct KernelArgs<'a> { - /// The executor-facing view, including the row count for this invocation. - pub execution: &'a dyn ExecutionArgs, - - /// The same inputs as concrete arrays for encoding-aware rewrites. + /// The input arrays for this kernel invocation. pub arrays: &'a [ArrayRef], + /// The number of rows in this kernel invocation. + pub row_count: usize, + /// The original input dtypes used to select the row implementation. pub dtypes: &'a [DType], diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index 1bdbaa4be44..62879e2d629 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -25,7 +25,6 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; -use crate::scalar_fn::BorrowedExecutionArgs; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::row::execute::RowExecution; @@ -42,13 +41,12 @@ enum ResolvedMask { } /// One batch of inputs and the metadata needed before its row kernel runs. -pub struct Batch<'a> { +pub struct Batch { /// The function being executed, named in the errors this raises. id: ScalarFnId, - /// The arguments as the execution layer handed them over. Every path but the filter strategy - /// gives the kernel these untouched, so it sees the original encodings. - args: &'a dyn ExecutionArgs, + /// The number of rows in the original execution scope. + row_count: usize, /// The input columns, collected once: constant folding inspects them and the filter strategy /// filters them. @@ -72,14 +70,14 @@ pub struct Batch<'a> { policy: RowPolicy, } -impl<'a> Batch<'a> { +impl Batch { /// Collect the inputs and derive their dtype, validity, and execution policy. /// /// **Not** for a nullary function: with no inputs there is no validity to propagate and no /// per-row work to fold, and the all-constant check below would vacuously pass. pub fn new( id: ScalarFnId, - args: &'a dyn ExecutionArgs, + args: &dyn ExecutionArgs, plan: impl FnOnce(&[DType]) -> VortexResult, ) -> VortexResult { let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) @@ -98,7 +96,7 @@ impl<'a> Batch<'a> { Ok(Self { id, - args, + row_count: args.row_count(), inputs, arg_dtypes, validity, @@ -135,7 +133,7 @@ impl<'a> Batch<'a> { // All inputs constant, and their conjoined validity proves every row non-null. This sees // through extension and masked wrappers just like argument decoding does. - if self.args.row_count() > 0 + if self.row_count > 0 && self.validity.definitely_no_nulls() && self .inputs @@ -168,11 +166,10 @@ impl<'a> Batch<'a> { .map(|input| input.slice(0..1)) .collect::>()?; - let args = BorrowedExecutionArgs::new(&one_row, 1); - let result = VortexResult::from(kernel(self.kernel_args(&args, &one_row), ctx)?)?; + let result = VortexResult::from(kernel(self.kernel_args(&one_row, 1), ctx)?)?; let scalar = self.finalize_output(result, 1)?.execute_scalar(0, ctx)?; - Ok(ConstantArray::new(scalar, self.args.row_count()).into_array()) + Ok(ConstantArray::new(scalar, self.row_count).into_array()) } /// Run the kernel over every row, including the rows behind nulls, then mask its result. @@ -191,13 +188,10 @@ impl<'a> Batch<'a> { return Ok(self.all_null()); } - let values = match kernel(self.kernel_args(self.args, &self.inputs), ctx)? { + let values = match kernel(self.kernel_args(&self.inputs, self.row_count), ctx)? { RowExecution::Output(values) => values, RowExecution::DeferredError(error) if retry_deferred_error => { - let valid = self - .validity - .clone() - .execute_mask(self.args.row_count(), ctx)?; + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; // Unlike `resolve_validity`, all-true preserves the deferred error and all-false // suppresses evidence that came entirely from null rows. An empty loop cannot @@ -218,11 +212,9 @@ impl<'a> Batch<'a> { match self.validity.clone() { Validity::NonNullable | Validity::AllValid => { - self.finalize_output(values, self.args.row_count()) - } - Validity::Array(valid) => { - self.finalize_output(values.mask(valid)?, self.args.row_count()) + self.finalize_output(values, self.row_count) } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), // Handled by the guard above, before the kernel ran. Validity::AllInvalid => Ok(self.all_null()), } @@ -235,18 +227,18 @@ impl<'a> Batch<'a> { kernel: &impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - let valid = self - .validity - .clone() - .execute_mask(self.args.row_count(), ctx)?; + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; // Check all-true before all-false: an empty mask is both, and must not be treated as // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). if valid.all_true() { return self .finalize_output( - VortexResult::from(kernel(self.kernel_args(self.args, &self.inputs), ctx)?)?, - self.args.row_count(), + VortexResult::from(kernel( + self.kernel_args(&self.inputs, self.row_count), + ctx, + )?)?, + self.row_count, ) .map(ResolvedMask::Decided); } @@ -293,7 +285,7 @@ impl<'a> Batch<'a> { ctx: &mut ExecutionCtx, ) -> VortexResult> { let Some(execution) = - try_unfiltered(self.kernel_args(self.args, &self.inputs), valid, ctx)? + try_unfiltered(self.kernel_args(&self.inputs, self.row_count), valid, ctx)? else { return Ok(None); }; @@ -318,30 +310,24 @@ impl<'a> Batch<'a> { .map(|input| input.filter(valid.clone())) .collect::>()?; - let args = BorrowedExecutionArgs::new(&filtered, valid.true_count()); - let values = VortexResult::from(kernel(self.kernel_args(&args, &filtered), ctx)?)?; + let values = VortexResult::from(kernel( + self.kernel_args(&filtered, valid.true_count()), + ctx, + )?)?; self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) } /// An all-null result of the function's declared return dtype. fn all_null(&self) -> ArrayRef { - ConstantArray::new( - Scalar::null(self.result_dtype.clone()), - self.args.row_count(), - ) - .into_array() + ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() } /// Pair an input view with this batch's planning metadata. - fn kernel_args<'b>( - &'b self, - execution: &'b dyn ExecutionArgs, - arrays: &'b [ArrayRef], - ) -> KernelArgs<'b> { + fn kernel_args<'b>(&'b self, arrays: &'b [ArrayRef], row_count: usize) -> KernelArgs<'b> { KernelArgs { - execution, arrays, + row_count, dtypes: &self.arg_dtypes, output_dtype: &self.output_dtype, } diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs index be3589a2335..bc63cf9ce21 100644 --- a/vortex-array/src/scalar_fn/row/batch/policy.rs +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -61,9 +61,9 @@ impl RowPolicy { /// The policy one concrete dispatch executes nullable rows under. /// /// This deliberately ignores [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. Batch execution always - /// tries [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) against the original arrays - /// before it tries the sink or filters the inputs. Skipping that probe can change the result of - /// an encoding-aware function. + /// tries [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) against the original + /// arrays before it tries the sink or filters the inputs. Skipping that probe can change the + /// result of an encoding-aware function. /// /// [`OutputSink::SUPPORTS_SKIPPED_ROWS`]: crate::scalar_fn::OutputSink::SUPPORTS_SKIPPED_ROWS pub const fn for_sink() -> Self { diff --git a/vortex-array/src/scalar_fn/row/execute/mod.rs b/vortex-array/src/scalar_fn/row/execute/mod.rs index 66f967be4b7..9c07363dd7c 100644 --- a/vortex-array/src/scalar_fn/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/row/execute/mod.rs @@ -53,7 +53,7 @@ impl From for VortexResult { } } -/// Ensure that every decoded varying column addresses the complete row loop. +/// Ensure that every decoded input addresses the complete row loop. pub(super) fn ensure_decoded_lengths( columns: &Args::Columns, varying: Option<&Args::VaryingColumns<'_>>, diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 272d48c464e..6374615a0d9 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -33,7 +33,7 @@ pub fn execute_sink( where Args: ElementTuple, Sink: OutputSink, - ApplyResult: SinkResult, + ApplyResult: SinkResult, { let row_count = args.row_count(); let mut sink = Sink::with_capacity(row_count, sink_dtype)?; @@ -91,7 +91,7 @@ pub fn execute_sink_valid_rows( where Args: ElementTuple, Sink: OutputSink, - ApplyResult: SinkResult, + ApplyResult: SinkResult, { // Batch execution needs a full-length result before applying the validity mask. Decline when // the sink cannot leave legal placeholders in positions this loop skips. diff --git a/vortex-array/src/scalar_fn/row/mod.rs b/vortex-array/src/scalar_fn/row/mod.rs index 3351c24d100..fda1dfdfa75 100644 --- a/vortex-array/src/scalar_fn/row/mod.rs +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -24,6 +24,7 @@ mod types; pub use types::DeferredError; pub use types::ElementTuple; pub use types::IndexedElementTuple; +pub use types::InitializedElement; pub use types::InputElement; pub use types::OutputElement; pub use types::OutputSink; diff --git a/vortex-array/src/scalar_fn/row/types/element/mod.rs b/vortex-array/src/scalar_fn/row/types/element/mod.rs index 82e177de606..176f7a00f0d 100644 --- a/vortex-array/src/scalar_fn/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -47,8 +47,8 @@ pub trait InputElement: 'static { /// Dense execution requires this of every argument; otherwise the row layer executes only /// valid rows. /// - /// A dense row closure can receive unspecified values from null rows. The closure must be - /// total over every stored value: it must not panic or have side effects. + /// Dense execution can pass unspecified values from null rows. The closure must be total over + /// every stored value: it cannot panic or cause side effects beyond its declared output. const DENSE_SAFE: bool = false; /// Whether [`decode`](Self::decode) can fail on _legal_ input data. diff --git a/vortex-array/src/scalar_fn/row/types/mod.rs b/vortex-array/src/scalar_fn/row/types/mod.rs index 2998e19ccbd..4032a7d25d4 100644 --- a/vortex-array/src/scalar_fn/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/row/types/mod.rs @@ -19,5 +19,6 @@ pub use result::DeferredError; pub use result::SinkResult; mod sink; +pub use sink::InitializedElement; pub use sink::OutputSink; pub use sink::UninitElementSink; diff --git a/vortex-array/src/scalar_fn/row/types/result.rs b/vortex-array/src/scalar_fn/row/types/result.rs index efd841cc969..4d36256b8df 100644 --- a/vortex-array/src/scalar_fn/row/types/result.rs +++ b/vortex-array/src/scalar_fn/row/types/result.rs @@ -7,6 +7,8 @@ use std::ops::BitOrAssign; use vortex_error::VortexResult; +use super::InitializedElement; + mod private { pub trait Sealed {} } @@ -45,6 +47,9 @@ impl BitOrAssign for DeferredError { /// should be no wider than the computed element so error tracking does not constrain vector width. /// This trait is sealed; row functions choose one of its supplied implementations. pub trait SinkResult: 'static + private::Sealed { + /// The [`OutputSink::WriteToken`](super::OutputSink::WriteToken) carried by a success. + type WriteToken: 'static; + /// The word this result reduces into, kept in a loop-local by the executor. type Accumulated: 'static + Copy + Default; @@ -64,6 +69,7 @@ pub trait SinkResult: 'static + private::Sealed { impl private::Sealed for () {} impl SinkResult for () { + type WriteToken = (); type Accumulated = (); const FALLIBLE: bool = false; @@ -81,6 +87,7 @@ impl SinkResult for () { impl private::Sealed for VortexResult<()> {} impl SinkResult for VortexResult<()> { + type WriteToken = (); type Accumulated = (); const FALLIBLE: bool = true; @@ -95,6 +102,24 @@ impl SinkResult for VortexResult<()> { } } +impl private::Sealed for VortexResult {} + +impl SinkResult for VortexResult { + type WriteToken = InitializedElement; + type Accumulated = (); + + const FALLIBLE: bool = true; + const DEFERRED: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + self.map(|_| ()) + } + + fn occurred(_accumulated: ()) -> bool { + false + } +} + /// The evidence widths a row closure may reduce. `bool` is the ordinary answer; the unsigned /// integers exist for a kernel whose per-row comparison would cost it its vectorization. macro_rules! impl_sink_result_word { @@ -103,6 +128,7 @@ macro_rules! impl_sink_result_word { impl private::Sealed for $word {} impl SinkResult for $word { + type WriteToken = (); type Accumulated = $word; const FALLIBLE: bool = false; diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index 712d209f7e9..b4d5a4c578c 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -47,6 +47,12 @@ pub trait OutputSink: 'static + Sized { where Self: 'a; + /// Proof that a successful row closure left its row handle initialized. + /// + /// Use `()` for initialized row handles. A sink exposing uninitialized storage uses an + /// unforgeable token returned after initialization. + type WriteToken: 'static; + /// The dtype of the column this sink builds, given the function's input dtypes. /// /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the @@ -81,8 +87,28 @@ pub trait OutputSink: 'static + Sized { fn finish(self, error: DeferredError) -> VortexResult; } +/// Proof that one uninitialized element row was initialized. +#[must_use = "return this token from the row closure to prove that it initialized the output"] +pub struct InitializedElement( + /// Private so safe code can only obtain this token by writing an uninitialized row. + (), +); + +impl InitializedElement { + /// Write `value` into an uninitialized row and return its proof token. + #[inline] + pub fn write(row: &mut MaybeUninit, value: T) -> Self { + row.write(value); + + Self(()) + } +} + /// An element sink that leaves dense output uninitialized before the row loop. /// +/// The row closure must return the [`InitializedElement`] from [`InitializedElement::write`] on +/// success. The token is zero-sized, so the proof adds no runtime row state. +/// /// Skip-invalid execution initializes placeholders before omitting rows. Immediate failures are /// safe because [`OutputSink::finish`] is not called after one. pub struct UninitElementSink { @@ -98,6 +124,7 @@ impl OutputSink for UninitElementSink { type Rows<'a> = &'a mut [MaybeUninit]; type Row<'a> = &'a mut MaybeUninit; + type WriteToken = InitializedElement; fn sink_dtype(_args: &[DType]) -> VortexResult { Ok(T::element_dtype()) @@ -129,9 +156,9 @@ impl OutputSink for UninitElementSink { } fn finish(mut self, _error: DeferredError) -> VortexResult { - // SAFETY: dense execution writes every row, while skip-invalid execution initializes every - // row before overwriting valid ones. The executor calls `finish` only after successful - // execution, and the allocation reserved every slot in `0..row_count`. + // SAFETY: dense execution reaches `finish` only after every row returned the token from + // `InitializedElement::write`. Skip-invalid execution initializes every row before + // overwriting valid ones. The allocation reserved every slot in `0..row_count`. unsafe { self.values.set_len(self.row_count) }; Ok(T::build(self.values)) diff --git a/vortex-array/src/scalar_fn/row/visitor/execute.rs b/vortex-array/src/scalar_fn/row/visitor/execute.rs index 36fd03af5d2..2cf4d2ac6f5 100644 --- a/vortex-array/src/scalar_fn/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/row/visitor/execute.rs @@ -90,7 +90,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { where Args: ElementTuple, Sink: OutputSink, - ApplyResult: SinkResult, + ApplyResult: SinkResult, { const { assert_sink_visit_contract::() }; @@ -193,7 +193,7 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { where Args: ElementTuple, Sink: OutputSink, - ApplyResult: SinkResult, + ApplyResult: SinkResult, { const { assert_sink_visit_contract::() }; diff --git a/vortex-array/src/scalar_fn/row/visitor/mod.rs b/vortex-array/src/scalar_fn/row/visitor/mod.rs index 7d7e615c8aa..aedc0912bba 100644 --- a/vortex-array/src/scalar_fn/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/row/visitor/mod.rs @@ -73,8 +73,9 @@ pub trait RowVisitor: private::Sealed + Sized { /// Visit a row computation that writes through a sink-provided row handle. /// - /// `apply` must be total over every stored element value: it must not panic or have side - /// effects. Dense execution can pass unspecified values from null rows. + /// `apply` must be total over every stored input value: it must not panic or cause side effects + /// other than writing the supplied row handle. Dense execution can pass unspecified values + /// from null rows. /// /// # Prerequisites /// @@ -93,7 +94,7 @@ pub trait RowVisitor: private::Sealed + Sized { where Args: ElementTuple, Sink: OutputSink, - ApplyResult: SinkResult, + ApplyResult: SinkResult, { self.visit_prepared_into::( |_| (), @@ -110,7 +111,7 @@ pub trait RowVisitor: private::Sealed + Sized { where Args: ElementTuple, Sink: OutputSink, - ApplyResult: SinkResult; + ApplyResult: SinkResult; /// Visit a row computation that returns an owned output and deferred failure evidence. /// diff --git a/vortex-array/src/scalar_fn/row/visitor/plan.rs b/vortex-array/src/scalar_fn/row/visitor/plan.rs index caa17ad651b..02acee208c1 100644 --- a/vortex-array/src/scalar_fn/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/row/visitor/plan.rs @@ -73,7 +73,7 @@ impl RowVisitor for PlanRows<'_, F> { where Args: ElementTuple, Sink: OutputSink, - ApplyResult: SinkResult, + ApplyResult: SinkResult, { const { assert_sink_visit_contract::() }; diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index 15d8242dfec..848ce47a612 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -18,6 +18,7 @@ use crate::dtype::DType; use crate::expr::Expression; use crate::expr::union_child_validities; use crate::scalar_fn::Arity; +use crate::scalar_fn::BorrowedExecutionArgs; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; @@ -70,8 +71,8 @@ impl ScalarFnVTable for F { if args.num_inputs() == 0 { let result_dtype = ScalarFnVTable::return_dtype(self, options, &[])?; let nullary_args = KernelArgs { - execution: args, arrays: &[], + row_count: args.row_count(), dtypes: &[], output_dtype: &result_dtype, }; @@ -125,10 +126,12 @@ fn execute_rows( return Ok(RowExecution::Output(reduced)); } + let execution = BorrowedExecutionArgs::new(args.arrays, args.row_count); + function.dispatch( options, args.dtypes, - ExecuteRows::::new(args.execution, args.output_dtype, ctx), + ExecuteRows::::new(&execution, args.output_dtype, ctx), ) } @@ -146,19 +149,21 @@ fn try_execute_rows_unfiltered( return Ok(Some(RowExecution::Output(reduced))); } + let execution = BorrowedExecutionArgs::new(args.arrays, args.row_count); + function.dispatch( options, args.dtypes, - ExecuteValidRows::::new(args.execution, args.output_dtype, valid, ctx), + ExecuteValidRows::::new(&execution, args.output_dtype, valid, ctx), ) } /// Prepare the batch inputs and execution plan for `function`. -fn prepare_batch<'args, F: RowFn>( +fn prepare_batch( function: &F, options: &F::Options, - args: &'args dyn ExecutionArgs, -) -> VortexResult> { + args: &dyn ExecutionArgs, +) -> VortexResult { Batch::new(RowFn::id(function), args, |arg_dtypes| { function.dispatch(options, arg_dtypes, PlanRows::::new(arg_dtypes)) }) diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index d8395bd35b8..30f38439dd5 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -196,8 +196,7 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// Returns whether this scalar function is strict. /// /// A strict function returns null for a row when any argument is null for that row. This - /// matches [PostgreSQL's `STRICT` convention](https://www.postgresql.org/docs/current/sql-createfunction.html) - /// for null propagation. + /// matches [PostgreSQL's `STRICT` convention][postgres-strict] for null propagation. /// /// Return `true` only when this holds for every argument. `add` is strict, but Kleene `AND` /// is not because `false AND null` returns `false`. `is_null` is also not strict. @@ -212,6 +211,8 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// /// This property applies only to the scalar function, not its child expressions. Nullary /// functions are vacuously strict. The default is conservatively `false`. + /// + /// [postgres-strict]: https://www.postgresql.org/docs/current/sql-createfunction.html fn is_strict(&self, options: &Self::Options) -> bool { _ = options; false @@ -365,7 +366,10 @@ impl ExecutionArgs for BorrowedExecutionArgs<'_> { /// A concrete [`ExecutionArgs`] backed by a `Vec`. pub struct VecExecutionArgs { + /// The owned arrays exposed through this execution view. inputs: Vec, + + /// The row count reported for this execution view. row_count: usize, } From 69607edb6eb135ea12f8c197c2ded006b1fd6507 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 19:19:15 -0400 Subject: [PATCH 018/160] Elide validated RowFn input bounds checks Signed-off-by: Connor Tsui --- .../src/scalar_fn/row/execute/sink.rs | 15 ++++++------ .../src/scalar_fn/row/types/element/bool.rs | 8 +++++++ .../src/scalar_fn/row/types/element/mod.rs | 15 ++++++++++++ .../scalar_fn/row/types/element/primitive.rs | 8 +++++++ .../src/scalar_fn/row/types/element/tuple.rs | 24 +++++++++++++++++++ 5 files changed, 63 insertions(+), 7 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 6374615a0d9..6b4750a9e66 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -56,12 +56,11 @@ where // mixed path instead reads collapsed batch constants at row zero. if let Some(varying) = varying { for index in 0..row_count { - apply( - &prepared, - Args::get_varying(&varying, index), - Sink::row(&mut rows, index), - ) - .accumulate(&mut accumulated)?; + // SAFETY: `ensure_decoded_lengths` proved every varying column has `row_count` + // rows before the loop. + let elements = unsafe { Args::get_varying_unchecked(&varying, index) }; + apply(&prepared, elements, Sink::row(&mut rows, index)) + .accumulate(&mut accumulated)?; } } else { for index in 0..row_count { @@ -139,7 +138,9 @@ where let result = match &varying { Some(varying) => apply( &prepared, - Args::get_varying(varying, index), + // SAFETY: `ensure_decoded_lengths` proved every varying column has + // `row_count` rows, and mask indices are below `row_count`. + unsafe { Args::get_varying_unchecked(varying, index) }, Sink::row(&mut rows, index), ), None => apply( diff --git a/vortex-array/src/scalar_fn/row/types/element/bool.rs b/vortex-array/src/scalar_fn/row/types/element/bool.rs index e5c29f756b5..bc966268e4d 100644 --- a/vortex-array/src/scalar_fn/row/types/element/bool.rs +++ b/vortex-array/src/scalar_fn/row/types/element/bool.rs @@ -54,6 +54,14 @@ impl InputElement for bool { { column.value(index) } + + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> bool + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { column.value_unchecked(index) } + } } impl OutputElement for bool { diff --git a/vortex-array/src/scalar_fn/row/types/element/mod.rs b/vortex-array/src/scalar_fn/row/types/element/mod.rs index 176f7a00f0d..1743e3f1db0 100644 --- a/vortex-array/src/scalar_fn/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -101,12 +101,27 @@ pub trait InputElement: 'static { fn varying(column: &Self::Column) -> Self::Varying<'_>; /// Number of rows addressable through a [`Varying`](Self::Varying) view. + /// + /// Every index below this length must be valid for + /// [`get_varying_unchecked`](Self::get_varying_unchecked). fn varying_len(column: &Self::Varying<'_>) -> usize; /// Read one row from a [`Varying`](Self::Varying) view. fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> Self::Elem<'a> where Self: 'a; + + /// Read one row without checking that `index` is in bounds. + /// + /// # Safety + /// + /// `index` must be less than [`varying_len`](Self::varying_len) for `column`. + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a, + { + Self::get_varying(column, index) + } } /// An owned row value that can be built into an all-valid column. diff --git a/vortex-array/src/scalar_fn/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/row/types/element/primitive.rs index 05fddbd25e4..071a7c55115 100644 --- a/vortex-array/src/scalar_fn/row/types/element/primitive.rs +++ b/vortex-array/src/scalar_fn/row/types/element/primitive.rs @@ -61,6 +61,14 @@ impl InputElement for T { { column[index] } + + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> T + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { *column.get_unchecked(index) } + } } impl OutputElement for T { diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple.rs b/vortex-array/src/scalar_fn/row/types/element/tuple.rs index 2a9b5f666a5..643d976f298 100644 --- a/vortex-array/src/scalar_fn/row/types/element/tuple.rs +++ b/vortex-array/src/scalar_fn/row/types/element/tuple.rs @@ -189,6 +189,16 @@ pub trait ElementTuple: 'static + private::Sealed { /// Read one row from columns already known to vary within the batch. fn get_varying<'a>(columns: &Self::VaryingColumns<'a>, index: usize) -> Self::Elems<'a>; + /// Read one row from varying columns without checking bounds. + /// + /// # Safety + /// + /// `index` must be in bounds for every column. + unsafe fn get_varying_unchecked<'a>( + columns: &Self::VaryingColumns<'a>, + index: usize, + ) -> Self::Elems<'a>; + /// Read the batch-constant elements out of the decoded columns. Called once per batch. fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; } @@ -281,6 +291,12 @@ impl ElementTuple for () { fn get_varying<'a>(_columns: &Self::VaryingColumns<'a>, _index: usize) -> Self::Elems<'a> {} + unsafe fn get_varying_unchecked<'a>( + _columns: &Self::VaryingColumns<'a>, + _index: usize, + ) -> Self::Elems<'a> { + } + fn constants(_columns: &Self::Columns) -> Self::ConstElems<'_> {} } @@ -356,6 +372,14 @@ macro_rules! element_tuple { ($($t::get_varying(&columns.$idx, index),)+) } + unsafe fn get_varying_unchecked<'a>( + columns: &Self::VaryingColumns<'a>, + index: usize, + ) -> Self::Elems<'a> { + // SAFETY: forwarded from this method's contract. + ($(unsafe { $t::get_varying_unchecked(&columns.$idx, index) },)+) + } + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { ($(columns.$idx.constant(),)+) } From 892717f304867d0761ec784173437a5320ac75a0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 19:19:29 -0400 Subject: [PATCH 019/160] Optimize RowFn tensor and spatial row access Signed-off-by: Connor Tsui --- vortex-spatial/src/scalar_fn/row.rs | 11 +++++++++++ vortex-tensor/src/scalar_fns/row.rs | 15 +++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index 496dc19c69d..f94a1a1aec7 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -66,6 +66,17 @@ impl InputElement for GeometryRow { &column[index] } + unsafe fn get_varying_unchecked<'a>( + column: &Self::Varying<'a>, + index: usize, + ) -> &'a Geometry + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { column.get_unchecked(index) } + } + /// Null rows decode to a placeholder geometry that the branch-and-skip row loop never reads. /// `Point` and `Polygon` columns are covered; other geometry types return `Ok(None)` and the /// batch falls back to the filter strategy. diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 3c02a1d2615..0d9bafb7740 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -108,6 +108,21 @@ impl InputElement for TensorRow { { Self::get(column, index) } + + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + let start = index * column.stride; + + // SAFETY: the caller guarantees that `index` addresses a complete row. + unsafe { + std::slice::from_raw_parts( + column.elements.as_slice().as_ptr().add(start), + column.list_size, + ) + } + } } /// Test-only probe recording which operands the last `prepare` step saw as batch-constant, so a From 4c936447a8bc416a4aceafd7db169e994d5400ac Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 21:26:36 -0400 Subject: [PATCH 020/160] Restore mixed-constant RowFn performance Signed-off-by: Connor Tsui --- .../rowfn-regressions-2026-08-08/README.md | 383 ++++++++++++++++++ .../src/scalar_fn/row/execute/owned.rs | 21 +- 2 files changed, 399 insertions(+), 5 deletions(-) create mode 100644 research/rowfn-regressions-2026-08-08/README.md diff --git a/research/rowfn-regressions-2026-08-08/README.md b/research/rowfn-regressions-2026-08-08/README.md new file mode 100644 index 00000000000..ce4661953ac --- /dev/null +++ b/research/rowfn-regressions-2026-08-08/README.md @@ -0,0 +1,383 @@ + + + +# RowFn regression and compiler-configuration research + +This document records the follow-up performance investigation for `ct/row-fn`. It covers the +benchmarks requested in the [original issue comment], comparison with the [CodSpeed report], four +compiler configurations, commit bisection, source ablations, and the selected optimization. + +The main result is narrow but important. Commit `5c02036a2` moved the `Args::varying` result and its +length check out of the branch that consumes the result. That source-only refactor made mixed +constant primitive operations about 4x slower with the default bench profile and more than 6x +slower with AVX2. Restoring the branch-local view and check recovers the performance. No algorithm +changed. + +The remaining spatial `envelope` regression is separate. It first appears when numeric RowFn code +is linked into the benchmark, even before the spatial functions use RowFn. The experiments below +show code-generation sensitivity, but they do not identify a specific compiler pass or source-level +cause. + +## Revisions and host + +- Candidate before the selected fix: `892717f30` (`ct/row-fn`). +- Develop baseline: `66d096b5d` (`origin/develop`). +- Last fast revision before the regression: `89fd28bc1`. +- First slow revision: `5c02036a2`. +- Rust: 1.91.0, LLVM 21.1.2. +- Host: AMD Ryzen 9 7950X, 16 physical cores and 32 hardware threads. +- Timed process: pinned to logical CPU 4. +- CPU governor: `powersave`; energy-performance preference: `power`. + +The governor could not be changed without elevated host privileges. Every comparison in a table +uses the same host and settings, so ratios are useful. Absolute times should not be compared +directly with the original performance-governor runs. + +The normal repository bench profile already matches two important CodSpeed settings: + +```toml +[profile.bench] +codegen-units = 16 +lto = false +``` + +CodSpeed also supplies `RUSTFLAGS=-C target-feature=+avx2`. Both the default target and this AVX2 +target were measured. + +## What `Args::varying` represents + +RowFn decodes each argument into an `ArgColumn`. An argument is either: + +- `Varying`, with one stored value for every logical row. +- `Constant`, with one stored value reused for every logical row. + +For a tuple, `Args::varying(&columns)` returns `Some` only when _every_ argument is varying. The +tuple implementation uses `?` for each column: + +```rust +fn varying(columns: &Self::Columns) -> Option> { + Some(($($t::varying(columns.$idx.varying()?),)+)) +} +``` + +One constant therefore makes the whole result `None`. This is not a statement about validity. It +classifies the physical row-addressing shape of the decoded arguments. + +The two results select different access mechanisms: + +1. `Some(varying)` contains a tuple of typed contiguous views. After one length check, + `indexed_source` and `map_checked_into` can use unchecked lane reads without per-row shape + dispatch. +2. `None` means at least one argument is constant. `Args::get(&columns, index)` then reads index + zero for each constant column and `index` for each varying column. + +The second mechanism sounds expensive, but it was already present in `89fd28bc1`, where constant +add and subtract took about 9.2 microseconds. The 4x regression was therefore not caused by +introducing the mixed-shape loop. + +The regression came from changing the optimizer-visible data flow around that loop. The slow form +first materialized `Option>`, passed `Option<&...>` to a separate generic +validation helper, and later consumed the original option in a branch: + +```rust +let varying = Args::varying(&columns); +ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; + +if let Some(varying) = varying { + // All-varying execution. +} else { + // Mixed constant and varying execution. +} +``` + +The fast form constructs and validates the typed view only in the selected branch: + +```rust +if let Some(varying) = Args::varying(&columns) { + vortex_ensure!(Args::varying_len_matches(&varying, row_count), ...); + // All-varying execution. +} else { + vortex_ensure!(Args::decoded_lens_match(&columns, row_count), ...); + // Mixed constant and varying execution. +} +``` + +On Rust 1.91.0 and LLVM 21.1.2, this placement determines whether the mixed-constant monomorphs are +well specialized. Source ablation and repeated benchmarks prove the relationship. They do not +prove which LLVM pass makes the poor decision. This should be treated as a measured compiler +workaround, not a general Rust rule. + +The code does need to retain this specific placement for the measured toolchain. The varying view, +its matching length proof, and its consumer should remain in one control-flow branch. Moving them +through the shared helper is semantically equivalent, but currently changes generated-code quality. +The sink executors still use the shared helper because moving their checks did not improve the +cosine or spatial benchmarks. + +## Commit bisection + +The large constant-input regression first appears in `5c02036a2`. + +| Revision | Add constant | Subtract constant | Multiply constant | Add varying | Multiply varying | +| --- | ---: | ---: | ---: | ---: | ---: | +| `89fd28bc1` | 9.219 us | 9.229 us | 18.94 us | 9.379 us | 26.68 us | +| `5c02036a2` | 30.46 us | 31.11 us | 37.73 us | 9.439 us | 26.61 us | + +That commit deduplicated five decoded-length checks into `ensure_decoded_lengths`. Reverting only +the owned executor to branch-local checks recovers constant inputs. Keeping the helper in the sink +executors preserves the useful deduplication where no regression was measured. + +Two other controls did not fix the regression: + +- Reverting the `BorrowedExecutionArgs` move and delegation. +- Adding `#[inline(never)]` to the spatial `box_corners` helper. + +## Selected optimization + +The selected change is confined to `row/execute/owned.rs`: + +- Call `Args::varying` in the `if let` condition. +- Validate `VaryingColumns` inside the all-varying branch. +- Validate the decoded `ArgColumn` tuple inside the mixed branch. +- Keep both validations before their loops so bounds-check elimination remains possible. +- Leave sink and valid-row execution unchanged. + +This is a control-flow and proof-placement change. It adds no per-row work and does not change +null, failure, constant, or output semantics. + +### Primitive binary results + +Default bench profile, median time: + +| Benchmark | Candidate | Fixed | Develop | Fixed/develop | +| --- | ---: | ---: | ---: | ---: | +| `add_i64_constant` | 35.4 us | 9.26 us | 8.38 us | 1.10x | +| `sub_i64_constant` | 36.2 us | 9.15 us | 8.23 us | 1.11x | +| `mul_i32_constant` | 41.9 us | 18.89 us | 26.45 us | 0.71x | +| `add_i64_nonnull` | 9.44 us | 9.44 us | approximately 9 us | approximately 1x | +| `mul_i32_nonnull` | 26.66 us | 26.66 us | approximately 26 us | approximately 1x | + +AVX2, median time: + +| Benchmark | Candidate | Fixed | Develop | Fixed/develop | +| --- | ---: | ---: | ---: | ---: | +| `add_i64_constant` | 29.70 us | 6.099 us | 4.919 us | 1.24x | +| `sub_i64_constant` | 30.65 us | 6.249 us | 4.959 us | 1.26x | +| `mul_i32_constant` | 37.97 us | 6.519 us | 5.689 us | 1.15x | + +The fix removes the major regression. Small constant add and subtract gaps remain, especially with +AVX2, but they are not the same failure mode. + +### RowFn executor microbenchmarks + +Default-profile medians before and after the selected fix: + +| Benchmark | Before | After | +| --- | ---: | ---: | +| Handwritten wrapping | approximately 127 ns | approximately 127 ns | +| RowFn sink wrapping | approximately 128 ns | approximately 128.5 ns | +| RowFn wrapping | approximately 128.5 ns | approximately 128.5 ns | +| RowFn checked | approximately 141.7 ns | approximately 141.7 ns | +| RowFn wrapping constant | approximately 62.1 ns | approximately 11.91 ns | +| RowFn checked constant | approximately 67.8 ns | approximately 35.69 ns | +| RowFn wrapping nullable | approximately 129.6 ns | approximately 130 ns | +| RowFn checked nullable | approximately 143.4 ns | approximately 143.6 ns | + +Only the mixed-constant cases move materially, which matches the source-level diagnosis. + +## Tensor results + +### Squared L2 distance + +Candidate and develop medians in microseconds: + +| Width | Candidate nonnull | Develop nonnull | Candidate nullable | Develop nullable | +| ---: | ---: | ---: | ---: | ---: | +| 2 | 17.29 | 31.77 | 18.47 | 32.45 | +| 32 | 6.77 | 7.26 | 7.95 | 8.01 | +| 256 | 10.15 | 10.00 | 11.28 | 10.72 | + +The candidate is about 1.83x faster at nonnull width 2 and 1.75x faster at nullable width 2. It is +about 7% and 1% faster at width 32. At width 256 it is about 1.5% slower for nonnull input and 5.2% +slower for nullable input. + +### Cosine similarity + +Candidate and develop medians in microseconds: + +| Shape and width | Candidate | Develop | Candidate speedup | +| --- | ---: | ---: | ---: | +| Column-column, 2 | 4.47 | 18.28 | 4.1x | +| Column-column, 32 | 2.44 | 4.97 | 2.0x | +| Column-column, 256 | 2.37 | 5.74 | 2.4x | +| Column-constant, 2 | 6.33 | 56.45 | 8.9x | +| Column-constant, 32 | 6.52 | 49.25 | 7.5x | +| Column-constant, 256 | 26.45 | 67.73 | 2.6x | +| Extension constant, 2 | 6.65 | 16.36 | 2.5x | +| Extension constant, 32 | 6.84 | 9.91 | 1.4x | +| Extension constant, 256 | 26.85 | 41.49 | 1.5x | + +The owned-executor optimization does not affect cosine similarity because that implementation uses +prepared sink execution. Moving the sink length proof into its selected branch was tested and did +not materially change these results. + +## Spatial results + +Most predicate benchmarks remain close to the handwritten kernels: + +- Column-column cases are generally 1% to 6% slower. +- Constant-input cases are generally 7% to 17% slower. +- Inputs with 90% nulls are about 4% faster. +- Dual-nullable inputs are about 2% slower. +- Polygon-column against constant-point cases are approximately equal. +- Constant-input `intersects` cases are about 4% to 9% slower. +- Exact and bounding-box diagnostic cases are approximately equal. +- The disjoint bounding-box diagnostic is slightly faster on the candidate. + +Moving the sink proof into its selected branch did not materially change these predicate results. + +### `envelope` + +`envelope` has a separate, reproducible regression. Default-profile multipolygon results in +microseconds were: + +| Input | Candidate before fix | Candidate after fix | Develop | +| --- | ---: | ---: | ---: | +| Mixed | 66.0 | 57.61 | 42.3 | +| Nonnull | 68.36 | 59.11 | 43.94 | +| Random | 54.28 | 48.40 | 33.63 | + +The owned-executor change removes part of the final branch's loss, but the remaining regression is +about 34% to 45%. + +Commit history isolates when it appears: + +| Revision | Mixed | Nonnull | Random | +| --- | ---: | ---: | ---: | +| Framework only, `fef191df5` | 42.52 us | 44.52 us | 33.73 us | +| Numeric RowFn port, `b324f3e26` | 58.02 us | 59.72 us | 49.11 us | +| Before geo RowFn, `aebe3caf7` | 58.43 us | 59.99 us | 49.50 us | + +The regression therefore predates the geo visitor conversion. The `envelope.rs` source is +unchanged. It appears when numeric RowFn code is linked into the benchmark binary. + +The generated candidate `envelope_array` function was smaller than develop, not larger: + +| Revision | Instructions | Calls | Jumps | +| --- | ---: | ---: | ---: | +| Candidate | 1,725 | 115 | 175 | +| Develop | 1,811 | 122 | 189 | + +This rules out the simple explanation that the candidate executes a visibly larger function. It +does not rule out placement, inlining, alignment, cache, or compiler phase-order effects elsewhere +in the linked binary. `perf` was unavailable on this host. LLVM-MCA was available, but no isolated +hot loop that retained the end-to-end regression was found. + +## `list_sum` and unrelated code-generation sensitivity + +`list_sum` does not call the RowFn owned executor, but it changed at the same source-shape commit. +This is evidence that generic code placement can perturb other monomorphs in the benchmark binary. + +Default-profile progression: + +| Revision | Large | Medium | +| --- | ---: | ---: | +| Framework only, `0a0ad0db1` | 13.84 ms | 59.71 us | +| Numeric RowFn, `89fd28bc1` | 13.49 ms | 61.8 us | +| Shared proof, `5c02036a2` | 14.99 ms | 77.82 us | +| Same revision with branch-local owned proof | 13.65 ms | 63.83 us | +| Final candidate with fix | 13.43 ms | 60.10 us | +| Develop | 13.59 ms | 60.48 us | + +With AVX2, the fixed candidate measured 12.96 ms and 61.75 us; develop measured 13.26 ms and +58.75 us. The large case is about 2% faster, while the medium case is about 5% slower. + +Because `list_sum` does not execute this RowFn path, the exact compiler mechanism remains an +inference. The commit bisection and one-change source ablation establish correlation and +reversibility, not a specific LLVM pass. + +## Compact-slice control + +The `compact_sliced(16384, 10)` benchmark did not reproduce the 26% CodSpeed loss: + +| Configuration | Candidate | Develop | Difference | +| --- | ---: | ---: | ---: | +| Default | 107.45 us | 105.7 us | Candidate 1.7% slower | +| One CGU | 105.7 us | 104.9 us | Candidate 0.8% slower | +| AVX2 | 64.21 us | 66.08 us | Candidate 2.8% faster | +| Thin LTO | 107.3 us | 107.5 us | Approximately equal | + +This result is consistent with simulation noise or linked-code layout sensitivity in CodSpeed. It +does not reproduce a durable algorithmic regression on this host. + +## Compiler-configuration matrix + +Changing codegen units, LTO, or AVX2 did not remove the two main regressions before the selected +fix. + +| Configuration | Constant operands | `list_sum` | `envelope` | Compact slice | +| --- | --- | --- | --- | --- | +| 16 CGUs, no LTO | About 4x slower | Medium 33% slower | 55% to 62% slower | 1.7% slower | +| 1 CGU, no LTO | Add/sub 3.9x; mul 1.33x | 13% / 25% slower | 53% to 60% slower | 0.8% slower | +| 16 CGUs, AVX2 | 6x to 6.7x slower | 9% / 26% slower | 58% to 63% slower | 2.8% faster | +| 16 CGUs, Thin LTO | Similar large loss | Large 9%; medium 32% slower | 52% to 58% slower | Equal | + +The repository's default of 16 CGUs and no LTO does not create the problem. One CGU and Thin LTO +also do not fix it. AVX2 amplifies the mixed-constant gap before the branch-local change. + +## Confirmed findings + +- `Args::varying` returns `Some` only when every decoded argument varies by row. +- Its `Some` value enables a typed indexed lane source; `None` selects mixed-shape row access. +- The mixed-shape loop itself was fast before `5c02036a2`. +- Hoisting the option and its proof through a generic helper causes the large mixed-constant loss on + Rust 1.91.0 and LLVM 21.1.2. +- Restoring branch-local construction and validation recovers the loss without new per-row work. +- All-varying numeric benchmarks are unchanged by the selected fix. +- Prepared-sink cosine and geo cases do not benefit from the analogous source change. +- `list_sum` tracks the source ablation even though it does not use owned RowFn execution. +- The `envelope` regression begins with the numeric RowFn port, before geo adopts RowFn. +- CGU count, Thin LTO, and AVX2 do not remove the unfixed regressions. +- The compact-slice CodSpeed regression does not reproduce materially on this host. + +## Inferences and unresolved questions + +- The mixed-constant result is likely an LLVM phase-order or specialization-quality problem. The + benchmark and source ablation do not identify the responsible pass. +- `list_sum` and `envelope` are likely sensitive to linked-code placement, inlining, alignment, or + another whole-program code-generation effect. No single mechanism has been proven. +- Smaller `envelope_array` assembly does not imply faster execution. The relevant difference may + be outside that symbol or may involve front-end behavior rather than instruction count. +- A compiler reduction should preserve both the timing delta and the production monomorph before + filing an LLVM or rustc issue. + +## Benchmark coverage and limitations + +The durable current-tree replacements for the original issue comment were run: + +- Primitive binary operations. +- RowFn executor microbenchmarks. +- Tensor L2 and cosine similarity. +- Geo predicates, bounding-box diagnostics, and envelope. +- `list_sum`. +- Compact sliced arrays. + +The old experimental `BytesLen` and forced null-strategy benchmarks no longer exist in the current +tree, so they could not be rerun. No substitute result is presented as if it were the removed +benchmark. + +Representative commands were: + +```bash +taskset -c 4 cargo bench -p vortex-array --bench binary_ops -- +taskset -c 4 cargo bench -p vortex-array --bench row_fn_executor -- +taskset -c 4 cargo bench -p vortex-array --bench list_sum -- +RUSTFLAGS='-C target-feature=+avx2' taskset -c 4 cargo bench ... +CARGO_PROFILE_BENCH_CODEGEN_UNITS=1 taskset -c 4 cargo bench ... +CARGO_PROFILE_BENCH_LTO=thin taskset -c 4 cargo bench ... +``` + +Compilations used separate target directories before timed runs when configurations differed. Timed +runs were serialized on one logical CPU. + +[original issue comment]: https://github.com/vortex-data/vortex/issues/9128#issuecomment-5151831802 +[CodSpeed report]: https://github.com/vortex-data/vortex/pull/9255#issuecomment-5211040550 diff --git a/vortex-array/src/scalar_fn/row/execute/owned.rs b/vortex-array/src/scalar_fn/row/execute/owned.rs index aa79e9c1076..8a9e27383ee 100644 --- a/vortex-array/src/scalar_fn/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/row/execute/owned.rs @@ -7,9 +7,9 @@ use std::ops::BitOrAssign; use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use super::RowExecution; -use super::ensure_decoded_lengths; use crate::ExecutionCtx; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::IndexedElementTuple; @@ -66,21 +66,32 @@ where let mut values = Vec::::with_capacity(row_count); let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); - let varying = Args::varying(&columns); - ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; let failure; { let output = &mut values.spare_capacity_mut()[..row_count]; // When every input varies, the indexed source removes argument-shape dispatch from the hot - // loop and lets the lane kernel optimize the traversal as one operation. - if let Some(varying) = varying { + // loop and lets the lane kernel optimize the traversal as one operation. Keep the varying + // view and its length proof in this branch: hoisting them through the shared validation + // helper produces slower mixed-constant code with LLVM 21.1.2. See + // `research/rowfn-regressions-2026-08-08/README.md`. + if let Some(varying) = Args::varying(&columns) { + vortex_ensure!( + Args::varying_len_matches(&varying, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + failure = Args::indexed_source(&varying) .map_checked_into(output, |elements| apply(&prepared, elements)); } else { // A batch-constant input was collapsed to one row during decoding. This path reads that // row repeatedly while indexing only the inputs that vary. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + let mut accumulated = Fail::default(); for index in 0..row_count { let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); From bdf95a77ecaab2e56e0d58c61b6a5ee7ede47bd8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 22:17:47 -0400 Subject: [PATCH 021/160] Document RowFn design and performance handoff Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/DESIGN.md | 495 ++++++++++++++++++ research/rowfn-reconstruction/HANDOFF.md | 125 +++++ research/rowfn-reconstruction/OPTIMIZATION.md | 329 ++++++++++++ research/rowfn-reconstruction/README.md | 95 ++++ research/rowfn-reconstruction/REPRODUCE.md | 380 ++++++++++++++ 5 files changed, 1424 insertions(+) create mode 100644 research/rowfn-reconstruction/DESIGN.md create mode 100644 research/rowfn-reconstruction/HANDOFF.md create mode 100644 research/rowfn-reconstruction/OPTIMIZATION.md create mode 100644 research/rowfn-reconstruction/README.md create mode 100644 research/rowfn-reconstruction/REPRODUCE.md diff --git a/research/rowfn-reconstruction/DESIGN.md b/research/rowfn-reconstruction/DESIGN.md new file mode 100644 index 00000000000..01297ba1e77 --- /dev/null +++ b/research/rowfn-reconstruction/DESIGN.md @@ -0,0 +1,495 @@ + + + +# RowFn design + +## Problem statement + +A scalar function receives arrays, but its mathematical definition often describes one row. For +example, checked addition has this row definition: + +```rust +fn checked_add(lhs: i64, rhs: i64) -> (i64, bool) { + lhs.overflowing_add(rhs) +} +``` + +A complete array implementation also needs to do this work: + +- Validate both dtypes. +- Decode both arrays into representations with cheap row access. +- Preserve or collapse batch constants. +- Combine input validity. +- Select dense or valid-only execution. +- Allocate output. +- Attribute failures only to valid rows. +- Build an array with the declared dtype and length. + +RowFn keeps the row definition small and implements the column concerns once. + +## Public declaration + +A row function declares its options, argument names, identity, fallibility, and dtype dispatch. +The essential trait has this shape: + +```rust +trait RowFn: Clone + Send + Sync + 'static { + type Options; + + const ARG_NAMES: &'static [&'static str]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId; + + fn dispatch( + &self, + options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult; + + fn reduce_encoded( + &self, + options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} +``` + +`dispatch` selects concrete Rust element types. Planning and execution call the same method with +different visitor types. Therefore, `dispatch` must select the same visit from only `options` and +the input dtypes. + +The `reduce_encoded` hook is optional. It gives an encoding-aware implementation the original +arrays before row decoding. `None` selects the row loop. + +## Why dispatch uses a visitor + +The return type of a generic visit depends on whether the caller plans or executes. Stable Rust +cannot return one closure with caller-selected generic types from a normal function. The visitor +reverses control: + +```text +ScalarFnVTable::return_dtype + -> RowFn::dispatch(PlanRows) + -> visitor.visit::(closure) + -> BatchPlan + +ScalarFnVTable::execute + -> RowFn::dispatch(ExecuteRows) + -> visitor.visit::(closure) + -> RowExecution +``` + +The function chooses `ConcreteArgs` and `ConcreteOutput`. The framework chooses what a visit does. +The compiler monomorphizes both paths for those concrete types. + +The planning visitor does not call the row closure. It validates the selected input and output +types, checks compile-time contracts, and selects a null policy. The execution visitor decodes the +arrays and runs the matching loop. + +## Visit capabilities + +The visitor has six entry points. Three unprepared methods delegate to three prepared methods. + +| Method | Output model | Row error model | Preparation | +| --- | --- | --- | --- | +| `visit` | Independent owned value | None | None | +| `visit_prepared` | Independent owned value | None | Once per batch | +| `visit_deferred` | Independent owned value | OR-reduced evidence | None | +| `visit_prepared_deferred` | Independent owned value | OR-reduced evidence | Once per batch | +| `visit_into` | Sink row handle | `SinkResult` | None | +| `visit_prepared_into` | Sink row handle | `SinkResult` | Once per batch | + +The unprepared methods exist for the common case: + +```rust +visitor.visit::<(i64, i64), i64>(|(lhs, rhs)| lhs.wrapping_add(rhs)) +``` + +Their default implementation supplies an empty prepared value: + +```rust +self.visit_prepared::( + |_| (), + move |&(), args| apply(args), +) +``` + +This delegation keeps planning and execution logic in the prepared methods only. + +## Input elements + +`InputElement` connects one logical Rust row value to one decoded array representation: + +```rust +trait InputElement { + type Column; + type Varying<'a>; + type Elem<'a>; + + const DENSE_SAFE: bool; + const DECODE_FALLIBLE: bool; + + fn validate(dtype: &DType) -> VortexResult<()>; + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>; + fn varying(column: &Self::Column) -> Self::Varying<'_>; + fn varying_len(column: &Self::Varying<'_>) -> usize; + unsafe fn get_varying_unchecked( + column: &Self::Varying<'_>, + index: usize, + ) -> Self::Elem<'_>; +} +``` + +`Column` owns the decoded batch representation. `Varying` is the cheaper view used by an +all-varying loop. `Elem` is the value that the row closure receives. + +For `i64`, these types are: + +```rust +type Column = Buffer; +type Varying<'a> = &'a [i64]; +type Elem<'a> = i64; +``` + +The decode step performs the array execution and ptype downcast once. The row loop sees a slice +and `i64` values. It does not see `ArrayRef`, a trait object, a ptype match, or an execution +context. + +For a tensor row of `f32`, these types are: + +```rust +type Column = TensorRows; +type Varying<'a> = &'a TensorRows; +type Elem<'a> = &'a [f32]; +``` + +`TensorRows` stores one typed flat buffer, the row count, the width, and a stride. The row access +computes one offset and returns a slice. This removes a ptype check and buffer downcast from every +row. + +## Concrete `Args::varying` examples + +`ElementTuple` combines input elements. It decodes each input into an `ArgColumn`: + +```rust +enum ArgColumnKind { + Varying(T::Column), + Constant(T::Column), +} +``` + +A constant column is decoded as one physical row. The logical batch length stays separate. + +### Example 1: column plus column + +Consider this logical input: + +```text +lhs = [10, 20, 30] +rhs = [ 1, 2, 3] +``` + +The decoded tuple is conceptually: + +```text +columns = ( + Varying(Buffer([10, 20, 30])), + Varying(Buffer([1, 2, 3])), +) +``` + +`Args::varying(&columns)` asks both arguments for direct varying views: + +```rust +Some(( + columns.0.varying()?, + columns.1.varying()?, +)) +``` + +Both calls return `Some`, so the result is: + +```text +Some((&[10, 20, 30], &[1, 2, 3])) +``` + +The executor validates both lengths once. It then creates a `LaneZip` source. The source yields: + +```text +index 0 -> (10, 1) +index 1 -> (20, 2) +index 2 -> (30, 3) +``` + +The hot loop does not inspect `ArgColumnKind`. + +### Example 2: column plus constant + +Now consider this logical input: + +```text +lhs = [10, 20, 30] +rhs = Constant(7, logical_len = 3) +``` + +The decoded tuple is conceptually: + +```text +columns = ( + Varying(Buffer([10, 20, 30])), + Constant(Buffer([7])), +) +``` + +The first `varying()?` succeeds. The second returns `None`. The `?` returns `None` from the tuple +method, so this is the result: + +```text +Args::varying(&columns) == None +``` + +`None` does not mean that no input varies. It means that the tuple is not _all varying_. The mixed +loop uses `Args::get`: + +```text +index 0 -> (columns.0[0], columns.1[0]) -> (10, 7) +index 1 -> (columns.0[1], columns.1[0]) -> (20, 7) +index 2 -> (columns.0[2], columns.1[0]) -> (30, 7) +``` + +This loop performs one `ArgColumnKind` match for each argument and row. It avoids allocating or +expanding `[7, 7, 7]`. + +The preparation input is independent from `Args::varying`: + +```text +Args::constants(&columns) == (None, Some(7)) +``` + +A prepared closure can precompute work from `7`. An ordinary closure can ignore the preparation +input and still use the mixed loop. + +### Example 3: constant plus constant + +If both inputs are non-null constants, batch execution takes a higher-level fast path. It executes +one row and broadcasts the result to the logical batch length. + +The row executor can still represent two constants. This representation matters for a masked +constant because the strict validity can prevent the all-constant broadcast path. + +## Why `Args::varying` exists + +The simplest loop can call `Args::get` for every input shape. That loop contains a branch for each +argument and row: + +```rust +for index in 0..row_count { + let lhs = match lhs_column { + Varying(values) => values[index], + Constant(value) => value[0], + }; + let rhs = match rhs_column { + Varying(values) => values[index], + Constant(value) => value[0], + }; + output[index] = apply(lhs, rhs); +} +``` + +For two varying arrays, these branches always choose the same arm. `Args::varying` selects that +shape once before the loop. The all-varying loop then contains only loads, arithmetic, failure +reduction, and stores. + +`VaryingColumns` also removes buffer descriptors from the row path. A primitive tuple becomes two +slices, and a `LaneZip` gives LLVM independent indexed loads. + +## Owned output + +`OutputElement` describes a Rust value that builds an all-valid array: + +```rust +trait OutputElement { + fn element_dtype() -> DType; + fn build(values: Vec) -> ArrayRef; +} +``` + +The dtype cannot depend on runtime input metadata. Primitive output fits this model. A tensor +output whose shape comes from an input dtype does not. + +The owned executor allocates `Vec` once. It exposes the spare capacity as +`[MaybeUninit]`. The loop writes each row directly into its final output slot. + +The vector length remains zero until the loop finishes. Therefore, an unwind does not drop +uninitialized slots. A compile-time assertion rejects output types that require drop glue. After +normal completion, the executor sets the length once and builds the array. + +## Output sinks + +An output sink supports runtime-shaped output and shared batch state: + +```rust +trait OutputSink { + type Rows<'a>; + type Row<'a>; + type WriteToken; + + fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; + fn rows(&mut self) -> Self::Rows<'_>; + fn row(rows: &mut Self::Rows<'_>, index: usize) -> Self::Row<'_>; + fn finish(self, error: DeferredError) -> VortexResult; +} +``` + +The executor borrows `Rows` once before the loop. This keeps the sink descriptor and shape as loop +invariants. The closure receives only the row handle. + +`UninitElementSink` avoids zero-initializing dense primitive output. Its row handle is +`&mut MaybeUninit`. Safe code must prove that it wrote the slot: + +```rust +let token = InitializedElement::write(output, value); +Ok(token) +``` + +`InitializedElement` is a zero-sized, unforgeable write token. The sink can call `Vec::set_len` +only after every successful row returns this token. A valid-only loop initializes placeholders +before it skips rows. + +## Failure models + +An immediate `VortexResult` leaves the loop on the first error. This model is appropriate when the +operation is expensive and scalar, such as integer division. + +Deferred failure separates cheap row evidence from expensive error construction: + +```rust +let mut failed = Fail::default(); +for index in 0..row_count { + let (value, row_failure) = apply(input[index]); + failed |= row_failure; + output[index].write(value); +} +finish_failure(failed) +``` + +The failure type must be no wider than the output type. A wide loop-carried reduction can limit +the vector width. The default failure value must mean success, including for an empty batch. + +The closure creates no `VortexError`. A cold function creates the rich error after the loop. + +## Why `RowExecution` exists + +Dense execution can evaluate stored payloads behind null rows. A checked operation can report a +failure from such a payload. That failure must not escape if the logical row is null. + +`RowExecution` preserves this distinction: + +```rust +enum RowExecution { + Output(ArrayRef), + DeferredError(VortexError), +} +``` + +An outer `VortexResult` carries immediate or structural errors. `DeferredError` means that the loop +finished and produced only retryable failure evidence. + +For mixed validity, batch execution filters to valid rows and repeats the dense loop. The second +result decides whether the error is observable. Once a path contains only valid rows, +`From for VortexResult` turns a deferred error into an ordinary error. + +## Null execution policies + +Planning derives one policy from the concrete input and result types. + +### `Dense` + +This policy applies when decoding and the closure tolerate all stored null payloads. The kernel +visits every row and batch execution masks the output. + +Primitive arithmetic uses this policy when it is infallible. A null primitive row still stores a +valid Rust primitive value, although that value is logically unspecified. + +### `DenseWithRetry` + +This policy applies to dense-safe inputs with deferred failure evidence. The first loop visits all +rows. If it reports failure, batch execution materializes validity and retries only valid rows. + +This policy preserves the fast dense loop for the common success case. It also prevents a null +payload from creating an observable error. + +### `ValidOnly` + +This policy applies when decoding or row access cannot tolerate null payloads. Batch execution +first asks the sink to skip invalid rows over the original arrays. If the input or sink cannot +support that path, batch execution filters every input and scatters the compact result. + +Geometry uses this policy. Some geometry encodings can decode a harmless placeholder for null +rows. The loop then reads only the valid indices. + +## Prepared constants + +A prepared visit receives `Option` for each argument before the row loop. `Some` means that +the argument is a batch constant. + +Cosine similarity uses this capability to compute a constant operand norm once: + +```rust +prepare((lhs, rhs)) -> ConstNorms { + lhs: lhs.map(l2_norm_row), + rhs: rhs.map(l2_norm_row), +} +``` + +Each row still computes its inner product. It reuses a prepared norm when an operand is constant. + +Spatial containment and intersection use the same pattern for constant geometry metadata and +bounding boxes. The preparation step removes repeated work without adding a specialized array +kernel. + +## Loop shape that LLVM receives + +For an all-varying primitive pair, monomorphization reduces the framework to this essential loop: + +```rust +let mut failed = Fail::default(); +for index in 0..len { + let lhs = unsafe { *lhs.get_unchecked(index) }; + let rhs = unsafe { *rhs.get_unchecked(index) }; + let (value, row_failure) = apply((lhs, rhs)); + failed |= row_failure; + unsafe { output.get_unchecked_mut(index).write(value) }; +} +``` + +The loop has these properties: + +- The input and output element types are concrete. +- The closure is concrete and inlineable. +- Input lengths are equal and validated before the loop. +- The output length equals the input length. +- Each iteration reads and writes an independent index. +- The failure reduction is associative bitwise OR. +- Rich errors, array construction, dtype dispatch, and validity logic are outside the loop. + +These properties make the loop suitable for LLVM autovectorization. They do not force LLVM to use +SIMD for every operation. + +## Compile-time contracts + +Const assertions reject these invalid declarations during compilation: + +- The element tuple arity differs from `RowFn::ARG_NAMES`. +- Input decoding can fail, but `RowFn::FALLIBLE` is false. +- A row result can fail, but `RowFn::FALLIBLE` is false. +- An owned output requires drop glue. +- Deferred failure evidence is wider than the output. +- A sink and its result disagree about deferred errors. + +Runtime planning validates input dtypes and output nullability. Batch finalization validates output +length and dtype. diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md new file mode 100644 index 00000000000..94244a2bc59 --- /dev/null +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -0,0 +1,125 @@ + + + +# RowFn investigation handoff + +This file records the exact state at the end of the 2026-08-09 investigation. Start with this +file, then read [`DESIGN.md`](DESIGN.md), [`OPTIMIZATION.md`](OPTIMIZATION.md), and +[`REPRODUCE.md`](REPRODUCE.md). + +## Branch state + +- Branch: `ct/row-fn`. +- Code head before this documentation commit: `4c936447a`. +- Comparison revision: develop at `66d096b5d`. +- No RowFn code changed during this final investigation. +- The only intended new branch changes are the files in `research/rowfn-reconstruction`. + +Two temporary remote refs exist for a future CodSpeed ablation: + +- `ct/row-fn-codspeed-framework` points to `0a0ad0db1`. +- `ct/row-fn-codspeed-numeric` points to `89fd28bc1`. + +The refs contain exact historical code. They do not contain the uncommitted focused-workflow edits +that were made only in temporary local worktrees. + +## Corrected CodSpeed history + +The latest push did not bring back the `take_filter_list_*` regressions. + +- The [CodSpeed check at `892717f30`] already reports the cases as about 15% to 16% slower. +- The [CodSpeed check at `4c936447a`] reports the same cases as about 14% to 16% slower. +- Most take/filter simulated times improve by less than 2% between those checks. +- `4c936447a` fixes the much larger constant add, subtract, and multiply regressions. This moves the + persistent take/filter entries higher in the ordered list of the 20 largest changes. +- Every retained RowFn CodSpeed summary from `0e5c19c00` through `4c936447a` that has a performance + table also contains take/filter regressions. + +The PR bot edits one current comment, and GitHub displays only the 20 largest changes. These two +details can make a persistent regression appear to leave and return. + +## What is known about take/filter + +The list, filter, and take source files are identical between develop and `4c936447a`. The +`take_filter_list` benchmark does not execute a RowFn operation. + +The linked AVX2 benchmark binaries still differ. Native inspection found: + +- The main filter-take function has the same `0x41cc` byte size on develop, `892717f30`, and + `4c936447a`. +- The main list `TakeExecute::take` function has the same `0x40ac` byte size. +- Normalized list-take disassembly has the same instructions. +- Function addresses, relative call targets, and linked layout differ. + +This evidence is consistent with a linked-layout effect or a changed callee outside the inspected +symbol. It does not prove which cache, branch, or callee causes the result. + +CodSpeed documents [function alignment] as a reason unchanged microbenchmarks can move after a +rebuild. Its differential flame graph is the correct next source of evidence. Inspect the +instruction, cache, and memory components separately. + +## Native measurements are separate evidence + +Pinned AVX2 wall-time runs on an AMD Ryzen 9 7950X found both `892717f30` and `4c936447a` about 25% +to 31% slower than develop for the tested take/filter list cases. The final push changes those +native medians by only 0% to 2%. + +Changing the bench profile from 16 codegen units to one did not remove that native gap. One +representative median pair was: + +| Profile | `4c936447a` | Develop | +| --- | ---: | ---: | +| 16 codegen units | 8.25 us | 6.41 us | +| One codegen unit | 7.86 us | 6.21 us | + +These measurements do not explain the CodSpeed simulation result. Do not use local wall time as a +proxy for CodSpeed CPU simulation. + +## Incomplete CodSpeed ablation + +Two `workflow_dispatch` runs were started and then canceled: + +- Framework only: [run `31289620637`]. +- Numeric RowFn: [run `31289622392`]. + +This approach was not sufficient. A workflow-dispatch run has no pull-request context, so it does +not update PR #9255's comment or create the PR comparison check needed for an inspectable result. +The framework array shard also reached an unrelated cancellation in +`take_slices_to_buffer_matrix`. Do not use either run as performance evidence. + +## Recommended next steps + +1. Open one affected `take_filter_list_*` benchmark in the existing `4c936447a` CodSpeed check. +2. Compare its differential flame graph with develop. Record executed instruction, cache, and + memory costs for the changed stack. +3. If the cost is extra instructions or a changed call path, follow that stack into assembly and + source. +4. If the cost is only instruction-cache placement, do not add arbitrary padding or unrelated + source edits. Determine whether a stable alignment or build-level remedy exists. +5. To locate the first bad revision, run focused `take_filter` simulations for `0a0ad0db1` and + `89fd28bc1` in a pull-request context. A dedicated temporary PR is less disruptive than moving + the head of PR #9255. Run only `cargo codspeed run --bench take_filter`. +6. If framework-only is clean and numeric RowFn is bad, compare those two profiles. If both are + clean, continue through `5c02036a2`, `a236e0b9d`, and `f4617a2b5`. +7. Recheck native wall time only after finding a CodSpeed cause. Keep the two result types labeled + separately. + +## Mixed-constant optimization + +Keep `4c936447a`. It fixes a real RowFn regression. + +For two varying inputs, `Args::varying` returns typed slices and selects the indexed lane source. +For an array plus a constant, one argument returns `None`, so the tuple returns `None`. Here, +`None` means "not every input varies," not "no input varies." The mixed loop reads the array at +`index` and the one-row constant at zero. + +The measured compiler requires the varying match and its length proof to remain inside the selected +owned-executor branch. Moving the proof through one shared `Option` helper made constant add and +subtract about 3.3 times slower. The branch-local form restored them. The semantic reason for the +source-placement sensitivity remains unknown. + +[CodSpeed check at `892717f30`]: https://github.com/vortex-data/vortex/runs/93169735961 +[CodSpeed check at `4c936447a`]: https://github.com/vortex-data/vortex/runs/93181527671 +[function alignment]: https://codspeed.io/docs/instruments/cpu/regression-causes#function-alignment +[run `31289620637`]: https://github.com/vortex-data/vortex/actions/runs/31289620637 +[run `31289622392`]: https://github.com/vortex-data/vortex/actions/runs/31289622392 diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md new file mode 100644 index 00000000000..19da17c9b39 --- /dev/null +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -0,0 +1,329 @@ + + + +# RowFn optimization guide + +## Performance model + +RowFn is fast when the hot loop contains only work that changes for each row. These operations can +stay outside the loop: + +- Array and dtype dispatch. +- Decoding and downcasts. +- Batch-constant detection. +- Input length validation. +- Output allocation and array construction. +- Validity policy. +- Rich error construction. +- Work derived only from constant operands. + +The design also gives LLVM concrete types and independent indexed lanes. A short row closure is not +enough by itself. The generic plumbing must disappear after monomorphization. + +## Optimization history + +### Stage 0: sink-only output + +The first shared executor required every row function to write through a sink. This model supported +runtime-shaped output, but it hid the independence of primitive output values. + +The checked primitive loop was slower for several wide integer types. Signed `i64` multiply was +about 29% slower than the baseline. Unsigned `u64` multiply was about 59% slower. + +### Stage 1: owned output + +The next design let the row closure return `(Output, Failure)`. Shared execution owned the final +store and reduced failure evidence. + +This change improved wide integer multiplication, but it did not give LLVM a simple input source. +For example, `i32` multiplication remained about 18% slower in the measured matrix. + +This stage proved that output ownership mattered. It also proved that output ownership alone was +not sufficient. + +### Stage 2: typed indexed input + +`IndexedElementTuple` added an all-varying source. A primitive pair becomes +`LaneZip<&[Left], &[Right]>`. Shared execution validates both lengths once and calls +`map_checked_into`. + +This stage restored varying and nullable multiplication to approximately baseline performance. It +also removed hot bounds checks from the inspected production monomorphs. + +The trait is separate from `ElementTuple`. Many element types do not have a contiguous source. +Stable Rust cannot combine a blanket fallback with a more specific primitive implementation +without specialization. + +### Stage 3: remove the `Output: Copy` bound + +The executor needs only one property from owned output: abandoning initialized spare capacity on +unwind must not leak a required destructor. `Output: Copy` was stronger than this property. + +On Rust 1.91.0 and LLVM 21.1.2, adding the public `Copy` bound changed the production `i32` checked +multiply monomorph from about 18.7 microseconds to about 29.9 microseconds. An inert marker bound +did not cause the loss. One codegen unit did not remove it. + +The selected design uses a compile-time `!needs_drop::()` assertion. It does not expose a +`Copy` bound that the executor does not need. + +The exact compiler mechanism remains unknown. Standalone reduced loops did not reproduce the +effect. The real trait, closure, vector, and monomorphization context was necessary. + +### Stage 4: preserve mixed-constant code placement + +Commit `5c02036a2` deduplicated length validation: + +```rust +let varying = Args::varying(&columns); +ensure_decoded_lengths(&columns, varying.as_ref(), row_count)?; + +if let Some(varying) = varying { + // All-varying loop. +} else { + // Mixed loop. +} +``` + +This source-only change made constant add and subtract about 3.3 times slower at that revision. It +did not change the all-varying cases. + +The selected form keeps the view and proof in the selected branch: + +```rust +if let Some(varying) = Args::varying(&columns) { + validate_varying_lengths(&varying, row_count)?; + // All-varying loop. +} else { + validate_mixed_lengths(&columns, row_count)?; + // Mixed loop. +} +``` + +This change restored constant add and subtract to about 9.2 microseconds. Constant `i32` multiply +returned to about 18.9 microseconds. The all-varying controls did not move. + +The source placement is a measured constraint for the current toolchain. Rust semantics do not +require it. The source ablation proves the performance relationship, but it does not identify the +LLVM pass that causes it. + +The sink executors retain the shared validator. Moving their proof into each branch did not improve +the cosine or spatial benchmarks. + +### Stage 5: typed tensor rows + +The old tensor row accessor repeated a ptype check and buffer downcast for every output row. The +new `TensorRows` representation performs these operations once during decode. + +Each row access uses a typed flat buffer, width, and stride. A constant-backed tensor uses stride +zero, so `index * stride` selects row zero without a branch. + +This representation makes the tensor inner loop ordinary slice arithmetic. It also keeps constant +input storage compact. + +### Stage 6: prepared tensor and spatial constants + +Prepared visits expose batch constants before the loop. Cosine similarity computes a constant norm +once. Spatial predicates compute constant bounding boxes and relation helpers once. + +This optimization does not require a new array kernel. The same row declaration handles both +constant and varying operands. + +## Source-placement constraints + +### Decode before the loop + +The `InputElement::decode` method must contain dtype checks, array execution, downcasts, and buffer +extraction. Calling these operations through `get` makes the loop pay batch work for every row. + +### Prepare before the loop + +`Args::constants` and the prepare closure run once after decode. The prepared value is borrowed by +the row closure. It must not be rebuilt for each row. + +### Validate lengths before the loop + +Unchecked input reads are sound only after each varying source proves that it contains +`row_count` rows. The output slice must also contain `row_count` slots. + +The validations must execute before the loop. A check in the loop keeps bounds control flow in the +hot path and can prevent bounds-check elimination. + +### Keep the owned varying proof in its branch + +The owned executor must not pass `Option<&VaryingColumns>` through the shared generic helper on the +measured toolchain. The option construction, proof, and consumer stay in one branch. + +This rule is intentionally narrow. Applying it to every executor adds duplication without measured +benefit. + +### Borrow sink rows once + +`sink.rows()` runs before the loop. The loop receives a stable row view instead of repeatedly +borrowing the sink object. This keeps the buffer descriptor and output shape invariant. + +### Keep rich errors cold + +The row closure computes a small failure word. A `#[cold]` and `#[inline(never)]` helper creates the +`VortexError` after the loop or on the immediate failure path. + +This arrangement prevents formatting, allocation, and error branches from entering successful +checked-arithmetic loops. + +### Use inlining evidence, not a blanket attribute + +The public wrappers use ordinary `#[inline]` only where a caller must see captured constants or a +small adapter. The implementation does not apply `#[inline(always)]` to checked arithmetic. + +The lane-kernel module contains small internal chunk helpers with stronger attributes. Those +helpers were measured as part of the pre-existing lane-kernel work. A new strong inlining attribute +requires separate assembly or benchmark evidence. + +## Why the loop can autovectorize + +The optimized all-varying primitive loop presents these facts to LLVM: + +1. The element types are concrete because `dispatch` selected `T` before execution. +2. The input sources are typed slices or a typed `LaneZip`. +3. Input and output lengths match. +4. Unchecked reads follow one pre-loop proof. +5. Each iteration reads and writes an independent row. +6. Failure combines with bitwise OR. +7. The closure is concrete and can inline into the loop. +8. Error construction and validity are outside the loop. + +The generated loop can use SIMD when LLVM has a legal and profitable lowering. Checked add and +small-width arithmetic often fit this model. + +The word _autovectorize_ must not describe every result. The inspected `i64` and `u64` widened +multiply loops remained scalar on x86. They recovered performance because RowFn matched the +handwritten scalar loop, not because LLVM found SIMD. + +The tensor outer loop returns one scalar for each tensor row. SIMD commonly appears in the inner +loop over each tensor slice. The outer RowFn loop does not need to vectorize across variable slice +references. + +## Rejected or incomplete alternatives + +### Keep every output behind a sink + +This model supports more output shapes, but it loses the independent owned-value contract that +primitive code generation needs. + +### Add a numeric `reduce_encoded` fast path + +This path recovered speed by duplicating shared null and constant policy inside the numeric +function. It made RowFn a slow fallback instead of making shared execution fast. + +### Add a numeric-specific visitor seam + +This design moved the same specialization into generic execution under a different name. It did +not establish a reusable capability for nonnumeric row functions. + +### Use safe zipped iterators + +The tested iterator forms caused 3x to 9x losses for narrow integer types. They did not preserve the +same indexed source shape across all monomorphs. + +### Depend on per-row bounds checks + +Unchecked access improved some cases, but it did not solve the original output and source-shape +problems. It also regressed some `u8` cases when applied without the final indexed design. + +### Scan output for failures + +The selected loop returns failure evidence directly. Scanning a finished output adds another pass +and cannot represent every error condition. + +### Use `Copy` as the no-drop proof + +`Copy` is stronger than required and triggered a measured compiler regression. The compile-time +no-drop assertion expresses the actual safety condition. + +### Apply branch-local validation to sinks + +This change did not improve cosine or spatial performance. The shared helper remains in those +paths. + +## Unrelated benchmark movement + +An unrelated benchmark can move after a RowFn source edit even when it never calls RowFn. The +source edit rebuilds `vortex-array` and the benchmark executable. This rebuild can change: + +- Codegen-unit partitioning. +- Inlining decisions in affected monomorphs. +- Function order and address alignment. +- Instruction-cache and decoded-instruction-cache set placement. +- Branch target placement. +- Linker layout of code that remains reachable through the shared session. + +These are code-generation dependencies, not semantic dependencies. + +[CodSpeed CPU simulation] measures executed instructions and models cache and memory access. It +can therefore report a different result when the instruction sequence or binary layout changes. +Local wall time can differ from the simulated ratio because it uses a real AMD processor instead +of the CodSpeed CPU model. + +CodSpeed documents [function alignment] as one reason an unchanged microbenchmark can move after +a rebuild. The correct diagnostic is the simulated instruction and cache counts in the +differential flame graph. + +An unrelated recovery does not prove that an algorithmic problem was fixed. The result is stable +only after source ablation, machine-code inspection, and repeated measurements agree on a cause. + +## Current `take_filter_list` evidence + +The [CodSpeed check at `4c936447a`] reports 31 regressions. Several `take_filter_list_*` +benchmarks are 14% to 16% slower than develop in CPU simulation. + +The [CodSpeed check at `892717f30`] already reported the same benchmarks as 15% to 16% slower. The +final mixed-constant fix did not bring them back. Most of their simulated times improved by less +than 2% between the two checks. The fix removed larger constant-arithmetic regressions, so the +unchanged take/filter entries became more prominent in the ordered report. + +Every retained RowFn CodSpeed summary from `0e5c19c00` through `4c936447a` that contains a +performance table also contains `take_filter_list_*` regressions. Some GitHub views show only the +20 largest changes, and the bot edits one current PR comment. Either behavior can make a persistent +regression appear to leave and return. + +The compared list, filter, and take source files are identical between develop and the branch. +The measured benchmark has no runtime call to RowFn. Therefore, the change is not an algorithmic +regression in list take or filter execution. + +AVX2 wall-time runs on CPU 4 provide a separate native-runtime observation: + +| Revision | Typical list/filter median | Difference from develop | +| --- | ---: | ---: | +| Develop `66d096b5d` | 6.2 to 7.0 us | Baseline | +| Before latest push `892717f30` | 7.9 to 8.8 us | About 25% to 31% slower | +| Latest push `4c936447a` | 8.0 to 8.9 us | About 25% to 31% slower | + +The latest push changes most local cases by only 0% to 2%. The branch already contains a native +wall-time gap before that push. This result does not explain the CodSpeed simulation result. + +Changing the bench profile from 16 codegen units to one did not remove the native gap. For one +representative case, the candidate and develop medians were 7.86 and 6.21 microseconds. The same +case measured 8.25 and 6.41 microseconds with 16 codegen units. + +The main filter-take and list-take function sizes are identical across the three AVX2 binaries. +Normalized disassembly of the list-take function has the same instructions. Relative addresses and +link layout differ. This native evidence points to linked-code layout or a called function outside +the compared symbol. It does not identify a specific cache or branch mechanism. The CodSpeed +differential flame graph and its instruction and cache counters are the correct evidence for the +simulation result. + +Do not fix this result with arbitrary padding or an unrelated source edit. Such a change can move +the report without removing the cause. + +## Current unresolved work + +- Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. +- Identify the linked-code cause of the list/filter wall-time gap. +- Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked + binary. +- Compare current CodSpeed flame graphs for list/filter and `envelope` against develop. +- Repeat the key results on a second compiler version before filing a compiler issue. + +[CodSpeed check at `4c936447a`]: https://github.com/vortex-data/vortex/runs/93181527671 +[CodSpeed check at `892717f30`]: https://github.com/vortex-data/vortex/runs/93169735961 +[CodSpeed CPU simulation]: https://codspeed.io/docs/instruments/cpu +[function alignment]: https://codspeed.io/docs/instruments/cpu/regression-causes#function-alignment diff --git a/research/rowfn-reconstruction/README.md b/research/rowfn-reconstruction/README.md new file mode 100644 index 00000000000..55d42c62cdf --- /dev/null +++ b/research/rowfn-reconstruction/README.md @@ -0,0 +1,95 @@ + + + +# RowFn reconstruction guide + +This guide explains the RowFn design without requiring access to its source. It records the type +model, execution model, performance constraints, implementation order, and benchmark procedure. +The goal is to let a new contributor reconstruct the branch and understand each unusual choice. + +The guide describes commit `4c936447a` on `ct/row-fn`. Its comparison revision is develop commit +`66d096b5d`. + +## Reading order + +1. Read [`HANDOFF.md`](HANDOFF.md) for the current branch state, corrected CodSpeed history, and + unfinished investigation. +2. Read [`DESIGN.md`](DESIGN.md) for the API, concrete input examples, null handling, failure + handling, and generated loop shape. +3. Read [`OPTIMIZATION.md`](OPTIMIZATION.md) for the performance history, source-placement + constraints, rejected designs, and current CodSpeed interpretation. +4. Read [`REPRODUCE.md`](REPRODUCE.md) to rebuild the implementation and repeat the experiments. + +These dated records contain the raw evidence behind this guide: + +- [`rowfn-x86-2026-08-07`](../rowfn-x86-2026-08-07/README.md) records the owned-output, indexed + source, `Copy`-bound, LLVM IR, assembly, and x86 experiments. +- [`rowfn-regressions-2026-08-08`](../rowfn-regressions-2026-08-08/README.md) records the branch + bisection, compiler-configuration matrix, and tensor, spatial, list, and compact benchmarks. +- [`NUMERIC_ROWFN_PLAN.md`](../../NUMERIC_ROWFN_PLAN.md) records the earlier Apple Silicon work and + the original numeric design alternatives. + +## Terms + +The guide uses these terms consistently: + +- A _batch_ is one invocation over zero or more equally sized arrays. +- A _row closure_ computes one logical result from one element of each input. +- A _varying input_ stores one decoded value for each logical row. +- A _batch constant_ stores one decoded value that every logical row reads. +- An _owned output_ returns one independent Rust value for each row. +- An _output sink_ gives the row closure a handle into batch-owned output state. +- A _dense loop_ visits all stored rows, including payloads behind nulls. +- A _valid-only loop_ visits only rows where every input is valid. +- _Failure evidence_ is a small value that the loop OR-reduces before it creates an error. +- A _semantic dependency_ means that the benchmark executes the changed code. +- A _code-generation dependency_ means that the rebuild changes machine code or layout without a + runtime call to the changed code. + +## Main conclusions + +- RowFn removes array dispatch, dtype dispatch, decoding, allocation, validity, and rich errors + from the hot row loop. +- Rust monomorphization gives the loop concrete input, output, closure, and failure types. +- Primitive all-varying inputs use a typed indexed source with one bounds proof before the loop. +- Mixed constant inputs use one branch per argument and row. Batch constants remain one-row + buffers and are not expanded. +- Prepared visits expose constant values once before the loop. Tensor norms and spatial bounding + boxes use this capability. +- Owned output and sink output are separate capabilities. One abstraction did not optimize both + use cases well. +- Deferred failure evidence keeps rich error construction outside the loop. It also lets batch + execution suppress failures that came only from null rows. +- Integer division uses immediate failure and an uninitialized sink. Division is expensive and + scalar, so deferred evidence does not preserve useful vectorization there. +- The mixed-constant owned loop is sensitive to one source placement with Rust 1.91.0 and LLVM + 21.1.2. The varying view and its length proof must remain in the selected branch. +- The current CodSpeed report still contains unrelated regressions. A changed result in an + unrelated benchmark is not evidence that RowFn changed its algorithm. + +## What “autovectorization” means here + +RowFn does not use explicit SIMD intrinsics. It presents LLVM with ordinary counted loops over +typed slices and independent output slots. This shape lets LLVM use SIMD when the operation and +target support it. + +Not every important result uses SIMD. The measured signed and unsigned 64-bit checked multiply +loops remain scalar on x86 because each lane needs a widened product. They still match the +handwritten baseline after the framework removes abstraction overhead. Tensor kernels often gain +SIMD inside each tensor row, rather than across RowFn output rows. + +The exact generated code is part of the contract for performance-sensitive paths. Benchmark +parity alone does not prove vectorization, and vector-shaped LLVM IR does not prove vector machine +instructions. + +## Future article structure + +The material supports two independent articles: + +1. The RowFn design: typed row declarations, planning through visitors, null policy, prepared + constants, and output capabilities. +2. The performance investigation: owned output, indexed sources, failure reduction, compiler + sensitivity, assembly inspection, and misleading unrelated benchmark movement. + +The dated records contain experiment details. This guide contains the stable explanatory model +that those articles can use. diff --git a/research/rowfn-reconstruction/REPRODUCE.md b/research/rowfn-reconstruction/REPRODUCE.md new file mode 100644 index 00000000000..2bd57c4e249 --- /dev/null +++ b/research/rowfn-reconstruction/REPRODUCE.md @@ -0,0 +1,380 @@ + + + +# RowFn reproduction guide + +This guide gives a new contributor enough information to rebuild RowFn and repeat its main +performance experiments. Read [`DESIGN.md`](DESIGN.md) before implementing the API. Read +[`OPTIMIZATION.md`](OPTIMIZATION.md) before changing a hot loop. + +## Recorded environment + +The final x86 measurements used this environment: + +- Candidate: `ct/row-fn` at `4c936447a`. +- Baseline: develop at `66d096b5d`. +- Rust: `rustc 1.91.0 (f8297e351 2025-10-28)`. +- LLVM: 21.1.2, as reported by `rustc -vV`. +- Host: AMD Ryzen 9 7950X, 16 cores and 32 hardware threads. +- Local benchmark CPU: hardware thread 4, selected with `taskset -c 4`. +- CodSpeed-compatible target feature: `RUSTFLAGS='-C target-feature=+avx2'`. +- Default bench profile: 16 codegen units and no LTO. + +Record the exact revisions, compiler, CPU, governor, and flags for every new run. A percentage +without this context is not reproducible. + +## Build order + +Implement the framework in this order. Each step has a correctness or performance control before +the next step adds another capability. + +### 1. Define decoded element types + +Create an `InputElement` trait with these associated types: + +- `Array`: the supported decoded array representation. +- `Value`: the value presented to a row closure. +- `Constant`: metadata extracted once for a batch constant. + +The trait decodes one array before execution and reads one logical row from that decoded form. It +also declares whether a dense loop is safe for values stored behind nulls. + +Start with primitive and Boolean elements. Do not add a hidden `scalar_at` call as a general +fallback. Such a call performs runtime dispatch in the hot loop. + +### 2. Compose elements into tuples + +Create an `ElementTuple` implementation for the arities that RowFn supports. Its decoded form must +distinguish two input shapes: + +```text +Varying(buffer with row_count values) +Constant(buffer with one value) +``` + +The tuple must provide: + +- Decoding for every input. +- Row lookup for mixed constant and varying inputs. +- Constant metadata for preparation. +- A validity mask for planning. + +Keep the one-value constant representation. Do not expand constants to `row_count` values. + +### 3. Add a typed all-varying source + +Add an indexed source capability for tuples whose values can be represented by contiguous typed +slices. For a primitive pair, its varying source is equivalent to: + +```rust +LaneZip<&[Left], &[Right]> +``` + +Validate every input length before the loop. The loop can then use unchecked indexed reads. The +single validation is both the safety proof and the condition that lets LLVM remove bounds checks. + +Keep this capability separate from the general tuple trait. Stable Rust cannot express a blanket +fallback plus a more specific primitive implementation without specialization. + +### 4. Define output capabilities + +Support two output models: + +1. An owned row value returned by the closure. +2. An output sink that lends a row handle to the closure. + +The owned executor allocates final storage and writes each returned value. It requires a +compile-time proof that abandoned initialized spare capacity does not contain a type with a +destructor. Use the existing no-drop assertion. Do not expose an unnecessary `Output: Copy` +bound. + +The uninitialized sink must make initialization a safe API invariant. Its row handle owns a +write-once token. Writing a value consumes the handle and returns a proof token. A successful +closure result must contain that token. This prevents safe code from reporting success without +initializing the output slot. + +### 5. Separate failure evidence from errors + +Represent common per-row failures with a small OR-reducible type. The loop returns failure +evidence, not a formatted `VortexError`. Convert the final evidence into an error outside the hot +loop with a cold, non-inlined helper. + +Keep immediate failure for operations such as integer division when that form measures better. +Do not assume that deferred failure always vectorizes or always wins. + +### 6. Add the visitor API + +Define visit methods for these independent capabilities: + +| Input preparation | Output | Failure | +| --- | --- | --- | +| None | Owned | None or deferred | +| None | Sink | None or immediate | +| Prepared constants | Owned | None or deferred | +| Prepared constants | Sink | None or immediate | + +The RowFn implementation declares one typed row operation. The execution visitor selects the loop +and null policy. A planning visitor obtains dtype and fallibility information without running the +row closure. + +### 7. Add batch planning and execution + +Planning records the output dtype, validity behavior, fallibility, and optional encoded rewrite. +Execution then: + +1. Decodes input arrays. +2. Computes conjoined validity. +3. Selects dense, dense-with-retry, valid-only, or filter-and-scatter execution. +4. Extracts constants and prepares batch state, when requested. +5. Runs the selected typed loop. +6. Builds the final array and validity. + +The closure used by a dense policy must be total for every stored lane value, including values +behind null rows. It must not panic or perform side effects for those values. + +### 8. Port primitive numeric functions first + +Primitive binary arithmetic gives the smallest useful performance matrix. Port wrapping, +checked, saturating, and division operations. Keep the previous implementation available as a +benchmark control until every shape is measured. + +Test at least these shapes: + +- Varying plus varying. +- Varying plus constant. +- Constant plus varying. +- Dense validity. +- Mixed validity. +- Checked success. +- Checked failure behind a null row. +- Checked visible failure. + +### 9. Add tensor and spatial row types + +Decode tensors into typed flat buffers with width and stride. Use stride zero for a constant +tensor. Do not repeat a ptype check or buffer downcast for every output row. + +Prepared tensor visits can compute a constant norm once. Prepared spatial visits can compute a +constant bounding box or relation helper once. These users prove that preparation is more than an +API placeholder. + +## Historical implementation map + +The branch history records useful intermediate designs. Recreate the final design from the steps +above, but use these commits to repeat an ablation or inspect why a design was rejected: + +| Commit | Purpose | +| --- | --- | +| `fef191df5` | Original RowFn framework | +| `ae099e890` | Initial executor and null-policy benchmarks | +| `b324f3e26` | First numeric RowFn port | +| `aebe3caf7` | First tensor port | +| `6c13e8516` | First spatial port | +| `0a0ad0db1` | Cleaned RowFn framework based on current develop | +| `89fd28bc1` | Owned primitive numeric execution | +| `59c4578ef` | Focused executor benchmarks | +| `5c02036a2` | Refined execution contracts and initial shared length check | +| `a236e0b9d` | Self-contained kernel arguments | +| `f4617a2b5` | Merge of the research and cleaned histories | +| `69607edb6` | Pre-loop bounds proofs for owned execution | +| `892717f30` | Typed tensor and spatial row access | +| `4c936447a` | Branch-local varying proof for mixed constants | + +The two histories before `f4617a2b5` are intentional. One preserves the original experiments. The +other preserves the cleaned implementation that was based on the latest develop revision. + +## Benchmark procedure + +### Choose the measurement before testing + +CodSpeed CPU simulation and local wall time answer different questions. Do not use one as a proxy +for the other. + +- Use the exact CodSpeed simulation workflow to reproduce a CodSpeed regression. Compare the + simulated instructions, cache costs, memory costs, and differential flame graph. +- Use a pinned local wall-time run to check native performance on that host. +- Treat agreement between the two as additional evidence. Do not require it. + +The repository workflow builds with AVX2 and runs `cargo codspeed run` in simulation mode. A +normal `cargo bench` invocation uses the wall-time compatibility runner and does not reproduce the +simulated metric. + +### Use isolated worktrees and target directories + +Build the baseline and candidate in separate worktrees. Give each build its own target directory. +This prevents one revision from reusing incompatible artifacts from another revision. + +```bash +git worktree add --detach /tmp/vortex-rowfn-base 66d096b5d +git worktree add --detach /tmp/vortex-rowfn-candidate 4c936447a + +RUSTFLAGS='-C target-feature=+avx2' \ + CARGO_TARGET_DIR=/tmp/rowfn-target-base \ + cargo bench -j 8 -p vortex-array --bench row_fn_executor --no-run + +RUSTFLAGS='-C target-feature=+avx2' \ + CARGO_TARGET_DIR=/tmp/rowfn-target-candidate \ + cargo bench -j 8 -p vortex-array --bench row_fn_executor --no-run +``` + +Build independent experiments in parallel. Run their benchmark binaries serially on the same +hardware thread. Parallel benchmark runs compete for caches and memory bandwidth. + +### Match CodSpeed compilation + +The repository bench profile uses the CodSpeed-relevant defaults: + +```text +codegen-units = 16 +lto = false +``` + +Set AVX2 explicitly for the local comparison: + +```bash +RUSTFLAGS='-C target-feature=+avx2' cargo bench -p vortex-array --bench take_filter --no-run +``` + +Test one codegen unit as a compiler ablation: + +```bash +RUSTFLAGS='-C target-feature=+avx2' \ + CARGO_PROFILE_BENCH_CODEGEN_UNITS=1 \ + cargo bench -p vortex-array --bench take_filter --no-run +``` + +The one-unit test does not emulate CodSpeed. It is only a compiler ablation. + +### Run CodSpeed simulation + +The CI workflow is the authoritative reproduction: + +```bash +RUSTFLAGS='-C target-feature=+avx2' \ + cargo codspeed build --features _test-harness -p vortex-array --profile bench +cargo codspeed run -m simulation +``` + +Local simulation requires `cargo-codspeed` and CodSpeed's Valgrind fork. A standard Valgrind +installation is not equivalent. If those tools are unavailable, dispatch the repository CodSpeed +workflow for the exact revision. Do not substitute a native timing run and label it CodSpeed. + +Use the CodSpeed benchmark page to compare the candidate with the same develop baseline. Inspect +the differential flame graph and record these values for the changed stack: + +- Simulated time. +- Executed instruction cost. +- Cache cost. +- Memory cost. +- Function self time and total time. + +### Pin a native benchmark process + +Find the generated executable under `target/release/deps`, then run it on one hardware thread: + +```bash +taskset -c 4 target/release/deps/row_fn_executor- \ + --bench --sample-count 100 --max-time 1 --color never +``` + +Run candidate and baseline in alternating order. Repeat a surprising result. Report medians and +the full range across repetitions. Label these results as native wall time. + +### Core benchmark set + +Use these commands to cover the framework and its migrated users: + +```bash +cargo bench -p vortex-array --bench row_fn_executor +cargo bench -p vortex-array --bench binary_ops +cargo bench -p vortex-array --bench take_filter +cargo bench -p vortex-array --bench compact +cargo bench -p vortex-tensor --bench cosine_similarity +cargo bench -p vortex-tensor --bench inner_product +cargo bench -p vortex-tensor --bench l2_norm +cargo bench -p vortex-spatial +``` + +Use benchmark name filters to keep each comparison focused. Record the exact filter with the +result. + +## Source ablation procedure + +When a small source edit causes a large result, do not infer a cause from the final diff. Use this +procedure: + +1. Keep compiler flags, target CPU, benchmark input, and toolchain fixed. +2. Change one source property. +3. Build into a new target directory. +4. Run the baseline and candidate serially on one CPU. +5. Inspect LLVM IR and final assembly for the production monomorph. +6. Revert the source property and confirm that the result returns. + +For the mixed-constant regression, the single property was the location of the varying-source +match and its length proof. Controls showed that all-varying execution did not move. + +Do not preserve a source edit only because an unrelated benchmark report improves. First prove +that the benchmark executes the changed path or that its machine-code change is stable and +understood. + +## Inspect generated code + +Build a focused crate with one codegen unit when you need readable LLVM IR or assembly: + +```bash +RUSTFLAGS='-C target-feature=+avx2' \ + CARGO_PROFILE_BENCH_CODEGEN_UNITS=1 \ + cargo rustc -p vortex-array --release --lib -- --emit=llvm-ir,asm +``` + +Search the emitted files for a concrete operation and type. Check these properties: + +- Array and dtype dispatch are outside the loop. +- The loop has no per-row bounds failure edge. +- The row closure is inlined. +- Failure evidence stays as a small value. +- Rich error construction is outside the loop. +- Vector instructions exist before claiming SIMD. + +For a linked benchmark binary, compare symbol sizes and disassembly: + +```bash +llvm-nm --demangle --print-size --size-sort target/release/deps/ > symbols.txt +llvm-objdump --demangle --disassemble-symbols='' \ + target/release/deps/ > symbol.asm +``` + +Normalize absolute addresses and relocation offsets before comparing instructions. Identical +instructions at different addresses still permit a layout-sensitive cache or branch result. + +## Correctness checks + +Run the narrow checks while iterating: + +```bash +cargo nextest run -p vortex-array +cargo test --doc -p vortex-array +cargo check -p vortex-array --benches +``` + +Run repository Rust checks before handing off code changes: + +```bash +cargo +nightly fmt --all +cargo clippy --all-targets --all-features +``` + +If cargo reports exactly `sccache: error: Operation not permitted`, rerun that command with +`RUSTC_WRAPPER=`. + +## Known limitations of the record + +- The host used a power-saving governor during some local runs. CPU pinning and repeated controls + reduce noise, but they do not replace a fixed-frequency benchmark host. +- `perf`, Samply, and local CodSpeed simulation were not available for the final take/filter + investigation. +- The current take/filter evidence identifies a linked-binary effect. It does not identify the + exact cache set, branch target, or called symbol that causes the wall-time gap. +- The exact cause of the public `Copy`-bound compiler regression remains unknown. +- Several early null-strategy and bytes-length benchmarks were research scaffolding and are not + part of the final API. From 61410ef211e6a793d4b6c7361403b2d44c617805 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 10:00:32 -0400 Subject: [PATCH 022/160] Avoid RowFn overhead when resetting list offsets Decode list offsets once and subtract the first offset in a typed loop. This removes the measured RowFn batch planning and decoding costs from small list conversions. Record the focused CodSpeed bisection and component counters. Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 128 +++++++++++++----- research/rowfn-reconstruction/OPTIMIZATION.md | 96 +++++++++++-- vortex-array/src/arrays/list/array.rs | 22 +-- .../src/arrays/listview/conversion.rs | 18 +++ 4 files changed, 212 insertions(+), 52 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 94244a2bc59..86dd46b71ce 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -3,25 +3,26 @@ # RowFn investigation handoff -This file records the exact state at the end of the 2026-08-09 investigation. Start with this -file, then read [`DESIGN.md`](DESIGN.md), [`OPTIMIZATION.md`](OPTIMIZATION.md), and +This file records the current state of the 2026-08-09 investigation. Start with this file, then +read [`DESIGN.md`](DESIGN.md), [`OPTIMIZATION.md`](OPTIMIZATION.md), and [`REPRODUCE.md`](REPRODUCE.md). ## Branch state - Branch: `ct/row-fn`. -- Code head before this documentation commit: `4c936447a`. +- Last RowFn code commit: `4c936447a`. +- Documentation head before the offsets fix: `bdf95a77e`. - Comparison revision: develop at `66d096b5d`. -- No RowFn code changed during this final investigation. -- The only intended new branch changes are the files in `research/rowfn-reconstruction`. +- The offsets fix does not change the RowFn API or implementation. -Two temporary remote refs exist for a future CodSpeed ablation: +Three temporary remote refs exist for the CodSpeed ablation: - `ct/row-fn-codspeed-framework` points to `0a0ad0db1`. - `ct/row-fn-codspeed-numeric` points to `89fd28bc1`. +- `ct/row-fn-codspeed-take-filter` is the head of temporary draft PR #9298. -The refs contain exact historical code. They do not contain the uncommitted focused-workflow edits -that were made only in temporary local worktrees. +The first two refs contain exact historical code. The third ref adds a PR-only workflow that runs +only `cargo codspeed run --bench take_filter`. ## Corrected CodSpeed history @@ -38,10 +39,55 @@ The latest push did not bring back the `take_filter_list_*` regressions. The PR bot edits one current comment, and GitHub displays only the 20 largest changes. These two details can make a persistent regression appear to leave and return. -## What is known about take/filter +## Verified take/filter cause The list, filter, and take source files are identical between develop and `4c936447a`. The -`take_filter_list` benchmark does not execute a RowFn operation. +benchmark still reaches RowFn through an indirect call: + +```text +take_filter + -> list_view_from_list + -> ListArrayExt::reset_offsets + -> binary(Sub) on offsets and the first offset + -> numeric RowFn +``` + +The differential profile therefore corrects the earlier claim that the benchmark does not execute +RowFn. `reset_offsets` creates a constant array and runs generic numeric subtraction. Numeric RowFn +adds batch planning, dispatch, argument decoding, and output reconciliation to this small operation. + +The representative benchmark is +`take_filter_list_small_uncached_random_mask_random_indices[256, 10]`. The current PR report gives +233.737 microseconds for develop and 280.793 microseconds for `bdf95a77e`. This is a 16.76% +regression. + +CodSpeed creates the downloadable callgraph in a separate profiling execution. Its total can +differ slightly from the aggregate report. The callgraph components are: + +| Revision | Instructions | Cache | Memory | Total | +| --- | ---: | ---: | ---: | ---: | +| Develop `66d096b5d` | 21.312 us | 83.443 us | 133.294 us | 238.050 us | +| RowFn `bdf95a77e` | 26.210 us | 104.293 us | 155.172 us | 285.675 us | +| Increase | 4.898 us | 20.850 us | 21.878 us | 47.626 us | + +The extra instructions and the changed stack rule out a cache-only layout explanation. Cache and +memory costs also increase, but they occur on newly executed RowFn work. + +The largest changed functions in the focused numeric profile are: + +| Function | Base self / total | Head self / total | +| --- | ---: | ---: | +| Old `execute_numeric_primitive` | 0.741 / 18.639 us | absent | +| RowFn `execute_numeric_primitive` | absent | 0.430 / 71.156 us | +| `Batch::execute` | absent | 1.033 / 49.972 us | +| `Batch::execute_dense` | absent | 0.634 / 45.781 us | +| `NumericBinary::dispatch` | absent | 1.316 / 45.736 us | +| `(A, B)::decode` | absent | 0.539 / 37.501 us | +| `ArgColumn::decode` | absent | 0.968 / 36.254 us | +| `list_view_from_list` | 3.543 / 79.144 us | 2.592 / 108.951 us | +| `Batch::new` | absent | 1.797 / 10.794 us | + +These totals are inclusive callgraph costs. A function can appear in more than one caller stack. The linked AVX2 benchmark binaries still differ. Native inspection found: @@ -51,12 +97,11 @@ The linked AVX2 benchmark binaries still differ. Native inspection found: - Normalized list-take disassembly has the same instructions. - Function addresses, relative call targets, and linked layout differ. -This evidence is consistent with a linked-layout effect or a changed callee outside the inspected -symbol. It does not prove which cache, branch, or callee causes the result. +That native inspection covered the large take and filter functions. It missed the changed numeric +callee reached during list offset normalization. CodSpeed documents [function alignment] as a reason unchanged microbenchmarks can move after a -rebuild. Its differential flame graph is the correct next source of evidence. Inspect the -instruction, cache, and memory components separately. +rebuild. That warning remains useful, but alignment is not the cause of this simulation regression. ## Native measurements are separate evidence @@ -75,7 +120,7 @@ representative median pair was: These measurements do not explain the CodSpeed simulation result. Do not use local wall time as a proxy for CodSpeed CPU simulation. -## Incomplete CodSpeed ablation +## Focused CodSpeed ablation Two `workflow_dispatch` runs were started and then canceled: @@ -83,26 +128,42 @@ Two `workflow_dispatch` runs were started and then canceled: - Numeric RowFn: [run `31289622392`]. This approach was not sufficient. A workflow-dispatch run has no pull-request context, so it does -not update PR #9255's comment or create the PR comparison check needed for an inspectable result. -The framework array shard also reached an unrelated cancellation in -`take_slices_to_buffer_matrix`. Do not use either run as performance evidence. +not create the needed comparison. Do not use either run as performance evidence. + +Draft PR [#9298] provides the required pull-request context. Its workflow builds and runs only the +`take_filter` benchmark. + +- [Focused framework check] at `0a0ad0db1`: 232.542 microseconds against 233.737 microseconds for + develop. This is a 0.51% improvement and CodSpeed classifies it as no change. +- [Focused numeric check] at `89fd28bc1`: 279.491 microseconds against 233.737 microseconds for + develop. This is a 16.37% regression. + +`89fd28bc1` is the first bad revision. It is the direct child of clean revision `0a0ad0db1`. + +The numeric revision's callgraph totals are 25.835 microseconds for instructions, 103.531 +microseconds for cache, and 154.728 microseconds for memory. Develop's totals are 21.312, 83.443, +and 133.294 microseconds. The total increases from 238.050 to 284.093 microseconds. + +## Focused fix + +`ListArrayExt::reset_offsets` now decodes offsets once and subtracts the first offset in a typed +loop. It no longer allocates a constant array or invokes the generic scalar-function path. + +The AVX2 release binary auto-vectorizes the benchmark's `u16` loop. The loop uses two packed +`psubw` instructions per iteration and processes 16 offsets. This is code-generation evidence, +not a local timing result. + +A new test covers nonzero `u16` offsets. The existing list and list-view tests cover other offset +types and conversion behavior. A push to PR #9255 is still required for CodSpeed validation. ## Recommended next steps -1. Open one affected `take_filter_list_*` benchmark in the existing `4c936447a` CodSpeed check. -2. Compare its differential flame graph with develop. Record executed instruction, cache, and - memory costs for the changed stack. -3. If the cost is extra instructions or a changed call path, follow that stack into assembly and - source. -4. If the cost is only instruction-cache placement, do not add arbitrary padding or unrelated - source edits. Determine whether a stable alignment or build-level remedy exists. -5. To locate the first bad revision, run focused `take_filter` simulations for `0a0ad0db1` and - `89fd28bc1` in a pull-request context. A dedicated temporary PR is less disruptive than moving - the head of PR #9255. Run only `cargo codspeed run --bench take_filter`. -6. If framework-only is clean and numeric RowFn is bad, compare those two profiles. If both are - clean, continue through `5c02036a2`, `a236e0b9d`, and `f4617a2b5`. -7. Recheck native wall time only after finding a CodSpeed cause. Keep the two result types labeled - separately. +1. Push the focused offsets fix to `ct/row-fn` so PR #9255 creates a CodSpeed comparison. +2. Verify the representative benchmark's report value and callgraph components. +3. Check the remaining `take_filter_list_*` cases for a consistent recovery. +4. Keep local wall time separate from CodSpeed CPU simulation. +5. Continue investigating the native wall-time gap only if it remains after the measured call path + is removed. ## Mixed-constant optimization @@ -121,5 +182,8 @@ source-placement sensitivity remains unknown. [CodSpeed check at `892717f30`]: https://github.com/vortex-data/vortex/runs/93169735961 [CodSpeed check at `4c936447a`]: https://github.com/vortex-data/vortex/runs/93181527671 [function alignment]: https://codspeed.io/docs/instruments/cpu/regression-causes#function-alignment +[#9298]: https://github.com/vortex-data/vortex/pull/9298 +[Focused framework check]: https://github.com/vortex-data/vortex/actions/runs/31316492455 +[Focused numeric check]: https://github.com/vortex-data/vortex/actions/runs/31316710479 [run `31289620637`]: https://github.com/vortex-data/vortex/actions/runs/31289620637 [run `31289622392`]: https://github.com/vortex-data/vortex/actions/runs/31289622392 diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 19da17c9b39..1a2757cb194 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -270,7 +270,7 @@ differential flame graph. An unrelated recovery does not prove that an algorithmic problem was fixed. The result is stable only after source ablation, machine-code inspection, and repeated measurements agree on a cause. -## Current `take_filter_list` evidence +## `take_filter_list` regression The [CodSpeed check at `4c936447a`] reports 31 regressions. Several `take_filter_list_*` benchmarks are 14% to 16% slower than develop in CPU simulation. @@ -286,8 +286,81 @@ performance table also contains `take_filter_list_*` regressions. Some GitHub vi regression appear to leave and return. The compared list, filter, and take source files are identical between develop and the branch. -The measured benchmark has no runtime call to RowFn. Therefore, the change is not an algorithmic -regression in list take or filter execution. +However, the benchmark reaches RowFn through code outside those files: + +```text +take_filter + -> list_view_from_list + -> ListArrayExt::reset_offsets + -> binary(Sub) on offsets and the first offset + -> numeric RowFn +``` + +The old implementation of `reset_offsets` used generic binary subtraction. It created a constant +array from the first offset. The numeric RowFn migration changed that generic call's implementation. + +### Differential simulation evidence + +For `take_filter_list_small_uncached_random_mask_random_indices[256, 10]`, the current PR report +measures 233.737 microseconds on develop and 280.793 microseconds on `bdf95a77e`. This is a 16.76% +regression. + +CodSpeed creates the downloadable callgraph during a separate profiling execution. Its absolute +total can differ slightly from the report aggregate. The component totals are: + +| Revision | Instructions | Cache | Memory | Total | +| --- | ---: | ---: | ---: | ---: | +| Develop `66d096b5d` | 21.312 us | 83.443 us | 133.294 us | 238.050 us | +| RowFn `bdf95a77e` | 26.210 us | 104.293 us | 155.172 us | 285.675 us | +| Increase | 4.898 us | 20.850 us | 21.878 us | 47.626 us | + +The profile contains extra executed instructions and a new call path. It does not support a +cache-only or alignment-only explanation. + +The focused numeric profile shows these self and inclusive function costs: + +| Function | Base self / total | Head self / total | +| --- | ---: | ---: | +| Old `execute_numeric_primitive` | 0.741 / 18.639 us | absent | +| RowFn `execute_numeric_primitive` | absent | 0.430 / 71.156 us | +| `Batch::execute` | absent | 1.033 / 49.972 us | +| `Batch::execute_dense` | absent | 0.634 / 45.781 us | +| `NumericBinary::dispatch` | absent | 1.316 / 45.736 us | +| `(A, B)::decode` | absent | 0.539 / 37.501 us | +| `ArgColumn::decode` | absent | 0.968 / 36.254 us | +| `list_view_from_list` | 3.543 / 79.144 us | 2.592 / 108.951 us | +| `Batch::new` | absent | 1.797 / 10.794 us | + +These inclusive costs overlap when functions call each other. They identify the changed stack. + +### First bad revision + +Temporary draft PR [#9298] runs only `cargo codspeed run --bench take_filter` in a pull-request +context. + +- The [focused framework check] at `0a0ad0db1` measures 232.542 microseconds. Develop measures + 233.737 microseconds, so CodSpeed classifies the 0.51% improvement as no change. +- The [focused numeric check] at `89fd28bc1` measures 279.491 microseconds. This is 16.37% slower + than develop. + +The two revisions are parent and child. Therefore, `89fd28bc1` is the first bad revision. + +The numeric revision's callgraph totals are 25.835 microseconds for instructions, 103.531 +microseconds for cache, and 154.728 microseconds for memory. Its total is 284.093 microseconds. + +### Focused remedy + +`ListArrayExt::reset_offsets` now decodes its offsets once. A typed loop subtracts the first offset +and builds the replacement primitive array. This removes the constant allocation, batch planning, +dispatch, argument decoding, and output reconciliation from this small internal operation. + +The AVX2 release binary auto-vectorizes the benchmark's `u16` subtraction. The generated loop has +two packed `psubw` operations and handles 16 offsets per iteration. No SIMD claim is made for the +other integer types without inspecting their machine code. + +This fix targets the measured changed call path. It does not add padding or unrelated structural +changes. Its local tests establish correctness only. A PR-context CodSpeed run must establish its +simulation effect. AVX2 wall-time runs on CPU 4 provide a separate native-runtime observation: @@ -306,24 +379,25 @@ case measured 8.25 and 6.41 microseconds with 16 codegen units. The main filter-take and list-take function sizes are identical across the three AVX2 binaries. Normalized disassembly of the list-take function has the same instructions. Relative addresses and -link layout differ. This native evidence points to linked-code layout or a called function outside -the compared symbol. It does not identify a specific cache or branch mechanism. The CodSpeed -differential flame graph and its instruction and cache counters are the correct evidence for the -simulation result. +link layout differ. The earlier inspection did not include the numeric callee in `reset_offsets`. -Do not fix this result with arbitrary padding or an unrelated source edit. Such a change can move -the report without removing the cause. +Do not fix unrelated movement with arbitrary padding or an unrelated source edit. Such a change can +move a report without removing a measured cause. ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. -- Identify the linked-code cause of the list/filter wall-time gap. +- Validate the offsets fix with PR-context CodSpeed simulation. +- Recheck the native list/filter wall-time gap after removing the measured call path. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked binary. -- Compare current CodSpeed flame graphs for list/filter and `envelope` against develop. +- Compare the current CodSpeed flame graph for `envelope` against develop. - Repeat the key results on a second compiler version before filing a compiler issue. [CodSpeed check at `4c936447a`]: https://github.com/vortex-data/vortex/runs/93181527671 [CodSpeed check at `892717f30`]: https://github.com/vortex-data/vortex/runs/93169735961 [CodSpeed CPU simulation]: https://codspeed.io/docs/instruments/cpu [function alignment]: https://codspeed.io/docs/instruments/cpu/regression-causes#function-alignment +[#9298]: https://github.com/vortex-data/vortex/pull/9298 +[focused framework check]: https://github.com/vortex-data/vortex/actions/runs/31316492455 +[focused numeric check]: https://github.com/vortex-data/vortex/actions/runs/31316710479 diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 419617c073c..f56e7a77bfc 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -6,6 +6,7 @@ use std::fmt::Formatter; use std::sync::Arc; use num_traits::AsPrimitive; +use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -26,17 +27,15 @@ use crate::array::TypedArrayRef; use crate::array::child_to_validity; use crate::array::validity_to_child; use crate::array_slots; -use crate::arrays::ConstantArray; use crate::arrays::List; use crate::arrays::ListArray; use crate::arrays::Primitive; -use crate::builtins::ArrayBuiltins; +use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::NativePType; use crate::legacy_session; use crate::match_each_integer_ptype; use crate::match_each_native_ptype; -use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; #[array_slots(List)] @@ -343,12 +342,17 @@ pub trait ListArrayExt: ListArraySlotsExt { .into_array(); } - let offsets = self.offsets(); - let first_offset = offsets.execute_scalar(0, ctx)?; - let adjusted_offsets = offsets.clone().binary( - ConstantArray::new(first_offset, offsets.len()).into_array(), - Operator::Sub, - )?; + let offsets = self.offsets().clone().execute::(ctx)?; + let adjusted_offsets = match_each_integer_ptype!(offsets.ptype(), |P| { + let offsets = offsets.as_slice::

(); + let first_offset = offsets[0]; + let adjusted = offsets + .iter() + .map(|offset| *offset - first_offset) + .collect::>(); + + PrimitiveArray::new(adjusted, Validity::NonNullable).into_array() + }); // SAFETY: By resetting the offsets we simply "shift" everything left and discard trailing garbage, so all invariants remain the same. Ok(unsafe { ListArray::new_unchecked(elements, adjusted_offsets, self.list_validity()) }) diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index f6b30b830c7..c3ca68236c5 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -350,6 +350,24 @@ mod tests { Ok(()) } + #[test] + fn test_list_to_listview_resets_nonzero_offsets() -> VortexResult<()> { + let elements = buffer![0i32, 1, 2, 3, 4].into_array(); + let offsets = buffer![2u16, 4, 5].into_array(); + let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?; + + let mut ctx = SESSION.create_execution_ctx(); + let list_view = list_view_from_list(list.clone(), &mut ctx)?; + + assert_arrays_eq!( + buffer![0u16, 2].into_array(), + list_view.offsets().clone(), + &mut ctx + ); + assert_arrays_eq!(list, list_view, &mut ctx); + Ok(()) + } + #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); From f9dfde730787c64071e11739550c7093470ce482 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 10:18:51 -0400 Subject: [PATCH 023/160] perf: reuse binary scalar function ID Reuse the registered Binary function ID for its internal numeric RowFn helper. Record the focused CodSpeed cause and the validated offsets result. Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 57 ++++++++++++++++--- research/rowfn-reconstruction/OPTIMIZATION.md | 41 +++++++++++-- research/rowfn-reconstruction/REPRODUCE.md | 7 ++- .../src/scalar_fn/fns/binary/numeric/row.rs | 5 +- 4 files changed, 90 insertions(+), 20 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 86dd46b71ce..13b740b5788 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -14,6 +14,7 @@ read [`DESIGN.md`](DESIGN.md), [`OPTIMIZATION.md`](OPTIMIZATION.md), and - Documentation head before the offsets fix: `bdf95a77e`. - Comparison revision: develop at `66d096b5d`. - The offsets fix does not change the RowFn API or implementation. +- `ct/row-fn-api` has one later optimization commit, `df8fcbe1a`. Three temporary remote refs exist for the CodSpeed ablation: @@ -149,20 +150,57 @@ and 133.294 microseconds. The total increases from 238.050 to 284.093 microsecon `ListArrayExt::reset_offsets` now decodes offsets once and subtracts the first offset in a typed loop. It no longer allocates a constant array or invokes the generic scalar-function path. -The AVX2 release binary auto-vectorizes the benchmark's `u16` loop. The loop uses two packed -`psubw` instructions per iteration and processes 16 offsets. This is code-generation evidence, -not a local timing result. +The AVX2 release binary auto-vectorizes every supported integer width. Each unrolled iteration has +two 128-bit packed subtracts: + +- `psubb` handles 32 `i8` or `u8` offsets. +- `psubw` handles 16 `i16` or `u16` offsets. +- `psubd` handles 8 `i32` or `u32` offsets. +- `psubq` handles 4 `i64` or `u64` offsets. + +Signed and unsigned monomorphs share machine code. This is code-generation evidence, not a local +timing result. A new test covers nonzero `u16` offsets. The existing list and list-view tests cover other offset -types and conversion behavior. A push to PR #9255 is still required for CodSpeed validation. +types and conversion behavior. + +The [offsets fix check] validates the change in CodSpeed CPU simulation. The representative case +measures 176.524 microseconds, compared with 233.737 microseconds on develop and 280.793 +microseconds before the fix. It changes from a 16.76% regression to a 32.41% improvement against +develop. All 14 `take_filter_list_*` cases improve by 25.61% to 35.54% against develop. + +The representative callgraph components after the fix are: + +| Revision | Instructions | Cache | Memory | Total | +| --- | ---: | ---: | ---: | ---: | +| Develop `66d096b5d` | 21.312 us | 83.443 us | 133.294 us | 238.050 us | +| Before fix `bdf95a77e` | 26.210 us | 104.293 us | 155.172 us | 285.675 us | +| Offsets fix `61410ef21` | 15.462 us | 59.031 us | 103.767 us | 178.259 us | + +The generic scalar-function stack is absent after the fix. The typed `reset_offsets` function +costs 0.933 microseconds self and 7.629 microseconds total. On develop, the old primitive numeric +function alone costs 0.741 microseconds self and 18.639 microseconds total. The larger reduction in +`list_view_from_list`, from 79.144 to 29.634 microseconds total, includes the lazy scalar-function +array and optimizer work removed by the direct operation. + +## Numeric helper ID + +The focused numeric profile also found 6.820 microseconds of new inclusive cost in +`CachedId::deref`. The new `vortex.numeric_binary` ID initializes during the measured call. +Develop's ID lookup costs 0.702 microseconds total. The numeric RowFn revision costs 7.522 +microseconds. + +`NumericBinary` is an internal helper for the registered `Binary` function. Commit `df8fcbe1a` on +`ct/row-fn-api` reuses `Binary`'s ID. This removes the second interner initialization and gives +errors the public function's name. It does not change the arithmetic loop or the public API. + +This is a first-execution cost, not a per-row cost. Its CodSpeed effect is not verified yet. ## Recommended next steps -1. Push the focused offsets fix to `ct/row-fn` so PR #9255 creates a CodSpeed comparison. -2. Verify the representative benchmark's report value and callgraph components. -3. Check the remaining `take_filter_list_*` cases for a consistent recovery. -4. Keep local wall time separate from CodSpeed CPU simulation. -5. Continue investigating the native wall-time gap only if it remains after the measured call path +1. Validate the internal numeric-helper ID change in a PR-context CodSpeed comparison. +2. Keep local wall time separate from CodSpeed CPU simulation. +3. Continue investigating the native wall-time gap only if it remains after the measured call path is removed. ## Mixed-constant optimization @@ -185,5 +223,6 @@ source-placement sensitivity remains unknown. [#9298]: https://github.com/vortex-data/vortex/pull/9298 [Focused framework check]: https://github.com/vortex-data/vortex/actions/runs/31316492455 [Focused numeric check]: https://github.com/vortex-data/vortex/actions/runs/31316710479 +[offsets fix check]: https://github.com/vortex-data/vortex/actions/runs/31317322594 [run `31289620637`]: https://github.com/vortex-data/vortex/actions/runs/31289620637 [run `31289622392`]: https://github.com/vortex-data/vortex/actions/runs/31289622392 diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 1a2757cb194..69d1c0206df 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -354,13 +354,41 @@ microseconds for cache, and 154.728 microseconds for memory. Its total is 284.09 and builds the replacement primitive array. This removes the constant allocation, batch planning, dispatch, argument decoding, and output reconciliation from this small internal operation. -The AVX2 release binary auto-vectorizes the benchmark's `u16` subtraction. The generated loop has -two packed `psubw` operations and handles 16 offsets per iteration. No SIMD claim is made for the -other integer types without inspecting their machine code. +The AVX2 release binary auto-vectorizes every integer width. Each unrolled iteration contains two +128-bit packed subtracts. `psubb` handles 32 offsets, `psubw` handles 16, `psubd` handles 8, and +`psubq` handles 4. Signed and unsigned monomorphs share their machine code. This fix targets the measured changed call path. It does not add padding or unrelated structural -changes. Its local tests establish correctness only. A PR-context CodSpeed run must establish its -simulation effect. +changes. + +The [offsets fix check] validates the result in CodSpeed CPU simulation. The representative case +measures 176.524 microseconds, compared with 233.737 microseconds on develop and 280.793 +microseconds before the fix. It changes from a 16.76% regression to a 32.41% improvement against +develop. All 14 `take_filter_list_*` cases improve by 25.61% to 35.54% against develop. + +The representative post-fix callgraph totals are 15.462 microseconds for instructions, 59.031 +microseconds for cache, and 103.767 microseconds for memory. Its total is 178.259 microseconds. +The generic scalar-function stack is absent. The typed `reset_offsets` path costs 0.933 +microseconds self and 7.629 microseconds total. `list_view_from_list` drops from 79.144 to 29.634 +microseconds total. + +This result is larger than a recovery to develop because develop also uses generic scalar-function +subtraction for this internal offset adjustment. The direct typed operation removes that older +overhead as well as the additional RowFn work. + +### Avoid a second ID for an internal helper + +The focused numeric profile shows another fixed cost. `CachedId::deref` increases from 0.702 to +7.522 microseconds inclusive. The new `vortex.numeric_binary` ID initializes inside the measured +call. + +`NumericBinary` is not registered. It executes the registered `Binary` operation's primitive path. +Commit `df8fcbe1a` on `ct/row-fn-api` therefore reuses `Binary`'s existing ID. This removes a second +interner initialization and makes internal errors name the public function. + +This change does not alter dispatch or the row loop. The cost occurs on first execution, so it is +separate from per-row vectorization. The 6.820-microsecond profile delta is evidence for the source +of the fixed cost. It is not a verified end-to-end improvement until CodSpeed measures the change. AVX2 wall-time runs on CPU 4 provide a separate native-runtime observation: @@ -387,7 +415,7 @@ move a report without removing a measured cause. ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. -- Validate the offsets fix with PR-context CodSpeed simulation. +- Validate the numeric-helper ID change with PR-context CodSpeed simulation. - Recheck the native list/filter wall-time gap after removing the measured call path. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked binary. @@ -401,3 +429,4 @@ move a report without removing a measured cause. [#9298]: https://github.com/vortex-data/vortex/pull/9298 [focused framework check]: https://github.com/vortex-data/vortex/actions/runs/31316492455 [focused numeric check]: https://github.com/vortex-data/vortex/actions/runs/31316710479 +[offsets fix check]: https://github.com/vortex-data/vortex/actions/runs/31317322594 diff --git a/research/rowfn-reconstruction/REPRODUCE.md b/research/rowfn-reconstruction/REPRODUCE.md index 2bd57c4e249..831cca3148e 100644 --- a/research/rowfn-reconstruction/REPRODUCE.md +++ b/research/rowfn-reconstruction/REPRODUCE.md @@ -256,8 +256,11 @@ cargo codspeed run -m simulation ``` Local simulation requires `cargo-codspeed` and CodSpeed's Valgrind fork. A standard Valgrind -installation is not equivalent. If those tools are unavailable, dispatch the repository CodSpeed -workflow for the exact revision. Do not substitute a native timing run and label it CodSpeed. +installation is not equivalent. If those tools are unavailable, push the exact revision to a +branch with an open pull request. That push gives CodSpeed the comparison context it needs. + +A plain `workflow_dispatch` run does not update a pull request's CodSpeed report. Do not use its +partial output as comparison evidence. Do not substitute a native timing run and label it CodSpeed. Use the CodSpeed benchmark page to compare the candidate with the same develop baseline. Inspect the differential flame graph and record these values for the changed stack: diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index e0b7a658d01..e286433d77f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -9,7 +9,6 @@ use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_err; -use vortex_session::registry::CachedId; use super::primitive::CheckedAdd; use super::primitive::CheckedArithmetic; @@ -29,6 +28,7 @@ use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::row::InitializedElement; use crate::scalar_fn::row::UninitElementSink; @@ -56,8 +56,7 @@ impl RowFn for NumericBinary { const FALLIBLE: bool = true; fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("vortex.numeric_binary"); - *ID + ScalarFnVTable::id(&Binary) } fn dispatch( From 7baa9fab743a57e68c1dd398c759ef91dac3bb28 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 10:34:44 -0400 Subject: [PATCH 024/160] perf: decode masked tensor values directly Dense RowFn execution owns input validity and restores it on the output. Decode a masked tensor from its child values so nullable tensor operations do not rebuild extension storage under the same mask. Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 34 ++++++++++++++++-- research/rowfn-reconstruction/OPTIMIZATION.md | 36 +++++++++++++++++-- vortex-tensor/src/scalar_fns/row.rs | 9 +++++ 3 files changed, 74 insertions(+), 5 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 13b740b5788..191a2b71eb1 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -194,11 +194,40 @@ microseconds. `ct/row-fn-api` reuses `Binary`'s ID. This removes the second interner initialization and gives errors the public function's name. It does not change the arithmetic loop or the public API. -This is a first-execution cost, not a per-row cost. Its CodSpeed effect is not verified yet. +This is a first-execution cost, not a per-row cost. The [numeric ID check] validates it: + +- `sub_i64_constant` improves from 675.849 to 670.968 microseconds. +- `CachedId::deref` drops from 5.327 to 0.376 microseconds total. +- `Id::new_static`, previously 3.723 microseconds total, disappears from the callgraph. +- CodSpeed still classifies the complete benchmark as no change against develop. The fixed 4.881 + microseconds is less than 1% of this operation. + +The take/filter control remains improved by 33.93% against develop. + +## Nullable tensor decode + +The current report has two remaining nullable tensor regressions at width 256. The differential +profile for `inner_product::nullable[256]` records these component increases: + +| Component | Develop | RowFn | Increase | +| --- | ---: | ---: | ---: | +| Instructions | 13.146 us | 14.378 us | 1.232 us | +| Cache | 62.165 us | 71.844 us | 9.679 us | +| Memory | 158.567 us | 186.622 us | 28.056 us | +| Total | 233.878 us | 272.845 us | 38.966 us | + +The floating-point row work is approximately unchanged. `TensorRow::decode` costs 33.553 +microseconds total, including 19.956 microseconds in `ArrayRef::mask`. The input is a `Masked` +tensor, but dense RowFn execution owns its validity and restores it on the output. Decoding the +child values directly avoids rebuilding extension storage under the same mask. + +The local `f64` inner-product loop is scalar-unrolled by four. It emits `mulsd` and `addsd` in the +source fold order, not packed floating-point SIMD. Reassociating this reduction could enable wider +SIMD, but it would change floating-point results. It is not a free RowFn code-generation change. ## Recommended next steps -1. Validate the internal numeric-helper ID change in a PR-context CodSpeed comparison. +1. Validate the masked tensor decode change in a PR-context CodSpeed comparison. 2. Keep local wall time separate from CodSpeed CPU simulation. 3. Continue investigating the native wall-time gap only if it remains after the measured call path is removed. @@ -224,5 +253,6 @@ source-placement sensitivity remains unknown. [Focused framework check]: https://github.com/vortex-data/vortex/actions/runs/31316492455 [Focused numeric check]: https://github.com/vortex-data/vortex/actions/runs/31316710479 [offsets fix check]: https://github.com/vortex-data/vortex/actions/runs/31317322594 +[numeric ID check]: https://github.com/vortex-data/vortex/actions/runs/31318131466 [run `31289620637`]: https://github.com/vortex-data/vortex/actions/runs/31289620637 [run `31289622392`]: https://github.com/vortex-data/vortex/actions/runs/31289622392 diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 69d1c0206df..971308b3303 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -387,8 +387,37 @@ Commit `df8fcbe1a` on `ct/row-fn-api` therefore reuses `Binary`'s existing ID. T interner initialization and makes internal errors name the public function. This change does not alter dispatch or the row loop. The cost occurs on first execution, so it is -separate from per-row vectorization. The 6.820-microsecond profile delta is evidence for the source -of the fixed cost. It is not a verified end-to-end improvement until CodSpeed measures the change. +separate from per-row vectorization. The [numeric ID check] validates the result: + +- `sub_i64_constant` improves from 675.849 to 670.968 microseconds. +- `CachedId::deref` drops from 5.327 to 0.376 microseconds total. +- `Id::new_static`, previously 3.723 microseconds total, disappears from the callgraph. +- CodSpeed still classifies the complete benchmark as no change against develop. The fixed 4.881 + microseconds is less than 1% of this operation. + +The take/filter control remains improved by 33.93% against develop. + +### Decode masked tensor values directly + +The report also shows 14.77% and 12.46% regressions for nullable width-256 inner product and L2 +norm. For `inner_product::nullable[256]`, the callgraph components are: + +| Component | Develop | RowFn | Increase | +| --- | ---: | ---: | ---: | +| Instructions | 13.146 us | 14.378 us | 1.232 us | +| Cache | 62.165 us | 71.844 us | 9.679 us | +| Memory | 158.567 us | 186.622 us | 28.056 us | +| Total | 233.878 us | 272.845 us | 38.966 us | + +The floating-point row work is approximately unchanged. `TensorRow::decode` costs 33.553 +microseconds total, including 19.956 microseconds in `ArrayRef::mask`. The benchmark passes a +`Masked` tensor. Dense RowFn execution owns that validity and restores it on the result, so the +tensor decoder can read the mask's child values directly. + +The linked `f64` inner-product loop is scalar-unrolled by four. It uses `mulsd` and `addsd` in the +source fold order, not packed floating-point SIMD. LLVM cannot reassociate the strict reduction. +Changing that order could enable wider SIMD, but it would change floating-point results and needs +an explicit numerical contract. AVX2 wall-time runs on CPU 4 provide a separate native-runtime observation: @@ -415,7 +444,7 @@ move a report without removing a measured cause. ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. -- Validate the numeric-helper ID change with PR-context CodSpeed simulation. +- Validate the masked tensor decode change with PR-context CodSpeed simulation. - Recheck the native list/filter wall-time gap after removing the measured call path. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked binary. @@ -430,3 +459,4 @@ move a report without removing a measured cause. [focused framework check]: https://github.com/vortex-data/vortex/actions/runs/31316492455 [focused numeric check]: https://github.com/vortex-data/vortex/actions/runs/31316710479 [offsets fix check]: https://github.com/vortex-data/vortex/actions/runs/31317322594 +[numeric ID check]: https://github.com/vortex-data/vortex/actions/runs/31318131466 diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 0d9bafb7740..55fc26ed3a2 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -10,7 +10,9 @@ use num_traits::Float; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::Masked; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::masked::MaskedArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; @@ -76,6 +78,13 @@ impl InputElement for TensorRow { } fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Dense batch execution owns the mask and restores it on the result. Decode the values + // directly so a nullable tensor does not rebuild its extension storage under that mask. + let array = match array.as_opt::() { + Some(masked) => masked.child().clone(), + None => array, + }; + let rows = array.len(); let list_size = validate_tensor_float_input(array.dtype())?.list_size() as usize; let ext: ExtensionArray = array.execute(ctx)?; From 7908685e3715f3ebd29700311987ac5a80cefa6c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 10:44:46 -0400 Subject: [PATCH 025/160] docs: record validated RowFn performance fixes Record the post-fix CodSpeed counters for take/filter, numeric ID initialization, and nullable tensor decoding. Document the remaining allocator-sensitive u8 multiplication result and floating-point reduction codegen. Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 40 +++++++++++++++---- research/rowfn-reconstruction/OPTIMIZATION.md | 35 +++++++++++++--- 2 files changed, 63 insertions(+), 12 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 191a2b71eb1..027933ba49e 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -13,8 +13,9 @@ read [`DESIGN.md`](DESIGN.md), [`OPTIMIZATION.md`](OPTIMIZATION.md), and - Last RowFn code commit: `4c936447a`. - Documentation head before the offsets fix: `bdf95a77e`. - Comparison revision: develop at `66d096b5d`. -- The offsets fix does not change the RowFn API or implementation. -- `ct/row-fn-api` has one later optimization commit, `df8fcbe1a`. +- Direct offsets fix: `61410ef21`. +- Numeric helper ID fix: `f9dfde730` on this branch and `df8fcbe1a` on `ct/row-fn-api`. +- Masked tensor decode fix: `7baa9fab7`. Three temporary remote refs exist for the CodSpeed ablation: @@ -216,18 +217,42 @@ profile for `inner_product::nullable[256]` records these component increases: | Memory | 158.567 us | 186.622 us | 28.056 us | | Total | 233.878 us | 272.845 us | 38.966 us | -The floating-point row work is approximately unchanged. `TensorRow::decode` costs 33.553 -microseconds total, including 19.956 microseconds in `ArrayRef::mask`. The input is a `Masked` -tensor, but dense RowFn execution owns its validity and restores it on the output. Decoding the -child values directly avoids rebuilding extension storage under the same mask. +The floating-point row work is approximately unchanged. Before the fix, `TensorRow::decode` costs +33.553 microseconds total. It spends 25.638 microseconds canonicalizing the masked extension. The +`ArrayRef::mask` node in this profile is Batch's expected output mask, not input decode. + +Dense RowFn execution owns input validity and restores it on the output. `TensorRow::decode` now +reads a `Masked` tensor's child values directly. The [masked tensor check] validates the change: + +- `inner_product::nullable[256]` improves from 270.674 to 247.710 microseconds. It changes from a + 14.77% regression to a 6.87% no-change result against develop. +- `l2_norm::nullable[256]` improves from 271.115 to 249.766 microseconds. It changes from a 12.46% + regression to a 4.98% no-change result against develop. +- `TensorRow::decode` drops from 33.553 to 6.801 microseconds total. +- Extension canonicalization under that decoder drops from 25.638 to 0.439 microseconds total. + +The post-fix inner-product callgraph totals are 12.537 microseconds for instructions, 64.096 +microseconds for cache, and 172.833 microseconds for memory. Its total is 249.466 microseconds. +The remaining difference from develop is memory cost, not extra executed instructions. The local `f64` inner-product loop is scalar-unrolled by four. It emits `mulsd` and `addsd` in the source fold order, not packed floating-point SIMD. Reassociating this reduction could enable wider SIMD, but it would change floating-point results. It is not a free RowFn code-generation change. +## Remaining `mul_u8_nonnull` regression + +The [numeric ID check] still reports `mul_u8_nonnull` as 12.74% slower than develop. Its callgraph +components increase by 1.149 microseconds for instructions, 6.147 microseconds for cache, and +18.411 microseconds for memory. The indexed loop's self cost is 69.973 microseconds on both sides. + +The RowFn run enters `mi_page_fresh_alloc`, which is absent on develop. Inclusive `__rust_alloc` +cost increases from 7.221 to 22.449 microseconds. The evidence points to allocator state or +benchmark-order sensitivity around the output allocation. It does not show a slower arithmetic +loop. Do not change the loop or add layout padding without an isolated allocator experiment. + ## Recommended next steps -1. Validate the masked tensor decode change in a PR-context CodSpeed comparison. +1. Isolate the allocator state before `mul_u8_nonnull` if its CodSpeed regression must be removed. 2. Keep local wall time separate from CodSpeed CPU simulation. 3. Continue investigating the native wall-time gap only if it remains after the measured call path is removed. @@ -254,5 +279,6 @@ source-placement sensitivity remains unknown. [Focused numeric check]: https://github.com/vortex-data/vortex/actions/runs/31316710479 [offsets fix check]: https://github.com/vortex-data/vortex/actions/runs/31317322594 [numeric ID check]: https://github.com/vortex-data/vortex/actions/runs/31318131466 +[masked tensor check]: https://github.com/vortex-data/vortex/actions/runs/31318825883 [run `31289620637`]: https://github.com/vortex-data/vortex/actions/runs/31289620637 [run `31289622392`]: https://github.com/vortex-data/vortex/actions/runs/31289622392 diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 971308b3303..8b24a50f698 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -409,16 +409,40 @@ norm. For `inner_product::nullable[256]`, the callgraph components are: | Memory | 158.567 us | 186.622 us | 28.056 us | | Total | 233.878 us | 272.845 us | 38.966 us | -The floating-point row work is approximately unchanged. `TensorRow::decode` costs 33.553 -microseconds total, including 19.956 microseconds in `ArrayRef::mask`. The benchmark passes a -`Masked` tensor. Dense RowFn execution owns that validity and restores it on the result, so the -tensor decoder can read the mask's child values directly. +The floating-point row work is approximately unchanged. Before the fix, `TensorRow::decode` costs +33.553 microseconds total. It spends 25.638 microseconds canonicalizing the masked extension. The +`ArrayRef::mask` node in this profile is Batch's expected output mask, not input decode. + +Dense RowFn execution owns input validity and restores it on the result. The tensor decoder now +reads a `Masked` tensor's child values directly. The [masked tensor check] validates the change: + +- `inner_product::nullable[256]` improves from 270.674 to 247.710 microseconds. It changes from a + 14.77% regression to a 6.87% no-change result against develop. +- `l2_norm::nullable[256]` improves from 271.115 to 249.766 microseconds. It changes from a 12.46% + regression to a 4.98% no-change result against develop. +- `TensorRow::decode` drops from 33.553 to 6.801 microseconds total. +- Extension canonicalization under that decoder drops from 25.638 to 0.439 microseconds total. + +The post-fix inner-product callgraph totals are 12.537 microseconds for instructions, 64.096 +microseconds for cache, and 172.833 microseconds for memory. Its total is 249.466 microseconds. +The remaining difference from develop is memory cost, not extra executed instructions. The linked `f64` inner-product loop is scalar-unrolled by four. It uses `mulsd` and `addsd` in the source fold order, not packed floating-point SIMD. LLVM cannot reassociate the strict reduction. Changing that order could enable wider SIMD, but it would change floating-point results and needs an explicit numerical contract. +### `mul_u8_nonnull` allocator path + +The [numeric ID check] still reports `mul_u8_nonnull` as 12.74% slower than develop. Its callgraph +components increase by 1.149 microseconds for instructions, 6.147 microseconds for cache, and +18.411 microseconds for memory. The indexed loop's self cost is 69.973 microseconds on both sides. + +The RowFn run enters `mi_page_fresh_alloc`, which is absent on develop. Inclusive `__rust_alloc` +cost increases from 7.221 to 22.449 microseconds. This points to allocator state or benchmark-order +sensitivity around the output allocation. It does not show a slower arithmetic loop. A focused +allocator-state experiment must precede any code or benchmark change. + AVX2 wall-time runs on CPU 4 provide a separate native-runtime observation: | Revision | Typical list/filter median | Difference from develop | @@ -444,7 +468,7 @@ move a report without removing a measured cause. ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. -- Validate the masked tensor decode change with PR-context CodSpeed simulation. +- Isolate allocator state before `mul_u8_nonnull` if its CodSpeed regression must be removed. - Recheck the native list/filter wall-time gap after removing the measured call path. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked binary. @@ -460,3 +484,4 @@ move a report without removing a measured cause. [focused numeric check]: https://github.com/vortex-data/vortex/actions/runs/31316710479 [offsets fix check]: https://github.com/vortex-data/vortex/actions/runs/31317322594 [numeric ID check]: https://github.com/vortex-data/vortex/actions/runs/31318131466 +[masked tensor check]: https://github.com/vortex-data/vortex/actions/runs/31318825883 From 309039fca4e426a1d8f49c4f512e66aa8c8e0582 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 15:59:12 -0400 Subject: [PATCH 026/160] docs: record native RowFn codegen findings Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/HANDOFF.md | 68 +++++++++++++++++-- research/rowfn-reconstruction/OPTIMIZATION.md | 38 ++++++++++- 2 files changed, 101 insertions(+), 5 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 027933ba49e..ffd2d043c0c 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -16,6 +16,15 @@ read [`DESIGN.md`](DESIGN.md), [`OPTIMIZATION.md`](OPTIMIZATION.md), and - Direct offsets fix: `61410ef21`. - Numeric helper ID fix: `f9dfde730` on this branch and `df8fcbe1a` on `ct/row-fn-api`. - Masked tensor decode fix: `7baa9fab7`. +- Cleaned `ct/row-fn-api` head: `6dd500f59`. + +The API branch was rewritten with an exact force-with-lease from seven commits to five: + +1. `71c3e7a58` adds the framework, refined contracts, and self-contained arguments. +2. `6e864bf8b` moves primitive numeric operators to RowFn and reuses `Binary`'s ID. +3. `41bb10143` adds focused executor benchmarks. +4. `266350488` removes validated input bounds checks. +5. `6dd500f59` restores mixed-constant performance. Three temporary remote refs exist for the CodSpeed ablation: @@ -107,6 +116,12 @@ rebuild. That warning remains useful, but alignment is not the cause of this sim ## Native measurements are separate evidence +For the rest of this investigation, pinned local x86 wall time is the primary acceptance signal. +CodSpeed remains useful for finding changed call paths and separating instruction, cache, and +memory costs, but a simulated microbenchmark movement is not by itself a reason to reject code +that has native parity or an improvement. Keep the two measurements labeled; neither predicts the +other. + Pinned AVX2 wall-time runs on an AMD Ryzen 9 7950X found both `892717f30` and `4c936447a` about 25% to 31% slower than develop for the tested take/filter list cases. The final push changes those native medians by only 0% to 2%. @@ -122,6 +137,50 @@ representative median pair was: These measurements do not explain the CodSpeed simulation result. Do not use local wall time as a proxy for CodSpeed CPU simulation. +### Primitive numeric matrix + +The cleaned API branch was compared with develop on an AMD Ryzen 9 7950X. Each Divan binary was +pinned to logical CPU 2 and used the TSC timer, 100 samples, and a 250-millisecond minimum time. +Five alternating runs covered 26 shared `binary_ops` cases. + +Before the mixed-constant fix, the varying cases were generally within 0% to 8.5% of develop. The +constant cases exposed a separate source-placement regression: + +| Benchmark | Develop | Before fix | Difference | +| --- | ---: | ---: | ---: | +| `add_i64_constant` | 8.369 us | 35.42 us | +323.2% | +| `sub_i64_constant` | 8.319 us | 36.19 us | +335.0% | +| `mul_i32_constant` | 26.43 us | 41.91 us | +58.6% | + +Commit `6dd500f59` keeps each length proof in the branch that consumes it. After the fix, +`add_i64_constant` measures 9.269 microseconds, `sub_i64_constant` measures 9.199 microseconds, and +`mul_i32_constant` measures 18.91 microseconds. The first two retain about 11% overhead; multiply +is 28.5% faster than develop. + +### `mul_u16_nonnull` code placement + +Ten one-second alternating runs isolate a stable native regression: + +| Binary | Median | Observed range | +| --- | ---: | ---: | +| Develop `66d096b5d` | 2.229 us | 2.229 to 2.239 us | +| Clean API `6dd500f59` | 2.809 us | 2.799 to 2.829 us | +| `-C llvm-args=-align-loops=64` diagnostic | 2.449 us | 2.439 to 2.499 us | + +The develop and RowFn steady-state loops have the same normalized instruction sequence: two +128-bit loads, `pmullw`, `pmulhuw`, failure accumulation, one store, and the loop branch. Both are +vectorized. Develop's loop starts 16 bytes into a cache line and fits in that line. The ordinary +RowFn loop starts 32 bytes into a line and crosses the boundary. + +The LLVM diagnostic did not force this loop to a 64-byte boundary. It changed the linked layout so +the loop starts 19 bytes into a line and fits. That recovers 0.360 microseconds of the 0.580 +microsecond gap, leaving the diagnostic binary 9.9% slower than develop. This is evidence that code +placement matters, but it is not a complete cause or a suitable global compiler flag. Do not add +padding or enable the hidden LLVM option as a production fix. + +Samply could not record this benchmark because `perf_event_paranoid` is 2 and the machine requires +1 or lower. The assembly comparison is available evidence; there is no sampled native profile. + ## Focused CodSpeed ablation Two `workflow_dispatch` runs were started and then canceled: @@ -252,10 +311,11 @@ loop. Do not change the loop or add layout padding without an isolated allocator ## Recommended next steps -1. Isolate the allocator state before `mul_u8_nonnull` if its CodSpeed regression must be removed. -2. Keep local wall time separate from CodSpeed CPU simulation. -3. Continue investigating the native wall-time gap only if it remains after the measured call path - is removed. +1. Use pinned, alternating local x86 runs for performance decisions and retain the raw per-run + medians. +2. Reduce the remaining `mul_u16_nonnull` native gap without relying on incidental padding. +3. Isolate allocator state before changing the `mul_u8_nonnull` loop. +4. Keep local wall time separate from CodSpeed CPU simulation. ## Mixed-constant optimization diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 8b24a50f698..8f0c457cc89 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -106,6 +106,11 @@ The source placement is a measured constraint for the current toolchain. Rust se require it. The source ablation proves the performance relationship, but it does not identify the LLVM pass that causes it. +Pinned local x86 measurements on an AMD Ryzen 9 7950X confirm that this is not only a CodSpeed +effect. Before the fix, constant `i64` add and subtract were 3.23 and 3.35 times slower than +develop. After the fix, they are about 11% slower. Constant `i32` multiply changes from 58.6% +slower than develop to 28.5% faster. + The sink executors retain the shared validator. Moving their proof into each branch did not improve the cosine or spatial benchmarks. @@ -270,6 +275,36 @@ differential flame graph. An unrelated recovery does not prove that an algorithmic problem was fixed. The result is stable only after source ablation, machine-code inspection, and repeated measurements agree on a cause. +### Native benchmark policy + +Pinned local x86 wall time is the primary performance acceptance signal for the remaining RowFn +work. Run separate copied binaries on the same logical CPU, alternate revision order, and report +the median of repeated run medians. Use enough minimum time to make a narrow result stable. + +CodSpeed simulation remains a diagnostic tool. Its instruction, cache, and memory components can +expose a changed stack that local wall time cannot explain. A CodSpeed-only movement does not +override native parity or improvement, and local wall time must not be presented as a prediction +of CodSpeed simulation. + +### Identical vector loops can retain a native gap + +`mul_u16_nonnull` is a useful counterexample to treating autovectorization as the end of the +investigation. Ten alternating one-second runs measure 2.229 microseconds on develop and 2.809 +microseconds on the cleaned API branch, a 26.0% native regression. + +Both hot loops contain the same normalized vector instructions. They load two 128-bit vectors, +execute `pmullw` and `pmulhuw`, combine the overflow evidence, store one vector, and branch. The +develop loop fits in one 64-byte cache line. The ordinary RowFn loop crosses a line boundary. + +A diagnostic build with `-C llvm-args=-align-loops=64` measures 2.449 microseconds. The option did +not align this loop to 64 bytes, but the resulting linked layout moved it wholly inside one cache +line. This recovers 62% of the gap while leaving a 9.9% difference from develop. + +This experiment supports front-end and code-placement sensitivity. It does not prove that line +crossing explains the complete regression. A hidden global LLVM option and source padding are not +stable remedies. The RowFn monomorph also contains all-varying and mixed shape branches in one +larger function, so entry and setup code remain candidates for the residual cost. + ## `take_filter_list` regression The [CodSpeed check at `4c936447a`] reports 31 regressions. Several `take_filter_list_*` @@ -468,7 +503,8 @@ move a report without removing a measured cause. ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. -- Isolate allocator state before `mul_u8_nonnull` if its CodSpeed regression must be removed. +- Explain the residual `mul_u16_nonnull` native gap after accounting for hot-loop placement. +- Isolate allocator state before changing the `mul_u8_nonnull` loop. - Recheck the native list/filter wall-time gap after removing the measured call path. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked binary. From 6b76c2ec4252d6f591f6cdc79fcc003b5bd0bfbb Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 16:15:35 -0400 Subject: [PATCH 027/160] docs: validate the offsets fix on native x86 Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/HANDOFF.md | 18 ++++++++++++++++++ research/rowfn-reconstruction/OPTIMIZATION.md | 13 +++++++++++++ 2 files changed, 31 insertions(+) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index ffd2d043c0c..f4586781e7f 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -181,6 +181,10 @@ padding or enable the hidden LLVM option as a production fix. Samply could not record this benchmark because `perf_event_paranoid` is 2 and the machine requires 1 or lower. The assembly comparison is available evidence; there is no sampled native profile. +Outlining the validated all-varying lane kernel behind `#[inline(never)]` did not change the result. +Ten runs measured 2.799 to 2.829 microseconds, the same range as the ordinary cleaned API binary. +Do not add this code movement; it does not isolate the residual cost. + ## Focused CodSpeed ablation Two `workflow_dispatch` runs were started and then canceled: @@ -243,6 +247,19 @@ function alone costs 0.741 microseconds self and 18.639 microseconds total. The `list_view_from_list`, from 79.144 to 29.634 microseconds total, includes the lazy scalar-function array and optimizer work removed by the direct operation. +PR [#9299] extracts this fix from RowFn. A pinned AVX2 comparison against its exact develop base +used separate binaries, logical CPU 2, the TSC timer, 100 samples, and a 500-millisecond minimum +time. Five alternating runs covered all 14 list benchmarks. Every median-of-run-medians improves: + +- The range is 27.9% to 33.3% faster. +- `take_filter_list_small_uncached_random_mask_random_indices[256, 10]` improves from 5.939 to + 4.059 microseconds, or 31.7%. +- The matching 768 case improves from 6.189 to 4.319 microseconds, or 30.2%. +- The smallest improvement is the nullable 768 case, from 6.419 to 4.629 microseconds, or 27.9%. + +This is native wall-time evidence that the extracted fix is worthwhile independently of the +CodSpeed result. + ## Numeric helper ID The focused numeric profile also found 6.820 microseconds of new inclusive cost in @@ -340,5 +357,6 @@ source-placement sensitivity remains unknown. [offsets fix check]: https://github.com/vortex-data/vortex/actions/runs/31317322594 [numeric ID check]: https://github.com/vortex-data/vortex/actions/runs/31318131466 [masked tensor check]: https://github.com/vortex-data/vortex/actions/runs/31318825883 +[#9299]: https://github.com/vortex-data/vortex/pull/9299 [run `31289620637`]: https://github.com/vortex-data/vortex/actions/runs/31289620637 [run `31289622392`]: https://github.com/vortex-data/vortex/actions/runs/31289622392 diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 8f0c457cc89..e303391455c 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -249,6 +249,13 @@ no-drop assertion expresses the actual safety condition. This change did not improve cosine or spatial performance. The shared helper remains in those paths. +### Outline the all-varying kernel + +Moving the validated all-varying lane kernel into a private `#[inline(never)]` helper does not +improve `mul_u16_nonnull`. Ten pinned runs remain between 2.799 and 2.829 microseconds, the same as +the ordinary cleaned API binary. The larger Rust function containing both argument shapes is not +by itself the residual cause. + ## Unrelated benchmark movement An unrelated benchmark can move after a RowFn source edit even when it never calls RowFn. The @@ -411,6 +418,11 @@ This result is larger than a recovery to develop because develop also uses gener subtraction for this internal offset adjustment. The direct typed operation removes that older overhead as well as the additional RowFn work. +PR [#9299] extracts the offset fix without RowFn. Five alternating native AVX2 runs against its +exact develop base improve all 14 list cases by 27.9% to 33.3%. The small uncached 256 case moves +from 5.939 to 4.059 microseconds, and its 768 counterpart moves from 6.189 to 4.319 microseconds. +The extraction is therefore a native win as well as a CodSpeed win. + ### Avoid a second ID for an internal helper The focused numeric profile shows another fixed cost. `CachedId::deref` increases from 0.702 to @@ -521,3 +533,4 @@ move a report without removing a measured cause. [offsets fix check]: https://github.com/vortex-data/vortex/actions/runs/31317322594 [numeric ID check]: https://github.com/vortex-data/vortex/actions/runs/31318131466 [masked tensor check]: https://github.com/vortex-data/vortex/actions/runs/31318825883 +[#9299]: https://github.com/vortex-data/vortex/pull/9299 From 8946803a21dc065abbc04e2c97095e48360db46c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 16:34:13 -0400 Subject: [PATCH 028/160] docs: compare RowFn with native CPU codegen Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/HANDOFF.md | 23 +++++++++++++++++++ research/rowfn-reconstruction/OPTIMIZATION.md | 17 ++++++++++++++ research/rowfn-reconstruction/REPRODUCE.md | 18 +++++++++++++++ 3 files changed, 58 insertions(+) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index f4586781e7f..ff67a57d175 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -185,6 +185,29 @@ Outlining the validated all-varying lane kernel behind `#[inline(never)]` did no Ten runs measured 2.799 to 2.829 microseconds, the same range as the ordinary cleaned API binary. Do not add this code movement; it does not isolate the residual cost. +Compiling both revisions with `-C target-cpu=native` reduces the gap. Ten alternating runs measure +2.259 to 2.269 microseconds on develop and 2.479 to 2.489 microseconds on the API branch. The +native difference is about 9.7%, not 26%. + +Both native loops use AVX-512. Develop handles 64 `u16` lanes per iteration with two ZMM vectors. +RowFn handles 128 lanes with four ZMM vectors. Both compute `vpmullw`, `vpmulhuw`, the failure OR, +and the output stores. The remaining difference is not lost autovectorization. + +Five alternating native runs across all 27 shared `binary_ops` cases give this shape: + +- Decimal arithmetic, integer division, comparisons, and nullable wide arithmetic are within 1%. +- Varying narrow integer operations are generally 4% to 12% slower. +- `mul_i64_nonnull` is 2.8% faster and `mul_u64_nonnull` is at parity. +- Constant `i64` add and subtract are 19.5% and 22.9% slower. +- Constant `i32` multiply is 22.5% slower. + +The mixed-constant native loops also use AVX-512 broadcasts and packed arithmetic. Their remaining +regressions are not scalar fallbacks. + +Replacing numeric dispatch's two-element `Vec` with a stack-backed borrowed view removes +an allocation but does not improve the repeated matrix. Do not keep that change without a smaller +benchmark that shows the allocation itself matters. + ## Focused CodSpeed ablation Two `workflow_dispatch` runs were started and then canceled: diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index e303391455c..a1e2eebf168 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -312,6 +312,23 @@ crossing explains the complete regression. A hidden global LLVM option and sourc stable remedies. The RowFn monomorph also contains all-varying and mixed shape branches in one larger function, so entry and setup code remain candidates for the residual cost. +With `-C target-cpu=native` on the Ryzen 9 7950X, develop measures 2.259 to 2.269 microseconds and +RowFn measures 2.479 to 2.489 microseconds. Native CPU targeting reduces the gap from 26.0% to +about 9.7%. + +Both native loops use AVX-512. Develop processes two ZMM vectors, or 64 `u16` lanes, per iteration. +RowFn processes four ZMM vectors, or 128 lanes. Both use packed low- and high-half multiply, +failure reduction, and packed stores. Autovectorization is intact; LLVM chose a different unroll +factor and the shared RowFn path retains additional batch setup. + +The complete native matrix shows the same distinction. Decimal arithmetic, integer division, and +most nullable wide cases are within 1%. Narrow varying integer cases are generally 4% to 12% +slower. Mixed-constant add, subtract, and multiply remain 19% to 23% slower even though their hot +loops use AVX-512 broadcasts and packed arithmetic. + +Changing numeric dispatch from a two-element `Vec` to a stack-backed borrowed argument +view does not improve repeated timings. Removing that allocation is not a measured remedy. + ## `take_filter_list` regression The [CodSpeed check at `4c936447a`] reports 31 regressions. Several `take_filter_list_*` diff --git a/research/rowfn-reconstruction/REPRODUCE.md b/research/rowfn-reconstruction/REPRODUCE.md index 831cca3148e..85fc2a7a101 100644 --- a/research/rowfn-reconstruction/REPRODUCE.md +++ b/research/rowfn-reconstruction/REPRODUCE.md @@ -220,6 +220,24 @@ RUSTFLAGS='-C target-feature=+avx2' \ Build independent experiments in parallel. Run their benchmark binaries serially on the same hardware thread. Parallel benchmark runs compete for caches and memory bandwidth. +### Match the native host + +Use the host CPU when native wall time is the acceptance signal: + +```bash +RUSTFLAGS='-C target-cpu=native' \ + CARGO_TARGET_DIR=/tmp/rowfn-native-base \ + cargo bench -j 8 -p vortex-array --bench binary_ops --no-run +``` + +Build the candidate into a different target directory with the same flags. Copy or retain both +executables, pin them to the same logical CPU, and alternate their run order. Record the compiler, +CPU model, flags, timer, sample count, minimum time, and every run median. + +This build answers how the code runs on that host. It does not match CodSpeed's AVX2 compilation. +For example, `target-cpu=native` enables AVX-512 on the Ryzen 9 7950X and reduces the measured +`mul_u16_nonnull` RowFn gap from 26.0% to about 9.7%. + ### Match CodSpeed compilation The repository bench profile uses the CodSpeed-relevant defaults: From 7bfbdde08d7f237d42bd1b93485273bf2ec6ac76 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 16:41:30 -0400 Subject: [PATCH 029/160] docs: separate RowFn setup from loop throughput Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/HANDOFF.md | 16 ++++++++++++++++ research/rowfn-reconstruction/OPTIMIZATION.md | 10 +++++++++- research/rowfn-reconstruction/REPRODUCE.md | 6 ++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index ff67a57d175..8c98a1a6247 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -208,6 +208,22 @@ Replacing numeric dispatch's two-element `Vec` with a stack-backed bor an allocation but does not improve the repeated matrix. Do not keep that change without a smaller benchmark that shows the allocation itself matters. +A 32-times-larger batch separates fixed setup from loop throughput. The benchmark-only ablation +changes `LEN` from 32,768 to 1,048,576 and keeps `target-cpu=native`: + +| Benchmark | Develop | RowFn | Difference | +| --- | ---: | ---: | ---: | +| `add_i64_constant` | 162.3 us | 163.0 us | +0.4% | +| `sub_i64_constant` | 162.5 us | 162.8 us | +0.2% | +| `mul_i32_constant` | 94.84 us | 92.99 us | -2.0% | +| `mul_u16_nonnull` | 61.04 us | 61.53 us | +0.8% | +| `add_i32_nonnull` | 121.8 us | 122.1 us | +0.2% | +| `mul_i64_nonnull` | 778.1 us | 749.4 us | -3.7% | + +The per-element loops have native parity or better at scale. The visible percentages at 32,768 +rows come primarily from fixed RowFn batch planning, dispatch, decode, and reconciliation costs. +Do not attribute them to failed autovectorization or slower arithmetic throughput. + ## Focused CodSpeed ablation Two `workflow_dispatch` runs were started and then canceled: diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index a1e2eebf168..0e3f3c62981 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -329,6 +329,14 @@ loops use AVX-512 broadcasts and packed arithmetic. Changing numeric dispatch from a two-element `Vec` to a stack-backed borrowed argument view does not improve repeated timings. Removing that allocation is not a measured remedy. +A benchmark-only 1,048,576-row ablation reduces the remaining differences to within 1% for +`mul_u16_nonnull`, `add_i32_nonnull`, and constant `i64` add and subtract. Constant `i32` multiply +is 2.0% faster than develop, and varying `i64` multiply is 3.7% faster. + +The large-batch result shows that RowFn preserves native per-element throughput. The percentages +in the 32,768-row microbenchmarks primarily measure fixed batch planning, dispatch, decode, and +output reconciliation. Optimize those costs as batch overhead; do not rewrite the vector loops. + ## `take_filter_list` regression The [CodSpeed check at `4c936447a`] reports 31 regressions. Several `take_filter_list_*` @@ -532,7 +540,7 @@ move a report without removing a measured cause. ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. -- Explain the residual `mul_u16_nonnull` native gap after accounting for hot-loop placement. +- Reduce fixed RowFn batch overhead if 32,768-row numeric calls are latency-critical. - Isolate allocator state before changing the `mul_u8_nonnull` loop. - Recheck the native list/filter wall-time gap after removing the measured call path. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked diff --git a/research/rowfn-reconstruction/REPRODUCE.md b/research/rowfn-reconstruction/REPRODUCE.md index 85fc2a7a101..a13549c3a38 100644 --- a/research/rowfn-reconstruction/REPRODUCE.md +++ b/research/rowfn-reconstruction/REPRODUCE.md @@ -334,6 +334,12 @@ procedure: For the mixed-constant regression, the single property was the location of the varying-source match and its length proof. Controls showed that all-varying execution did not move. +To distinguish fixed setup from per-row throughput, repeat a focused case with a much larger +`LEN`. Keep every other source property and build flag fixed. Compare both the percentage and the +absolute time difference. If a 32-times-larger batch reaches parity while the small batch moves, +investigate planning, dispatch, decode, allocation, and output construction before changing the +loop. + Do not preserve a source edit only because an unrelated benchmark report improves. First prove that the benchmark executes the changed path or that its machine-code change is stable and understood. From 2b2079dce2ba8da0d7eb684e7e231951760f7f81 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 16:50:24 -0400 Subject: [PATCH 030/160] docs: record current list offset benchmark Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 32 ++++++++++++------- research/rowfn-reconstruction/OPTIMIZATION.md | 14 +++++--- research/rowfn-reconstruction/REPRODUCE.md | 7 ++++ 3 files changed, 37 insertions(+), 16 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 8c98a1a6247..e31998d5edf 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -17,6 +17,7 @@ read [`DESIGN.md`](DESIGN.md), [`OPTIMIZATION.md`](OPTIMIZATION.md), and - Numeric helper ID fix: `f9dfde730` on this branch and `df8fcbe1a` on `ct/row-fn-api`. - Masked tensor decode fix: `7baa9fab7`. - Cleaned `ct/row-fn-api` head: `6dd500f59`. +- PR #9299 head measured locally: `d97e53e66`. The API branch was rewritten with an exact force-with-lease from seven commits to five: @@ -286,18 +287,25 @@ function alone costs 0.741 microseconds self and 18.639 microseconds total. The `list_view_from_list`, from 79.144 to 29.634 microseconds total, includes the lazy scalar-function array and optimizer work removed by the direct operation. -PR [#9299] extracts this fix from RowFn. A pinned AVX2 comparison against its exact develop base -used separate binaries, logical CPU 2, the TSC timer, 100 samples, and a 500-millisecond minimum -time. Five alternating runs covered all 14 list benchmarks. Every median-of-run-medians improves: - -- The range is 27.9% to 33.3% faster. -- `take_filter_list_small_uncached_random_mask_random_indices[256, 10]` improves from 5.939 to - 4.059 microseconds, or 31.7%. -- The matching 768 case improves from 6.189 to 4.319 microseconds, or 30.2%. -- The smallest improvement is the nullable 768 case, from 6.419 to 4.629 microseconds, or 27.9%. - -This is native wall-time evidence that the extracted fix is worthwhile independently of the -CodSpeed result. +PR [#9299] originally extracted this direct typed subtraction at `fa54891b`. Five alternating +native AVX2 runs found that superseded revision 27.9% to 33.3% faster than its exact develop base. +Do not attribute those numbers to the current PR implementation. + +The current PR head, `d97e53e66`, keeps the generic lazy subtraction in `reset_offsets`. It executes +the normalized offsets once in `list_view_from_list`, then uses the same primitive array to build +sizes and output offsets. A fresh pinned AVX2 comparison used separate binaries, logical CPU 2, the +TSC timer, 100 samples, and a 500-millisecond minimum time. Five alternating runs covered all 14 +list benchmarks. Every median-of-run-medians improves: + +- The range is 17.2% to 19.3% faster. +- `take_filter_list_small_uncached_random_mask_random_indices[256, 10]` improves from 5.909 to + 4.879 microseconds, or 17.4%. +- The matching 768 case improves from 6.169 to 5.109 microseconds, or 17.2%. +- The largest improvement is the small random 256 case, from 5.659 to 4.569 microseconds, or 19.3%. + +This is native wall-time evidence that executing and reusing the normalized offsets is worthwhile +independently of the CodSpeed result. It does not measure the same implementation as the direct +typed fix on `ct/row-fn`. ## Numeric helper ID diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 0e3f3c62981..e64ba652212 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -443,10 +443,16 @@ This result is larger than a recovery to develop because develop also uses gener subtraction for this internal offset adjustment. The direct typed operation removes that older overhead as well as the additional RowFn work. -PR [#9299] extracts the offset fix without RowFn. Five alternating native AVX2 runs against its -exact develop base improve all 14 list cases by 27.9% to 33.3%. The small uncached 256 case moves -from 5.939 to 4.059 microseconds, and its 768 counterpart moves from 6.189 to 4.319 microseconds. -The extraction is therefore a native win as well as a CodSpeed win. +PR [#9299] first extracted the direct typed offset fix at `fa54891b`. Five alternating native AVX2 +runs against its exact develop base improved all 14 list cases by 27.9% to 33.3%. That commit is no +longer the PR head, so those results describe only the superseded implementation. + +The current PR head, `d97e53e66`, leaves the generic lazy subtraction in `reset_offsets`. It +materializes that result once in `list_view_from_list`, then reuses the primitive offsets for both +sizes and output offsets. Five fresh alternating runs improve all 14 cases by 17.2% to 19.3%. The +small uncached 256 case moves from 5.909 to 4.879 microseconds, and its 768 counterpart moves from +6.169 to 5.109 microseconds. This implementation is also a native win, but it is distinct from the +direct typed fix measured in CodSpeed and retained on `ct/row-fn`. ### Avoid a second ID for an internal helper diff --git a/research/rowfn-reconstruction/REPRODUCE.md b/research/rowfn-reconstruction/REPRODUCE.md index a13549c3a38..c9c76f617d1 100644 --- a/research/rowfn-reconstruction/REPRODUCE.md +++ b/research/rowfn-reconstruction/REPRODUCE.md @@ -298,6 +298,13 @@ taskset -c 4 target/release/deps/row_fn_executor- \ --bench --sample-count 100 --max-time 1 --color never ``` +For the `take_filter` comparison in this record, the exact runner options were: + +```bash +taskset -c 2 target/release/deps/take_filter- \ + --bench take_filter_list --timer tsc --sample-count 100 --min-time 0.5 --color never +``` + Run candidate and baseline in alternating order. Repeat a surprising result. Report medians and the full range across repetitions. Label these results as native wall time. From ec44dbd41777e5063c36cc38d3fefd95c80f4e79 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 17:03:10 -0400 Subject: [PATCH 031/160] docs: validate native RowFn throughput Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 11 +++++++++++ research/rowfn-reconstruction/OPTIMIZATION.md | 10 ++++++++++ 2 files changed, 21 insertions(+) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index e31998d5edf..6f0810489f3 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -225,6 +225,13 @@ The per-element loops have native parity or better at scale. The visible percent rows come primarily from fixed RowFn batch planning, dispatch, decode, and reconciliation costs. Do not attribute them to failed autovectorization or slower arithmetic throughput. +The framework control reaches the same conclusion without the numeric wrapper. Five +`target-cpu=native` runs of `row_fn_executor` compare 65,536-row loops in one linked binary. The +hand-written sink median is 137.4 microseconds. Infallible owned RowFn execution is 138.8 +microseconds, and sink RowFn execution is 138.5 microseconds, both within 1%. Checked owned +execution is 141.9 microseconds, or 3.3% slower. The shared executor does not impose a large +steady-state throughput cost. + ## Focused CodSpeed ablation Two `workflow_dispatch` runs were started and then canceled: @@ -307,6 +314,10 @@ This is native wall-time evidence that executing and reusing the normalized offs independently of the CodSpeed result. It does not measure the same implementation as the direct typed fix on `ct/row-fn`. +The same five-run comparison with `-C target-cpu=native` improves every case by 16.0% to 19.6%. +The small uncached cases move from 6.159 to 4.979 microseconds and from 6.389 to 5.209 +microseconds. The improvement therefore survives the host's AVX-512 code generation. + ## Numeric helper ID The focused numeric profile also found 6.820 microseconds of new inclusive cost in diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index e64ba652212..d2cadb78ae1 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -337,6 +337,12 @@ The large-batch result shows that RowFn preserves native per-element throughput. in the 32,768-row microbenchmarks primarily measure fixed batch planning, dispatch, decode, and output reconciliation. Optimize those costs as batch overhead; do not rewrite the vector loops. +The `row_fn_executor` control isolates the framework in one linked binary. Across five +`target-cpu=native` runs at 65,536 rows, the hand-written sink median is 137.4 microseconds. +Infallible owned RowFn execution is 138.8 microseconds, and sink RowFn execution is 138.5 +microseconds. Checked owned execution is 141.9 microseconds. The infallible executor variants are +within 1% of the hand-written loop, while deferred overflow reduction retains about 3.3% overhead. + ## `take_filter_list` regression The [CodSpeed check at `4c936447a`] reports 31 regressions. Several `take_filter_list_*` @@ -454,6 +460,10 @@ small uncached 256 case moves from 5.909 to 4.879 microseconds, and its 768 coun 6.169 to 5.109 microseconds. This implementation is also a native win, but it is distinct from the direct typed fix measured in CodSpeed and retained on `ct/row-fn`. +With `-C target-cpu=native`, five more alternating runs improve every case by 16.0% to 19.6%. The +small uncached cases move from 6.159 to 4.979 microseconds and from 6.389 to 5.209 microseconds. +The optimization therefore remains effective under this host's AVX-512 code generation. + ### Avoid a second ID for an internal helper The focused numeric profile shows another fixed cost. `CachedId::deref` increases from 0.702 to From 07dbbf1456869b632afe25a4fc184b011537030c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 17:09:36 -0400 Subject: [PATCH 032/160] docs: record rejected validity fast path Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 5 +++++ research/rowfn-reconstruction/OPTIMIZATION.md | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 6f0810489f3..02beecf9ebd 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -209,6 +209,11 @@ Replacing numeric dispatch's two-element `Vec` with a stack-backed bor an allocation but does not improve the repeated matrix. Do not keep that change without a smaller benchmark that shows the allocation itself matters. +Skipping `Array::validity` for inputs whose dtype is non-nullable is also not a measured fast path. +Five focused native comparisons move non-nullable and constant cases by less than 1%. Nullable +controls move by a similar amount even though their executed logic is unchanged. Treat those +differences as linked-layout noise and keep the uniform validity fold. + A 32-times-larger batch separates fixed setup from loop throughput. The benchmark-only ablation changes `LEN` from 32,768 to 1,048,576 and keeps `target-cpu=native`: diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index d2cadb78ae1..48fb71c3d9a 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -329,6 +329,10 @@ loops use AVX-512 broadcasts and packed arithmetic. Changing numeric dispatch from a two-element `Vec` to a stack-backed borrowed argument view does not improve repeated timings. Removing that allocation is not a measured remedy. +Skipping each encoding's validity function when its dtype is non-nullable also moves focused +native cases by less than 1%. Nullable controls move by a similar amount without a call-path +change. This is linked-layout noise, not evidence for a second batch-planning path. + A benchmark-only 1,048,576-row ablation reduces the remaining differences to within 1% for `mul_u16_nonnull`, `add_i32_nonnull`, and constant `i64` add and subtract. Constant `i32` multiply is 2.0% faster than develop, and varying `i64` multiply is 3.7% faster. From 001d8934090a324f95fc25062d2a363709823d13 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 17:11:12 -0400 Subject: [PATCH 033/160] docs: update focused ablation branch state Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 02beecf9ebd..24c4ccf30e6 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -27,14 +27,14 @@ The API branch was rewritten with an exact force-with-lease from seven commits t 4. `266350488` removes validated input bounds checks. 5. `6dd500f59` restores mixed-constant performance. -Three temporary remote refs exist for the CodSpeed ablation: +Two temporary remote refs remain for the CodSpeed ablation: - `ct/row-fn-codspeed-framework` points to `0a0ad0db1`. - `ct/row-fn-codspeed-numeric` points to `89fd28bc1`. -- `ct/row-fn-codspeed-take-filter` is the head of temporary draft PR #9298. -The first two refs contain exact historical code. The third ref adds a PR-only workflow that runs -only `cargo codspeed run --bench take_filter`. +Both refs contain exact historical code. Temporary draft PR #9298 supplied the pull-request context +for the focused comparisons. It is now closed, and its `ct/row-fn-codspeed-take-filter` head branch +has been deleted. ## Corrected CodSpeed history From 232b51424961d3707c8cff7674c774e675b076c5 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 17:11:51 -0400 Subject: [PATCH 034/160] docs: focus remaining work on native evidence Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/OPTIMIZATION.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 48fb71c3d9a..642506a89ee 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -562,11 +562,12 @@ move a report without removing a measured cause. - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. - Reduce fixed RowFn batch overhead if 32,768-row numeric calls are latency-critical. - Isolate allocator state before changing the `mul_u8_nonnull` loop. -- Recheck the native list/filter wall-time gap after removing the measured call path. +- Choose between the direct typed offset fix and PR #9299's materialize-once design based on API + maintenance and correctness. Both are native wins, but they are different implementations. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked binary. -- Compare the current CodSpeed flame graph for `envelope` against develop. -- Repeat the key results on a second compiler version before filing a compiler issue. +- Repeat the key local results on a second x86 machine and compiler version before filing a + compiler issue. [CodSpeed check at `4c936447a`]: https://github.com/vortex-data/vortex/runs/93181527671 [CodSpeed check at `892717f30`]: https://github.com/vortex-data/vortex/runs/93169735961 From 13afa1d81a5a5ef1f8645c6f648754783832b352 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 17:20:49 -0400 Subject: [PATCH 035/160] docs: isolate the native u8 multiply gap Signed-off-by: "Connor Tsui" --- research/rowfn-reconstruction/HANDOFF.md | 9 ++++++++- research/rowfn-reconstruction/OPTIMIZATION.md | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 24c4ccf30e6..d8a98d19198 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -389,12 +389,19 @@ cost increases from 7.221 to 22.449 microseconds. The evidence points to allocat benchmark-order sensitivity around the output allocation. It does not show a slower arithmetic loop. Do not change the loop or add layout padding without an isolated allocator experiment. +The isolated native benchmark does not reproduce that allocator-order explanation. Ten +alternating `target-cpu=native` runs measure a 1.939-microsecond develop median and a +2.149-microsecond RowFn median, a stable 10.8% gap. Both hot loops execute the same normalized +64-lane AVX-512 sequence. Develop's loop target is 64-byte aligned; RowFn's is seven bytes into a +line. A global `-align-loops=64` diagnostic neither aligned this loop nor changed its timing, so it +does not prove an alignment cause. Native counters remain unavailable on this host. + ## Recommended next steps 1. Use pinned, alternating local x86 runs for performance decisions and retain the raw per-run medians. 2. Reduce the remaining `mul_u16_nonnull` native gap without relying on incidental padding. -3. Isolate allocator state before changing the `mul_u8_nonnull` loop. +3. Profile the isolated `mul_u8_nonnull` case on a host that permits native performance counters. 4. Keep local wall time separate from CodSpeed CPU simulation. ## Mixed-constant optimization diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 642506a89ee..3677d9cc16d 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -326,6 +326,13 @@ most nullable wide cases are within 1%. Narrow varying integer cases are general slower. Mixed-constant add, subtract, and multiply remain 19% to 23% slower even though their hot loops use AVX-512 broadcasts and packed arithmetic. +`mul_u8_nonnull` retains a stable 10.8% gap when run alone: 1.939 microseconds on develop and 2.149 +microseconds with RowFn across ten alternating runs. Both hot loops process 64 lanes with the same +normalized AVX-512 instructions. Develop's loop target is 64-byte aligned, while RowFn's is seven +bytes into a line. The global `-align-loops=64` diagnostic did not align this loop and did not +change the timing. This rules out local benchmark-order allocator state, but it does not establish +an alignment cause. + Changing numeric dispatch from a two-element `Vec` to a stack-backed borrowed argument view does not improve repeated timings. Removing that allocation is not a measured remedy. @@ -561,7 +568,7 @@ move a report without removing a measured cause. - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. - Reduce fixed RowFn batch overhead if 32,768-row numeric calls are latency-critical. -- Isolate allocator state before changing the `mul_u8_nonnull` loop. +- Profile the isolated `mul_u8_nonnull` case with native performance counters. - Choose between the direct typed offset fix and PR #9299's materialize-once design based on API maintenance and correctness. Both are native wins, but they are different implementations. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked From 1f822e7dbfa287badeb04994d028c766c5f3f25d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 18:10:14 -0400 Subject: [PATCH 036/160] Address RowFn review findings Signed-off-by: Connor Tsui --- vortex-array/benches/like.rs | 7 +++-- vortex-array/src/arrays/list/array.rs | 27 ++++++++++--------- .../src/scalar_fn/row/batch/execution.rs | 2 +- .../src/scalar_fn/row/batch/policy.rs | 3 ++- .../src/scalar_fn/row/execute/sink.rs | 4 +++ vortex-array/src/scalar_fn/row/types/sink.rs | 4 +-- vortex-spatial/src/scalar_fn/contains.rs | 4 +-- vortex-spatial/src/scalar_fn/execute.rs | 3 +-- vortex-tensor/src/scalar_fns/l2_norm.rs | 6 +++++ vortex-tensor/src/scalar_fns/mod.rs | 2 +- vortex-tensor/src/scalar_fns/row.rs | 18 ++++++++++--- 11 files changed, 51 insertions(+), 29 deletions(-) diff --git a/vortex-array/benches/like.rs b/vortex-array/benches/like.rs index e83fae69b28..657f44a9c51 100644 --- a/vortex-array/benches/like.rs +++ b/vortex-array/benches/like.rs @@ -136,11 +136,10 @@ fn like_per_row_distinct_patterns(bencher: Bencher) { bench_per_row_patterns(bencher, patterns); } -/// A distinct three-letter lowercase infix per row, so `ARRAY_SIZE` rows never repeat a pattern -/// while every pattern keeps the same shape and compiles the same way. +/// A distinct three-letter lowercase infix for each index below 26³. fn distinct_trigram(i: usize) -> String { - let letter = |shift: usize| char::from(b'a' + u8::try_from((i >> shift) % 26).unwrap()); - [letter(0), letter(5), letter(10)].iter().collect() + let letter = |place: usize| char::from(b'a' + u8::try_from((i / place) % 26).unwrap()); + [letter(1), letter(26), letter(26 * 26)].iter().collect() } #[divan::bench] diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index f56e7a77bfc..ed28761ab46 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -6,6 +6,7 @@ use std::fmt::Formatter; use std::sync::Arc; use num_traits::AsPrimitive; +use num_traits::Zero; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -269,10 +270,6 @@ impl ListData { Ok(()) } - // TODO(connor)[ListView]: Create 2 functions `reset_offsets` and `recursive_reset_offsets`, - // where `reset_offsets` is infallible. - // Also, `reset_offsets` can be made more efficient by replacing `sub_scalar` with a match on - // the offset type and manual subtraction and fast path where `offsets[0] == 0`. } pub trait ListArrayExt: ListArraySlotsExt { @@ -344,14 +341,20 @@ pub trait ListArrayExt: ListArraySlotsExt { let offsets = self.offsets().clone().execute::(ctx)?; let adjusted_offsets = match_each_integer_ptype!(offsets.ptype(), |P| { - let offsets = offsets.as_slice::

(); - let first_offset = offsets[0]; - let adjusted = offsets - .iter() - .map(|offset| *offset - first_offset) - .collect::>(); - - PrimitiveArray::new(adjusted, Validity::NonNullable).into_array() + let offset_values = offsets.as_slice::

(); + let first_offset = offset_values[0]; + if first_offset == P::zero() { + offsets.clone().into_array() + } else { + // ListData validation requires sorted offsets, so every offset is at least the + // first offset. + let adjusted = offset_values + .iter() + .map(|offset| *offset - first_offset) + .collect::>(); + + PrimitiveArray::new(adjusted, Validity::NonNullable).into_array() + } }); // SAFETY: By resetting the offsets we simply "shift" everything left and discard trailing garbage, so all invariants remain the same. diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index 62879e2d629..4c5eb51ff41 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -250,7 +250,7 @@ impl Batch { Ok(ResolvedMask::Mixed(valid)) } - /// Resolve validity, try unfiltered execution when worthwhile, then fall back to filtering. + /// Resolve validity, try unfiltered execution, then fall back to filtering. fn execute_valid_only( &self, kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs index bc63cf9ce21..1ea3a500baa 100644 --- a/vortex-array/src/scalar_fn/row/batch/policy.rs +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -4,6 +4,7 @@ //! Nullable execution strategies derived from a concrete row dispatch. use crate::dtype::DType; +use crate::dtype::Nullability; use crate::scalar_fn::ElementTuple; use crate::scalar_fn::SinkResult; @@ -20,7 +21,7 @@ impl BatchPlan { /// Return the output dtype widened with strict input nullability. pub fn result_dtype(&self, args: &[DType]) -> DType { let nullability = self.output_dtype.nullability() - | crate::dtype::Nullability::from(args.iter().any(DType::is_nullable)); + | Nullability::from(args.iter().any(DType::is_nullable)); self.output_dtype.with_nullability(nullability) } diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 6b4750a9e66..312dd2d2db4 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -112,6 +112,10 @@ where let AllOr::Some(valid) = valid.bit_buffer() else { vortex_bail!("execute_sink_valid_rows requires a mixed mask"); }; + vortex_ensure!( + valid.len() == row_count, + "the validity mask does not address exactly {row_count} rows", + ); { let mut rows = sink.rows(); diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index b4d5a4c578c..7cefccf3a4f 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -49,8 +49,8 @@ pub trait OutputSink: 'static + Sized { /// Proof that a successful row closure left its row handle initialized. /// - /// Use `()` for initialized row handles. A sink exposing uninitialized storage uses an - /// unforgeable token returned after initialization. + /// Use `()` for initialized row handles. A sink exposing uninitialized storage uses a distinct + /// token returned after initialization. type WriteToken: 'static; /// The dtype of the column this sink builds, given the function's input dtypes. diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index becea3f094a..3a6e54a7e4c 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -642,7 +642,7 @@ mod tests { /// Nullable geometry operands conjoin their validity before computing containment. #[test] - fn test_contains_nullable_geometries_conjoins_validity() -> VortexResult<()> { + fn contains_nullable_geometries_conjoins_validity() -> VortexResult<()> { let session = vortex_array::array_session(); let mut ctx = session.create_execution_ctx(); @@ -669,7 +669,7 @@ mod tests { /// Geometry types without a null-tolerant decode fall back to filtering valid rows. #[test] - fn test_contains_unsupported_geometry_falls_back_to_filter() -> VortexResult<()> { + fn contains_unsupported_geometry_falls_back_to_filter() -> VortexResult<()> { let session = vortex_array::array_session(); let mut ctx = session.create_execution_ctx(); diff --git a/vortex-spatial/src/scalar_fn/execute.rs b/vortex-spatial/src/scalar_fn/execute.rs index 836577ec26e..e1836e3ad65 100644 --- a/vortex-spatial/src/scalar_fn/execute.rs +++ b/vortex-spatial/src/scalar_fn/execute.rs @@ -8,7 +8,6 @@ mod unary; pub(crate) use unary::dispatch_unary; use vortex_array::ArrayRef; use vortex_array::scalar::Scalar; -use vortex_mask::Mask; /// A non-null operand presented to a geometry kernel. pub(crate) enum Operand { @@ -19,7 +18,7 @@ pub(crate) enum Operand { } /// Shared batch state presented to a null-propagating geometry kernel with `N` operands. -pub(crate) struct Execution { +pub(crate) struct Execution { /// Constant/column shape of each operand. pub(crate) operands: [Operand; N], /// Validity state required by the kernel. diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 8df7f670c0a..ef36b5dd39f 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -22,6 +22,7 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::UninitElementSink; use vortex_array::serde::ArrayChildren; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -100,6 +101,11 @@ impl RowFn for L2Norm { } let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); let (_, norms) = extract_normalized_children(input); + vortex_ensure!( + norms.dtype().is_primitive(), + "normalized norms must be primitive, got {}", + norms.dtype(), + ); vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); Ok(Some(norms)) } diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 706392d3b25..da9b8950e7a 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -6,7 +6,7 @@ pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; -pub mod row; +pub(crate) mod row; #[cfg(test)] mod tests; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 55fc26ed3a2..340cdc1ed92 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -89,12 +89,21 @@ impl InputElement for TensorRow { let list_size = validate_tensor_float_input(array.dtype())?.list_size() as usize; let ext: ExtensionArray = array.execute(ctx)?; let flat = extract_flat_elements(ext.storage_array(), list_size, ctx)?; + let list_size = flat.list_size(); + let stride = flat.row_stride(); + let elements = flat.into_buffer::(); + + debug_assert!(if stride == 0 { + elements.len() == list_size + } else { + stride == list_size && rows.checked_mul(stride) == Some(elements.len()) + }); Ok(TensorRows { + elements, rows, - list_size: flat.list_size(), - stride: flat.row_stride(), - elements: flat.into_buffer::(), + list_size, + stride, }) } @@ -124,7 +133,8 @@ impl InputElement for TensorRow { { let start = index * column.stride; - // SAFETY: the caller guarantees that `index` addresses a complete row. + // SAFETY: decode established one complete stored row for stride 0, or `rows` contiguous + // `list_size`-element rows otherwise. The caller guarantees `index < rows`. unsafe { std::slice::from_raw_parts( column.elements.as_slice().as_ptr().add(start), From 065a727b9648708c63ff2cd82ebd30b1a220a5f0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 18:10:36 -0400 Subject: [PATCH 037/160] docs: update RowFn review handoff Signed-off-by: Connor Tsui --- SCALAR_FN_HANDOFF.md | 40 +++++-------------- docs/strictness-and-validity-pushdown.typ | 2 +- research/rowfn-reconstruction/HANDOFF.md | 15 +++++++ research/rowfn-reconstruction/OPTIMIZATION.md | 15 +++++++ 4 files changed, 40 insertions(+), 32 deletions(-) diff --git a/SCALAR_FN_HANDOFF.md b/SCALAR_FN_HANDOFF.md index 10224704829..acce1c5c003 100644 --- a/SCALAR_FN_HANDOFF.md +++ b/SCALAR_FN_HANDOFF.md @@ -61,15 +61,6 @@ cargo bench -p vortex-spatial --bench envelope cargo bench -p vortex-spatial --bench predicate_bbox ``` -For the spatial PR, also run the branch-only `vortex-spatial` `null_strategies` diagnostic. It -forces branch-and-skip and filter-and-scatter for the measured nullable geometry shapes. Confirm -that automatic selection uses the faster mechanism for one costly decode at 50% survivors and for -two costly decodes at about 81% survivors. - -```bash -cargo bench -p vortex-spatial --bench null_strategies -``` - The public benchmark names are shared with `develop`, so cross-revision comparisons do not need a frozen benchmark-local implementation as their primary control. @@ -215,36 +206,25 @@ types: - `Dense` may execute over garbage behind nulls and masks afterward; - `DenseWithRetry` may execute densely, then retry valid rows when deferred evidence reports an error; and -- `ValidOnly { filtered_decode_cost }` guarantees that the row closure sees only valid rows. +- `ValidOnly` guarantees that the row closure sees only valid rows. An early-failing row or a decoder that is not dense-safe must use valid-only execution. A deferred kernel may use dense execution because it writes a legal provisional value for every row. If only garbage behind nulls reports an error, the valid-row retry discards it. -Valid-only execution has two mechanisms. Filter-and-scatter shrinks inputs before decoding. -Branch-and-skip decodes the original batch and visits set bits from the conjoined validity mask. A -sink that does not support skipped rows automatically falls back to filter-and-scatter. - -The selector needs more than a boolean "decode shrinks" flag. Every `InputElement` declares an -additive `FILTERED_DECODE_COST`, defaulting to zero. `ElementTuple` sums the costs across arguments: - -- cost 0 always prefers branch-and-skip; -- cost 1 prefers branch-and-skip at 50% or more surviving rows; and -- cost 2 or greater prefers branch-and-skip at 85% or more surviving rows. - -This distinction comes from the x86 measurement in #9128. One nullable geometry input at 50% nulls -favored branching, while two independently nullable geometry inputs at 10% nulls each, about 81% -survivors, favored filtering. OR-ing a per-argument flag loses exactly that distinction. +Valid-only execution first calls `reduce_encoded` on the original arrays. If reduction declines, +the executor tries branch-and-skip on the original batch. This path decodes values behind nulls and +visits the set bits from the conjoined validity mask. It requires null-tolerant input decoding and a +sink that supports skipped rows. -The values are still a coarse heuristic. There is no evidence yet to separate cost 2 from cost 3, -and the batch-size crossover has not been measured. `NullStrategy` remains only as a test-harness -seam for forcing a mechanism. Do not expose the private row policy as an author contract. +If branch-and-skip declines, filter-and-scatter shrinks the inputs before decoding. It then scatters +the output into a full-length nullable array. Authors declare local safety through their input and +result types. They do not select the mechanism or provide a decode-cost estimate. ## Performance and generated-code evidence The older Ryzen 9 7950X AVX-512 measurements remain the production-performance record in the -[#9128 follow-up](https://github.com/vortex-data/vortex/issues/9128#issuecomment-5151831802). They -also supplied the per-argument null-selection evidence above. +[#9128 follow-up](https://github.com/vortex-data/vortex/issues/9128#issuecomment-5151831802). The final API cleanup was checked separately against its parent, `53c51d803c`, by cross-compiling the optimized `row_fn_executor` benchmark for `x86_64-apple-darwin` with `target-cpu=x86-64-v3`. @@ -442,8 +422,6 @@ Use the IR gate for loop shape and a focused microbenchmark for anything the loo ## Remaining boundaries -- Complete the required x86 production and forced-null-strategy benchmark run above before treating - the thresholds or overall performance as settled. - Keep nullable outputs separate until the first real function can define the validity contract. - Do not add another sink composition abstraction. Put multiple builders in one custom sink. - Do not add a general runtime-shaped sink until a production function needs one. diff --git a/docs/strictness-and-validity-pushdown.typ b/docs/strictness-and-validity-pushdown.typ index d45d87a37c9..d15805d7653 100644 --- a/docs/strictness-and-validity-pushdown.typ +++ b/docs/strictness-and-validity-pushdown.typ @@ -226,7 +226,7 @@ dictionary values safe to evaluate. [infallible], [no legal evaluation errors], [speculative evaluation], [dense-safe], [bytes behind nulls may be read safely], - [`NullHandling::Dense`], + [`RowPolicy::Dense`], ) Representability is a type-level obligation: a strict `cast` with a pinned non-nullable return type diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index d8a98d19198..0529416b5f3 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -396,6 +396,21 @@ alternating `target-cpu=native` runs measure a 1.939-microsecond develop median line. A global `-align-loops=64` diagnostic neither aligned this loop nor changed its timing, so it does not prove an alignment cause. Native counters remain unavailable on this host. +## Zero-based list offsets + +The review follow-up adds an early return when `ListArray::reset_offsets` receives primitive +offsets that already start at zero. This reuses the executed offsets instead of copying and +subtracting zero from the complete buffer. + +An isolated `target-cpu=native` A/B used the same review edits on both sides. The control removed +only this early return. Three alternating runs on CPU 2 used the TSC timer, 100 samples, and a +0.5-second minimum per case. All 14 `take_filter_list_*` cases improve by 1.66% to 3.29%. +The small uncached cases move from 4.149 to 4.049 microseconds at 256 rows and from 4.379 to +4.299 microseconds at 768 rows. + +This result is native wall-time evidence for the early return. It is not CodSpeed simulation +evidence and does not explain earlier CodSpeed movement. + ## Recommended next steps 1. Use pinned, alternating local x86 runs for performance decisions and retain the raw per-run diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 3677d9cc16d..f0a703485a9 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -564,6 +564,21 @@ link layout differ. The earlier inspection did not include the numeric callee in Do not fix unrelated movement with arbitrary padding or an unrelated source edit. Such a change can move a report without removing a measured cause. +### Reuse zero-based list offsets + +The review follow-up adds the remaining fast path from the old `reset_offsets` TODO. When the +executed primitive offsets start at zero, `reset_offsets` now reuses that array. It does not copy +the complete offsets buffer to subtract zero. + +The native control contains every other review edit and removes only the early return. Three +alternating runs used `-C target-cpu=native`, CPU 2, the TSC timer, 100 samples, and a 0.5-second +minimum per case. The early return improves all 14 `take_filter_list_*` cases by 1.66% to 3.29%. +The small uncached 256 case moves from 4.149 to 4.049 microseconds. The matching 768 case moves +from 4.379 to 4.299 microseconds. + +This isolated result supports the code change, but it remains native wall-time evidence. It does +not provide CodSpeed instruction, cache, or memory counters. + ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. From 3dbe279bf945fd9891c55c2b38c24fca76fb0d8e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 18:59:22 -0400 Subject: [PATCH 038/160] Benchmark primitive comparison shapes Cover lane widths, equality, nullability, and both constant operand positions before routing primitive comparisons through RowFn. Signed-off-by: Connor Tsui --- vortex-array/benches/compare.rs | 115 ++++++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 4a399760dc2..9e6dd3e4e5b 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -4,6 +4,7 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use divan::counter::ItemsCount; use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; @@ -38,6 +39,7 @@ const ARRAY_SIZE: usize = 65_536; fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); bencher + .counter(ItemsCount::new(ARRAY_SIZE)) .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input @@ -49,6 +51,31 @@ fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { }); } +fn u8_array(offset: u8) -> ArrayRef { + (0u8..=u8::MAX) + .cycle() + .take(ARRAY_SIZE) + .map(|value| value.wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn i32_array(offset: i32) -> ArrayRef { + (0i32..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn u64_array(offset: u64) -> ArrayRef { + (0u64..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + fn bool_array(rng: &mut StdRng) -> ArrayRef { BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.5))).into_array() } @@ -87,6 +114,13 @@ fn float_array(rng: &mut StdRng) -> ArrayRef { .into_array() } +fn f32_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f32..1.0)) + .collect::>() + .into_array() +} + fn string_array(rng: &mut StdRng) -> ArrayRef { VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { let len = rng.random_range(1usize..24); @@ -153,6 +187,14 @@ fn compare_int_constant(bencher: Bencher) { bench_compare(bencher, arr, constant, Operator::Gte); } +#[divan::bench] +fn compare_int_constant_lhs(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let constant = ConstantArray::new(50_000_000i64, ARRAY_SIZE).into_array(); + let arr = int_array(&mut rng); + bench_compare(bencher, constant, arr, Operator::Gte); +} + #[divan::bench] fn compare_int_eq(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -161,6 +203,55 @@ fn compare_int_eq(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Eq); } +#[divan::bench] +fn compare_i32(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = i32_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_i32_constant(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = ConstantArray::new(1_000_000i32, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = u8_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8_constant(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = ConstantArray::new(127u8, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_constant(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = ConstantArray::new(1_000_000u64, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_eq(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Eq); +} + #[divan::bench] fn compare_float(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -169,6 +260,30 @@ fn compare_float(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Gte); } +#[divan::bench] +fn compare_float_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = float_array(&mut rng); + let arr2 = float_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_f32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_f32_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + #[divan::bench] fn compare_decimal(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); From 8128137cc63956098a744fd202e59ae23d618fff Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 18:59:29 -0400 Subject: [PATCH 039/160] Execute primitive comparisons with RowFn Use RowFn for primitive comparisons while retaining fused x86 bit-packing for the measured wide ordered cases where LLVM generates faster code. Signed-off-by: Connor Tsui --- .../src/scalar_fn/fns/binary/compare/mod.rs | 2 +- .../scalar_fn/fns/binary/compare/primitive.rs | 164 ++++++++---------- .../fns/binary/compare/primitive/columnar.rs | 120 +++++++++++++ .../primitive/operand.rs} | 19 +- vortex-array/src/scalar_fn/fns/binary/mod.rs | 1 - vortex-array/src/test_harness/trace/tests.rs | 8 + 6 files changed, 219 insertions(+), 95 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs rename vortex-array/src/scalar_fn/fns/binary/{primitive_operand.rs => compare/primitive/operand.rs} (78%) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index 0452f4a3156..2ef3d7b424e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -211,7 +211,7 @@ fn compare_arrays( ) .into_array()), DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx), - DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), + DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, ctx), DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 1247358dce1..3e6ee8023df 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,27 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays via bit-packing lane kernels. +//! Primitive comparison execution through [`RowFn`]. + +#[cfg(target_arch = "x86_64")] +mod columnar; +#[cfg(target_arch = "x86_64")] +mod operand; -use vortex_buffer::BitBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_bail; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::dtype::NativePType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::compare::collect_bits; -use crate::scalar_fn::fns::binary::compare::collect_zip_bits; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::RowVisitor; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; /// Compare two primitive arrays of the same [`PType`]. @@ -32,99 +33,78 @@ pub(super) fn compare_primitive( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - match_each_native_ptype!(ptype, |T| { - compare_primitive_typed::(lhs, rhs, op, nullability, ctx) - }) + #[cfg(target_arch = "x86_64")] + if use_columnar_comparison(lhs, rhs, op)? { + return columnar::compare_primitive(lhs, rhs, op, ctx); + } + + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + ScalarFnVTable::execute(&PrimitiveCompare, &op, &args, ctx) } -fn compare_primitive_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: CompareOperator, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let len = lhs.len(); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - if lhs.len() != rhs.len() { - vortex_bail!( - "compare operator requires equal lengths, got {} and {}", - lhs.len(), - rhs.len() - ); +/// Internal row execution for primitive comparison operators. +#[derive(Clone)] +struct PrimitiveCompare; + +impl RowFn for PrimitiveCompare { + type Options = CompareOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + ScalarFnVTable::id(&Binary) } - let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; - - let bits = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slices(lhs, rhs, op), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => compare_slice_constant(lhs, *rhs, op), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slice_constant(rhs, *lhs, op.swap()), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Unreachable through `execute_compare` (constant-constant is folded there), but - // cheap to answer anyway. - BitBuffer::full(apply_op(*lhs, *rhs, op), len) - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { - return Ok( - ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) - .into_array(), - ); - } - }; - - Ok(BoolArray::try_new(bits, validity)?.into_array()) -} + fn dispatch( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = + PType::try_from(args.first().ok_or_else(|| { + vortex_err!("a comparison operator takes two operands, got none") + })?)?; -#[inline(always)] -fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { - match op { - CompareOperator::Eq => lhs.is_eq(rhs), - CompareOperator::NotEq => !lhs.is_eq(rhs), - CompareOperator::Gt => lhs.is_gt(rhs), - CompareOperator::Gte => lhs.is_ge(rhs), - CompareOperator::Lt => lhs.is_lt(rhs), - CompareOperator::Lte => lhs.is_le(rhs), + match_each_native_ptype!(ptype, |T| { visit_compare::(*op, visitor) }) } } -fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { - // Dispatch the operator outside the lane loop so each instantiation vectorizes a single - // branch-free predicate. - match op { - CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)), - CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)), - CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), - CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), - CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), - CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), +#[cfg(target_arch = "x86_64")] +fn use_columnar_comparison( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, +) -> VortexResult { + if matches!(op, CompareOperator::Eq | CompareOperator::NotEq) { + return Ok(false); } + + let ptype = PType::try_from(lhs.dtype())?; + Ok(match ptype { + // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit + // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. + PType::I64 | PType::F64 => true, + // LLVM vectorizes varying u64 inputs, but not the mixed-constant RowFn loop. + PType::U64 => lhs.as_constant().is_some() || rhs.as_constant().is_some(), + _ => false, + }) } -fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { +fn visit_compare(op: CompareOperator, visitor: V) -> VortexResult +where + T: NativePType, + V: RowVisitor, +{ match op { - CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), - CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), - CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)), - CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)), - CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)), - CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), + CompareOperator::Eq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)), + CompareOperator::NotEq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| !lhs.is_eq(rhs)), + CompareOperator::Gt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_gt(rhs)), + CompareOperator::Gte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_ge(rhs)), + CompareOperator::Lt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_lt(rhs)), + CompareOperator::Lte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_le(rhs)), } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs new file mode 100644 index 00000000000..ccd7ffb8719 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fused comparison and bit-packing for wide x86 lanes. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use super::operand::PrimitiveOperand; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// Compare primitive operands with one fused comparison and bit-packing loop. +pub(super) fn compare_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match PType::try_from(lhs.dtype())? { + PType::I64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::U64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::F64 => compare_primitive_typed::(lhs, rhs, op, ctx), + ptype => vortex_bail!("columnar comparison is not selected for {ptype}"), + } +} + +fn compare_primitive_typed( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable()); + let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; + let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + let bits = match (&lhs, &rhs) { + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slices(lhs, rhs, op), + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => compare_slice_constant(lhs, *rhs, op), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slice_constant(rhs, *lhs, op.swap()), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => BitBuffer::full(apply_op(*lhs, *rhs, op), len), + (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) + .into_array(), + ); + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +#[inline(always)] +fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { + match op { + CompareOperator::Eq => lhs.is_eq(rhs), + CompareOperator::NotEq => !lhs.is_eq(rhs), + CompareOperator::Gt => lhs.is_gt(rhs), + CompareOperator::Gte => lhs.is_ge(rhs), + CompareOperator::Lt => lhs.is_lt(rhs), + CompareOperator::Lte => lhs.is_le(rhs), + } +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + } +} + +fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(lhs, |lhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_bits(lhs, |lhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_bits(lhs, |lhs: T| lhs.is_gt(rhs)), + CompareOperator::Gte => collect_bits(lhs, |lhs: T| lhs.is_ge(rhs)), + CompareOperator::Lt => collect_bits(lhs, |lhs: T| lhs.is_lt(rhs)), + CompareOperator::Lte => collect_bits(lhs, |lhs: T| lhs.is_le(rhs)), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs similarity index 78% rename from vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs rename to vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs index 71d1122fc79..55d81153b1f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Decoding shared by primitive binary operators. +//! Operand decoding for the fused primitive comparison path. use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -15,19 +15,33 @@ use crate::validity::Validity; /// A materialized primitive column, a non-null constant, or an all-null constant. pub(super) enum PrimitiveOperand { + /// A varying primitive column and its validity. Array { + /// The materialized values. values: Buffer, + + /// The validity of the values. validity: Validity, }, + + /// A non-null value repeated for every row. Constant { + /// The repeated value. value: T, + + /// The number of repeated rows. len: usize, + + /// The validity implied by the constant's dtype. validity: Validity, }, + + /// An all-null constant with this row count. Null(usize), } impl PrimitiveOperand { + /// Decode an operand once for the fused comparison loop. pub(super) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { if let Some(constant) = array.as_opt::() { return Ok( @@ -49,9 +63,11 @@ impl PrimitiveOperand { let array = array.clone().execute::(ctx)?; let validity = array.validity()?; let values = array.into_buffer::(); + Ok(Self::Array { values, validity }) } + /// Return the logical row count. pub(super) fn len(&self) -> usize { match self { Self::Array { values, .. } => values.len(), @@ -59,6 +75,7 @@ impl PrimitiveOperand { } } + /// Return the operand validity. pub(super) fn validity(&self) -> Validity { match self { Self::Array { validity, .. } => validity.clone(), diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 80faff20e0a..a5b9fe70539 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -43,7 +43,6 @@ mod compare; pub use compare::*; mod numeric; pub(crate) use numeric::*; -mod primitive_operand; use crate::scalar::NumericOperator; use crate::scalar::Scalar; diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index 98043caec13..fd9685dcb35 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -685,6 +685,14 @@ fn trace_compare_on_dict() -> VortexResult<()> { iter 0 current=vortex.dict(bool, len=5) builder_active=false ExecuteSlot slot=1 parent=vortex.dict(bool, len=5) child=vortex.binary(bool, len=3) iter 1 current=vortex.binary(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=5) From f9b223d266a3d3283373c10ba10d7d491bf03e32 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 19:07:08 -0400 Subject: [PATCH 040/160] Document primitive comparison RowFn results Record the local wall-time matrix, the wide ordered fallback, and the linked-layout sensitivity without treating the results as CodSpeed simulation evidence. Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/HANDOFF.md | 55 +++++++++++++++++++ research/rowfn-reconstruction/OPTIMIZATION.md | 31 +++++++++++ 2 files changed, 86 insertions(+) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 0529416b5f3..fc57f50a47a 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -411,6 +411,61 @@ The small uncached cases move from 4.149 to 4.049 microseconds at 256 rows and f This result is native wall-time evidence for the early return. It is not CodSpeed simulation evidence and does not explain earlier CodSpeed movement. +## Primitive comparison RowFn + +The primitive comparison port has two commits on `ct/row-fn`. The first expands `compare` with +lane-width, equality, nullability, and constant-operand cases. The second routes the primitive +comparison loop through RowFn. + +The local A/B used the default bench profile on an AMD Ryzen 9 7950X. Each run used the OS timer, +100 samples, a one-second minimum per case, and 65,536-row inputs. No benchmark measurements ran in +parallel. The baseline and final measurements used the same benchmark source. + +```bash +cargo bench -p vortex-array --bench compare -- \ + compare_i32 compare_u8 compare_int compare_float compare_u64 compare_f32 \ + --timer os --sample-count 100 --min-time 1 --color never +``` + +Representative medians are: + +| Case | Columnar baseline | Final | Change | +| --- | ---: | ---: | ---: | +| `compare_u8` | 39.07 us | 3.419 us | 91.2% faster | +| `compare_u8_constant` | 31.35 us | 3.349 us | 89.3% faster | +| `compare_i32` | 19.07 us | 7.419 us | 61.1% faster | +| `compare_f32` | 33.16 us | 14.40 us | 56.6% faster | +| `compare_int_eq` | 21.68 us | 19.47 us | 10.2% faster | +| `compare_u64` | 27.22 us | 24.33 us | 10.6% faster | +| `compare_float_eq` | 21.71 us | 19.40 us | 10.6% faster | +| `compare_int` | 27.14 us | 27.12 us | parity | +| `compare_float` | 49.31 us | 49.66 us | parity | +| `compare_u64_constant` | 23.18 us | 23.25 us | parity | + +Two baseline runs and two final runs covered the original matrix. Their medians remained within +1%. The extended `u64` and floating-point baseline used one run. The final extended matrix used +two runs. + +The direct RowFn experiment did not keep all cases. It made ordered `i64` 25% slower, nullable +ordered `i64` 28% slower, and ordered `f64` 11% slower. Constant ordered `u64` was 34% slower. +Equality remained faster at each measured wide type, and varying ordered `u64` improved by 11%. + +Packing 65,536 materialized `bool` values into a `BitBuffer` has a 570-nanosecond median. This is +less than 2% of the direct RowFn `i64` time. The wide ordered regression therefore comes from the +generated comparison loop, not the separate packing pass. + +The final x86 path keeps fused comparison and bit-packing for ordered `i64`, ordered `f64`, and +constant ordered `u64`. It uses RowFn for the other primitive shapes. The fallback only +instantiates columnar kernels for `i64`, `u64`, and `f64`. + +Pruning the eight unreachable fallback type instantiations moved the `compare_u8` median from +approximately 3.06 to 3.42 microseconds. The selected source path did not change. This is +consistent with native linked-layout sensitivity, but no normalized machine-code comparison was +performed for these two binaries. + +These measurements are local wall-time evidence. They contain no CodSpeed instruction, cache, or +memory counters and do not predict a CodSpeed simulation result. + ## Recommended next steps 1. Use pinned, alternating local x86 runs for performance decisions and retain the raw per-run diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index f0a703485a9..66f44800ee7 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -579,6 +579,37 @@ from 4.379 to 4.299 microseconds. This isolated result supports the code change, but it remains native wall-time evidence. It does not provide CodSpeed instruction, cache, or memory counters. +### Select primitive comparison output by measured code generation + +Primitive comparisons expose a second output trade-off. The owned RowFn path writes one `bool` per +row, then `OutputElement for bool` packs the values into a `BitBuffer`. The old columnar path fuses +the predicate and bit-packing loop. + +The separate pack is cheap on the current x86 host. Packing 65,536 values takes 570 nanoseconds. +The comparison loop determines the larger differences: + +- RowFn improves measured `u8`, `i32`, `f32`, equality, and varying `u64` cases by 10% to 92%. +- The fused path remains faster for ordered `i64`, ordered `f64`, and constant ordered `u64`. +- A direct RowFn port regresses those cases by 11% to 34%. + +Dispatch each operator to a separate RowFn closure. This keeps the operator match outside the row +loop and gives LLVM one predicate per monomorph. Do not move the operator match into the closure. + +On x86, select the fused path before RowFn planning for the measured wide ordered cases. A +`reduce_encoded` prototype recovered the loop but repeated planning and validity work. Nullable +`i64` remained 5.7% slower. Selecting at the primitive entry point restores parity. + +Keep the fallback instantiation set narrow. Only `i64`, `u64`, and `f64` can reach it, so a full +`match_each_native_ptype!` adds unused columnar monomorphs. Explicit dispatch avoids that code-size +cost. + +This pruning moved the local `u8` median from approximately 3.06 to 3.42 microseconds without +changing its selected source path. Treat this as layout sensitivity, not a loop regression, until +a normalized machine-code comparison shows otherwise. + +The benchmark source, commands, and representative medians are in `HANDOFF.md`. These results use +local wall time, not CodSpeed CPU simulation. + ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. From 9bed9c902f35657bb73e17a90c53b1f12619c40d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 19:09:43 -0400 Subject: [PATCH 041/160] Fix RowFn research spell check Use the valid seven-character tensor-port revision because Typos parses the longer hash suffix as a misspelled word. Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/REPRODUCE.md | 2 +- research/rowfn-regressions-2026-08-08/README.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/research/rowfn-reconstruction/REPRODUCE.md b/research/rowfn-reconstruction/REPRODUCE.md index c9c76f617d1..4eb0419486d 100644 --- a/research/rowfn-reconstruction/REPRODUCE.md +++ b/research/rowfn-reconstruction/REPRODUCE.md @@ -168,7 +168,7 @@ above, but use these commits to repeat an ablation or inspect why a design was r | `fef191df5` | Original RowFn framework | | `ae099e890` | Initial executor and null-policy benchmarks | | `b324f3e26` | First numeric RowFn port | -| `aebe3caf7` | First tensor port | +| `aebe3ca` | First tensor port | | `6c13e8516` | First spatial port | | `0a0ad0db1` | Cleaned RowFn framework based on current develop | | `89fd28bc1` | Owned primitive numeric execution | diff --git a/research/rowfn-regressions-2026-08-08/README.md b/research/rowfn-regressions-2026-08-08/README.md index ce4661953ac..79672241f62 100644 --- a/research/rowfn-regressions-2026-08-08/README.md +++ b/research/rowfn-regressions-2026-08-08/README.md @@ -255,7 +255,7 @@ Commit history isolates when it appears: | --- | ---: | ---: | ---: | | Framework only, `fef191df5` | 42.52 us | 44.52 us | 33.73 us | | Numeric RowFn port, `b324f3e26` | 58.02 us | 59.72 us | 49.11 us | -| Before geo RowFn, `aebe3caf7` | 58.43 us | 59.99 us | 49.50 us | +| Before geo RowFn, `aebe3ca` | 58.43 us | 59.99 us | 49.50 us | The regression therefore predates the geo visitor conversion. The `envelope.rs` source is unchanged. It appears when numeric RowFn code is linked into the benchmark binary. From ef3fc1c2b1ec0b282e555c1269fd9420788090cc Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 19:19:40 -0400 Subject: [PATCH 042/160] Document comparison benchmark simulation results Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/HANDOFF.md | 11 +++++++++++ research/rowfn-reconstruction/OPTIMIZATION.md | 6 ++++++ 2 files changed, 17 insertions(+) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index fc57f50a47a..493c5d71b56 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -466,6 +466,17 @@ performed for these two binaries. These measurements are local wall-time evidence. They contain no CodSpeed instruction, cache, or memory counters and do not predict a CodSpeed simulation result. +The full CodSpeed workflow for `9bed9c9` completed successfully on all nine CPU shards. Its PR +report compares against `66d096b`, because CodSpeed had no successful run for the newer develop +head. It reports 36 improvements, 45 regressions, and 30 new benchmarks. The regressions include +unrelated expression, FastLanes, compact, and file benchmarks, while the local comparison A/B above +is at parity or faster for every selected production path. This disagreement is CodSpeed simulation +evidence, not native wall-time evidence. The report does not expose instruction, cache, or memory +counters in the PR comment, so it does not establish a cause for those movements. + +- [CodSpeed workflow](https://github.com/vortex-data/vortex/actions/runs/31341241599) +- [CodSpeed PR report](https://github.com/vortex-data/vortex/pull/9255#issuecomment-5211040550) + ## Recommended next steps 1. Use pinned, alternating local x86 runs for performance decisions and retain the raw per-run diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 66f44800ee7..5f54b45d60e 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -610,6 +610,12 @@ a normalized machine-code comparison shows otherwise. The benchmark source, commands, and representative medians are in `HANDOFF.md`. These results use local wall time, not CodSpeed CPU simulation. +The completed CodSpeed run for `9bed9c9` moved many benchmarks outside this comparison path. Its PR +report has 36 improvements and 45 regressions, including expression, FastLanes, compact, and file +benchmarks. It also fell back to `66d096b` rather than the newer develop head. Without the simulated +instruction, cache, and memory counters, this broad movement cannot distinguish changed work from +linked-layout costs. Do not use it to override the focused native A/B above. + ## Current unresolved work - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. From 443aed0b99d67378f1f562024faf2625c7b8eac3 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 20:56:06 -0400 Subject: [PATCH 043/160] Make uninitialized RowFn writes explicit Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/DESIGN.md | 22 ++++++++++++------- .../src/scalar_fn/fns/binary/numeric/row.rs | 3 ++- vortex-array/src/scalar_fn/row/types/sink.rs | 21 +++++++++++++----- vortex-spatial/src/scalar_fn/contains.rs | 3 ++- vortex-spatial/src/scalar_fn/distance.rs | 5 ++++- vortex-spatial/src/scalar_fn/intersects.rs | 3 ++- .../src/scalar_fns/cosine_similarity.rs | 11 ++++++---- vortex-tensor/src/scalar_fns/inner_product.rs | 5 ++++- vortex-tensor/src/scalar_fns/l2_norm.rs | 3 ++- vortex-tensor/src/scalar_fns/tests/row.rs | 3 ++- 10 files changed, 54 insertions(+), 25 deletions(-) diff --git a/research/rowfn-reconstruction/DESIGN.md b/research/rowfn-reconstruction/DESIGN.md index 01297ba1e77..5a073bce363 100644 --- a/research/rowfn-reconstruction/DESIGN.md +++ b/research/rowfn-reconstruction/DESIGN.md @@ -91,7 +91,7 @@ arrays and runs the matching loop. ## Visit capabilities -The visitor has six entry points. Three unprepared methods delegate to three prepared methods. +The visitor has six entry points. | Method | Output model | Row error model | Preparation | | --- | --- | --- | --- | @@ -348,17 +348,23 @@ trait OutputSink { The executor borrows `Rows` once before the loop. This keeps the sink descriptor and shape as loop invariants. The closure receives only the row handle. -`UninitElementSink` avoids zero-initializing dense primitive output. Its row handle is -`&mut MaybeUninit`. Safe code must prove that it wrote the slot: +`OutputSink::WriteToken` ties each sink to the result from its row closure. Initialized sinks use +`()`. `UninitElementSink` requires `InitializedElement` and exposes each row as +`&mut MaybeUninit`: ```rust -let token = InitializedElement::write(output, value); -Ok(token) +visitor.visit_into::, _>(|args, output| { + let value = apply(args); + + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, value) } +}) ``` -`InitializedElement` is a zero-sized, unforgeable write token. The sink can call `Vec::set_len` -only after every successful row returns this token. A valid-only loop initializes placeholders -before it skips rows. +`InitializedElement` is zero-sized write evidence. Only unsafe code can construct it. The caller +must write the current callback's row and return the token from that callback. The sink calls +`Vec::set_len` only after every successful row returns this evidence. A valid-only loop initializes +placeholders before it skips rows. ## Failure models diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index e286433d77f..e3201356f4f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -118,7 +118,8 @@ where return Err(numeric_error(>::ERROR)); } - Ok(InitializedElement::write(output, value)) + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) }) } diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index 7cefccf3a4f..bf8dfab68b3 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -50,7 +50,9 @@ pub trait OutputSink: 'static + Sized { /// Proof that a successful row closure left its row handle initialized. /// /// Use `()` for initialized row handles. A sink exposing uninitialized storage uses a distinct - /// token returned after initialization. + /// token returned after initialization. A sink that uses this token to justify unsafe code + /// **must** prevent safe construction that does not establish the invariant. Make construction + /// unsafe when Rust cannot tie the token to the supplied row handle. type WriteToken: 'static; /// The dtype of the column this sink builds, given the function's input dtypes. @@ -90,14 +92,20 @@ pub trait OutputSink: 'static + Sized { /// Proof that one uninitialized element row was initialized. #[must_use = "return this token from the row closure to prove that it initialized the output"] pub struct InitializedElement( - /// Private so safe code can only obtain this token by writing an uninitialized row. + /// Private so constructing initialization evidence requires an unsafe operation. (), ); impl InitializedElement { /// Write `value` into an uninitialized row and return its proof token. + /// + /// # Safety + /// + /// `row` must be the [`UninitElementSink`] row supplied to the current callback. The caller must + /// return the token from that callback. Using another row or returning the token from another + /// callback can cause undefined behavior. #[inline] - pub fn write(row: &mut MaybeUninit, value: T) -> Self { + pub unsafe fn write(row: &mut MaybeUninit, value: T) -> Self { row.write(value); Self(()) @@ -156,9 +164,10 @@ impl OutputSink for UninitElementSink { } fn finish(mut self, _error: DeferredError) -> VortexResult { - // SAFETY: dense execution reaches `finish` only after every row returned the token from - // `InitializedElement::write`. Skip-invalid execution initializes every row before - // overwriting valid ones. The allocation reserved every slot in `0..row_count`. + // SAFETY: the `WriteToken` equality requires each successful dense callback to return an + // `InitializedElement`. Its unsafe constructor requires initialization of that callback's + // row. Skip-invalid execution initializes every row before overwriting valid ones. + // `with_capacity` reserved every slot in `0..row_count`. unsafe { self.values.set_len(self.row_count) }; Ok(T::build(self.values)) diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 3a6e54a7e4c..7fed9106d6e 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -86,7 +86,8 @@ impl RowFn for SpatialContains { } }, |operands, (a, b), output| { - InitializedElement::write(output, contains_row_prepared(operands, a, b)) + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, contains_row_prepared(operands, a, b)) } }, ) } diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index a5ac84e1e67..59226381bc7 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -67,7 +67,10 @@ impl RowFn for SpatialDistance { visitor: V, ) -> VortexResult { visitor.visit_into::<(GeometryRow, GeometryRow), UninitElementSink, _>( - |(a, b), output| InitializedElement::write(output, Euclidean.distance(a, b)), + |(a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, Euclidean.distance(a, b)) } + }, ) } } diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index b84506cd918..a7d3aa17d45 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -77,7 +77,8 @@ impl RowFn for SpatialIntersects { ConstBboxes::new(a, b) }, |bboxes, (a, b), output| { - InitializedElement::write(output, intersects_row_prepared(bboxes, a, b)) + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, intersects_row_prepared(bboxes, a, b)) } }, ) } diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index da39b35cbad..18930b89009 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -99,10 +99,13 @@ impl RowFn for CosineSimilarity { } }, |norms, (lhs, rhs), output| { - InitializedElement::write( - output, - cosine_similarity_row_prepared(norms, lhs, rhs), - ) + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { + InitializedElement::write( + output, + cosine_similarity_row_prepared(norms, lhs, rhs), + ) + } }, ) }) diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 0e44fa0aa4f..b972fe54b96 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -76,7 +76,10 @@ impl RowFn for InnerProduct { ) -> VortexResult { match_each_float_ptype!(tensor_element_ptype(args)?, |T| { visitor.visit_into::<(TensorRow, TensorRow), UninitElementSink, _>( - |(lhs, rhs), output| InitializedElement::write(output, inner_product_row(lhs, rhs)), + |(lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, inner_product_row(lhs, rhs)) } + }, ) }) } diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index ef36b5dd39f..dbc2e27d33c 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -81,7 +81,8 @@ impl RowFn for L2Norm { ) -> VortexResult { match_each_float_ptype!(tensor_element_ptype(args)?, |T| { visitor.visit_into::<(TensorRow,), UninitElementSink, _>(|(row,), output| { - InitializedElement::write(output, l2_norm_row(row)) + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l2_norm_row(row)) } }) }) } diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index 3d8a784435a..ef86cdc80e1 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -50,7 +50,8 @@ impl RowFn for L1Norm { ) -> VortexResult { match_each_float_ptype!(tensor_element_ptype(args)?, |T| { visitor.visit_into::<(TensorRow,), UninitElementSink, _>(|(row,), output| { - InitializedElement::write(output, l1_norm_row(row)) + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l1_norm_row(row)) } }) }) } From 833632aaa3cab59fb4a7d4f001df26975b2267a1 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 21:32:12 -0400 Subject: [PATCH 044/160] Update RowFn session handoff Signed-off-by: Connor Tsui --- research/rowfn-reconstruction/HANDOFF.md | 60 ++++++++++++++----- research/rowfn-reconstruction/OPTIMIZATION.md | 11 ++-- research/rowfn-reconstruction/README.md | 4 +- 3 files changed, 55 insertions(+), 20 deletions(-) diff --git a/research/rowfn-reconstruction/HANDOFF.md b/research/rowfn-reconstruction/HANDOFF.md index 493c5d71b56..704e1487f75 100644 --- a/research/rowfn-reconstruction/HANDOFF.md +++ b/research/rowfn-reconstruction/HANDOFF.md @@ -10,22 +10,26 @@ read [`DESIGN.md`](DESIGN.md), [`OPTIMIZATION.md`](OPTIMIZATION.md), and ## Branch state - Branch: `ct/row-fn`. -- Last RowFn code commit: `4c936447a`. +- Current RowFn code head: `443aed0b9`. - Documentation head before the offsets fix: `bdf95a77e`. - Comparison revision: develop at `66d096b5d`. - Direct offsets fix: `61410ef21`. -- Numeric helper ID fix: `f9dfde730` on this branch and `df8fcbe1a` on `ct/row-fn-api`. +- Numeric helper ID fix: `f9dfde730`. - Masked tensor decode fix: `7baa9fab7`. -- Cleaned `ct/row-fn-api` head: `6dd500f59`. +- Primitive comparison implementation: `8128137cc`. +- Cleaned `ct/row-fn-api` head: `29e3db1b8`. +- Cleaned `ct/row-fn-numeric` head: `2aae5992d`. - PR #9299 head measured locally: `d97e53e66`. -The API branch was rewritten with an exact force-with-lease from seven commits to five: +The focused branches now separate the framework from primitive numeric arithmetic: -1. `71c3e7a58` adds the framework, refined contracts, and self-contained arguments. -2. `6e864bf8b` moves primitive numeric operators to RowFn and reuses `Binary`'s ID. -3. `41bb10143` adds focused executor benchmarks. -4. `266350488` removes validated input bounds checks. -5. `6dd500f59` restores mixed-constant performance. +1. `7b9cf51ea` adds the cleaned framework to `ct/row-fn-api`. +2. `29e3db1b8` adds the focused executor benchmarks to `ct/row-fn-api`. +3. `2aae5992d` adds primitive numeric RowFn execution on `ct/row-fn-numeric`. + +Both focused branches use develop commit `7ec7ffbae` as their base. The API branch contains no +primitive numeric implementation. The numeric branch differs from it in seven numeric source and +benchmark files. All three local branch tips match their `origin` refs. Two temporary remote refs remain for the CodSpeed ablation: @@ -36,6 +40,32 @@ Both refs contain exact historical code. Temporary draft PR #9298 supplied the p for the focused comparisons. It is now closed, and its `ct/row-fn-codspeed-take-filter` head branch has been deleted. +## Final output-sink safety contract + +`443aed0b9` keeps `RowVisitor::visit_into` and `RowVisitor::visit_prepared_into` safe. The selected +`SinkResult::WriteToken` must match `OutputSink::WriteToken`, so +`UninitElementSink` requires an `InitializedElement` for every successful row. + +`InitializedElement::write` is the unsafe boundary. Its caller must write the +`UninitElementSink` row from the current callback and return that token from the same callback. +The token has no safe constructor. Ordinary initialized sinks use `()` and require no unsafe code. + +The final API has no `visit_uninit`, `try_visit_uninit`, or `visit_prepared_uninit` wrappers. +`UninitElementSink` remains public and uses the generic `visit_into` path. This keeps the unsafe +operation inside each uninitialized-output closure without making the visitor API unsafe. + +The final validation completed these commands: + +```bash +cargo +nightly fmt --all +cargo nextest run -p vortex-array -p vortex-tensor -p vortex-spatial +cargo test --doc -p vortex-array -p vortex-tensor -p vortex-spatial +cargo clippy --all-targets --all-features +``` + +The targeted run passed 3,884 tests. The cleaned numeric branch also passed all 3,460 +`vortex-array` tests and `cargo clippy --all-targets --all-features -- -D warnings`. + ## Corrected CodSpeed history The latest push did not bring back the `take_filter_list_*` regressions. @@ -153,7 +183,8 @@ constant cases exposed a separate source-placement regression: | `sub_i64_constant` | 8.319 us | 36.19 us | +335.0% | | `mul_i32_constant` | 26.43 us | 41.91 us | +58.6% | -Commit `6dd500f59` keeps each length proof in the branch that consumes it. After the fix, +The measured API revision `6dd500f59` keeps each length proof in the branch that consumes it. After +the fix, `add_i64_constant` measures 9.269 microseconds, `sub_i64_constant` measures 9.199 microseconds, and `mul_i32_constant` measures 18.91 microseconds. The first two retain about 11% overhead; multiply is 28.5% faster than develop. @@ -165,7 +196,7 @@ Ten one-second alternating runs isolate a stable native regression: | Binary | Median | Observed range | | --- | ---: | ---: | | Develop `66d096b5d` | 2.229 us | 2.229 to 2.239 us | -| Clean API `6dd500f59` | 2.809 us | 2.799 to 2.829 us | +| Measured API revision `6dd500f59` | 2.809 us | 2.799 to 2.829 us | | `-C llvm-args=-align-loops=64` diagnostic | 2.449 us | 2.439 to 2.499 us | The develop and RowFn steady-state loops have the same normalized instruction sequence: two @@ -330,9 +361,10 @@ The focused numeric profile also found 6.820 microseconds of new inclusive cost Develop's ID lookup costs 0.702 microseconds total. The numeric RowFn revision costs 7.522 microseconds. -`NumericBinary` is an internal helper for the registered `Binary` function. Commit `df8fcbe1a` on -`ct/row-fn-api` reuses `Binary`'s ID. This removes the second interner initialization and gives -errors the public function's name. It does not change the arithmetic loop or the public API. +`NumericBinary` is an internal helper for the registered `Binary` function. Commit `f9dfde730` on +this branch reuses `Binary`'s ID. The cleaned focused implementation is commit `2aae5992d` on +`ct/row-fn-numeric`. This removes the second interner initialization and gives errors the public +function's name. It does not change the arithmetic loop or the public API. This is a first-execution cost, not a per-row cost. The [numeric ID check] validates it: diff --git a/research/rowfn-reconstruction/OPTIMIZATION.md b/research/rowfn-reconstruction/OPTIMIZATION.md index 5f54b45d60e..a5ee44824e9 100644 --- a/research/rowfn-reconstruction/OPTIMIZATION.md +++ b/research/rowfn-reconstruction/OPTIMIZATION.md @@ -482,8 +482,9 @@ The focused numeric profile shows another fixed cost. `CachedId::deref` increase call. `NumericBinary` is not registered. It executes the registered `Binary` operation's primitive path. -Commit `df8fcbe1a` on `ct/row-fn-api` therefore reuses `Binary`'s existing ID. This removes a second -interner initialization and makes internal errors name the public function. +Commit `f9dfde730` on the monolithic branch therefore reuses `Binary`'s existing ID. The cleaned +focused implementation is commit `2aae5992d` on `ct/row-fn-numeric`. This removes a second interner +initialization and makes internal errors name the public function. This change does not alter dispatch or the row loop. The cost occurs on first execution, so it is separate from per-row vectorization. The [numeric ID check] validates the result: @@ -621,8 +622,10 @@ linked-layout costs. Do not use it to override the focused native A/B above. - Reduce the mixed-constant LLVM sensitivity while preserving the production monomorph. - Reduce fixed RowFn batch overhead if 32,768-row numeric calls are latency-critical. - Profile the isolated `mul_u8_nonnull` case with native performance counters. -- Choose between the direct typed offset fix and PR #9299's materialize-once design based on API - maintenance and correctness. Both are native wins, but they are different implementations. +- Reconcile `61410ef21` with PR #9299 before merging the monolithic branch. PR #9299 identifies the + double execution of lazy reset offsets and materializes them once in `list_view_from_list`. + `61410ef21` makes `reset_offsets` eager, which also prevents the second execution. Remove the + direct fix if PR #9299 makes it redundant, then repeat the focused native comparison. - Identify the spatial `envelope` regression that begins when numeric RowFn code enters the linked binary. - Repeat the key local results on a second x86 machine and compiler version before filing a diff --git a/research/rowfn-reconstruction/README.md b/research/rowfn-reconstruction/README.md index 55d42c62cdf..07f7b2cd4dd 100644 --- a/research/rowfn-reconstruction/README.md +++ b/research/rowfn-reconstruction/README.md @@ -7,8 +7,8 @@ This guide explains the RowFn design without requiring access to its source. It model, execution model, performance constraints, implementation order, and benchmark procedure. The goal is to let a new contributor reconstruct the branch and understand each unusual choice. -The guide describes commit `4c936447a` on `ct/row-fn`. Its comparison revision is develop commit -`66d096b5d`. +The guide describes the implementation through `443aed0b9` on `ct/row-fn`. Historical CodSpeed +comparisons use develop commit `66d096b5d`. ## Reading order From 44457b76546b40987c5328eea7759cd095a4761e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 10:47:24 -0400 Subject: [PATCH 045/160] Harden RowFn element and sink contracts Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 2 +- .../src/scalar_fn/fns/binary/numeric/row.rs | 5 ++ .../src/scalar_fn/row/batch/policy.rs | 9 ++-- .../src/scalar_fn/row/execute/sink.rs | 16 +++---- .../src/scalar_fn/row/types/element/bool.rs | 3 +- .../src/scalar_fn/row/types/element/mod.rs | 9 +++- .../scalar_fn/row/types/element/primitive.rs | 3 +- vortex-array/src/scalar_fn/row/types/sink.rs | 48 ++++++++++--------- vortex-spatial/src/scalar_fn/row.rs | 4 +- vortex-tensor/src/scalar_fns/row.rs | 4 +- 10 files changed, 58 insertions(+), 45 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index 4c42aaa37f4..d53afd57770 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -136,7 +136,7 @@ impl OutputSink for I64Sink { &mut rows[index] } - fn finish(self, _error: DeferredError) -> VortexResult { + unsafe fn finish(self, _error: DeferredError) -> VortexResult { Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) } } diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index e3201356f4f..c65209b3220 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -56,6 +56,11 @@ impl RowFn for NumericBinary { const FALLIBLE: bool = true; fn id(&self) -> ScalarFnId { + // This private helper is never registered or serialized. It executes the registered + // `Binary` operation's primitive path directly, so reusing that public ID keeps errors and + // first-use interning attributed to the function the caller invoked. Privacy is the guard: + // making this type registrable requires giving it an independent ID and persistence + // contract first. ScalarFnVTable::id(&Binary) } diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs index 1ea3a500baa..d017b5278d6 100644 --- a/vortex-array/src/scalar_fn/row/batch/policy.rs +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -61,12 +61,9 @@ impl RowPolicy { /// The policy one concrete dispatch executes nullable rows under. /// - /// This deliberately ignores [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. Batch execution always - /// tries [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) against the original - /// arrays before it tries the sink or filters the inputs. Skipping that probe can change the - /// result of an encoding-aware function. - /// - /// [`OutputSink::SUPPORTS_SKIPPED_ROWS`]: crate::scalar_fn::OutputSink::SUPPORTS_SKIPPED_ROWS + /// Batch execution always tries [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) + /// against the original arrays before it tries the sink or filters the inputs. Skipping that + /// probe can change the result of an encoding-aware function. pub const fn for_sink() -> Self { if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { if ApplyResult::DEFERRED { diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 312dd2d2db4..64fc96624ca 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -42,7 +42,6 @@ where let varying = Args::varying(&columns); ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; let mut accumulated = ApplyResult::Accumulated::default(); - { // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This // scope releases the borrow before `finish_sink` consumes the sink. @@ -92,12 +91,6 @@ where Sink: OutputSink, ApplyResult: SinkResult, { - // Batch execution needs a full-length result before applying the validity mask. Decline when - // the sink cannot leave legal placeholders in positions this loop skips. - if !Sink::SUPPORTS_SKIPPED_ROWS { - return Ok(None); - } - // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An // element representation may decline when it cannot provide those values safely. let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { @@ -107,7 +100,6 @@ where let row_count = args.row_count(); let mut sink = Sink::with_capacity(row_count, sink_dtype)?; let mut accumulated = ApplyResult::Accumulated::default(); - // Batch execution resolves all-valid and all-null inputs before selecting this path. let AllOr::Some(valid) = valid.bit_buffer() else { vortex_bail!("execute_sink_valid_rows requires a mixed mask"); @@ -129,7 +121,9 @@ where // The loop writes only valid indices, but the sink still finishes a full-length output. // Initialize placeholders now; batch execution masks them before the result escapes. - Sink::initialize_skipped_rows(&mut rows); + if !Sink::initialize_skipped_rows(&mut rows) { + return Ok(None); + } // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first // immediate error, turn later callbacks into no-ops, and return before finishing the sink. @@ -175,7 +169,9 @@ fn finish_sink( sink: S, deferred_error: DeferredError, ) -> VortexResult { - match sink.finish(deferred_error) { + // SAFETY: callers reach this helper only after successful traversal. Dense traversal visited + // every addressable row; skipped-row traversal initialized every row before visiting its mask. + match unsafe { sink.finish(deferred_error) } { Ok(output) => Ok(RowExecution::Output(output)), Err(error) if deferred_error.occurred() => Ok(RowExecution::DeferredError(error)), Err(error) => Err(error), diff --git a/vortex-array/src/scalar_fn/row/types/element/bool.rs b/vortex-array/src/scalar_fn/row/types/element/bool.rs index bc966268e4d..3b68dfb1e80 100644 --- a/vortex-array/src/scalar_fn/row/types/element/bool.rs +++ b/vortex-array/src/scalar_fn/row/types/element/bool.rs @@ -15,7 +15,8 @@ use crate::scalar_fn::InputElement; use crate::scalar_fn::OutputElement; use crate::validity::Validity; -impl InputElement for bool { +// SAFETY: the varying view is a bit buffer, and its reported length is the buffer length. +unsafe impl InputElement for bool { type Column = BitBuffer; type Varying<'a> = &'a BitBuffer; type Elem<'a> = bool; diff --git a/vortex-array/src/scalar_fn/row/types/element/mod.rs b/vortex-array/src/scalar_fn/row/types/element/mod.rs index 1743e3f1db0..232040cde12 100644 --- a/vortex-array/src/scalar_fn/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -23,7 +23,14 @@ pub use tuple::IndexedElementTuple; pub use tuple::batch_constant; /// An element type that can be read row-wise out of an input column. -pub trait InputElement: 'static { +/// +/// # Safety +/// +/// For every view returned by [`varying`](Self::varying), every index below +/// [`varying_len`](Self::varying_len) **must** satisfy the safety contract of +/// [`get_varying_unchecked`](Self::get_varying_unchecked). Shared execution relies on this proof to +/// perform unchecked reads after one pre-loop length check. +pub unsafe trait InputElement: 'static { /// The decoded column representation supporting `O(1)` row access. type Column; diff --git a/vortex-array/src/scalar_fn/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/row/types/element/primitive.rs index 071a7c55115..9e80d937ef5 100644 --- a/vortex-array/src/scalar_fn/row/types/element/primitive.rs +++ b/vortex-array/src/scalar_fn/row/types/element/primitive.rs @@ -17,7 +17,8 @@ use crate::scalar_fn::InputElement; use crate::scalar_fn::OutputElement; use crate::validity::Validity; -impl InputElement for T { +// SAFETY: the varying view is a native slice, and its reported length is the slice length. +unsafe impl InputElement for T { type Column = Buffer; type Varying<'a> = &'a [T]; type Elem<'a> = T; diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index bf8dfab68b3..ebf1053d842 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -19,7 +19,7 @@ use crate::scalar_fn::OutputElement; /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; /// skip-invalid execution can omit invalid rows when -/// [`SUPPORTS_SKIPPED_ROWS`](Self::SUPPORTS_SKIPPED_ROWS) is `true`. +/// [`initialize_skipped_rows`](Self::initialize_skipped_rows) returns `true`. pub trait OutputSink: 'static + Sized { /// Whether this sink accepts [`DeferredError`] from its row closure instead of requiring a /// per-row [`VortexResult`]. @@ -28,12 +28,6 @@ pub trait OutputSink: 'static + Sized { /// argument occurred. const ERRORS_ARE_DEFERRED: bool = false; - /// Whether this sink can finish a full-length output when some rows were never visited. - /// - /// A supporting sink must use [`initialize_skipped_rows`](Self::initialize_skipped_rows) to - /// leave a legal arbitrary value at every skipped row. Batch execution masks those values. - const SUPPORTS_SKIPPED_ROWS: bool = false; - /// A loop-local view of all output rows. /// /// Borrowed once before execution so the sink's buffer descriptor and shape become loop @@ -62,7 +56,8 @@ pub trait OutputSink: 'static + Sized { fn sink_dtype(args: &[DType]) -> VortexResult; /// Allocate a sink for `rows` rows of `dtype`, which is this sink's own - /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + /// [`sink_dtype`](Self::sink_dtype). Called once per batch with whether any deferred row error + /// occurred. fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; /// Borrow all output rows for the hot loop. @@ -74,19 +69,28 @@ pub trait OutputSink: 'static + Sized { /// optimizer the output bounds it needs to remove the bounds check hidden in each row accessor. fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool; - /// Initialize output positions that skip-invalid execution can omit. + /// Initialize every output position so skip-invalid execution may omit rows. /// - /// Called only when [`SUPPORTS_SKIPPED_ROWS`](Self::SUPPORTS_SKIPPED_ROWS) is `true`. The - /// default is for sinks whose allocation already contains legal values. - fn initialize_skipped_rows(_rows: &mut Self::Rows<'_>) {} + /// Return `true` after leaving a legal arbitrary value in every row. Return `false` without + /// changing `rows` when the sink cannot support skipped rows. Combining the capability probe + /// and initialization prevents a separate support flag from disagreeing with a no-op method. + fn initialize_skipped_rows(_rows: &mut Self::Rows<'_>) -> bool { + false + } /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; /// Finish into the built column, whose dtype **must** be this sink's - /// [`sink_dtype`](Self::sink_dtype). Called once per batch with whether any deferred row error - /// occurred. - fn finish(self, error: DeferredError) -> VortexResult; + /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + /// + /// # Safety + /// + /// The executor must have completed every row callback successfully, and each callback must + /// have returned this sink's [`WriteToken`](Self::WriteToken). When skipped rows are allowed, + /// [`initialize_skipped_rows`](Self::initialize_skipped_rows) must have returned `true` before + /// traversal. + unsafe fn finish(self, error: DeferredError) -> VortexResult; } /// Proof that one uninitialized element row was initialized. @@ -128,8 +132,6 @@ pub struct UninitElementSink { } impl OutputSink for UninitElementSink { - const SUPPORTS_SKIPPED_ROWS: bool = true; - type Rows<'a> = &'a mut [MaybeUninit]; type Row<'a> = &'a mut MaybeUninit; type WriteToken = InitializedElement; @@ -153,21 +155,21 @@ impl OutputSink for UninitElementSink { rows.len() == row_count } - fn initialize_skipped_rows(rows: &mut Self::Rows<'_>) { + fn initialize_skipped_rows(rows: &mut Self::Rows<'_>) -> bool { for row in rows.iter_mut() { row.write(T::default()); } + + true } fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { &mut rows[index] } - fn finish(mut self, _error: DeferredError) -> VortexResult { - // SAFETY: the `WriteToken` equality requires each successful dense callback to return an - // `InitializedElement`. Its unsafe constructor requires initialization of that callback's - // row. Skip-invalid execution initializes every row before overwriting valid ones. - // `with_capacity` reserved every slot in `0..row_count`. + unsafe fn finish(mut self, _error: DeferredError) -> VortexResult { + // SAFETY: the caller guarantees every slot in `0..row_count` was initialized, and + // `with_capacity` reserved every slot in that range. unsafe { self.values.set_len(self.row_count) }; Ok(T::build(self.values)) diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index f94a1a1aec7..52d3818f73f 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -24,7 +24,9 @@ use crate::extension::is_native_geometry; /// column is *some* native geometry. pub struct GeometryRow; -impl InputElement for GeometryRow { +// SAFETY: the varying view is the decoded geometry slice, and its reported length is that slice's +// length. +unsafe impl InputElement for GeometryRow { type Column = Vec>; type Varying<'a> = &'a [Geometry]; type Elem<'a> = &'a Geometry; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 340cdc1ed92..68f7f44d1c3 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -55,7 +55,9 @@ pub struct TensorRows { stride: usize, } -impl InputElement for TensorRow { +// SAFETY: `TensorRows` records the row count validated during decode, and both checked and +// unchecked access use the same stride and row width. +unsafe impl InputElement for TensorRow { type Column = TensorRows; type Varying<'a> = &'a TensorRows; type Elem<'a> = &'a [T]; From 5204fb2be5f069fb33eee1f14f3fca2b9760f894 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 10:51:38 -0400 Subject: [PATCH 046/160] Fix RowFn batch boundary and retry behavior Signed-off-by: Connor Tsui --- .../src/scalar_fn/row/batch/execution.rs | 42 +++- vortex-array/src/scalar_fn/row/batch/mod.rs | 3 + vortex-array/src/scalar_fn/row/batch/tests.rs | 214 ++++++++++++++++++ vortex-array/src/scalar_fn/row/row_fn.rs | 7 +- vortex-array/src/scalar_fn/row/vtable.rs | 13 +- 5 files changed, 263 insertions(+), 16 deletions(-) create mode 100644 vortex-array/src/scalar_fn/row/batch/tests.rs diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index 4c5eb51ff41..c227c656697 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -80,9 +80,22 @@ impl Batch { args: &dyn ExecutionArgs, plan: impl FnOnce(&[DType]) -> VortexResult, ) -> VortexResult { + let row_count = args.row_count(); let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) .map(|index| args.get(index)) .collect::>()?; + if let Some((index, input)) = inputs + .iter() + .enumerate() + .find(|(_, input)| input.len() != row_count) + { + vortex_ensure_eq!( + input.len(), + row_count, + "the {id} input {index} has {} rows but execution declares {row_count}", + input.len(), + ); + } let arg_dtypes: SmallVec<[DType; 4]> = inputs.iter().map(|input| input.dtype().clone()).collect(); @@ -96,7 +109,7 @@ impl Batch { Ok(Self { id, - row_count: args.row_count(), + row_count, inputs, arg_dtypes, validity, @@ -109,10 +122,12 @@ impl Batch { /// Add null propagation, constant folding, and strategy selection around `kernel`. /// /// The kernel may ignore input validity. It receives valid-only rows when required, and its - /// output **must** match the planned dtype up to nullability. `try_unfiltered` receives the - /// original inputs plus a mixed validity mask. `Ok(None)` selects filter-and-scatter. + /// output **must** match the planned dtype up to nullability. `reduce` receives the original + /// inputs exactly once. `try_unfiltered` receives the originals plus a mixed validity mask; + /// `Ok(None)` selects filter-and-scatter. pub fn execute( &self, + reduce: impl FnOnce(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult>, kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_unfiltered: impl FnOnce( KernelArgs<'_>, @@ -131,6 +146,15 @@ impl Batch { return Ok(self.all_null()); } + // An all-null batch has no observable row work. Other batches offer the encoding-aware + // hook the original inputs once, before slicing or filtering can change their encodings. + if matches!(self.validity, Validity::AllInvalid) { + return Ok(self.all_null()); + } + if let Some(values) = reduce(self.kernel_args(&self.inputs, self.row_count), ctx)? { + return self.finalize_reduced(values); + } + // All inputs constant, and their conjoined validity proves every row non-null. This sees // through extension and masked wrappers just like argument decoding does. if self.row_count > 0 @@ -323,6 +347,18 @@ impl Batch { ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() } + /// Reconcile an encoding-aware result and apply the batch's strict input validity. + fn finalize_reduced(&self, values: ArrayRef) -> VortexResult { + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.row_count) + } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // Handled before the encoding-aware hook runs. + Validity::AllInvalid => Ok(self.all_null()), + } + } + /// Pair an input view with this batch's planning metadata. fn kernel_args<'b>(&'b self, arrays: &'b [ArrayRef], row_count: usize) -> KernelArgs<'b> { KernelArgs { diff --git a/vortex-array/src/scalar_fn/row/batch/mod.rs b/vortex-array/src/scalar_fn/row/batch/mod.rs index 1b492f1ff6f..3dd7b08603f 100644 --- a/vortex-array/src/scalar_fn/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/row/batch/mod.rs @@ -20,3 +20,6 @@ pub(super) use execution::finalize_kernel_output; mod policy; pub(super) use policy::BatchPlan; pub(super) use policy::RowPolicy; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs new file mode 100644 index 00000000000..f31b0f578a0 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::Batch; +use super::BatchPlan; +use super::RowPolicy; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::scalar_fn::DeferredError; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::RowVisitor; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::validity::Validity; + +#[derive(Clone)] +struct RetryConstantAdd; + +#[derive(Clone)] +struct NullarySeven; + +struct I64Sink(BufferMut); + +impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self, _error: DeferredError) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +impl RowFn for NullarySeven { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.nullary_seven"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), I64Sink, _>(|(), output| { + *output = 7; + }) + } +} + +impl RowFn for RetryConstantAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.retry_constant_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(u8, u8), u8, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "checked add overflowed")); + } + + Ok(()) + }, + ) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 1 { + return Ok(Some(ConstantArray::new(0u8, args[0].len()).into_array())); + } + + Ok(None) + } +} + +#[test] +fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.row_batch"); + + let input = PrimitiveArray::new(vec![1i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let result = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::Dense, + }) + }); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let result = ScalarFnVTable::execute(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx); + + assert!(result.is_err()); + Ok(()) +} + +#[rstest] +#[case::dense(RowPolicy::Dense)] +#[case::dense_with_retry(RowPolicy::DenseWithRetry)] +#[case::valid_only(RowPolicy::ValidOnly)] +fn test_strategy_matrix(#[case] policy: RowPolicy) -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.row_strategy"); + + let input = PrimitiveArray::new(vec![1i64, 2, 3], Validity::from_iter([true, false, true])) + .into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 3); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |_args, _ctx| Ok(None), + |args, _ctx| { + Ok(super::super::execute::RowExecution::Output( + args.arrays[0].clone(), + )) + }, + |args, _valid, _ctx| { + Ok(Some(super::super::execute::RowExecution::Output( + args.arrays[0].clone(), + ))) + }, + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_nullary_row_function_broadcasts() -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([7i64, 7, 7]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index 68d92c192ad..7fd8b3c08fb 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -69,7 +69,12 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// Try an encoding-aware implementation before decoding the inputs into row elements. /// /// `None` continues to the dispatched row loop. `Some(output)` skips that loop. The output can - /// remain encoded or lazy. Filter-and-scatter execution can pass compacted inputs. + /// remain encoded or lazy. For non-nullary functions, batch execution calls this hook at most + /// once with the original, unfiltered arrays; slices and compacted retries do not reach it. + /// + /// Like a dense row closure, this hook must be total over every stored payload, including + /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or + /// retried through the row layer. /// /// # Requirements /// diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index 848ce47a612..8ed9163866f 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -90,6 +90,7 @@ impl ScalarFnVTable for F { let batch = prepare_batch(self, options, args)?; batch.execute( + |args, ctx| self.reduce_encoded(options, args.arrays, ctx), |args, ctx| execute_rows(self, options, args, ctx), |args, valid, ctx| try_execute_rows_unfiltered(self, options, args, valid, ctx), ctx, @@ -120,12 +121,6 @@ fn execute_rows( args: KernelArgs<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { - if !args.arrays.is_empty() - && let Some(reduced) = function.reduce_encoded(options, args.arrays, ctx)? - { - return Ok(RowExecution::Output(reduced)); - } - let execution = BorrowedExecutionArgs::new(args.arrays, args.row_count); function.dispatch( @@ -143,12 +138,6 @@ fn try_execute_rows_unfiltered( valid: &Mask, ctx: &mut ExecutionCtx, ) -> VortexResult> { - // Try the encoding-aware path before filtering changes the inputs. The caller masks its - // full-length result with `valid` before returning it. - if let Some(reduced) = function.reduce_encoded(options, args.arrays, ctx)? { - return Ok(Some(RowExecution::Output(reduced))); - } - let execution = BorrowedExecutionArgs::new(args.arrays, args.row_count); function.dispatch( From 38814c7fb5c25920a8bd927c18d10d52894a89f7 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 10:57:43 -0400 Subject: [PATCH 047/160] Document RowFn review follow-up Signed-off-by: Connor Tsui --- research/rowfn-review-followup/README.md | 168 ++++++++++++++++++ .../codegen/api-contract-summary.md | 54 ++++++ .../codegen/final-batch-summary.md | 39 ++++ 3 files changed, 261 insertions(+) create mode 100644 research/rowfn-review-followup/README.md create mode 100644 research/rowfn-review-followup/codegen/api-contract-summary.md create mode 100644 research/rowfn-review-followup/codegen/final-batch-summary.md diff --git a/research/rowfn-review-followup/README.md b/research/rowfn-review-followup/README.md new file mode 100644 index 00000000000..2e701266e16 --- /dev/null +++ b/research/rowfn-review-followup/README.md @@ -0,0 +1,168 @@ + + + +# RowFn review follow-up + +Platform: Apple Silicon `arm64`, macOS 15.7.3, Rust 1.91.0, LLVM 21.1.2. This machine uses +128-bit NEON and cannot reproduce or compare the pinned Ryzen wall-clock results. This +investigation uses optimized LLVM IR and correctness tests only. + +Baseline: `833632aaa3cab59fb4a7d4f001df26975b2267a1` on `ct/row-fn`. + +## Result + +| # | Verdict | Evidence and status | +| ---: | --- | --- | +| 1 | Fixed, not upstreamed | `44457b7654` makes `OutputSink::finish` unsafe. The zero-unsafe external reproduction changes from reading allocator contents to `E0133`. All gate items pass, but the local `ct/row-fn-api` ref contains divergent unpushed history and was not overwritten. | +| 2 | Fixed, not upstreamed | `InputElement` is an unsafe trait with local safety proofs. `ElementTuple` and `IndexedElementTuple` remain sealed framework traits. Same upstream status as #1. | +| 3 | Fixed, not upstreamed | The existing unsafe `InitializedElement::write` token remains unchanged. Unsafe publication now belongs to `OutputSink::finish`, and the generic executor owns the proof. Same upstream status as #1. | +| 4 | Fixed, not upstreamed | `initialize_skipped_rows` now returns its capability result. A separate support constant cannot disagree with a no-op default. Same upstream status as #1. | +| 5 | Fixed, not upstreamed | `5204fb2be5` documents totality and probes original inputs once. Gate item 3 fails because the batch refactor adds an owned mixed-path bounds edge. | +| 6 | Not started | The requested redesign was out of scope. A temporary manual implementation reproduced `E0119`. The options analysis is below. | +| 7 | Fixed, not upstreamed | `44457b7654` documents why private `NumericBinary` borrows the registered `Binary` ID and makes privacy the registration guard. Same upstream status as #1. | +| 8 | Fixed, not upstreamed | `5204fb2be5` validates every input length in `Batch::new`. The regression test fails before the fix. Gate item 3 fails as described for #5. | +| 9 | Refuted | `MaskedArray::try_new` enforces an all-valid child, and null `ConstantArray` validity is `AllInvalid`. Both couplings are constructor invariants. | +| 10 | Fixed, not upstreamed | `5204fb2be5` probes once before retry. The regression test fails before the fix and passes after it. Gate item 3 fails as described for #5. | +| 11 | Investigated, needs x86 | The path is unreachable because no sink sets `ERRORS_ARE_DEFERRED`. Deleting it perturbed arithmetic IR, so no change remains. | +| 12 | Refuted | Filtering preserves every constant shape recognized by `batch_constant`: literal `Constant`, constant `Masked` child, and constant `Extension` storage. `ConstElems` stays consistent. | +| 13 | Fixed, not upstreamed | `5204fb2be5` runs the encoding probe before one-row broadcast, against the original arrays. Gate item 3 fails as described for #5. | +| 14 | Investigated, needs x86 | The five deferred word implementations are unreachable. Their deletion changed codegen-unit placement and owned-loop IR, so the machinery remains. | +| 15 | Refuted | Mixed `i64` add/sub and `i32` multiply contain broadcast vector loops on `arm64`. They do not fall back to scalar row execution. | +| 16 | Investigated, needs x86 | Full-row skipped initialization and non-breaking mask traversal are independent costs. The proposed API and early-exit split are below. | +| 17 | Investigated, needs x86 | `scatter_valid` allocates one `u64` per original row. The cited `vortex-spatial/src/scalar_fn/execute/geo_types.rs` path does not exist at this revision. | +| 18 | Refuted | No consumer requires row output to retain 256-byte physical alignment. Alignment-sensitive consumers call `ensure_aligned`. Any performance change still needs x86 evidence. | + +## Reproductions + +### Uninitialized sink publication + +A separate crate used no unsafe code. It freed a `Vec` filled with +`0xd1775eedd1775eed`, created `UninitElementSink::::with_capacity(4096, ...)`, called +`finish`, and read the result. The baseline returned the recognizable value in all 4,096 rows. + +After `44457b7654`, the unchanged call fails to compile: + +```text +error[E0133]: call to unsafe function `OutputSink::finish` is unsafe and requires unsafe block +``` + +The sink keeps the existing write-token mechanism. The generic executor makes one documented +unsafe `finish` call after successful traversal or successful skipped-row initialization and +traversal. + +### Blanket vtable implementation + +A temporary `impl ScalarFnVTable for LazyDouble` in `strict_validity.rs` fails with `E0119` because +the blanket `impl ScalarFnVTable for F` already applies. The temporary change was removed. + +### Batch failures + +`test_batch_rejects_input_length_mismatch` fails before the length check. The baseline reaches a +strategy-specific panic, error, or wrong-length result instead of rejecting the batch boundary. + +`test_dense_retry_does_not_reduce_filtered_inputs` uses a reducer that answers only for compacted +one-row inputs. The baseline retries after a deferred error, re-probes the filtered inputs, and +suppresses the valid-row failure. The candidate probes the originals once and preserves the error. + +The first proposed reproduction required every filtered input to become `ConstantArray`. Filtering +a one-row `PrimitiveArray` keeps it primitive, so that version did not exercise the defect. The +one-row encoding probe establishes the same retry violation without assuming a filter encoding. + +## API options for the blanket vtable + +### `RowFnAdaptor` + +An adaptor can own the blanket `ScalarFnVTable` implementation while `F` remains a `RowFn`. +Registration and expression construction must wrap every function. Existing call sites that name +the concrete function type also change. Arithmetic can delegate to the same generic row executor, +but the wrapper changes monomorph identities and requires the full IR gate. + +### Public `execute_rows` + +A public free function lets each adopter implement `ScalarFnVTable` and delegate only execution. +It preserves all five vtable hooks but repeats arity, child naming, return dtype, strictness, and +fallibility boilerplate in each implementation. The arithmetic row loop can remain the same helper +monomorph. The surrounding vtable implementation still requires IR verification. + +### Hooks on `RowFn` + +`RowFn` can redeclare `coerce_args`, `simplify`, `simplify_untyped`, `reduce`, and `fmt_sql`. +The blanket implementation can forward them. This keeps call sites unchanged but duplicates the +`ScalarFnVTable` surface and creates two contracts for each hook. Default hook forwarding remains +outside the arithmetic loop, but the public trait change still requires the full IR gate. + +No option is implemented here. + +## Performance findings + +### Mixed constants + +The optimized `arm64` IR contains fixed-width vector loops for mixed `i64` add/sub and `i32` +multiply. The constant operand is broadcast with `insertelement` and `shufflevector`. Failure is a +loop-carried vector phi and rich error construction remains outside the loop. See +`codegen/final-batch-summary.md`. + +### Skipped rows + +`UninitElementSink::initialize_skipped_rows` writes `T::default()` to every row because its API has +no mask. A mask-aware API can accept `&Mask` and initialize only unset positions. This is a public +sink API change and needs x86 measurement for sparse and dense masks. + +The early-exit problem is separable. A fallible or breakable mask iterator can stop after the first +row error without changing `OutputSink`. `Mask::indices()` is not an equivalent production fix +because it can materialize indices. No change is made here. + +### Scatter + +`scatter_valid` allocates `vec![0u64; valid.len()]`, fills ranks for set-bit runs, performs `take`, +and applies validity. A run-based scatter can copy dense value ranges directly into a pre-sized +output. That design is encoding-sensitive and needs a focused x86 benchmark. The spatial precedent +named in the finding is absent from this revision. + +### Alignment + +`OutputElement::build(Vec)` reports `Alignment::of::()` and does not retain the 256-byte +physical over-alignment from `BufferMut`. Repository consumers that need a stronger alignment use +`BufferHandle::ensure_aligned` in IO, serialization, Zstd, and benchmark paths. No numeric consumer +assumes 256-byte alignment. This is not a correctness requirement. + +## Code generation gate + +The API-only commit preserves the six owned arithmetic monomorphs and both `i64` division sink +paths. No new closure call, loop bounds check, vector loss, failure spill, or in-loop error +construction appears. See `codegen/api-contract-summary.md`. + +The batch correctness commit adds one `panic_bounds_check` site for `output[index]` in the mixed +owned branch even though `owned.rs` is unchanged. The vector loops remain, but gate item 3 fails. +The commit stays on `ct/row-fn` and needs an x86 rerun before any upstream attempt. + +## Do not do this + +Do not delete the unreachable deferred sink machinery as semantically inert cleanup. The deletion +changed codegen-unit placement and the optimized owned arithmetic IR. It was reverted. + +Do not move the `reduce_encoded` probe without checking the owned mixed-path bounds edge. The +correctness fix is valid, but the current source shape does not pass the arithmetic IR gate. + +The investigation did not challenge any guardrail in the task. `owned.rs`, numeric primitive +inlining policy, the numeric `Vec`, `BorrowedExecutionArgs`, and LLVM loop flags remain +unchanged. + +## Verification + +```text +cargo nextest run -p vortex-array -p vortex-spatial -p vortex-tensor + 3805 passed, 1 skipped + +cargo test --doc -p vortex-array + 73 passed, 13 ignored + +cargo +nightly fmt --all + passed + +PYO3_PYTHON=.venv/bin/python cargo clippy --all-targets --all-features + passed +``` + +The first clippy run selected `/usr/bin/python3` 3.9 and stopped because `abi3-py311` requires +Python 3.11. Rerunning with the repository virtual environment completed cleanly. diff --git a/research/rowfn-review-followup/codegen/api-contract-summary.md b/research/rowfn-review-followup/codegen/api-contract-summary.md new file mode 100644 index 00000000000..1f445537351 --- /dev/null +++ b/research/rowfn-review-followup/codegen/api-contract-summary.md @@ -0,0 +1,54 @@ + + + +# RowFn API contract code generation + +Platform: Apple Silicon `arm64`, macOS 15.7.3, Rust 1.91.0, LLVM 21.1.2. + +Baseline: `833632aaa3cab59fb4a7d4f001df26975b2267a1` + +Candidate: `44457b76546b40987c5328eea7759cd095a4761e` + +The candidate contains only the `InputElement` and `OutputSink` contract changes. The later batch +execution fixes are not present. + +## Command + +Both revisions used the repository bench profile with 16 codegen units and no LTO: + +```text +RUSTFLAGS='-C symbol-mangling-version=v0' \ + cargo rustc --profile bench -p vortex-array --lib -- \ + --emit=llvm-ir -C debuginfo=0 +``` + +`symbol-mangling-version=v0` identifies each generic monomorph. It does not change optimization. +The ordinary mangling build used the same profile and produced the same structural conclusions. + +## Structural comparison + +| Monomorph | Vector IR | Bounds sites | Closure calls | Failure reduction | +| --- | --- | ---: | ---: | --- | +| `i64` add | `<16 x i64>` plus remainder | 2 -> 2 | 0 -> 0 | vector and scalar phi | +| `i64` sub | `<16 x i64>` plus remainder | 2 -> 2 | 0 -> 0 | vector and scalar phi | +| `i64` mul | scalar unroll plus `<2 x i64>` support | 2 -> 2 | 0 -> 0 | scalar phi | +| `i32` mul | `<4 x i32>` and `<16 x i32>` | 2 -> 2 | 0 -> 0 | vector and scalar phi | +| `u64` mul | scalar unroll plus `<2 x i64>` support | 2 -> 2 | 0 -> 0 | scalar phi | +| `u16` mul | `<8 x i16>` plus wider unroll groups | 2 -> 2 | 0 -> 0 | vector and scalar phi | +| `i64` div, dense sink | scalar | 2 -> 2 | 0 -> 0 | immediate error | +| `i64` div, valid-row sink | scalar mask walk | 1 -> 1 | 0 -> 0 | immediate error | + +The bounds sites are outside the arithmetic vector bodies. No `panic_bounds_check` edge appears in +a vector body. The candidate preserves every branch, vector factor, broadcast, load, arithmetic +operation, store, and loop-carried failure phi in the six owned arithmetic monomorphs. The raw diff +contains only metadata numbering, attribute numbering, and allocation-location symbol hashes. + +Signed add and subtract use the existing XOR/AND sign test instead of LLVM overflow intrinsics. +Multiply uses the existing widening or high-half checks. Rich `VortexError` construction remains +outside every row loop. + +## Sink publication + +`OutputSink::finish` is now unsafe. The generic sink executor contains one documented unsafe call +after successful dense traversal or successful skipped-row initialization and traversal. LLVM +removes the unsafe boundary. Both `i64` division sink monomorphs retain their prior loop structure. diff --git a/research/rowfn-review-followup/codegen/final-batch-summary.md b/research/rowfn-review-followup/codegen/final-batch-summary.md new file mode 100644 index 00000000000..305aaea8714 --- /dev/null +++ b/research/rowfn-review-followup/codegen/final-batch-summary.md @@ -0,0 +1,39 @@ + + + +# RowFn batch correctness code generation + +Platform: Apple Silicon `arm64`, macOS 15.7.3, Rust 1.91.0, LLVM 21.1.2. + +Baseline: `833632aaa3cab59fb4a7d4f001df26975b2267a1` + +Candidate: `5204fb2be5f069fb33eee1f14f3fca2b9760f894` + +The candidate includes the API contract commit and the batch length, single-probe, and retry fixes. +It used the command from `api-contract-summary.md`. + +## Gate result + +The batch correctness commit does not pass the arithmetic IR gate. The mixed-input branch of each +owned arithmetic monomorph gains a `panic_bounds_check` edge for +`row/execute/owned.rs:98`: + +```rust +output[index].write(value); +``` + +The source line itself did not change. Moving the `reduce_encoded` probe out of the generic row +executor changed codegen-unit placement and inlining decisions. The all-varying and mixed-constant +vector loops remain present, the row closure remains inlined, failure remains loop-carried, and +error construction remains outside the loop. The new output bounds edge is still a structural +regression, so this commit stays on `ct/row-fn` and is not upstreamed. + +## Mixed constants on `arm64` + +The mixed-constant `i64` add and subtract branches contain `insertelement` and `shufflevector` +broadcasts, `<16 x i64>` arithmetic, `<16 x i1>` loop-carried failure phis, and vector stores. The +mixed-constant `i32` multiply branch also contains fixed-width vector arithmetic and broadcasts. + +Apple Silicon exposes 128-bit NEON. The LLVM vector factors span groups of NEON registers, but the +optimized IR is not a scalar fallback. The claim that mixed constants fall off the vectorized lane +kernel is false on this platform. From e85d65cf89d014c7f2998288b52b9130e166eede Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 11:17:12 -0400 Subject: [PATCH 048/160] Polish RowFn contracts and tests Signed-off-by: Connor Tsui --- .../src/scalar_fn/fns/binary/numeric/row.rs | 8 +++----- vortex-array/src/scalar_fn/row/batch/tests.rs | 13 +++---------- vortex-array/src/scalar_fn/row/execute/sink.rs | 2 ++ vortex-array/src/scalar_fn/row/types/sink.rs | 6 +++--- 4 files changed, 11 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index c65209b3220..089adcbcbde 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -56,11 +56,9 @@ impl RowFn for NumericBinary { const FALLIBLE: bool = true; fn id(&self) -> ScalarFnId { - // This private helper is never registered or serialized. It executes the registered - // `Binary` operation's primitive path directly, so reusing that public ID keeps errors and - // first-use interning attributed to the function the caller invoked. Privacy is the guard: - // making this type registrable requires giving it an independent ID and persistence - // contract first. + // `NumericBinary` is a private implementation detail of `Binary`: it is never registered or + // serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. ScalarFnVTable::id(&Binary) } diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index f31b0f578a0..7c708582c9e 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -7,6 +7,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::registry::CachedId; +use super::super::execute::RowExecution; use super::Batch; use super::BatchPlan; use super::RowPolicy; @@ -184,16 +185,8 @@ fn test_strategy_matrix(#[case] policy: RowPolicy) -> VortexResult<()> { let actual = batch.execute( |_args, _ctx| Ok(None), - |args, _ctx| { - Ok(super::super::execute::RowExecution::Output( - args.arrays[0].clone(), - )) - }, - |args, _valid, _ctx| { - Ok(Some(super::super::execute::RowExecution::Output( - args.arrays[0].clone(), - ))) - }, + |args, _ctx| Ok(RowExecution::Output(args.arrays[0].clone())), + |args, _valid, _ctx| Ok(Some(RowExecution::Output(args.arrays[0].clone()))), &mut ctx, )?; diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 64fc96624ca..4805dc86e1f 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -42,6 +42,7 @@ where let varying = Args::varying(&columns); ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; let mut accumulated = ApplyResult::Accumulated::default(); + { // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This // scope releases the borrow before `finish_sink` consumes the sink. @@ -100,6 +101,7 @@ where let row_count = args.row_count(); let mut sink = Sink::with_capacity(row_count, sink_dtype)?; let mut accumulated = ApplyResult::Accumulated::default(); + // Batch execution resolves all-valid and all-null inputs before selecting this path. let AllOr::Some(valid) = valid.bit_buffer() else { vortex_bail!("execute_sink_valid_rows requires a mixed mask"); diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index ebf1053d842..2bfb6968445 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -56,8 +56,7 @@ pub trait OutputSink: 'static + Sized { fn sink_dtype(args: &[DType]) -> VortexResult; /// Allocate a sink for `rows` rows of `dtype`, which is this sink's own - /// [`sink_dtype`](Self::sink_dtype). Called once per batch with whether any deferred row error - /// occurred. + /// [`sink_dtype`](Self::sink_dtype). Called once per batch. fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; /// Borrow all output rows for the hot loop. @@ -82,7 +81,8 @@ pub trait OutputSink: 'static + Sized { fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; /// Finish into the built column, whose dtype **must** be this sink's - /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + /// [`sink_dtype`](Self::sink_dtype). Called once per batch with whether any deferred row error + /// occurred. /// /// # Safety /// From a2709dca5389169e84a42684840467c5a340e030 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 11:29:34 -0400 Subject: [PATCH 049/160] Restore early RowFn sink decline Signed-off-by: Connor Tsui --- .../src/scalar_fn/row/execute/sink.rs | 13 +++- .../src/scalar_fn/row/execute/sink/tests.rs | 69 +++++++++++++++++++ vortex-array/src/scalar_fn/row/types/sink.rs | 37 +++++----- 3 files changed, 96 insertions(+), 23 deletions(-) create mode 100644 vortex-array/src/scalar_fn/row/execute/sink/tests.rs diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 4805dc86e1f..7e53b686653 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -92,6 +92,12 @@ where Sink: OutputSink, ApplyResult: SinkResult, { + // Decline before input decoding or sink allocation when this sink cannot initialize rows that + // the mask skips. The capability and the operation are the same function pointer. + let Some(initialize_skipped_rows) = Sink::SKIPPED_ROWS_INITIALIZER else { + return Ok(None); + }; + // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An // element representation may decline when it cannot provide those values safely. let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { @@ -123,9 +129,7 @@ where // The loop writes only valid indices, but the sink still finishes a full-length output. // Initialize placeholders now; batch execution masks them before the result escapes. - if !Sink::initialize_skipped_rows(&mut rows) { - return Ok(None); - } + initialize_skipped_rows(&mut rows); // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first // immediate error, turn later callbacks into no-ops, and return before finishing the sink. @@ -179,3 +183,6 @@ fn finish_sink( Err(error) => Err(error), } } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/row/execute/sink/tests.rs b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs new file mode 100644 index 00000000000..bf036fe0d63 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use super::execute_sink_valid_rows; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::scalar_fn::DeferredError; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::VecExecutionArgs; +use crate::validity::Validity; + +struct NonSkippingSink; + +impl OutputSink for NonSkippingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + Err(vortex_err!( + "a non-skipping sink must decline before allocation" + )) + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { + true + } + + fn row<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} + + unsafe fn finish(self, _error: DeferredError) -> VortexResult { + Err(vortex_err!("a non-skipping sink must not finish")) + } +} + +#[test] +fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([true, false]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, ()>( + &args, + &DType::from(i64::PTYPE), + &valid, + &mut ctx, + |_| (), + |_, _, _| (), + )?; + + assert!(execution.is_none()); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index 2bfb6968445..c0f1c1e2cd7 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -19,7 +19,7 @@ use crate::scalar_fn::OutputElement; /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; /// skip-invalid execution can omit invalid rows when -/// [`initialize_skipped_rows`](Self::initialize_skipped_rows) returns `true`. +/// [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) is present. pub trait OutputSink: 'static + Sized { /// Whether this sink accepts [`DeferredError`] from its row closure instead of requiring a /// per-row [`VortexResult`]. @@ -36,6 +36,14 @@ pub trait OutputSink: 'static + Sized { where Self: 'a; + /// The operation that initializes every output position before skip-invalid execution. + /// + /// `None` declines skip-invalid execution before input decoding or sink allocation. A present + /// initializer **must** leave a legal arbitrary value in every row. Encoding support as the + /// initializer's presence prevents a separate capability flag from disagreeing with a no-op + /// method. + const SKIPPED_ROWS_INITIALIZER: Option fn(&mut Self::Rows<'a>)> = None; + /// The place a row closure writes one row through, borrowed from the sink. type Row<'a> where @@ -68,15 +76,6 @@ pub trait OutputSink: 'static + Sized { /// optimizer the output bounds it needs to remove the bounds check hidden in each row accessor. fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool; - /// Initialize every output position so skip-invalid execution may omit rows. - /// - /// Return `true` after leaving a legal arbitrary value in every row. Return `false` without - /// changing `rows` when the sink cannot support skipped rows. Combining the capability probe - /// and initialization prevents a separate support flag from disagreeing with a no-op method. - fn initialize_skipped_rows(_rows: &mut Self::Rows<'_>) -> bool { - false - } - /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; @@ -88,8 +87,7 @@ pub trait OutputSink: 'static + Sized { /// /// The executor must have completed every row callback successfully, and each callback must /// have returned this sink's [`WriteToken`](Self::WriteToken). When skipped rows are allowed, - /// [`initialize_skipped_rows`](Self::initialize_skipped_rows) must have returned `true` before - /// traversal. + /// [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) must have run before traversal. unsafe fn finish(self, error: DeferredError) -> VortexResult; } @@ -133,6 +131,13 @@ pub struct UninitElementSink { impl OutputSink for UninitElementSink { type Rows<'a> = &'a mut [MaybeUninit]; + + const SKIPPED_ROWS_INITIALIZER: Option fn(&mut Self::Rows<'a>)> = Some(|rows| { + for row in rows.iter_mut() { + row.write(T::default()); + } + }); + type Row<'a> = &'a mut MaybeUninit; type WriteToken = InitializedElement; @@ -155,14 +160,6 @@ impl OutputSink for UninitElementSink { rows.len() == row_count } - fn initialize_skipped_rows(rows: &mut Self::Rows<'_>) -> bool { - for row in rows.iter_mut() { - row.write(T::default()); - } - - true - } - fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { &mut rows[index] } From 1993f225eeb574dd435d69e2ed719d22a2e22942 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 11:29:43 -0400 Subject: [PATCH 050/160] Tighten RowFn batch fast paths Signed-off-by: Connor Tsui --- .../src/scalar_fn/row/batch/execution.rs | 39 ++++++--------- vortex-array/src/scalar_fn/row/batch/tests.rs | 49 +++++++++++++++++++ 2 files changed, 64 insertions(+), 24 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index c227c656697..fb485b0c2d8 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -84,15 +84,12 @@ impl Batch { let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) .map(|index| args.get(index)) .collect::>()?; - if let Some((index, input)) = inputs - .iter() - .enumerate() - .find(|(_, input)| input.len() != row_count) - { + + for (index, input) in inputs.iter().enumerate() { vortex_ensure_eq!( input.len(), row_count, - "the {id} input {index} has {} rows but execution declares {row_count}", + "the {id} input {index} must have {row_count} rows, got {}", input.len(), ); } @@ -123,8 +120,9 @@ impl Batch { /// /// The kernel may ignore input validity. It receives valid-only rows when required, and its /// output **must** match the planned dtype up to nullability. `reduce` receives the original - /// inputs exactly once. `try_unfiltered` receives the originals plus a mixed validity mask; - /// `Ok(None)` selects filter-and-scatter. + /// inputs exactly once, before the generic all-constant broadcast, so a function-owned encoded + /// implementation takes precedence. `try_unfiltered` receives the originals plus a mixed + /// validity mask; `Ok(None)` selects filter-and-scatter. pub fn execute( &self, reduce: impl FnOnce(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult>, @@ -136,21 +134,19 @@ impl Batch { ) -> VortexResult>, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Strictness: any null-constant input forces an all-null result without evaluating the - // kernel. - if self - .inputs - .iter() - .any(|input| input.as_constant().is_some_and(|scalar| scalar.is_null())) + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.validity, Validity::AllInvalid) + || self + .inputs + .iter() + .any(|input| input.as_constant().is_some_and(|scalar| scalar.is_null())) { return Ok(self.all_null()); } - // An all-null batch has no observable row work. Other batches offer the encoding-aware - // hook the original inputs once, before slicing or filtering can change their encodings. - if matches!(self.validity, Validity::AllInvalid) { - return Ok(self.all_null()); - } + // The function-owned encoded path takes precedence over the generic all-constant + // broadcast and sees the original inputs before slicing or filtering changes them. if let Some(values) = reduce(self.kernel_args(&self.inputs, self.row_count), ctx)? { return self.finalize_reduced(values); } @@ -207,11 +203,6 @@ impl Batch { retry_deferred_error: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Every row is null, so the kernel has nothing to contribute. - if matches!(self.validity, Validity::AllInvalid) { - return Ok(self.all_null()); - } - let values = match kernel(self.kernel_args(&self.inputs, self.row_count), ctx)? { RowExecution::Output(values) => values, RowExecution::DeferredError(error) if retry_deferred_error => { diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 7c708582c9e..69d0a145528 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -37,6 +37,9 @@ struct RetryConstantAdd; #[derive(Clone)] struct NullarySeven; +#[derive(Clone)] +struct OriginalInputReducer; + struct I64Sink(BufferMut); impl OutputSink for I64Sink { @@ -134,6 +137,39 @@ impl RowFn for RetryConstantAdd { } } +impl RowFn for OriginalInputReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.original_input_reducer"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 3 { + return Ok(Some(ConstantArray::new(42_i64, 3).into_array())); + } + + Ok(None) + } +} + #[test] fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { static ID: CachedId = CachedId::new("test.row_batch"); @@ -165,6 +201,19 @@ fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { Ok(()) } +#[test] +fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + let expected = ConstantArray::new(42_i64, 3).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + #[rstest] #[case::dense(RowPolicy::Dense)] #[case::dense_with_retry(RowPolicy::DenseWithRetry)] From 7814a2adb735784dcbcd403c057d178292bd876e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 11:32:14 -0400 Subject: [PATCH 051/160] Defer encoded RowFn errors through valid rows Signed-off-by: Connor Tsui --- vortex-array/benches/strict_validity.rs | 5 +- .../src/scalar_fn/row/batch/execution.rs | 40 +++++++++- vortex-array/src/scalar_fn/row/batch/tests.rs | 75 ++++++++++++++++++- vortex-array/src/scalar_fn/row/mod.rs | 1 + vortex-array/src/scalar_fn/row/row_fn.rs | 11 ++- .../src/scalar_fns/cosine_similarity.rs | 11 +-- vortex-tensor/src/scalar_fns/inner_product.rs | 6 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 5 +- 8 files changed, 132 insertions(+), 22 deletions(-) diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs index bd2d7bbef10..b2fa7a3a824 100644 --- a/vortex-array/benches/strict_validity.rs +++ b/vortex-array/benches/strict_validity.rs @@ -38,6 +38,7 @@ use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowExecution; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; @@ -100,8 +101,8 @@ impl RowFn for LazyDouble { _options: &Self::Options, args: &[ArrayRef], ctx: &mut ExecutionCtx, - ) -> VortexResult> { - doubled(&args[0], ctx).map(Some) + ) -> VortexResult> { + doubled(&args[0], ctx).map(|output| Some(RowExecution::Output(output))) } } diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index fb485b0c2d8..fd4b0401fb3 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -4,6 +4,7 @@ //! Null propagation, constant folding, and strategy execution for one columnar batch. use smallvec::SmallVec; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -125,7 +126,7 @@ impl Batch { /// validity mask; `Ok(None)` selects filter-and-scatter. pub fn execute( &self, - reduce: impl FnOnce(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult>, + reduce: impl FnOnce(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult>, kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_unfiltered: impl FnOnce( KernelArgs<'_>, @@ -147,8 +148,13 @@ impl Batch { // The function-owned encoded path takes precedence over the generic all-constant // broadcast and sees the original inputs before slicing or filtering changes them. - if let Some(values) = reduce(self.kernel_args(&self.inputs, self.row_count), ctx)? { - return self.finalize_reduced(values); + if let Some(execution) = reduce(self.kernel_args(&self.inputs, self.row_count), ctx)? { + match execution { + RowExecution::Output(values) => return self.finalize_reduced(values), + RowExecution::DeferredError(error) => { + return self.resolve_reduced_error(error, kernel, try_unfiltered, ctx); + } + } } // All inputs constant, and their conjoined validity proves every row non-null. This sees @@ -350,6 +356,34 @@ impl Batch { } } + /// Resolve deferred evidence from the encoded path by executing only observable rows. + fn resolve_reduced_error( + &self, + error: VortexError, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + /// Pair an input view with this batch's planning metadata. fn kernel_args<'b>(&'b self, arrays: &'b [ArrayRef], row_count: usize) -> KernelArgs<'b> { KernelArgs { diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 69d0a145528..6a59be49ef0 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -40,6 +40,9 @@ struct NullarySeven; #[derive(Clone)] struct OriginalInputReducer; +#[derive(Clone)] +struct DeferredOriginalReducer; + struct I64Sink(BufferMut); impl OutputSink for I64Sink { @@ -128,9 +131,11 @@ impl RowFn for RetryConstantAdd { _options: &Self::Options, args: &[ArrayRef], _ctx: &mut ExecutionCtx, - ) -> VortexResult> { + ) -> VortexResult> { if args[0].len() == 1 { - return Ok(Some(ConstantArray::new(0u8, args[0].len()).into_array())); + return Ok(Some(RowExecution::Output( + ConstantArray::new(0u8, args[0].len()).into_array(), + ))); } Ok(None) @@ -161,15 +166,49 @@ impl RowFn for OriginalInputReducer { _options: &Self::Options, args: &[ArrayRef], _ctx: &mut ExecutionCtx, - ) -> VortexResult> { + ) -> VortexResult> { if args[0].len() == 3 { - return Ok(Some(ConstantArray::new(42_i64, 3).into_array())); + return Ok(Some(RowExecution::Output( + ConstantArray::new(42_i64, 3).into_array(), + ))); } Ok(None) } } +impl RowFn for DeferredOriginalReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.deferred_original_reducer"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::DeferredError(vortex_err!( + InvalidArgument: "encoded payload failed" + )))) + } +} + #[test] fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { static ID: CachedId = CachedId::new("test.row_batch"); @@ -201,6 +240,34 @@ fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { Ok(()) } +#[test] +fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![1, u8::MAX], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1_u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::new(vec![2_u8, 0], Validity::from_iter([true, false])); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} + +#[test] +fn test_reduce_encoded_defers_errors_behind_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + #[test] fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { let input = ConstantArray::new(7_i64, 3).into_array(); diff --git a/vortex-array/src/scalar_fn/row/mod.rs b/vortex-array/src/scalar_fn/row/mod.rs index fda1dfdfa75..50c005fabe6 100644 --- a/vortex-array/src/scalar_fn/row/mod.rs +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -14,6 +14,7 @@ //! batch. mod execute; +pub use execute::RowExecution; mod batch; diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index 7fd8b3c08fb..99b544e44aa 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -15,6 +15,7 @@ use super::visitor::RowVisitor; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; +use crate::scalar_fn::RowExecution; use crate::scalar_fn::ScalarFnId; /// A scalar function computed one row at a time. @@ -68,9 +69,11 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// Try an encoding-aware implementation before decoding the inputs into row elements. /// - /// `None` continues to the dispatched row loop. `Some(output)` skips that loop. The output can - /// remain encoded or lazy. For non-nullary functions, batch execution calls this hook at most - /// once with the original, unfiltered arrays; slices and compacted retries do not reach it. + /// `None` continues to the dispatched row loop. [`Output`](RowExecution::Output) skips that + /// loop and may remain encoded or lazy. [`DeferredError`](RowExecution::DeferredError) reruns + /// only valid rows when null payloads may have caused the failure. For non-nullary functions, + /// batch execution calls this hook at most once with the original, unfiltered arrays; slices + /// and compacted retries do not reach it. /// /// Like a dense row closure, this hook must be total over every stored payload, including /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or @@ -88,7 +91,7 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { options: &Self::Options, args: &[ArrayRef], ctx: &mut ExecutionCtx, - ) -> VortexResult> { + ) -> VortexResult> { _ = (options, args, ctx); Ok(None) } diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 18930b89009..4b957528eab 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -18,6 +18,7 @@ use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowExecution; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; @@ -122,18 +123,18 @@ impl RowFn for CosineSimilarity { _options: &Self::Options, args: &[ArrayRef], ctx: &mut ExecutionCtx, - ) -> VortexResult> { + ) -> VortexResult> { let lhs = args[0].clone(); let rhs = args[1].clone(); match NormalizedOrientation::classify(&lhs, &rhs) { - NormalizedOrientation::Both { lhs, rhs } => { - cosine_both_normalized(lhs, rhs, ctx).map(Some) - } + NormalizedOrientation::Both { lhs, rhs } => cosine_both_normalized(lhs, rhs, ctx) + .map(|output| Some(RowExecution::Output(output))), NormalizedOrientation::One { normalized_array, plain, - } => cosine_one_normalized(normalized_array, plain, ctx).map(Some), + } => cosine_one_normalized(normalized_array, plain, ctx) + .map(|output| Some(RowExecution::Output(output))), NormalizedOrientation::Neither => Ok(None), } } diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index b972fe54b96..e2b6def5595 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -16,6 +16,7 @@ use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowExecution; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; @@ -95,7 +96,7 @@ impl RowFn for InnerProduct { _options: &Self::Options, args: &[ArrayRef], _ctx: &mut ExecutionCtx, - ) -> VortexResult> { + ) -> VortexResult> { let len = args[0].len(); Ok(match NormalizedOrientation::classify(&args[0], &args[1]) { @@ -119,7 +120,8 @@ impl RowFn for InnerProduct { Some(dot.binary(norms, Operator::Mul)?) } NormalizedOrientation::Neither => None, - }) + } + .map(RowExecution::Output)) } } diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index dbc2e27d33c..50fa7819d66 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -16,6 +16,7 @@ use vortex_array::dtype::proto::dtype as pb; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowExecution; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; @@ -95,7 +96,7 @@ impl RowFn for L2Norm { _options: &Self::Options, args: &[ArrayRef], _ctx: &mut ExecutionCtx, - ) -> VortexResult> { + ) -> VortexResult> { let input = &args[0]; if !input.is::() { return Ok(None); @@ -108,7 +109,7 @@ impl RowFn for L2Norm { norms.dtype(), ); vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); - Ok(Some(norms)) + Ok(Some(RowExecution::Output(norms))) } } From 6714a420e6a579dc0285ee5f09c2eb748714b206 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 11:40:36 -0400 Subject: [PATCH 052/160] Cover RowFn batch strategy contracts Signed-off-by: Connor Tsui --- .../src/scalar_fn/row/batch/policy.rs | 76 +++++++ vortex-array/src/scalar_fn/row/batch/tests.rs | 195 ++++++++++++++++++ .../src/scalar_fn/row/types/element/tuple.rs | 42 ++++ 3 files changed, 313 insertions(+) diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs index d017b5278d6..5cdbee5af55 100644 --- a/vortex-array/src/scalar_fn/row/batch/policy.rs +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -76,3 +76,79 @@ impl RowPolicy { } } } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::RowPolicy; + use crate::ArrayRef; + use crate::ExecutionCtx; + use crate::dtype::DType; + use crate::scalar_fn::InputElement; + + struct SparseFallibleElement; + + // SAFETY: the varying view reports length zero, so no index satisfies the unchecked-read + // precondition. + unsafe impl InputElement for SparseFallibleElement { + type Column = (); + type Varying<'a> = (); + type Elem<'a> = (); + + const DENSE_SAFE: bool = false; + const DECODE_FALLIBLE: bool = true; + + fn validate(_dtype: &DType) -> VortexResult<()> { + Ok(()) + } + + fn decode(_array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(()) + } + + fn get(_column: &Self::Column, _index: usize) -> Self::Elem<'_> {} + + fn varying(_column: &Self::Column) -> Self::Varying<'_> {} + + fn varying_len(_column: &Self::Varying<'_>) -> usize { + 0 + } + + fn get_varying<'a>(_column: &Self::Varying<'a>, _index: usize) -> Self::Elem<'a> {} + } + + #[test] + fn test_owned_output_policy() { + assert_eq!(RowPolicy::for_owned_output::<(i64,)>(), RowPolicy::Dense); + assert_eq!( + RowPolicy::for_owned_output::<(SparseFallibleElement,)>(), + RowPolicy::ValidOnly, + ); + } + + #[test] + fn test_deferred_output_policy() { + assert_eq!( + RowPolicy::for_deferred_output::<(i64,)>(), + RowPolicy::DenseWithRetry, + ); + assert_eq!( + RowPolicy::for_deferred_output::<(SparseFallibleElement,)>(), + RowPolicy::ValidOnly, + ); + } + + #[test] + fn test_sink_policy() { + assert_eq!(RowPolicy::for_sink::<(i64,), ()>(), RowPolicy::Dense); + assert_eq!( + RowPolicy::for_sink::<(i64,), VortexResult<()>>(), + RowPolicy::ValidOnly, + ); + assert_eq!( + RowPolicy::for_sink::<(SparseFallibleElement,), ()>(), + RowPolicy::ValidOnly, + ); + } +} diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 6a59be49ef0..be0b80bcc22 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -1,6 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + use rstest::rstest; use vortex_buffer::BufferMut; use vortex_error::VortexResult; @@ -11,16 +15,19 @@ use super::super::execute::RowExecution; use super::Batch; use super::BatchPlan; use super::RowPolicy; +use super::finalize_kernel_output; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::NativePType; +use crate::dtype::Nullability; use crate::scalar_fn::DeferredError; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::OutputSink; @@ -43,6 +50,19 @@ struct OriginalInputReducer; #[derive(Clone)] struct DeferredOriginalReducer; +#[derive(Clone)] +struct PreparedAdd { + visit: PreparedVisit, + prepares: Arc, +} + +#[derive(Clone, Copy)] +enum PreparedVisit { + Owned, + Sink, + Deferred, +} + struct I64Sink(BufferMut); impl OutputSink for I64Sink { @@ -209,6 +229,55 @@ impl RowFn for DeferredOriginalReducer { } } +impl RowFn for PreparedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.prepared_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + let prepares = Arc::clone(&self.prepares); + let prepare = move |(_lhs, rhs): (Option, Option)| { + prepares.fetch_add(1, Ordering::Relaxed); + rhs + }; + + match self.visit { + PreparedVisit::Owned => visitor + .visit_prepared::<(i64, i64), i64, _>(prepare, |constant_rhs, (lhs, rhs)| { + lhs.wrapping_add(constant_rhs.unwrap_or(rhs)) + }), + PreparedVisit::Sink => visitor.visit_prepared_into::<(i64, i64), I64Sink, _, ()>( + prepare, + |constant_rhs, (lhs, rhs), output| { + *output = lhs.wrapping_add(constant_rhs.unwrap_or(rhs)); + }, + ), + PreparedVisit::Deferred => visitor.visit_prepared_deferred::<(i64, i64), i64, _, bool>( + prepare, + |constant_rhs, (lhs, rhs)| lhs.overflowing_add(constant_rhs.unwrap_or(rhs)), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "prepared add overflowed")); + } + + Ok(()) + }, + ), + } + } +} + #[test] fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { static ID: CachedId = CachedId::new("test.row_batch"); @@ -281,6 +350,132 @@ fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { Ok(()) } +#[test] +fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid([true, true])] +#[case::all_invalid([false, false])] +fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.resolve_validity"); + + let validity = Validity::Array(BoolArray::from_iter(validity).into_array()); + let input = PrimitiveArray::new(vec![4_i64, 5], validity).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_filters_and_scatters() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.filter_and_scatter"); + + let input = PrimitiveArray::new( + vec![10_i64, 20, 30, 40], + Validity::from_iter([true, false, true, false]), + ) + .into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 4); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.finalize_kernel_output"); + + let values = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let result_dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); + let mut ctx = array_session().create_execution_ctx(); + + let actual = finalize_kernel_output(*ID, &result_dtype, 2, values.clone())?; + let expected = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + assert_eq!(actual.dtype(), &result_dtype); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + assert!(finalize_kernel_output(*ID, &result_dtype, 3, values).is_err()); + + let bools = BoolArray::from_iter([true, false]).into_array(); + assert!(finalize_kernel_output(*ID, &result_dtype, 2, bools).is_err()); + Ok(()) +} + +#[rstest] +#[case::owned_constant(PreparedVisit::Owned, true)] +#[case::owned_varying(PreparedVisit::Owned, false)] +#[case::sink_constant(PreparedVisit::Sink, true)] +#[case::sink_varying(PreparedVisit::Sink, false)] +#[case::deferred_constant(PreparedVisit::Deferred, true)] +#[case::deferred_varying(PreparedVisit::Deferred, false)] +fn test_prepared_visits( + #[case] visit: PreparedVisit, + #[case] constant_rhs: bool, +) -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let rhs = if constant_rhs { + ConstantArray::new(3_i64, 2).into_array() + } else { + PrimitiveArray::from_iter([3_i64, 4]).into_array() + }; + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let prepares = Arc::new(AtomicUsize::new(0)); + let function = PreparedAdd { + visit, + prepares: Arc::clone(&prepares), + }; + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&function, &EmptyOptions, &args, &mut ctx)?; + let expected = if constant_rhs { + PrimitiveArray::from_iter([4_i64, 5]).into_array() + } else { + PrimitiveArray::from_iter([4_i64, 6]).into_array() + }; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(prepares.load(Ordering::Relaxed), 1); + Ok(()) +} + #[rstest] #[case::dense(RowPolicy::Dense)] #[case::dense_with_retry(RowPolicy::DenseWithRetry)] diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple.rs b/vortex-array/src/scalar_fn/row/types/element/tuple.rs index 643d976f298..680ae0b4e9c 100644 --- a/vortex-array/src/scalar_fn/row/types/element/tuple.rs +++ b/vortex-array/src/scalar_fn/row/types/element/tuple.rs @@ -419,8 +419,20 @@ impl IndexedElementTuple for (Left, Right #[cfg(test)] mod tests { use vortex_compute::lane_kernels::IndexedSource; + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_mask::Mask; use super::UnaryTupleSource; + use super::batch_constant; + use crate::IntoArray; + use crate::arrays::ConstantArray; + use crate::arrays::ExtensionArray; + use crate::arrays::MaskedArray; + use crate::dtype::Nullability; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::validity::Validity; #[test] fn test_unary_tuple_source_reads_one_tuple_per_row() { @@ -430,4 +442,34 @@ mod tests { // SAFETY: index one is within the three-element source. assert_eq!(unsafe { source.get_unchecked(1) }, (20,)); } + + #[test] + fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> { + let child = ConstantArray::new(7_i64, 3).into_array(); + let masked = + MaskedArray::try_new(child, Validity::from_iter([true, false, true]))?.into_array(); + let filtered = masked.filter(Mask::from_iter([true, true, false]))?; + + let Some(constant) = batch_constant(&filtered) else { + vortex_bail!("filtered masked constant must remain batch-constant"); + }; + + assert!(constant.as_constant().is_some()); + Ok(()) + } + + #[test] + fn test_batch_constant_preserves_filtered_extension() -> VortexResult<()> { + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); + let extension = + ExtensionArray::new(ext_dtype, ConstantArray::new(7_i64, 3).into_array()).into_array(); + let filtered = extension.filter(Mask::from_iter([true, false, true]))?; + + let Some(constant) = batch_constant(&filtered) else { + vortex_bail!("filtered extension storage must remain batch-constant"); + }; + + assert_eq!(constant.dtype(), extension.dtype()); + Ok(()) + } } From 443ef2d3bb82ff3d27754c03767ff0e13bd4016e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 11:41:08 -0400 Subject: [PATCH 053/160] Clarify the RowFn sink finish proof Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/row/execute/sink.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 7e53b686653..4cb987e282c 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -175,8 +175,9 @@ fn finish_sink( sink: S, deferred_error: DeferredError, ) -> VortexResult { - // SAFETY: callers reach this helper only after successful traversal. Dense traversal visited - // every addressable row; skipped-row traversal initialized every row before visiting its mask. + // SAFETY: callers reach this helper only after every completed callback returned the sink's + // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. + // The sink contract defines how that evidence establishes initialization of its row storage. match unsafe { sink.finish(deferred_error) } { Ok(output) => Ok(RowExecution::Output(output)), Err(error) if deferred_error.occurred() => Ok(RowExecution::DeferredError(error)), From d20e898c7c3fcead6947e47e9945a7005458514c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 12:30:10 -0400 Subject: [PATCH 054/160] Document RowFn target branch gate Signed-off-by: Connor Tsui --- research/rowfn-review-followup/README.md | 175 +++++++++++++----- .../codegen/target-branch-round2-summary.md | 94 ++++++++++ 2 files changed, 225 insertions(+), 44 deletions(-) create mode 100644 research/rowfn-review-followup/codegen/target-branch-round2-summary.md diff --git a/research/rowfn-review-followup/README.md b/research/rowfn-review-followup/README.md index 2e701266e16..015a34b798f 100644 --- a/research/rowfn-review-followup/README.md +++ b/research/rowfn-review-followup/README.md @@ -7,30 +7,33 @@ Platform: Apple Silicon `arm64`, macOS 15.7.3, Rust 1.91.0, LLVM 21.1.2. This ma 128-bit NEON and cannot reproduce or compare the pinned Ryzen wall-clock results. This investigation uses optimized LLVM IR and correctness tests only. -Baseline: `833632aaa3cab59fb4a7d4f001df26975b2267a1` on `ct/row-fn`. +Round 1 baseline: `833632aaa3cab59fb4a7d4f001df26975b2267a1` on `ct/row-fn`. + +Round 2 target baselines: `84837ad36f` on `origin/ct/row-fn-api` and `32ad0bf3b7` on +`origin/ct/row-fn-numeric`. ## Result -| # | Verdict | Evidence and status | -| ---: | --- | --- | -| 1 | Fixed, not upstreamed | `44457b7654` makes `OutputSink::finish` unsafe. The zero-unsafe external reproduction changes from reading allocator contents to `E0133`. All gate items pass, but the local `ct/row-fn-api` ref contains divergent unpushed history and was not overwritten. | -| 2 | Fixed, not upstreamed | `InputElement` is an unsafe trait with local safety proofs. `ElementTuple` and `IndexedElementTuple` remain sealed framework traits. Same upstream status as #1. | -| 3 | Fixed, not upstreamed | The existing unsafe `InitializedElement::write` token remains unchanged. Unsafe publication now belongs to `OutputSink::finish`, and the generic executor owns the proof. Same upstream status as #1. | -| 4 | Fixed, not upstreamed | `initialize_skipped_rows` now returns its capability result. A separate support constant cannot disagree with a no-op default. Same upstream status as #1. | -| 5 | Fixed, not upstreamed | `5204fb2be5` documents totality and probes original inputs once. Gate item 3 fails because the batch refactor adds an owned mixed-path bounds edge. | -| 6 | Not started | The requested redesign was out of scope. A temporary manual implementation reproduced `E0119`. The options analysis is below. | -| 7 | Fixed, not upstreamed | `44457b7654` documents why private `NumericBinary` borrows the registered `Binary` ID and makes privacy the registration guard. Same upstream status as #1. | -| 8 | Fixed, not upstreamed | `5204fb2be5` validates every input length in `Batch::new`. The regression test fails before the fix. Gate item 3 fails as described for #5. | -| 9 | Refuted | `MaskedArray::try_new` enforces an all-valid child, and null `ConstantArray` validity is `AllInvalid`. Both couplings are constructor invariants. | -| 10 | Fixed, not upstreamed | `5204fb2be5` probes once before retry. The regression test fails before the fix and passes after it. Gate item 3 fails as described for #5. | -| 11 | Investigated, needs x86 | The path is unreachable because no sink sets `ERRORS_ARE_DEFERRED`. Deleting it perturbed arithmetic IR, so no change remains. | -| 12 | Refuted | Filtering preserves every constant shape recognized by `batch_constant`: literal `Constant`, constant `Masked` child, and constant `Extension` storage. `ConstElems` stays consistent. | -| 13 | Fixed, not upstreamed | `5204fb2be5` runs the encoding probe before one-row broadcast, against the original arrays. Gate item 3 fails as described for #5. | -| 14 | Investigated, needs x86 | The five deferred word implementations are unreachable. Their deletion changed codegen-unit placement and owned-loop IR, so the machinery remains. | -| 15 | Refuted | Mixed `i64` add/sub and `i32` multiply contain broadcast vector loops on `arm64`. They do not fall back to scalar row execution. | -| 16 | Investigated, needs x86 | Full-row skipped initialization and non-breaking mask traversal are independent costs. The proposed API and early-exit split are below. | -| 17 | Investigated, needs x86 | `scatter_valid` allocates one `u64` per original row. The cited `vortex-spatial/src/scalar_fn/execute/geo_types.rs` path does not exist at this revision. | -| 18 | Refuted | No consumer requires row output to retain 256-byte physical alignment. Alignment-sensitive consumers call `ensure_aligned`. Any performance change still needs x86 evidence. | +| # | Verdict | Gate branch | Evidence and status | +| ---: | --- | --- | --- | +| 1 | Fixed, not upstreamed | API target plus downstream numeric target | `OutputSink::finish(self)` is unsafe. The zero-unsafe external reproduction changes from reading allocator contents to `E0133`. Publication awaits approval for the combined API commit. | +| 2 | Fixed, not upstreamed | API target plus downstream numeric target | `InputElement` is an unsafe trait with local safety proofs. `ElementTuple` and `IndexedElementTuple` remain sealed framework traits. Publication awaits approval for the combined API commit. | +| 3 | Mitigated | API target plus downstream numeric target | Unsafe publication belongs to `OutputSink::finish`, and each successful callback must return the sink's write token. The token is not type-tied to the exact row handle; misuse now requires violating an unsafe sink/input contract rather than safe client code. | +| 4 | Fixed, not upstreamed | API target plus downstream numeric target | `SKIPPED_ROWS_INITIALIZER: Option` makes capability and operation one fact while restoring decline before decode and allocation. Publication awaits approval for the combined API commit. | +| 5 | Fixed, not upstreamed | API target plus downstream numeric target | `reduce_encoded` probes original inputs once and returns `RowExecution`, so encoded reductions can defer errors behind nulls through the same validity path as row execution. Publication awaits approval for the combined API commit. | +| 6 | Investigated | Design only | A temporary manual implementation still reproduces `E0119`. A public `execute_rows` free function is the recommended redesign; it is not implemented in this pass. | +| 7 | Fixed, not upstreamed | Numeric target | Private `NumericBinary` deliberately borrows the registered `Binary` ID and is guarded by privacy. Publication awaits the API commit. | +| 8 | Fixed, not upstreamed | API target plus downstream numeric target | `Batch::new` validates each input length directly. The regression test fails before the fix. Publication awaits approval for the combined API commit. | +| 9 | Refuted | API target spot-check | `MaskedArray::try_new` enforces an all-valid child, and null `ConstantArray` validity is `AllInvalid`. Both couplings remain constructor invariants. | +| 10 | Fixed, not upstreamed | API target plus downstream numeric target | The encoding probe runs once before retry. The regression test fails before the fix and passes after it. Publication awaits approval for the combined API commit. | +| 11 | Fixed, not upstreamed | API target plus downstream numeric target | No sink can defer errors. The unreachable `finish_sink` retry classification and deferred finish argument are deleted. Publication awaits explicit approval for this deletion. | +| 12 | Refuted | API target tests and source spot-check | Filtering preserves literal, masked-child, and extension-storage constants recognized by `batch_constant`; new masked and extension tests pin the behavior. | +| 13 | Fixed, not upstreamed | API target plus downstream numeric target | The encoding probe runs before one-row broadcast and sees the original arrays. Publication awaits approval for the combined API commit. | +| 14 | Fixed, not upstreamed | API target plus downstream numeric target | The five unreachable deferred `SinkResult` word implementations and their capability pairing are deleted. Publication awaits explicit approval for this deletion. | +| 15 | Refuted | Downstream numeric target IR | Mixed `i64` add/sub and `i32` multiply contain broadcast vector loops on `arm64`; they are not scalar fallbacks. | +| 16 | Investigated, needs x86 | Analysis only | Full-row skipped initialization and non-breaking mask traversal remain independent costs. No performance change is made. | +| 17 | Investigated, needs x86 | API target source | `scatter_valid` still allocates one `u64` per original row. Whether replacing the gather pays for itself is a standalone benchmark question. | +| 18 | Refuted | API target source spot-check | No consumer requires row output to retain 256-byte physical alignment. Alignment-sensitive consumers call `ensure_aligned`; a performance change would still need x86 evidence. | ## Reproductions @@ -68,6 +71,14 @@ The first proposed reproduction required every filtered input to become `Constan a one-row `PrimitiveArray` keeps it primitive, so that version did not exercise the defect. The one-row encoding probe establishes the same retry violation without assuming a filter encoding. +Round 2 adds `test_reduce_encoded_defers_errors_behind_nulls`. An encoded reducer reports a +`RowExecution::DeferredError`; mixed validity reruns only observable rows, all-valid validity makes +the error fatal, and all-invalid validity produces the declared all-null result. + +The early-decline regression uses a sink whose `with_capacity` returns an error. Before the round 2 +fix, `execute_sink_valid_rows` reaches that allocation before declining. The candidate observes the +absent initializer and returns `None` before decoding inputs or constructing the sink. + ## API options for the blanket vtable ### `RowFnAdaptor` @@ -77,12 +88,15 @@ Registration and expression construction must wrap every function. Existing call the concrete function type also change. Arithmetic can delegate to the same generic row executor, but the wrapper changes monomorph identities and requires the full IR gate. -### Public `execute_rows` +### Public `execute_rows` (recommended) -A public free function lets each adopter implement `ScalarFnVTable` and delegate only execution. -It preserves all five vtable hooks but repeats arity, child naming, return dtype, strictness, and -fallibility boilerplate in each implementation. The arithmetic row loop can remain the same helper -monomorph. The surrounding vtable implementation still requires IR verification. +A public free function lets each adopter keep its `ScalarFnVTable` implementation and delegate only +execution. `Binary` keeps `coerce_args`, both simplifiers, and `fmt_sql`; `Between` and `Like` keep +their SQL formatting; the other concrete functions keep their existing encoded and validity hooks. +This is the least disruptive path for functions that already have a vtable. Its call-site cost is +one explicit `execute` delegation per adopter. Simple RowFn-only functions lose the blanket +one-line adoption story unless a separate opt-in adaptor is also provided. The arithmetic row loop +can remain the same helper monomorph, but the surrounding delegation still needs the full IR gate. ### Hooks on `RowFn` @@ -91,7 +105,18 @@ The blanket implementation can forward them. This keeps call sites unchanged but `ScalarFnVTable` surface and creates two contracts for each hook. Default hook forwarding remains outside the arithmetic loop, but the public trait change still requires the full IR gate. -No option is implemented here. +`RowFnAdaptor` changes registration and expression construction to name a wrapper and changes +monomorph identities. Re-declaring the hooks on `RowFn` duplicates the `ScalarFnVTable` contract. +Neither cost is justified merely to preserve the blanket implementation. No redesign is implemented +here. + +## Coverage added in round 2 + +The new tests cover array-backed validity resolution, filter/scatter finalization, one-row constant +broadcast, output dtype and length validation, masked and extension constant unwrapping, all three +`RowPolicy` constructors, early decline for non-skipping sinks, dense-retry suppression, and all +three prepared visitor forms with both constant and varying inputs. `prepare` is asserted to run +once per batch. ## Performance findings @@ -104,9 +129,14 @@ loop-carried vector phi and rich error construction remains outside the loop. Se ### Skipped rows -`UninitElementSink::initialize_skipped_rows` writes `T::default()` to every row because its API has -no mask. A mask-aware API can accept `&Mask` and initialize only unset positions. This is a public -sink API change and needs x86 measurement for sparse and dense masks. +`UninitElementSink::SKIPPED_ROWS_INITIALIZER` writes `T::default()` to every row because the +initializer receives no mask. A mask-aware API can accept `&Mask` and initialize only unset +positions. This is a public sink API change and needs x86 measurement for sparse and dense masks. + +The function-pointer option restores a compile-time capability fact without creating a second +boolean that can disagree with the operation. `RowPolicy::for_sink` does not need it: policy decides +whether dense execution is semantically legal, while skip support decides whether the later +valid-row sink attempt proceeds or falls back to filter-and-scatter. The early-exit problem is separable. A fallible or breakable mask iterator can stop after the first row error without changing `OutputSink`. `Mask::indices()` is not an equivalent production fix @@ -115,9 +145,11 @@ because it can materialize indices. No change is made here. ### Scatter `scatter_valid` allocates `vec![0u64; valid.len()]`, fills ranks for set-bit runs, performs `take`, -and applies validity. A run-based scatter can copy dense value ranges directly into a pre-sized -output. That design is encoding-sensitive and needs a focused x86 benchmark. The spatial precedent -named in the finding is absent from this revision. +and applies validity. The cited spatial helper does exist on the target branches, but its typed +primitive/bool output construction is not a reusable implementation for arbitrary `ArrayRef` +encodings. A run-based scatter can copy dense value ranges directly into a pre-sized output, but +that design is encoding-sensitive. Whether avoiding the eight-byte-per-row gather index matters +needs a focused x86 benchmark. ### Alignment @@ -128,27 +160,47 @@ assumes 256-byte alignment. This is not a correctness requirement. ## Code generation gate -The API-only commit preserves the six owned arithmetic monomorphs and both `i64` division sink -paths. No new closure call, loop bounds check, vector loss, failure spill, or in-loop error -construction appears. See `codegen/api-contract-summary.md`. +Round 1's gate was measured on `ct/row-fn`, not on either branch that would receive the change. +Round 2 starts from exact `origin/ct/row-fn-api` and `origin/ct/row-fn-numeric` tips. The API branch +does not instantiate the production arithmetic monomorphs, so its candidate is compiled through +the numeric branch after rebasing the two numeric commits onto it. + +On that target lineage, the `5204fb2` `output[index]` bounds regression does not reproduce. Owned +arithmetic has two bounds sites after the candidate rather than three at the numeric baseline. The +candidate and the separate deferred-sink deletion preserve inlining, vector factors, loop-carried +failure phis, overflow checks, and out-of-loop error construction. The deletion changes codegen-unit +placement but none of the stated loop properties. See +`codegen/target-branch-round2-summary.md`. + +Staged capture found that the `RowExecution` reducer signature alone restores the third bounds +edge. The required coverage and final safety comment change placement again and remove it in the +final combined tree. Do not upstream or benchmark the reducer-signature commit in isolation. + +The exact tested trees were consolidated without changing their contents: + +- API candidate `f759998dcefa69b11d11b281e3bbebb6b88584e4`, based on `84837ad36f`. +- Numeric candidate `b4900072c1300238d246d395d986466db21583f6`, based on the API candidate. -The batch correctness commit adds one `panic_bounds_check` site for `output[index]` in the mixed -owned branch even though `owned.rs` is unchanged. The vector loops remain, but gate item 3 fails. -The commit stays on `ct/row-fn` and needs an x86 rerun before any upstream attempt. +Both remote tips still matched the recorded baselines after the final fetch. The candidates remain +unpublished because publishing the combined API commit, including the broad deferred-sink +deletion, requires explicit approval. ## Do not do this -Do not delete the unreachable deferred sink machinery as semantically inert cleanup. The deletion -changed codegen-unit placement and the optimized owned arithmetic IR. It was reverted. +Do not use a `ct/row-fn` IR result as the upstream gate. Every hot-path file and the linked numeric +code differ on the target branches; round 1's bounds regression does not reproduce there. -Do not move the `reduce_encoded` probe without checking the owned mixed-path bounds edge. The -correctness fix is valid, but the current source shape does not pass the arithmetic IR gate. +Do not call the deferred-sink deletion byte-identical. It changes codegen-unit placement on the +numeric target even though the arithmetic loop structure passes the gate. + +Do not cherry-pick the `RowExecution` reducer-signature change without its required coverage and +the final target-tree gate. Its isolated target build reintroduces the owned mixed-path bounds edge. The investigation did not challenge any guardrail in the task. `owned.rs`, numeric primitive inlining policy, the numeric `Vec`, `BorrowedExecutionArgs`, and LLVM loop flags remain unchanged. -## Verification +## Round 1 verification ```text cargo nextest run -p vortex-array -p vortex-spatial -p vortex-tensor @@ -166,3 +218,38 @@ PYO3_PYTHON=.venv/bin/python cargo clippy --all-targets --all-features The first clippy run selected `/usr/bin/python3` 3.9 and stopped because `abi3-py311` requires Python 3.11. Rerunning with the repository virtual environment completed cleanly. + +## Round 2 verification + +API candidate: + +```text +cargo nextest run -p vortex-array -p vortex-spatial -p vortex-tensor + 3757 passed, 1 skipped + +cargo test --doc -p vortex-array + 73 passed, 13 ignored + +cargo +nightly fmt --all + passed + +PYO3_PYTHON=/Users/connor/spiral/vortex-data/vortex1/.venv/bin/python \ + cargo clippy --all-targets --all-features + passed +``` + +Downstream numeric candidate: + +```text +cargo nextest run -p vortex-array + 3401 passed, 1 skipped + +cargo test --doc -p vortex-array + 73 passed, 13 ignored + +cargo +nightly fmt --all + passed + +PYO3_PYTHON=.venv/bin/python cargo clippy -p vortex-array --all-targets --all-features + passed +``` diff --git a/research/rowfn-review-followup/codegen/target-branch-round2-summary.md b/research/rowfn-review-followup/codegen/target-branch-round2-summary.md new file mode 100644 index 00000000000..1538c1cc68e --- /dev/null +++ b/research/rowfn-review-followup/codegen/target-branch-round2-summary.md @@ -0,0 +1,94 @@ + + + +# RowFn target-branch code generation, round 2 + +Platform: Apple Silicon `arm64`, macOS 15.7.3, Rust 1.91.0, LLVM 21.1.2. + +API baseline: `84837ad36f9c8e7c2cec76920ec07f303d60be11` on +`origin/ct/row-fn-api`. + +Numeric baseline: `32ad0bf3b7` on `origin/ct/row-fn-numeric`, whose parent is the API baseline. + +The API crate does not instantiate the production numeric `execute_owned` and `execute_sink` +monomorphs. The arithmetic gate therefore uses the numeric branch twice: first at the exact numeric +baseline, then with its two commits rebased onto the candidate API tree. This measures the code that +will actually receive the framework changes. + +The final tested source trees were consolidated into API candidate `f759998dce` and downstream +numeric candidate `b4900072c`. Consolidation did not change either tree. + +## Command + +All captures use the repository bench profile with 16 codegen units and no LTO: + +```text +RUSTFLAGS='-C symbol-mangling-version=v0' \ + cargo rustc --profile bench -p vortex-array --lib -- \ + --emit=llvm-ir -C debuginfo=0 +``` + +## Framework candidate + +The candidate includes the round 1 contract and batch commits, the round 2 early-decline and batch +corrections, `RowExecution` from `reduce_encoded`, and the added coverage. + +| Monomorph | Vector IR at candidate | Bounds sites, baseline -> candidate | Closure calls | Failure reduction | +| --- | --- | ---: | ---: | --- | +| `i64` add | `<16 x i64>` plus remainder | 3 -> 2 | 0 -> 0 | vector and scalar phi | +| `i64` sub | `<16 x i64>` plus remainder | 3 -> 2 | 0 -> 0 | vector and scalar phi | +| `i64` mul | scalar unroll plus `<2 x i64>` support | 3 -> 2 | 0 -> 0 | scalar phi | +| `i32` mul | `<4 x i32>` and `<16 x i32>` | 3 -> 2 | 0 -> 0 | vector and scalar phi | +| `u64` mul | scalar unroll plus `<2 x i64>` support | 3 -> 2 | 0 -> 0 | scalar phi | +| `u16` mul | `<8 x i16>` plus wider unroll groups | 3 -> 2 | 0 -> 0 | vector and scalar phi | +| `i64` div, dense sink | scalar | 2 -> 2 | 0 -> 0 | immediate error | +| `i64` div, valid-row sink | scalar mask walk | 1 -> 1 | 0 -> 0 | immediate error | + +The extra `output[index]` bounds edge observed for `5204fb2` on `ct/row-fn` does not reproduce on +the target lineage. The target candidate instead removes that edge from every owned arithmetic +monomorph. The remaining bounds sites are outside the arithmetic vector bodies. + +The nullable `i64` add path reuses the same owned monomorph after validity resolution or filtered +retry. Nullable division uses the valid-row sink monomorph shown separately. + +Signed add and subtract retain their XOR/AND overflow test. Multiply retains its widening or +high-half checks. Failure values remain loop-carried phis rather than stack loads and stores, and +rich error construction remains outside every row loop. + +## Staged attribution + +The production changes were also compiled one stage at a time on the downstream numeric lineage. +Bounds counts below are per owned arithmetic monomorph; the dense and valid-row division sinks stay +at two and one throughout. + +| Stage | Owned bounds sites | Result | +| --- | ---: | --- | +| Round 1 API/batch candidate | 3 | Starting point for round 2 | +| Restore early skipped-row decline | 3 | No call, vector, failure-phi, or bounds regression | +| Fix length check, all-invalid branch, and reducer precedence | 2 | Removes the `output[index]` edge; other loop properties remain | +| Change `reduce_encoded` to `RowExecution`, alone | 3 | Fails the gate as a standalone change | +| Add the required strategy/prepared coverage and final safety comment | 2 | Final combined API tree passes | +| Delete unreachable deferred sink results | 2 | Final deletion also passes | + +The reducer-signature change must not be cherry-picked alone. Test-only source and comments should +not normally affect optimized library IR, but in this framework they change codegen-unit placement +and restore the desired bounds proof. Upstreaming is therefore gated on the combined API commit and +its exact final tree, not on an intermediate commit or on `ct/row-fn`. + +## Deferred sink deletion + +Findings 11 and 14 were tested as a separate commit after the framework candidate. The commit +removes `OutputSink::ERRORS_ARE_DEFERRED`, the five unreachable deferred `SinkResult` word +implementations, the unused deferred evidence passed to `OutputSink::finish`, and the unreachable +retry classification in `finish_sink`. + +The deletion moves the numeric monomorphs from codegen unit 10 to codegen unit 9, so the raw IR is +not byte-identical. Structurally, all eight monomorphs retain the candidate's vector factors, +broadcasts, two/one bounds sites, inlined closures, loop-carried failure phis, overflow checks, and +out-of-loop error construction. It passes the stated IR gate on the target lineage. + +## Mixed constants + +Mixed `i64` add and subtract retain `insertelement`/`shufflevector` broadcasts and `<16 x i64>` +vector arithmetic. Mixed `i32` multiply also retains fixed-width broadcast vector loops. LLVM's +vector factors span groups of 128-bit NEON registers; these paths are not scalar fallbacks. From 1ce3f1ee78cd49c4de4fdd00d85f5097780bff07 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 11:51:24 -0400 Subject: [PATCH 055/160] Remove unreachable deferred sink results Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 3 +- .../src/scalar_fn/row/batch/policy.rs | 6 +- vortex-array/src/scalar_fn/row/batch/tests.rs | 3 +- .../src/scalar_fn/row/execute/sink.rs | 22 +---- .../src/scalar_fn/row/execute/sink/tests.rs | 3 +- vortex-array/src/scalar_fn/row/mod.rs | 5 +- vortex-array/src/scalar_fn/row/types/mod.rs | 1 - .../src/scalar_fn/row/types/result.rs | 91 +------------------ vortex-array/src/scalar_fn/row/types/sink.rs | 15 +-- .../src/scalar_fn/row/visitor/check.rs | 11 +-- .../src/scalar_fn/row/visitor/execute.rs | 4 +- vortex-array/src/scalar_fn/row/visitor/mod.rs | 2 - .../src/scalar_fn/row/visitor/plan.rs | 2 +- 13 files changed, 19 insertions(+), 149 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index d53afd57770..2ce9323f09e 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -19,7 +19,6 @@ use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::DeferredError; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::OutputSink; use vortex_array::scalar_fn::RowFn; @@ -136,7 +135,7 @@ impl OutputSink for I64Sink { &mut rows[index] } - unsafe fn finish(self, _error: DeferredError) -> VortexResult { + unsafe fn finish(self) -> VortexResult { Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) } } diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs index 5cdbee5af55..d7e024ce763 100644 --- a/vortex-array/src/scalar_fn/row/batch/policy.rs +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -66,11 +66,7 @@ impl RowPolicy { /// probe can change the result of an encoding-aware function. pub const fn for_sink() -> Self { if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { - if ApplyResult::DEFERRED { - Self::DenseWithRetry - } else { - Self::Dense - } + Self::Dense } else { Self::ValidOnly } diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index be0b80bcc22..a2da250e957 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -28,7 +28,6 @@ use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; -use crate::scalar_fn::DeferredError; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::OutputSink; use crate::scalar_fn::RowFn; @@ -90,7 +89,7 @@ impl OutputSink for I64Sink { &mut rows[index] } - unsafe fn finish(self, _error: DeferredError) -> VortexResult { + unsafe fn finish(self) -> VortexResult { Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) } } diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 4cb987e282c..42d2e3991e9 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -13,7 +13,6 @@ use super::RowExecution; use super::ensure_decoded_lengths; use crate::ExecutionCtx; use crate::dtype::DType; -use crate::scalar_fn::DeferredError; use crate::scalar_fn::ElementTuple; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::OutputSink; @@ -74,8 +73,7 @@ where } } - // Immediate failures returned above. Only reduced, deferred failure evidence reaches finish. - finish_sink(sink, DeferredError::new(ApplyResult::occurred(accumulated))) + finish_sink(sink) } /// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. @@ -163,26 +161,14 @@ where } } - // Immediate failures returned above. Only reduced, deferred failure evidence reaches finish. - finish_sink(sink, DeferredError::new(ApplyResult::occurred(accumulated))).map(Some) + finish_sink(sink).map(Some) } -/// Classify a sink error as retryable only when row accumulation recorded a deferred failure. -/// -/// The sink contract requires [`OutputSink::finish`] to surface recorded failure evidence. Without -/// that evidence, its error is structural and retrying over a different set of rows cannot help. -fn finish_sink( - sink: S, - deferred_error: DeferredError, -) -> VortexResult { +fn finish_sink(sink: S) -> VortexResult { // SAFETY: callers reach this helper only after every completed callback returned the sink's // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. // The sink contract defines how that evidence establishes initialization of its row storage. - match unsafe { sink.finish(deferred_error) } { - Ok(output) => Ok(RowExecution::Output(output)), - Err(error) if deferred_error.occurred() => Ok(RowExecution::DeferredError(error)), - Err(error) => Err(error), - } + unsafe { sink.finish() }.map(RowExecution::Output) } #[cfg(test)] diff --git a/vortex-array/src/scalar_fn/row/execute/sink/tests.rs b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs index bf036fe0d63..8395aa69833 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink/tests.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs @@ -13,7 +13,6 @@ use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::NativePType; -use crate::scalar_fn::DeferredError; use crate::scalar_fn::OutputSink; use crate::scalar_fn::VecExecutionArgs; use crate::validity::Validity; @@ -43,7 +42,7 @@ impl OutputSink for NonSkippingSink { fn row<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} - unsafe fn finish(self, _error: DeferredError) -> VortexResult { + unsafe fn finish(self) -> VortexResult { Err(vortex_err!("a non-skipping sink must not finish")) } } diff --git a/vortex-array/src/scalar_fn/row/mod.rs b/vortex-array/src/scalar_fn/row/mod.rs index 50c005fabe6..1ad4e98b15a 100644 --- a/vortex-array/src/scalar_fn/row/mod.rs +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -6,8 +6,8 @@ //! Start with [`RowFn`] to define an operation and [`RowVisitor`] to select one of its output //! capabilities. [`InputElement`] and [`ElementTuple`] describe how columns become row values. //! [`OutputElement`] covers independent owned values, while [`OutputSink`] supports outputs that -//! need row handles or shared batch state. [`SinkResult`] and [`DeferredError`] describe how a -//! sink-writing closure reports errors. +//! need row handles or shared batch state. [`SinkResult`] describes how a sink-writing closure +//! reports errors. //! //! The internal executor owns decoding, batch constants, null propagation, allocation, and //! validity. A visitor's prepare closure may derive shared state from constant operands once per @@ -22,7 +22,6 @@ mod row_fn; pub use row_fn::RowFn; mod types; -pub use types::DeferredError; pub use types::ElementTuple; pub use types::IndexedElementTuple; pub use types::InitializedElement; diff --git a/vortex-array/src/scalar_fn/row/types/mod.rs b/vortex-array/src/scalar_fn/row/types/mod.rs index 4032a7d25d4..e47f195410c 100644 --- a/vortex-array/src/scalar_fn/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/row/types/mod.rs @@ -15,7 +15,6 @@ pub use element::OutputElement; pub(super) use element::batch_constant; mod result; -pub use result::DeferredError; pub use result::SinkResult; mod sink; diff --git a/vortex-array/src/scalar_fn/row/types/result.rs b/vortex-array/src/scalar_fn/row/types/result.rs index 6a3284604d0..cda44802c40 100644 --- a/vortex-array/src/scalar_fn/row/types/result.rs +++ b/vortex-array/src/scalar_fn/row/types/result.rs @@ -3,8 +3,6 @@ //! What a sink-writing row closure may return. -use std::ops::BitOrAssign; - use vortex_error::VortexResult; use super::InitializedElement; @@ -13,57 +11,21 @@ mod private { pub trait Sealed {} } -/// A value-dependent failure bit reduced across the row loop and handed to the output sink. -/// -/// Unlike [`VortexResult`], this never exits the loop. Use it when every row can write a safe -/// provisional value and report failure once at the end. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub struct DeferredError( - /// Whether this value records a deferred error. - bool, -); - -impl DeferredError { - /// Record whether this row encountered an error. - pub const fn new(failed: bool) -> Self { - Self(failed) - } - - /// Whether any row accumulated into this value failed. - pub const fn occurred(self) -> bool { - self.0 - } -} - -impl BitOrAssign for DeferredError { - fn bitor_assign(&mut self, rhs: Self) { - self.0 |= rhs.0; - } -} - -/// The result of writing one row: success, an immediate error, or deferred error evidence. +/// The result of writing one row: success or an immediate error. /// -/// The executor OR-reduces [`Accumulated`](Self::Accumulated) in a loop-local. The accumulated word -/// should be no wider than the computed element so error tracking does not constrain vector width. /// This trait is sealed; row functions choose one of its supplied implementations. pub trait SinkResult: 'static + private::Sealed { /// The [`OutputSink::WriteToken`](super::OutputSink::WriteToken) carried by a success. type WriteToken: 'static; - /// The word this result reduces into, kept in a loop-local by the executor. + /// Loop-local state used while accumulating row results. type Accumulated: 'static + Copy + Default; /// Whether this return type can carry an error. const FALLIBLE: bool; - /// Whether this result defers failure reporting until the sink finishes. - const DEFERRED: bool; - /// Merge this row's outcome into the batch-wide reduction. fn accumulate(self, accumulated: &mut Self::Accumulated) -> VortexResult<()>; - - /// Whether the finished reduction means some row failed. - fn occurred(accumulated: Self::Accumulated) -> bool; } impl private::Sealed for () {} @@ -73,15 +35,10 @@ impl SinkResult for () { type Accumulated = (); const FALLIBLE: bool = false; - const DEFERRED: bool = false; fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { Ok(()) } - - fn occurred(_accumulated: ()) -> bool { - false - } } impl private::Sealed for InitializedElement {} @@ -91,15 +48,10 @@ impl SinkResult for InitializedElement { type Accumulated = (); const FALLIBLE: bool = false; - const DEFERRED: bool = false; fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { Ok(()) } - - fn occurred(_accumulated: ()) -> bool { - false - } } impl private::Sealed for VortexResult<()> {} @@ -109,15 +61,10 @@ impl SinkResult for VortexResult<()> { type Accumulated = (); const FALLIBLE: bool = true; - const DEFERRED: bool = false; fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { self } - - fn occurred(_accumulated: ()) -> bool { - false - } } impl private::Sealed for VortexResult {} @@ -127,42 +74,8 @@ impl SinkResult for VortexResult { type Accumulated = (); const FALLIBLE: bool = true; - const DEFERRED: bool = false; fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { self.map(|_| ()) } - - fn occurred(_accumulated: ()) -> bool { - false - } } - -/// The evidence widths a row closure may reduce. `bool` is the ordinary answer; the unsigned -/// integers exist for a kernel whose per-row comparison would cost it its vectorization. -macro_rules! impl_sink_result_word { - ($($word:ty),+ $(,)?) => { - $( - impl private::Sealed for $word {} - - impl SinkResult for $word { - type WriteToken = (); - type Accumulated = $word; - - const FALLIBLE: bool = false; - const DEFERRED: bool = true; - - fn accumulate(self, accumulated: &mut $word) -> VortexResult<()> { - *accumulated |= self; - Ok(()) - } - - fn occurred(accumulated: $word) -> bool { - accumulated != <$word>::default() - } - } - )+ - }; -} - -impl_sink_result_word!(bool, u8, u16, u32, u64); diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index c0f1c1e2cd7..adbfdbfa878 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -9,7 +9,6 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::dtype::DType; -use crate::scalar_fn::DeferredError; use crate::scalar_fn::OutputElement; /// A column allocated once per batch that a row closure writes into, one row at a time. @@ -21,13 +20,6 @@ use crate::scalar_fn::OutputElement; /// skip-invalid execution can omit invalid rows when /// [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) is present. pub trait OutputSink: 'static + Sized { - /// Whether this sink accepts [`DeferredError`] from its row closure instead of requiring a - /// per-row [`VortexResult`]. - /// - /// A supporting sink must return an error from [`finish`](Self::finish) when its deferred error - /// argument occurred. - const ERRORS_ARE_DEFERRED: bool = false; - /// A loop-local view of all output rows. /// /// Borrowed once before execution so the sink's buffer descriptor and shape become loop @@ -80,15 +72,14 @@ pub trait OutputSink: 'static + Sized { fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; /// Finish into the built column, whose dtype **must** be this sink's - /// [`sink_dtype`](Self::sink_dtype). Called once per batch with whether any deferred row error - /// occurred. + /// [`sink_dtype`](Self::sink_dtype). Called once per batch. /// /// # Safety /// /// The executor must have completed every row callback successfully, and each callback must /// have returned this sink's [`WriteToken`](Self::WriteToken). When skipped rows are allowed, /// [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) must have run before traversal. - unsafe fn finish(self, error: DeferredError) -> VortexResult; + unsafe fn finish(self) -> VortexResult; } /// Proof that one uninitialized element row was initialized. @@ -164,7 +155,7 @@ impl OutputSink for UninitElementSink { &mut rows[index] } - unsafe fn finish(mut self, _error: DeferredError) -> VortexResult { + unsafe fn finish(mut self) -> VortexResult { // SAFETY: the caller guarantees every slot in `0..row_count` was initialized, and // `with_capacity` reserved every slot in that range. unsafe { self.values.set_len(self.row_count) }; diff --git a/vortex-array/src/scalar_fn/row/visitor/check.rs b/vortex-array/src/scalar_fn/row/visitor/check.rs index a000ce38369..189aa49b75b 100644 --- a/vortex-array/src/scalar_fn/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/row/visitor/check.rs @@ -54,11 +54,10 @@ where } /// Assert that a sink visit obeys the input, fallibility, and deferred-error contracts. -pub(super) const fn assert_sink_visit_contract() +pub(super) const fn assert_sink_visit_contract() where Function: RowFn, Args: ElementTuple, - Sink: OutputSink, ApplyResult: SinkResult, { assert_input_visit_contract::(); @@ -66,14 +65,6 @@ where !ApplyResult::FALLIBLE || Function::FALLIBLE, "RowFn::FALLIBLE must be true when a row result can fail", ); - assert!( - !ApplyResult::DEFERRED || Function::FALLIBLE, - "RowFn::FALLIBLE must be true when a row result defers failure evidence", - ); - assert!( - Sink::ERRORS_ARE_DEFERRED == ApplyResult::DEFERRED, - "OutputSink::ERRORS_ARE_DEFERRED must match SinkResult::DEFERRED", - ); } /// Assert the owned-output contract, fallibility declaration, and failure-evidence width bound. diff --git a/vortex-array/src/scalar_fn/row/visitor/execute.rs b/vortex-array/src/scalar_fn/row/visitor/execute.rs index 2cf4d2ac6f5..73314df9172 100644 --- a/vortex-array/src/scalar_fn/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/row/visitor/execute.rs @@ -92,7 +92,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { Sink: OutputSink, ApplyResult: SinkResult, { - const { assert_sink_visit_contract::() }; + const { assert_sink_visit_contract::() }; execute_sink::( self.args, @@ -195,7 +195,7 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { Sink: OutputSink, ApplyResult: SinkResult, { - const { assert_sink_visit_contract::() }; + const { assert_sink_visit_contract::() }; execute_sink_valid_rows::( self.args, diff --git a/vortex-array/src/scalar_fn/row/visitor/mod.rs b/vortex-array/src/scalar_fn/row/visitor/mod.rs index aedc0912bba..0ef085268b7 100644 --- a/vortex-array/src/scalar_fn/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/row/visitor/mod.rs @@ -85,8 +85,6 @@ pub trait RowVisitor: private::Sealed + Sized { /// [`RowFn::ARG_NAMES`](crate::scalar_fn::RowFn::ARG_NAMES). /// - [`RowFn::FALLIBLE`](crate::scalar_fn::RowFn::FALLIBLE) **must** be `true` when decoding /// `Args` or computing the result can fail. - /// - [`OutputSink::ERRORS_ARE_DEFERRED`] **must** match [`SinkResult::DEFERRED`] for the - /// selected `Sink` and `ApplyResult`. fn visit_into( self, apply: impl Fn(Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, diff --git a/vortex-array/src/scalar_fn/row/visitor/plan.rs b/vortex-array/src/scalar_fn/row/visitor/plan.rs index 02acee208c1..894acf33082 100644 --- a/vortex-array/src/scalar_fn/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/row/visitor/plan.rs @@ -75,7 +75,7 @@ impl RowVisitor for PlanRows<'_, F> { Sink: OutputSink, ApplyResult: SinkResult, { - const { assert_sink_visit_contract::() }; + const { assert_sink_visit_contract::() }; Ok(BatchPlan { output_dtype: validate_sink_visit::(self.dtypes)?, From d209fcfba23c135456fae7d76425cc6317646a34 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 14:13:16 -0400 Subject: [PATCH 056/160] Record RowFn upstream gate results Signed-off-by: Connor Tsui --- research/rowfn-review-followup/README.md | 47 ++++++++++++++++-------- 1 file changed, 31 insertions(+), 16 deletions(-) diff --git a/research/rowfn-review-followup/README.md b/research/rowfn-review-followup/README.md index 015a34b798f..c0ddba7044e 100644 --- a/research/rowfn-review-followup/README.md +++ b/research/rowfn-review-followup/README.md @@ -16,20 +16,20 @@ Round 2 target baselines: `84837ad36f` on `origin/ct/row-fn-api` and `32ad0bf3b7 | # | Verdict | Gate branch | Evidence and status | | ---: | --- | --- | --- | -| 1 | Fixed, not upstreamed | API target plus downstream numeric target | `OutputSink::finish(self)` is unsafe. The zero-unsafe external reproduction changes from reading allocator contents to `E0133`. Publication awaits approval for the combined API commit. | -| 2 | Fixed, not upstreamed | API target plus downstream numeric target | `InputElement` is an unsafe trait with local safety proofs. `ElementTuple` and `IndexedElementTuple` remain sealed framework traits. Publication awaits approval for the combined API commit. | +| 1 | Resolved | `ct/row-fn-api` at `f759998dce` | `OutputSink::finish(self)` is unsafe. The zero-unsafe external reproduction changes from reading allocator contents to `E0133`. | +| 2 | Resolved | `ct/row-fn-api` at `f759998dce` | `InputElement` is an unsafe trait with local safety proofs. `ElementTuple` and `IndexedElementTuple` remain sealed framework traits. | | 3 | Mitigated | API target plus downstream numeric target | Unsafe publication belongs to `OutputSink::finish`, and each successful callback must return the sink's write token. The token is not type-tied to the exact row handle; misuse now requires violating an unsafe sink/input contract rather than safe client code. | -| 4 | Fixed, not upstreamed | API target plus downstream numeric target | `SKIPPED_ROWS_INITIALIZER: Option` makes capability and operation one fact while restoring decline before decode and allocation. Publication awaits approval for the combined API commit. | -| 5 | Fixed, not upstreamed | API target plus downstream numeric target | `reduce_encoded` probes original inputs once and returns `RowExecution`, so encoded reductions can defer errors behind nulls through the same validity path as row execution. Publication awaits approval for the combined API commit. | +| 4 | Resolved | `ct/row-fn-api` at `f759998dce` | `SKIPPED_ROWS_INITIALIZER: Option` makes capability and operation one fact while restoring decline before decode and allocation. | +| 5 | Resolved | `ct/row-fn-api` at `f759998dce` | `reduce_encoded` probes original inputs once and returns `RowExecution`, so encoded reductions can defer errors behind nulls through the same validity path as row execution. | | 6 | Investigated | Design only | A temporary manual implementation still reproduces `E0119`. A public `execute_rows` free function is the recommended redesign; it is not implemented in this pass. | -| 7 | Fixed, not upstreamed | Numeric target | Private `NumericBinary` deliberately borrows the registered `Binary` ID and is guarded by privacy. Publication awaits the API commit. | -| 8 | Fixed, not upstreamed | API target plus downstream numeric target | `Batch::new` validates each input length directly. The regression test fails before the fix. Publication awaits approval for the combined API commit. | +| 7 | Resolved | `ct/row-fn-numeric` at `b4900072c1` | Private `NumericBinary` deliberately borrows the registered `Binary` ID and is guarded by privacy. | +| 8 | Resolved | `ct/row-fn-api` at `f759998dce` | `Batch::new` validates each input length directly. The regression test fails before the fix. | | 9 | Refuted | API target spot-check | `MaskedArray::try_new` enforces an all-valid child, and null `ConstantArray` validity is `AllInvalid`. Both couplings remain constructor invariants. | -| 10 | Fixed, not upstreamed | API target plus downstream numeric target | The encoding probe runs once before retry. The regression test fails before the fix and passes after it. Publication awaits approval for the combined API commit. | -| 11 | Fixed, not upstreamed | API target plus downstream numeric target | No sink can defer errors. The unreachable `finish_sink` retry classification and deferred finish argument are deleted. Publication awaits explicit approval for this deletion. | +| 10 | Resolved | `ct/row-fn-api` at `f759998dce` | The encoding probe runs once before retry. The regression test fails before the fix and passes after it. | +| 11 | Resolved | `ct/row-fn-api` at `f759998dce` | No sink can defer errors. The unreachable `finish_sink` retry classification and deferred finish argument are deleted. | | 12 | Refuted | API target tests and source spot-check | Filtering preserves literal, masked-child, and extension-storage constants recognized by `batch_constant`; new masked and extension tests pin the behavior. | -| 13 | Fixed, not upstreamed | API target plus downstream numeric target | The encoding probe runs before one-row broadcast and sees the original arrays. Publication awaits approval for the combined API commit. | -| 14 | Fixed, not upstreamed | API target plus downstream numeric target | The five unreachable deferred `SinkResult` word implementations and their capability pairing are deleted. Publication awaits explicit approval for this deletion. | +| 13 | Resolved | `ct/row-fn-api` at `f759998dce` | The encoding probe runs before one-row broadcast and sees the original arrays. | +| 14 | Resolved | `ct/row-fn-api` at `f759998dce` | The five unreachable deferred `SinkResult` word implementations and their capability pairing are deleted. | | 15 | Refuted | Downstream numeric target IR | Mixed `i64` add/sub and `i32` multiply contain broadcast vector loops on `arm64`; they are not scalar fallbacks. | | 16 | Investigated, needs x86 | Analysis only | Full-row skipped initialization and non-breaking mask traversal remain independent costs. No performance change is made. | | 17 | Investigated, needs x86 | API target source | `scatter_valid` still allocates one `u64` per original row. Whether replacing the gather pays for itself is a standalone benchmark question. | @@ -176,14 +176,13 @@ Staged capture found that the `RowExecution` reducer signature alone restores th edge. The required coverage and final safety comment change placement again and remove it in the final combined tree. Do not upstream or benchmark the reducer-signature commit in isolation. -The exact tested trees were consolidated without changing their contents: +The exact tested trees were consolidated without changing their contents and upstreamed: -- API candidate `f759998dcefa69b11d11b281e3bbebb6b88584e4`, based on `84837ad36f`. -- Numeric candidate `b4900072c1300238d246d395d986466db21583f6`, based on the API candidate. +- `ct/row-fn-api` at `f759998dcefa69b11d11b281e3bbebb6b88584e4`, based on `84837ad36f`. +- `ct/row-fn-numeric` at `b4900072c1300238d246d395d986466db21583f6`, based on the API commit. -Both remote tips still matched the recorded baselines after the final fetch. The candidates remain -unpublished because publishing the combined API commit, including the broad deferred-sink -deletion, requires explicit approval. +Both remote tips still matched the recorded baselines before publication. The numeric branch was +updated with an exact lease on `32ad0bf3b7`. ## Do not do this @@ -253,3 +252,19 @@ cargo +nightly fmt --all PYO3_PYTHON=.venv/bin/python cargo clippy -p vortex-array --all-targets --all-features passed ``` + +Final `ct/row-fn` after applying the deferred-sink deletion: + +```text +cargo nextest run -p vortex-array -p vortex-spatial -p vortex-tensor + 3825 passed, 1 skipped + +cargo test --doc -p vortex-array + 73 passed, 13 ignored + +cargo +nightly fmt --all + passed + +PYO3_PYTHON=.venv/bin/python cargo clippy --all-targets --all-features + passed +``` From 562109b3dcf968b3273f7eeeb7e066d2c5dd734b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:41:34 -0400 Subject: [PATCH 057/160] Preserve RowFn performance evidence Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/row/execute/owned.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/execute/owned.rs b/vortex-array/src/scalar_fn/row/execute/owned.rs index 8a9e27383ee..97519408937 100644 --- a/vortex-array/src/scalar_fn/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/row/execute/owned.rs @@ -74,8 +74,9 @@ where // When every input varies, the indexed source removes argument-shape dispatch from the hot // loop and lets the lane kernel optimize the traversal as one operation. Keep the varying // view and its length proof in this branch: hoisting them through the shared validation - // helper produces slower mixed-constant code with LLVM 21.1.2. See - // `research/rowfn-regressions-2026-08-08/README.md`. + // helper changed mixed-constant add, subtract, and multiply from 9.219, 9.229, and 18.94 us + // to 30.46, 31.11, and 37.73 us on a Ryzen 9 7950X with rustc 1.91.0 and LLVM 21.1.2. + // Restoring this placement recovered the fast code under the 16-CGU, no-LTO bench profile. if let Some(varying) = Args::varying(&columns) { vortex_ensure!( Args::varying_len_matches(&varying, row_count), From b45e9d72423c7f48f784b5105c0f4bd4e3d951b8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:42:22 -0400 Subject: [PATCH 058/160] Add RowFn boundary regression tests Signed-off-by: Connor Tsui --- .../scalar_fn/fns/binary/compare/primitive.rs | 48 +++++++++-- .../src/scalar_fn/fns/binary/compare/tests.rs | 70 ++++++++++++++++ vortex-array/src/scalar_fn/row/batch/tests.rs | 80 +++++++++++++++++++ 3 files changed, 193 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 3e6ee8023df..ffd03c38536 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -3,9 +3,7 @@ //! Primitive comparison execution through [`RowFn`]. -#[cfg(target_arch = "x86_64")] mod columnar; -#[cfg(target_arch = "x86_64")] mod operand; use vortex_error::VortexResult; @@ -35,8 +33,49 @@ pub(super) fn compare_primitive( op: CompareOperator, ctx: &mut ExecutionCtx, ) -> VortexResult { - #[cfg(target_arch = "x86_64")] - if use_columnar_comparison(lhs, rhs, op)? { + compare_primitive_with_path(lhs, rhs, op, PrimitiveComparisonPath::Auto, ctx) +} + +/// Selects automatic production dispatch or a forced implementation in tests. +#[derive(Clone, Copy)] +pub(super) enum PrimitiveComparisonPath { + /// Use the architecture and operand-specific production policy. + Auto, + + /// Force row execution. + #[cfg(test)] + Row, + + /// Force fused columnar execution. + #[cfg(test)] + Columnar, +} + +/// Compare primitives through the selected implementation. +pub(super) fn compare_primitive_with_path( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + path: PrimitiveComparisonPath, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let use_columnar = match path { + PrimitiveComparisonPath::Auto => { + #[cfg(target_arch = "x86_64")] + { + use_columnar_comparison(lhs, rhs, op)? + } + #[cfg(not(target_arch = "x86_64"))] + { + false + } + } + #[cfg(test)] + PrimitiveComparisonPath::Row => false, + #[cfg(test)] + PrimitiveComparisonPath::Columnar => true, + }; + if use_columnar { return columnar::compare_primitive(lhs, rhs, op, ctx); } @@ -73,7 +112,6 @@ impl RowFn for PrimitiveCompare { } } -#[cfg(target_arch = "x86_64")] fn use_columnar_comparison( lhs: &ArrayRef, rhs: &ArrayRef, diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index c2af83561a4..06d8e20ba11 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -7,11 +7,14 @@ use rstest::rstest; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::VTable; use crate::array_session; +use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; @@ -20,6 +23,7 @@ use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::ScalarFn; use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; @@ -36,6 +40,8 @@ use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; use crate::scalar::DecimalValue; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::primitive::PrimitiveComparisonPath; +use crate::scalar_fn::fns::binary::compare::primitive::compare_primitive_with_path; use crate::scalar_fn::fns::binary::scalar_cmp; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -426,6 +432,70 @@ fn float_total_order() { ); } +#[rstest] +#[case::row(PrimitiveComparisonPath::Row)] +#[case::columnar(PrimitiveComparisonPath::Columnar)] +fn test_primitive_comparison_paths_preserve_semantics_and_encoding( + #[case] path: PrimitiveComparisonPath, +) -> VortexResult<()> { + let lhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::NAN, // Null on the left. + -0.0, // Signed zero ordering. + 1.0, // A finite value below NaN. + f64::NAN, // Null on the right. + ], + Validity::from_iter([ + true, // + false, // + true, // + true, // + true, // + ]), + ) + .into_array(); + let rhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::INFINITY, // Null on the left. + 0.0, // Signed zero ordering. + f64::NAN, // A finite value below NaN. + 1.0, // Null on the right. + ], + Validity::from_iter([ + true, // + true, // + true, // + true, // + false, // + ]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let actual = compare_primitive_with_path(&lhs, &rhs, CompareOperator::Lt, path, &mut ctx)?; + let expected = BoolArray::from_iter([ + Some(false), // + None, // + Some(true), // + Some(true), // + None, // + ]); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(actual.dtype(), &DType::Bool(Nullability::Nullable)); + + // This encoding difference is intentional: the fused path materializes bits and validity + // together, while the RowFn path keeps masking lazy. + match path { + PrimitiveComparisonPath::Columnar => assert_eq!(actual.encoding_id(), Bool.id()), + PrimitiveComparisonPath::Row => assert!(actual.as_opt::().is_some()), + PrimitiveComparisonPath::Auto => unreachable!(), + } + Ok(()) +} + #[rstest] #[case(Operator::Eq, [true, false, true, true])] #[case(Operator::Lt, [false, true, false, false])] diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index a2da250e957..3ccf2114525 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -8,6 +8,7 @@ use std::sync::atomic::Ordering; use rstest::rstest; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_session::registry::CachedId; @@ -29,6 +30,7 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::OutputElement; use crate::scalar_fn::OutputSink; use crate::scalar_fn::RowFn; use crate::scalar_fn::RowVisitor; @@ -49,6 +51,12 @@ struct OriginalInputReducer; #[derive(Clone)] struct DeferredOriginalReducer; +#[derive(Clone)] +struct InvalidKernelOutput; + +/// Deliberately violates [`OutputElement::build`] to test validation at the public boundary. +struct NullProducingI64(i64); + #[derive(Clone)] struct PreparedAdd { visit: PreparedVisit, @@ -62,6 +70,19 @@ enum PreparedVisit { Deferred, } +impl OutputElement for NullProducingI64 { + fn element_dtype() -> DType { + DType::from(i64::PTYPE) + } + + fn build(values: Vec) -> ArrayRef { + let values: Vec<_> = values.into_iter().map(|value| value.0).collect(); + let validity = Validity::from_iter((0..values.len()).map(|index| index != 0)); + + PrimitiveArray::new(values, validity).into_array() + } +} + struct I64Sink(BufferMut); impl OutputSink for I64Sink { @@ -228,6 +249,26 @@ impl RowFn for DeferredOriginalReducer { } } +impl RowFn for InvalidKernelOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_kernel_output"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), NullProducingI64>(|(value,)| NullProducingI64(value)) + } +} + impl RowFn for PreparedAdd { type Options = EmptyOptions; @@ -438,6 +479,45 @@ fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { Ok(()) } +#[test] +fn test_nonnullable_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + + assert_invalid_kernel_output(input) +} + +#[test] +fn test_all_valid_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + + assert_invalid_kernel_output(input) +} + +#[track_caller] +fn assert_invalid_kernel_output(input: ArrayRef) -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + let execution = ScalarFnVTable::execute(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); + let error = match execution { + Err(error) => error, + Ok(output) => match output.execute::(&mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an invalid row kernel output passed boundary validation"), + }, + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_kernel_output"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("row kernel produced nulls for valid rows"), + "the boundary error must identify invalid row output, got {error}", + ); + Ok(()) +} + #[rstest] #[case::owned_constant(PreparedVisit::Owned, true)] #[case::owned_varying(PreparedVisit::Owned, false)] From 6fc726a304adda90715d345bf3909fc6a44987ec Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:49:40 -0400 Subject: [PATCH 059/160] Validate RowFn kernel outputs at the boundary Signed-off-by: Connor Tsui --- .../src/scalar_fn/row/batch/execution.rs | 74 +++++++++++++++---- vortex-array/src/scalar_fn/row/batch/tests.rs | 11 +-- vortex-array/src/scalar_fn/row/vtable.rs | 1 + 3 files changed, 67 insertions(+), 19 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index fd4b0401fb3..83df651d4f4 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -193,7 +193,9 @@ impl Batch { .collect::>()?; let result = VortexResult::from(kernel(self.kernel_args(&one_row, 1), ctx)?)?; - let scalar = self.finalize_output(result, 1)?.execute_scalar(0, ctx)?; + let result = self.validate_kernel_output(result, 1, ctx)?; + let result = self.finalize_output(result, 1)?; + let scalar = result.execute_scalar(0, ctx)?; Ok(ConstantArray::new(scalar, self.row_count).into_array()) } @@ -230,6 +232,7 @@ impl Batch { } RowExecution::DeferredError(error) => return Err(error), }; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; match self.validity.clone() { Validity::NonNullable | Validity::AllValid => { @@ -253,15 +256,12 @@ impl Batch { // Check all-true before all-false: an empty mask is both, and must not be treated as // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). if valid.all_true() { - return self - .finalize_output( - VortexResult::from(kernel( - self.kernel_args(&self.inputs, self.row_count), - ctx, - )?)?, - self.row_count, - ) - .map(ResolvedMask::Decided); + let values = + VortexResult::from(kernel(self.kernel_args(&self.inputs, self.row_count), ctx)?)?; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = self.finalize_output(values, self.row_count)?; + + return Ok(ResolvedMask::Decided(values)); } if valid.all_false() { @@ -311,6 +311,7 @@ impl Batch { return Ok(None); }; let values = VortexResult::from(execution)?; + let values = self.validate_kernel_output(values, valid.len(), ctx)?; let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); self.finalize_output(values.mask(mask)?, valid.len()) @@ -335,6 +336,7 @@ impl Batch { self.kernel_args(&filtered, valid.true_count()), ctx, )?)?; + let values = self.validate_kernel_output(values, valid.true_count(), ctx)?; self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) } @@ -396,7 +398,17 @@ impl Batch { /// Finalize an output against this batch's expected length and declared return dtype. fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { - finalize_kernel_output(self.id, &self.result_dtype, expected_len, values) + reconcile_output(self.id, &self.result_dtype, expected_len, values) + } + + /// Validate the output from a row kernel before batch validity is attached. + fn validate_kernel_output( + &self, + values: ArrayRef, + expected_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx) } /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set @@ -448,17 +460,46 @@ impl Batch { } } -/// Validate a kernel output, then cast it to the row function's declared nullability. +/// Validate the output produced directly by a row kernel. /// /// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring -/// nullability. The kernel may omit nullability because batch execution owns strict null -/// propagation, so a nullability-only difference is cast to `result_dtype`. +/// nullability, and every produced row **must** be valid. Batch execution owns strict null +/// propagation and attaches input-derived validity only after this boundary. pub fn finalize_kernel_output( id: ScalarFnId, result_dtype: &DType, expected_len: usize, values: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + vortex_ensure!( + values.all_valid(ctx)?, + "the {id} row kernel produced nulls for valid rows", + ); + + cast_output_nullability(result_dtype, values) +} + +/// Reconcile an output with the function's declared shape and nullability. +fn reconcile_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, ) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + + cast_output_nullability(result_dtype, values) +} + +/// Validate an output's shape and logical dtype without executing a nullability cast. +fn validate_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: &ArrayRef, +) -> VortexResult<()> { vortex_ensure_eq!( values.len(), expected_len, @@ -471,6 +512,11 @@ pub fn finalize_kernel_output( values.dtype(), ); + Ok(()) +} + +/// Cast only the output nullability after its shape, dtype, and validity are accepted. +fn cast_output_nullability(result_dtype: &DType, values: ArrayRef) -> VortexResult { if values.dtype() == result_dtype { Ok(values) } else { diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 3ccf2114525..21fa9cd9eca 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -26,6 +26,7 @@ use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; @@ -467,15 +468,15 @@ fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { let result_dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); let mut ctx = array_session().create_execution_ctx(); - let actual = finalize_kernel_output(*ID, &result_dtype, 2, values.clone())?; + let actual = finalize_kernel_output(*ID, &result_dtype, 2, values.clone(), &mut ctx)?; let expected = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); assert_eq!(actual.dtype(), &result_dtype); assert_arrays_eq!(&actual, &expected, &mut ctx); - assert!(finalize_kernel_output(*ID, &result_dtype, 3, values).is_err()); + assert!(finalize_kernel_output(*ID, &result_dtype, 3, values, &mut ctx).is_err()); let bools = BoolArray::from_iter([true, false]).into_array(); - assert!(finalize_kernel_output(*ID, &result_dtype, 2, bools).is_err()); + assert!(finalize_kernel_output(*ID, &result_dtype, 2, bools, &mut ctx).is_err()); Ok(()) } @@ -575,8 +576,8 @@ fn test_strategy_matrix(#[case] policy: RowPolicy) -> VortexResult<()> { let actual = batch.execute( |_args, _ctx| Ok(None), - |args, _ctx| Ok(RowExecution::Output(args.arrays[0].clone())), - |args, _valid, _ctx| Ok(Some(RowExecution::Output(args.arrays[0].clone()))), + |args, _ctx| Ok(RowExecution::Output(args.arrays[0].fill_null(0_i64)?)), + |args, _valid, _ctx| Ok(Some(RowExecution::Output(args.arrays[0].fill_null(0_i64)?))), &mut ctx, )?; diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index 8ed9163866f..6c96959b788 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -85,6 +85,7 @@ impl ScalarFnVTable for F { &result_dtype, args.row_count(), values, + ctx, ); } From c027963c37d3fb5ee31b5a3963612c44cfa7d481 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:51:42 -0400 Subject: [PATCH 060/160] Finalize RowFn API decisions Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 17 +- vortex-array/benches/strict_validity.rs | 4 +- .../scalar_fn/fns/binary/compare/primitive.rs | 7 +- .../src/scalar_fn/fns/binary/numeric/row.rs | 9 +- vortex-array/src/scalar_fn/mod.rs | 7 + vortex-array/src/scalar_fn/row/batch/tests.rs | 100 ++++++-- .../src/scalar_fn/row/execute/sink.rs | 56 +++-- .../src/scalar_fn/row/execute/sink/tests.rs | 7 +- vortex-array/src/scalar_fn/row/mod.rs | 2 + vortex-array/src/scalar_fn/row/row_fn.rs | 2 +- .../src/scalar_fn/row/types/element/mod.rs | 3 +- vortex-array/src/scalar_fn/row/types/sink.rs | 15 +- .../src/scalar_fn/row/visitor/check.rs | 11 +- .../src/scalar_fn/row/visitor/execute.rs | 28 ++- vortex-array/src/scalar_fn/row/visitor/mod.rs | 23 +- .../src/scalar_fn/row/visitor/plan.rs | 26 ++- vortex-array/src/scalar_fn/row/vtable.rs | 217 +++++++++++------- vortex-spatial/src/scalar_fn/contains.rs | 4 +- vortex-spatial/src/scalar_fn/distance.rs | 4 +- vortex-spatial/src/scalar_fn/intersects.rs | 4 +- .../src/scalar_fns/cosine_similarity.rs | 4 +- vortex-tensor/src/scalar_fns/inner_product.rs | 4 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 4 +- vortex-tensor/src/scalar_fns/tests/row.rs | 4 +- 24 files changed, 373 insertions(+), 189 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index 2ce9323f09e..056fe2007ff 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -24,6 +24,7 @@ use vortex_array::scalar_fn::OutputSink; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -55,7 +56,7 @@ impl RowFn for RowWrappingAdd { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -79,7 +80,7 @@ impl RowFn for RowCheckedAdd { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -110,12 +111,12 @@ struct I64Sink( BufferMut, ); -impl OutputSink for I64Sink { +impl OutputSink for I64Sink { type Rows<'a> = &'a mut [i64]; type Row<'a> = &'a mut i64; type WriteToken = (); - fn sink_dtype(_args: &[DType]) -> VortexResult { + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } @@ -153,7 +154,7 @@ impl RowFn for RowSinkWrappingAdd { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -165,6 +166,10 @@ impl RowFn for RowSinkWrappingAdd { } } +vortex_array::impl_row_fn_vtable!(RowWrappingAdd); +vortex_array::impl_row_fn_vtable!(RowCheckedAdd); +vortex_array::impl_row_fn_vtable!(RowSinkWrappingAdd); + fn inputs() -> (ArrayRef, ArrayRef) { let lhs = (0..ROWS) .map(|index| index as i64) @@ -201,7 +206,7 @@ fn nullable_inputs() -> (ArrayRef, ArrayRef) { fn bench_row_fn(bencher: Bencher, function: F, make_inputs: fn() -> (ArrayRef, ArrayRef)) where - F: RowFn, + F: RowFn + ScalarFnVTable, { bencher .with_inputs(make_inputs) diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs index b2fa7a3a824..e0a8b6a565e 100644 --- a/vortex-array/benches/strict_validity.rs +++ b/vortex-array/benches/strict_validity.rs @@ -87,7 +87,7 @@ impl RowFn for LazyDouble { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -106,6 +106,8 @@ impl RowFn for LazyDouble { } } +vortex_array::impl_row_fn_vtable!(LazyDouble); + /// The same function, applying validity the way the adapter used to: materialize a mask first. #[derive(Clone)] struct EagerDouble; diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index ffd03c38536..53ab5448585 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -20,6 +20,7 @@ use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::execute_rows; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; @@ -81,7 +82,7 @@ pub(super) fn compare_primitive_with_path( let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); - ScalarFnVTable::execute(&PrimitiveCompare, &op, &args, ctx) + execute_rows(&PrimitiveCompare, &op, &args, ctx) } /// Internal row execution for primitive comparison operators. @@ -97,7 +98,7 @@ impl RowFn for PrimitiveCompare { ScalarFnVTable::id(&Binary) } - fn dispatch( + fn dispatch>( &self, op: &Self::Options, args: &[DType], @@ -135,7 +136,7 @@ fn use_columnar_comparison( fn visit_compare(op: CompareOperator, visitor: V) -> VortexResult where T: NativePType, - V: RowVisitor, + V: RowVisitor, { match op { CompareOperator::Eq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)), diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index 089adcbcbde..bc38dd0e5eb 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -28,6 +28,7 @@ use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::execute_rows; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::row::InitializedElement; use crate::scalar_fn::row::UninitElementSink; @@ -40,7 +41,7 @@ pub(super) fn execute_numeric_primitive( ) -> VortexResult { let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); - ScalarFnVTable::execute(&NumericBinary, &op, &args, ctx) + execute_rows(&NumericBinary, &op, &args, ctx) } /// Internal row execution for the primitive arithmetic operators. @@ -62,7 +63,7 @@ impl RowFn for NumericBinary { ScalarFnVTable::id(&Binary) } - fn dispatch( + fn dispatch>( &self, op: &Self::Options, args: &[DType], @@ -88,7 +89,7 @@ fn visit_checked(visitor: V) -> VortexResult where T: NativePType, Op: CheckedPrimitiveOp, - V: RowVisitor, + V: RowVisitor, { visitor.visit_deferred::<(T, T), T, Op::Fail>( |(lhs, rhs)| Op::apply(lhs, rhs), @@ -105,7 +106,7 @@ where fn visit_div(visitor: V) -> VortexResult where T: CheckedArithmetic, - V: RowVisitor, + V: RowVisitor, { if T::PTYPE.is_float() { return visit_checked::(visitor); diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 5e73caefdfa..70a8443314d 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -19,6 +19,13 @@ use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::literal::Literal; +/// Reexports used by [`impl_row_fn_vtable!`] without requiring downstream transitive dependencies. +#[doc(hidden)] +pub mod row_fn_macro_support { + pub use vortex_error::VortexResult; + pub use vortex_session::VortexSession; +} + mod vtable; pub use vtable::*; diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 21fa9cd9eca..9760db84f68 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -36,8 +36,9 @@ use crate::scalar_fn::OutputSink; use crate::scalar_fn::RowFn; use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::execute_rows; +use crate::scalar_fn::row_fn_return_dtype; use crate::validity::Validity; #[derive(Clone)] @@ -52,6 +53,11 @@ struct OriginalInputReducer; #[derive(Clone)] struct DeferredOriginalReducer; +#[derive(Clone)] +struct SinkOptions; + +struct OptionsCheckingSink; + #[derive(Clone)] struct InvalidKernelOutput; @@ -71,6 +77,36 @@ enum PreparedVisit { Deferred, } +impl OutputSink for OptionsCheckingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn sink_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { + if !enabled { + vortex_bail!(InvalidArgument: "the test sink is disabled"); + } + + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + vortex_bail!("the planning-only test sink must not be allocated") + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { + true + } + + fn row<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} + + unsafe fn finish(self) -> VortexResult { + vortex_bail!("the planning-only test sink must not finish") + } +} + impl OutputElement for NullProducingI64 { fn element_dtype() -> DType { DType::from(i64::PTYPE) @@ -86,12 +122,12 @@ impl OutputElement for NullProducingI64 { struct I64Sink(BufferMut); -impl OutputSink for I64Sink { +impl OutputSink for I64Sink { type Rows<'a> = &'a mut [i64]; type Row<'a> = &'a mut i64; type WriteToken = (); - fn sink_dtype(_args: &[DType]) -> VortexResult { + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } @@ -126,7 +162,7 @@ impl RowFn for NullarySeven { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -149,7 +185,7 @@ impl RowFn for RetryConstantAdd { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -193,7 +229,7 @@ impl RowFn for OriginalInputReducer { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -229,7 +265,7 @@ impl RowFn for DeferredOriginalReducer { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -250,6 +286,26 @@ impl RowFn for DeferredOriginalReducer { } } +impl RowFn for SinkOptions { + type Options = bool; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.sink_options"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), OptionsCheckingSink, _>(|(), ()| ()) + } +} + impl RowFn for InvalidKernelOutput { type Options = EmptyOptions; @@ -260,7 +316,7 @@ impl RowFn for InvalidKernelOutput { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -281,7 +337,7 @@ impl RowFn for PreparedAdd { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -344,7 +400,7 @@ fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { let args = VecExecutionArgs::new(vec![lhs, rhs], 2); let mut ctx = array_session().create_execution_ctx(); - let result = ScalarFnVTable::execute(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx); + let result = execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx); assert!(result.is_err()); Ok(()) @@ -358,7 +414,7 @@ fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { let args = VecExecutionArgs::new(vec![lhs, rhs], 2); let mut ctx = array_session().create_execution_ctx(); - let actual = ScalarFnVTable::execute(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx)?; let expected = PrimitiveArray::new(vec![2_u8, 0], Validity::from_iter([true, false])); assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); @@ -372,7 +428,7 @@ fn test_reduce_encoded_defers_errors_behind_nulls() -> VortexResult<()> { let args = VecExecutionArgs::new(vec![input.clone()], 2); let mut ctx = array_session().create_execution_ctx(); - let actual = ScalarFnVTable::execute(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; assert_arrays_eq!(&actual, &input, &mut ctx); Ok(()) @@ -384,7 +440,7 @@ fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { let args = VecExecutionArgs::new(vec![input], 3); let mut ctx = array_session().create_execution_ctx(); - let actual = ScalarFnVTable::execute(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; let expected = ConstantArray::new(42_i64, 3).into_array(); assert_arrays_eq!(&actual, &expected, &mut ctx); @@ -397,7 +453,7 @@ fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { let args = VecExecutionArgs::new(vec![input.clone()], 2); let mut ctx = array_session().create_execution_ctx(); - let actual = ScalarFnVTable::execute(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; assert_arrays_eq!(&actual, &input, &mut ctx); Ok(()) @@ -480,6 +536,16 @@ fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { Ok(()) } +#[test] +fn test_sink_dtype_receives_function_options() -> VortexResult<()> { + assert_eq!( + row_fn_return_dtype(&SinkOptions, &true, &[])?, + DType::from(i64::PTYPE) + ); + assert!(row_fn_return_dtype(&SinkOptions, &false, &[]).is_err()); + Ok(()) +} + #[test] fn test_nonnullable_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { let input = PrimitiveArray::from_iter([1_i64, 2]).into_array(); @@ -498,7 +564,7 @@ fn test_all_valid_kernel_output_rejects_nulls_at_function_boundary() -> VortexRe fn assert_invalid_kernel_output(input: ArrayRef) -> VortexResult<()> { let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = array_session().create_execution_ctx(); - let execution = ScalarFnVTable::execute(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); + let execution = execute_rows(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); let error = match execution { Err(error) => error, Ok(output) => match output.execute::(&mut ctx) { @@ -544,7 +610,7 @@ fn test_prepared_visits( }; let mut ctx = array_session().create_execution_ctx(); - let actual = ScalarFnVTable::execute(&function, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; let expected = if constant_rhs { PrimitiveArray::from_iter([4_i64, 5]).into_array() } else { @@ -590,7 +656,7 @@ fn test_nullary_row_function_broadcasts() -> VortexResult<()> { let args = VecExecutionArgs::new(vec![], 3); let mut ctx = array_session().create_execution_ctx(); - let actual = ScalarFnVTable::execute(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; let expected = PrimitiveArray::from_iter([7i64, 7, 7]).into_array(); assert_arrays_eq!(&actual, &expected, &mut ctx); diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index 42d2e3991e9..cbf0b536b89 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -22,20 +22,20 @@ use crate::scalar_fn::SinkResult; /// /// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state /// does not need to be captured by the closure. -pub fn execute_sink( +pub fn execute_sink( args: &dyn ExecutionArgs, sink_dtype: &DType, ctx: &mut ExecutionCtx, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, { let row_count = args.row_count(); - let mut sink = Sink::with_capacity(row_count, sink_dtype)?; + let mut sink = >::with_capacity(row_count, sink_dtype)?; let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); let varying = Args::varying(&columns); @@ -45,9 +45,9 @@ where { // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This // scope releases the borrow before `finish_sink` consumes the sink. - let mut rows = sink.rows(); + let mut rows = >::rows(&mut sink); vortex_ensure!( - Sink::row_count_matches(&rows, row_count), + >::row_count_matches(&rows, row_count), "the output sink does not address exactly {row_count} rows", ); @@ -58,41 +58,46 @@ where // SAFETY: `ensure_decoded_lengths` proved every varying column has `row_count` // rows before the loop. let elements = unsafe { Args::get_varying_unchecked(&varying, index) }; - apply(&prepared, elements, Sink::row(&mut rows, index)) - .accumulate(&mut accumulated)?; + apply( + &prepared, + elements, + >::row(&mut rows, index), + ) + .accumulate(&mut accumulated)?; } } else { for index in 0..row_count { apply( &prepared, Args::get(&columns, index), - Sink::row(&mut rows, index), + >::row(&mut rows, index), ) .accumulate(&mut accumulated)?; } } } - finish_sink(sink) + finish_sink::(sink) } /// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. -pub fn execute_sink_valid_rows( +pub fn execute_sink_valid_rows( args: &dyn ExecutionArgs, sink_dtype: &DType, valid: &Mask, ctx: &mut ExecutionCtx, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, ) -> VortexResult> where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, { // Decline before input decoding or sink allocation when this sink cannot initialize rows that // the mask skips. The capability and the operation are the same function pointer. - let Some(initialize_skipped_rows) = Sink::SKIPPED_ROWS_INITIALIZER else { + let Some(initialize_skipped_rows) = >::SKIPPED_ROWS_INITIALIZER + else { return Ok(None); }; @@ -103,7 +108,7 @@ where }; let prepared = prepare(Args::constants(&columns)); let row_count = args.row_count(); - let mut sink = Sink::with_capacity(row_count, sink_dtype)?; + let mut sink = >::with_capacity(row_count, sink_dtype)?; let mut accumulated = ApplyResult::Accumulated::default(); // Batch execution resolves all-valid and all-null inputs before selecting this path. @@ -116,9 +121,9 @@ where ); { - let mut rows = sink.rows(); + let mut rows = >::rows(&mut sink); vortex_ensure!( - Sink::row_count_matches(&rows, row_count), + >::row_count_matches(&rows, row_count), "the output sink does not address exactly {row_count} rows", ); @@ -143,12 +148,12 @@ where // SAFETY: `ensure_decoded_lengths` proved every varying column has // `row_count` rows, and mask indices are below `row_count`. unsafe { Args::get_varying_unchecked(varying, index) }, - Sink::row(&mut rows, index), + >::row(&mut rows, index), ), None => apply( &prepared, Args::get(&columns, index), - Sink::row(&mut rows, index), + >::row(&mut rows, index), ), }; if let Err(err) = result.accumulate(&mut accumulated) { @@ -161,14 +166,17 @@ where } } - finish_sink(sink).map(Some) + finish_sink::(sink).map(Some) } -fn finish_sink(sink: S) -> VortexResult { +fn finish_sink(sink: S) -> VortexResult +where + S: OutputSink, +{ // SAFETY: callers reach this helper only after every completed callback returned the sink's // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. // The sink contract defines how that evidence establishes initialization of its row storage. - unsafe { sink.finish() }.map(RowExecution::Output) + unsafe { >::finish(sink) }.map(RowExecution::Output) } #[cfg(test)] diff --git a/vortex-array/src/scalar_fn/row/execute/sink/tests.rs b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs index 8395aa69833..01f6f0b2bd3 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink/tests.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs @@ -13,18 +13,19 @@ use crate::array_session; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::NativePType; +use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::OutputSink; use crate::scalar_fn::VecExecutionArgs; use crate::validity::Validity; struct NonSkippingSink; -impl OutputSink for NonSkippingSink { +impl OutputSink for NonSkippingSink { type Rows<'a> = (); type Row<'a> = (); type WriteToken = (); - fn sink_dtype(_args: &[DType]) -> VortexResult { + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } @@ -54,7 +55,7 @@ fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { let valid = Mask::from_iter([true, false]); let mut ctx = array_session().create_execution_ctx(); - let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, ()>( + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( &args, &DType::from(i64::PTYPE), &valid, diff --git a/vortex-array/src/scalar_fn/row/mod.rs b/vortex-array/src/scalar_fn/row/mod.rs index 1ad4e98b15a..754097a91ca 100644 --- a/vortex-array/src/scalar_fn/row/mod.rs +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -35,3 +35,5 @@ mod visitor; pub use visitor::RowVisitor; mod vtable; +pub use vtable::execute_rows; +pub use vtable::row_fn_return_dtype; diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index 99b544e44aa..4105a3690dc 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -60,7 +60,7 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// /// Plan time and run time both call this method, so the choice **must** be a pure function of /// `options` and `args`. Cross-argument dtype validation belongs here. - fn dispatch( + fn dispatch>( &self, options: &Self::Options, args: &[DType], diff --git a/vortex-array/src/scalar_fn/row/types/element/mod.rs b/vortex-array/src/scalar_fn/row/types/element/mod.rs index 232040cde12..304e52b7d00 100644 --- a/vortex-array/src/scalar_fn/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -139,7 +139,8 @@ pub trait OutputElement: 'static + Sized { /// Taking no arguments confines an element's dtype to a property of its Rust type, so an output /// whose dtype depends on runtime data (a tensor, whose dtype carries its shape) cannot be an /// element. Such an output uses an [`OutputSink`](crate::scalar_fn::OutputSink), whose - /// [`sink_dtype`](crate::scalar_fn::OutputSink::sink_dtype) does see the input dtypes. + /// [`sink_dtype`](crate::scalar_fn::OutputSink::sink_dtype) sees the function options and input + /// dtypes. fn element_dtype() -> DType; /// Build a column from one value per row. Called once per batch. diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index adbfdbfa878..a75da4c44ca 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -13,13 +13,14 @@ use crate::scalar_fn::OutputElement; /// A column allocated once per batch that a row closure writes into, one row at a time. /// -/// A sink may use the input dtypes to build a runtime-shaped output or own shared batch state. The -/// executor passes each row slot into an [`Fn`] closure, keeping mutable state out of its capture. +/// A sink may use the function's `Options` and input dtypes to build a runtime-shaped output or own +/// shared batch state. The executor passes each row slot into an [`Fn`] closure, keeping mutable +/// state out of its capture. /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; /// skip-invalid execution can omit invalid rows when /// [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) is present. -pub trait OutputSink: 'static + Sized { +pub trait OutputSink: 'static + Sized { /// A loop-local view of all output rows. /// /// Borrowed once before execution so the sink's buffer descriptor and shape become loop @@ -49,11 +50,11 @@ pub trait OutputSink: 'static + Sized { /// unsafe when Rust cannot tie the token to the supplied row handle. type WriteToken: 'static; - /// The dtype of the column this sink builds, given the function's input dtypes. + /// The dtype of the column this sink builds, given the function options and input dtypes. /// /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the /// result, and masks the null rows. - fn sink_dtype(args: &[DType]) -> VortexResult; + fn sink_dtype(options: &Options, args: &[DType]) -> VortexResult; /// Allocate a sink for `rows` rows of `dtype`, which is this sink's own /// [`sink_dtype`](Self::sink_dtype). Called once per batch. @@ -120,7 +121,7 @@ pub struct UninitElementSink { row_count: usize, } -impl OutputSink for UninitElementSink { +impl OutputSink for UninitElementSink { type Rows<'a> = &'a mut [MaybeUninit]; const SKIPPED_ROWS_INITIALIZER: Option fn(&mut Self::Rows<'a>)> = Some(|rows| { @@ -132,7 +133,7 @@ impl OutputSink for UninitElementSink { type Row<'a> = &'a mut MaybeUninit; type WriteToken = InitializedElement; - fn sink_dtype(_args: &[DType]) -> VortexResult { + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(T::element_dtype()) } diff --git a/vortex-array/src/scalar_fn/row/visitor/check.rs b/vortex-array/src/scalar_fn/row/visitor/check.rs index 189aa49b75b..6caf7f903a0 100644 --- a/vortex-array/src/scalar_fn/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/row/visitor/check.rs @@ -102,12 +102,17 @@ pub(super) fn validate_owned_visit( } /// Validate the input dtypes and return the non-nullable dtype built by `Sink`. -pub(super) fn validate_sink_visit( +pub(super) fn validate_sink_visit( + options: &Options, dtypes: &[DType], -) -> VortexResult { +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, +{ Args::validate(dtypes)?; - let dtype = Sink::sink_dtype(dtypes)?; + let dtype = Sink::sink_dtype(options, dtypes)?; vortex_ensure!( !dtype.is_nullable(), "row output sinks must declare a non-nullable dtype, got {dtype}", diff --git a/vortex-array/src/scalar_fn/row/visitor/execute.rs b/vortex-array/src/scalar_fn/row/visitor/execute.rs index 73314df9172..61cae08de72 100644 --- a/vortex-array/src/scalar_fn/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/row/visitor/execute.rs @@ -65,7 +65,7 @@ impl<'args, 'ctx, F> ExecuteRows<'args, 'ctx, F> { impl private::Sealed for ExecuteRows<'_, '_, F> {} -impl RowVisitor for ExecuteRows<'_, '_, F> { +impl RowVisitor for ExecuteRows<'_, '_, F> { type VisitResult = RowExecution; fn visit_prepared( @@ -85,16 +85,20 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { fn visit_prepared_into( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, { const { assert_sink_visit_contract::() }; - execute_sink::( + execute_sink::( self.args, self.output_dtype, self.ctx, @@ -166,7 +170,7 @@ impl<'args, 'ctx, F> ExecuteValidRows<'args, 'ctx, F> { impl private::Sealed for ExecuteValidRows<'_, '_, F> {} -impl RowVisitor for ExecuteValidRows<'_, '_, F> { +impl RowVisitor for ExecuteValidRows<'_, '_, F> { type VisitResult = Option; fn visit_prepared( @@ -188,16 +192,20 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { fn visit_prepared_into( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, { const { assert_sink_visit_contract::() }; - execute_sink_valid_rows::( + execute_sink_valid_rows::( self.args, self.output_dtype, self.valid, diff --git a/vortex-array/src/scalar_fn/row/visitor/mod.rs b/vortex-array/src/scalar_fn/row/visitor/mod.rs index 0ef085268b7..f8634f02081 100644 --- a/vortex-array/src/scalar_fn/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/row/visitor/mod.rs @@ -29,7 +29,7 @@ pub(super) use plan::PlanRows; /// /// Only the framework implements this trait. The `visit_prepared*` methods derive shared state /// from constant arguments before visiting any rows. -pub trait RowVisitor: private::Sealed + Sized { +pub trait RowVisitor: private::Sealed + Sized { /// The framework result of visiting one concrete row signature. /// /// This is a batch plan or execution result, not the per-row `Out` returned by [`visit`] and @@ -77,6 +77,11 @@ pub trait RowVisitor: private::Sealed + Sized { /// other than writing the supplied row handle. Dense execution can pass unspecified values /// from null rows. /// + /// On success, `apply` must return the write token produced by writing the `Sink::Row` supplied + /// to that same invocation. It must not return evidence produced for another row, sink, or + /// unrelated local cell. Violating this requirement can make the unsafe + /// [`OutputSink::finish`] precondition false. + /// /// # Prerequisites /// /// The framework checks these at compile time: @@ -87,12 +92,12 @@ pub trait RowVisitor: private::Sealed + Sized { /// `Args` or computing the result can fail. fn visit_into( self, - apply: impl Fn(Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + apply: impl Fn(Args::Elems<'_>, >::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, { self.visit_prepared_into::( |_| (), @@ -104,12 +109,16 @@ pub trait RowVisitor: private::Sealed + Sized { fn visit_prepared_into( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult; + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>; /// Visit a row computation that returns an owned output and deferred failure evidence. /// diff --git a/vortex-array/src/scalar_fn/row/visitor/plan.rs b/vortex-array/src/scalar_fn/row/visitor/plan.rs index 894acf33082..bf6301b1b5a 100644 --- a/vortex-array/src/scalar_fn/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/row/visitor/plan.rs @@ -26,26 +26,30 @@ use crate::scalar_fn::row::batch::BatchPlan; use crate::scalar_fn::row::batch::RowPolicy; /// The plan-time visit that validates dtypes and derives the nullable execution policy. -pub struct PlanRows<'a, F> { +pub struct PlanRows<'a, F: RowFn> { /// The input dtypes for this plan. dtypes: &'a [DType], + /// The function options used to derive a sink's runtime dtype. + options: &'a F::Options, + /// The visited function, carried only so the dispatch check can name its contract. function: PhantomData, } -impl<'a, F> PlanRows<'a, F> { - pub fn new(dtypes: &'a [DType]) -> Self { +impl<'a, F: RowFn> PlanRows<'a, F> { + pub fn new(dtypes: &'a [DType], options: &'a F::Options) -> Self { Self { dtypes, + options, function: PhantomData, } } } -impl private::Sealed for PlanRows<'_, F> {} +impl private::Sealed for PlanRows<'_, F> {} -impl RowVisitor for PlanRows<'_, F> { +impl RowVisitor for PlanRows<'_, F> { type VisitResult = BatchPlan; fn visit_prepared( @@ -68,17 +72,21 @@ impl RowVisitor for PlanRows<'_, F> { fn visit_prepared_into( self, _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - _apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + _apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, { const { assert_sink_visit_contract::() }; Ok(BatchPlan { - output_dtype: validate_sink_visit::(self.dtypes)?, + output_dtype: validate_sink_visit::(self.options, self.dtypes)?, policy: RowPolicy::for_sink::(), }) } diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index 6c96959b788..d31fde89cf2 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The [`ScalarFnVTable`] adapter shared by every [`RowFn`]. +//! The explicit [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) adapter for [`RowFn`]. //! //! The [`visitor`](super::visitor) module validates and executes the concrete row signature //! selected by dispatch. This module connects those visits to batch execution and exposes the @@ -9,20 +9,13 @@ use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_session::VortexSession; use super::row_fn::RowFn; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; -use crate::expr::Expression; -use crate::expr::union_child_validities; -use crate::scalar_fn::Arity; use crate::scalar_fn::BorrowedExecutionArgs; -use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::row::batch::Batch; use crate::scalar_fn::row::batch::KernelArgs; use crate::scalar_fn::row::batch::finalize_kernel_output; @@ -31,92 +24,144 @@ use crate::scalar_fn::row::visitor::ExecuteRows; use crate::scalar_fn::row::visitor::ExecuteValidRows; use crate::scalar_fn::row::visitor::PlanRows; -/// Implement [`ScalarFnVTable`] for every [`RowFn`]. -impl ScalarFnVTable for F { - type Options = F::Options; - - fn id(&self) -> ScalarFnId { - RowFn::id(self) - } - - fn serialize(&self, options: &Self::Options) -> VortexResult>> { - RowFn::serialize(self, options) - } - - fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { - RowFn::deserialize(self, metadata, session) - } - - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(F::ARG_NAMES.len()) - } - - fn child_name(&self, _options: &Self::Options, child_index: usize) -> ChildName { - ChildName::from(F::ARG_NAMES[child_index]) - } - - fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { - let plan = self.dispatch(options, args, PlanRows::::new(args))?; - - Ok(plan.result_dtype(args)) - } - - fn execute( - &self, - options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - // Nullary functions have no input validity to propagate, so they skip batch execution. - if args.num_inputs() == 0 { - let result_dtype = ScalarFnVTable::return_dtype(self, options, &[])?; - let nullary_args = KernelArgs { - arrays: &[], - row_count: args.row_count(), - dtypes: &[], - output_dtype: &result_dtype, - }; - - let execution = execute_rows(self, options, nullary_args, ctx)?; - let values = VortexResult::from(execution)?; - - return finalize_kernel_output( - RowFn::id(self), - &result_dtype, - args.row_count(), - values, - ctx, - ); +/// Implement [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) for one concrete [`RowFn`]. +/// +/// Adoption is explicit so a function that needs custom coercion, simplification, reduction, or +/// formatting hooks can implement the vtable itself and delegate only execution to +/// [`execute_rows`](crate::scalar_fn::execute_rows). +#[macro_export] +macro_rules! impl_row_fn_vtable { + ($function:ty) => { + impl $crate::scalar_fn::ScalarFnVTable for $function { + type Options = <$function as $crate::scalar_fn::RowFn>::Options; + + fn id(&self) -> $crate::scalar_fn::ScalarFnId { + $crate::scalar_fn::RowFn::id(self) + } + + fn serialize( + &self, + options: &Self::Options, + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult>> { + $crate::scalar_fn::RowFn::serialize(self, options) + } + + fn deserialize( + &self, + metadata: &[u8], + session: &$crate::scalar_fn::row_fn_macro_support::VortexSession, + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult { + $crate::scalar_fn::RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> $crate::scalar_fn::Arity { + $crate::scalar_fn::Arity::Exact( + <$function as $crate::scalar_fn::RowFn>::ARG_NAMES.len(), + ) + } + + fn child_name( + &self, + _options: &Self::Options, + child_index: usize, + ) -> $crate::scalar_fn::ChildName { + $crate::scalar_fn::ChildName::from( + <$function as $crate::scalar_fn::RowFn>::ARG_NAMES[child_index], + ) + } + + fn return_dtype( + &self, + options: &Self::Options, + args: &[$crate::dtype::DType], + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult<$crate::dtype::DType> { + $crate::scalar_fn::row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn $crate::scalar_fn::ExecutionArgs, + ctx: &mut $crate::ExecutionCtx, + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult<$crate::ArrayRef> { + $crate::scalar_fn::execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &$crate::expr::Expression, + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult< + Option<$crate::expr::Expression>, + > { + $crate::expr::union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + <$function as $crate::scalar_fn::RowFn>::FALLIBLE + } } + }; +} - let batch = prepare_batch(self, options, args)?; - batch.execute( - |args, ctx| self.reduce_encoded(options, args.arrays, ctx), - |args, ctx| execute_rows(self, options, args, ctx), - |args, valid, ctx| try_execute_rows_unfiltered(self, options, args, valid, ctx), - ctx, - ) - } +/// Compute the return dtype for a [`RowFn`] without adopting its complete scalar-function vtable. +pub fn row_fn_return_dtype( + function: &F, + options: &F::Options, + args: &[DType], +) -> VortexResult { + let plan = function.dispatch(options, args, PlanRows::::new(args, options))?; - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } + Ok(plan.result_dtype(args)) +} - fn is_strict(&self, _options: &Self::Options) -> bool { - true +/// Execute a [`RowFn`] while preserving a caller-owned scalar-function vtable. +/// +/// Existing vtables delegate here when they need row execution but retain custom hooks for other +/// capabilities. A simple RowFn-only function can use [`impl_row_fn_vtable`] instead. +pub fn execute_rows( + function: &F, + options: &F::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, +) -> VortexResult { + // Nullary functions have no input validity to propagate, so they skip batch execution. + if args.num_inputs() == 0 { + let result_dtype = row_fn_return_dtype(function, options, &[])?; + let nullary_args = KernelArgs { + arrays: &[], + row_count: args.row_count(), + dtypes: &[], + output_dtype: &result_dtype, + }; + + let execution = execute_row_kernel(function, options, nullary_args, ctx)?; + let values = VortexResult::from(execution)?; + + return finalize_kernel_output( + RowFn::id(function), + &result_dtype, + args.row_count(), + values, + ctx, + ); } - fn is_fallible(&self, _options: &Self::Options) -> bool { - F::FALLIBLE - } + let batch = prepare_batch(function, options, args)?; + batch.execute( + |args, ctx| function.reduce_encoded(options, args.arrays, ctx), + |args, ctx| execute_row_kernel(function, options, args, ctx), + |args, valid, ctx| try_execute_rows_unfiltered(function, options, args, valid, ctx), + ctx, + ) } /// Run the encoding-aware rewrite when available, or execute the selected row loop. -fn execute_rows( +fn execute_row_kernel( function: &F, options: &F::Options, args: KernelArgs<'_>, @@ -155,6 +200,6 @@ fn prepare_batch( args: &dyn ExecutionArgs, ) -> VortexResult { Batch::new(RowFn::id(function), args, |arg_dtypes| { - function.dispatch(options, arg_dtypes, PlanRows::::new(arg_dtypes)) + function.dispatch(options, arg_dtypes, PlanRows::::new(arg_dtypes, options)) }) } diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 7fed9106d6e..88854f24da0 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -70,7 +70,7 @@ impl RowFn for SpatialContains { } /// Containment is not symmetric, so `a` is always the container and `b` the contained. - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -93,6 +93,8 @@ impl RowFn for SpatialContains { } } +vortex_array::impl_row_fn_vtable!(SpatialContains); + /// Per-batch state for the contains row kernel: the prepared form of whichever operand is /// constant for the batch. `None` marks an operand that varies by row. struct ConstOperands { diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index 59226381bc7..751e402f118 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -60,7 +60,7 @@ impl RowFn for SpatialDistance { Ok(EmptyOptions) } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -75,6 +75,8 @@ impl RowFn for SpatialDistance { } } +vortex_array::impl_row_fn_vtable!(SpatialDistance); + #[cfg(test)] mod tests { use vortex_array::ArrayRef; diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index a7d3aa17d45..6f04e86b14a 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -64,7 +64,7 @@ impl RowFn for SpatialIntersects { Ok(EmptyOptions) } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, _args: &[DType], @@ -84,6 +84,8 @@ impl RowFn for SpatialIntersects { } } +vortex_array::impl_row_fn_vtable!(SpatialIntersects); + /// Per-batch state for the intersects row kernel: the bounding rect of each operand that is /// constant for the batch. /// diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 4b957528eab..91e4e837d64 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -83,7 +83,7 @@ impl RowFn for CosineSimilarity { Ok(EmptyOptions) } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, args: &[DType], @@ -140,6 +140,8 @@ impl RowFn for CosineSimilarity { } } +vortex_array::impl_row_fn_vtable!(CosineSimilarity); + impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index e2b6def5595..c0968d30474 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -69,7 +69,7 @@ impl RowFn for InnerProduct { Ok(EmptyOptions) } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, args: &[DType], @@ -125,6 +125,8 @@ impl RowFn for InnerProduct { } } +vortex_array::impl_row_fn_vtable!(InnerProduct); + impl ScalarFnArrayVTable for InnerProduct { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 50fa7819d66..3bad1d6d1d9 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -74,7 +74,7 @@ impl RowFn for L2Norm { Ok(EmptyOptions) } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, args: &[DType], @@ -122,6 +122,8 @@ pub(super) struct L2NormMetadata { input_dtype: Option, } +vortex_array::impl_row_fn_vtable!(L2Norm); + impl ScalarFnArrayVTable for L2Norm { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index ef86cdc80e1..77248329ae5 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -42,7 +42,7 @@ impl RowFn for L1Norm { *ID } - fn dispatch( + fn dispatch>( &self, _options: &Self::Options, args: &[DType], @@ -57,6 +57,8 @@ impl RowFn for L1Norm { } } +vortex_array::impl_row_fn_vtable!(L1Norm); + fn l1_norm_row(row: &[T]) -> T { row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) } From 18c7093d8724fd6fab8708fcfc8b0c495fd3d984 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:47:23 -0400 Subject: [PATCH 061/160] Harden RowFn sink and adapter contracts Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 66 +++++++++++++- vortex-array/benches/strict_validity.rs | 45 +++++++++- vortex-array/src/scalar_fn/mod.rs | 7 -- vortex-array/src/scalar_fn/row/batch/tests.rs | 8 +- .../src/scalar_fn/row/execute/sink/tests.rs | 4 +- vortex-array/src/scalar_fn/row/row_fn.rs | 11 ++- vortex-array/src/scalar_fn/row/types/sink.rs | 45 +++++++++- vortex-array/src/scalar_fn/row/vtable.rs | 87 +------------------ vortex-spatial/src/scalar_fn/contains.rs | 2 +- vortex-spatial/src/scalar_fn/distance.rs | 2 +- vortex-spatial/src/scalar_fn/intersects.rs | 2 +- vortex-spatial/src/scalar_fn/mod.rs | 77 ++++++++++++++++ .../src/scalar_fns/cosine_similarity.rs | 2 +- vortex-tensor/src/scalar_fns/inner_product.rs | 2 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 2 +- vortex-tensor/src/scalar_fns/mod.rs | 77 ++++++++++++++++ vortex-tensor/src/scalar_fns/tests/row.rs | 2 +- 17 files changed, 330 insertions(+), 111 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index 056fe2007ff..77baf03ea83 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -38,6 +38,62 @@ const ROWS: usize = 65_536; static SESSION: LazyLock = LazyLock::new(array_session); +/// Adopt the standard scalar-function behavior for a row function in this benchmark. +macro_rules! impl_row_fn_scalar_vtable { + ($function:ty) => { + impl ScalarFnVTable for $function { + type Options = <$function as RowFn>::Options; + + fn id(&self) -> ScalarFnId { + RowFn::id(self) + } + + fn arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::Arity { + vortex_array::scalar_fn::Arity::Exact(<$function as RowFn>::ARG_NAMES.len()) + } + + fn child_name( + &self, + _options: &Self::Options, + child_index: usize, + ) -> vortex_array::scalar_fn::ChildName { + vortex_array::scalar_fn::ChildName::from( + <$function as RowFn>::ARG_NAMES[child_index], + ) + } + + fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { + vortex_array::scalar_fn::row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn vortex_array::scalar_fn::ExecutionArgs, + ctx: &mut vortex_array::ExecutionCtx, + ) -> VortexResult { + vortex_array::scalar_fn::execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &vortex_array::expr::Expression, + ) -> VortexResult> { + vortex_array::expr::union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + <$function as RowFn>::FALLIBLE + } + } + }; +} + fn main() { LazyLock::force(&SESSION); divan::main(); @@ -111,7 +167,9 @@ struct I64Sink( BufferMut, ); -impl OutputSink for I64Sink { +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { type Rows<'a> = &'a mut [i64]; type Row<'a> = &'a mut i64; type WriteToken = (); @@ -166,9 +224,9 @@ impl RowFn for RowSinkWrappingAdd { } } -vortex_array::impl_row_fn_vtable!(RowWrappingAdd); -vortex_array::impl_row_fn_vtable!(RowCheckedAdd); -vortex_array::impl_row_fn_vtable!(RowSinkWrappingAdd); +impl_row_fn_scalar_vtable!(RowWrappingAdd); +impl_row_fn_scalar_vtable!(RowCheckedAdd); +impl_row_fn_scalar_vtable!(RowSinkWrappingAdd); fn inputs() -> (ArrayRef, ArrayRef) { let lhs = (0..ROWS) diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs index e0a8b6a565e..e765889cbd4 100644 --- a/vortex-array/benches/strict_validity.rs +++ b/vortex-array/benches/strict_validity.rs @@ -106,7 +106,50 @@ impl RowFn for LazyDouble { } } -vortex_array::impl_row_fn_vtable!(LazyDouble); +impl ScalarFnVTable for LazyDouble { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + RowFn::id(self) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(Self::ARG_NAMES.len()) + } + + fn child_name(&self, _options: &Self::Options, child_index: usize) -> ChildName { + ChildName::from(Self::ARG_NAMES[child_index]) + } + + fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { + vortex_array::scalar_fn::row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + vortex_array::scalar_fn::execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + Self::FALLIBLE + } +} /// The same function, applying validity the way the adapter used to: materialize a mask first. #[derive(Clone)] diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 70a8443314d..5e73caefdfa 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -19,13 +19,6 @@ use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::literal::Literal; -/// Reexports used by [`impl_row_fn_vtable!`] without requiring downstream transitive dependencies. -#[doc(hidden)] -pub mod row_fn_macro_support { - pub use vortex_error::VortexResult; - pub use vortex_session::VortexSession; -} - mod vtable; pub use vtable::*; diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 9760db84f68..06539114638 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -77,7 +77,9 @@ enum PreparedVisit { Deferred, } -impl OutputSink for OptionsCheckingSink { +// SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or +// `finish` through the executor. The row-initialization requirements are therefore vacuous. +unsafe impl OutputSink for OptionsCheckingSink { type Rows<'a> = (); type Row<'a> = (); type WriteToken = (); @@ -122,7 +124,9 @@ impl OutputElement for NullProducingI64 { struct I64Sink(BufferMut); -impl OutputSink for I64Sink { +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { type Rows<'a> = &'a mut [i64]; type Row<'a> = &'a mut i64; type WriteToken = (); diff --git a/vortex-array/src/scalar_fn/row/execute/sink/tests.rs b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs index 01f6f0b2bd3..fe23c07c445 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink/tests.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs @@ -20,7 +20,9 @@ use crate::validity::Validity; struct NonSkippingSink; -impl OutputSink for NonSkippingSink { +// SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or +// `finish` through the executor. The row-initialization requirements are therefore vacuous. +unsafe impl OutputSink for NonSkippingSink { type Rows<'a> = (); type Row<'a> = (); type WriteToken = (); diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index 4105a3690dc..c84c07b62e2 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -21,8 +21,15 @@ use crate::scalar_fn::ScalarFnId; /// A scalar function computed one row at a time. /// /// Declare the argument names and use [`dispatch`](Self::dispatch) to choose concrete element and -/// sink types for each accepted dtype combination. Implement -/// [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) directly for columnar kernels. +/// sink types for each accepted dtype combination. +/// +/// A `RowFn` does not automatically implement +/// [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable). A function adopted by the expression +/// system implements that trait explicitly, delegates its `return_dtype` method to +/// [`row_fn_return_dtype`](crate::scalar_fn::row_fn_return_dtype), and delegates `execute` to +/// [`execute_rows`](crate::scalar_fn::execute_rows). Explicit adoption leaves the function free to +/// provide custom coercion, simplification, reduction, or formatting hooks. Implement only +/// `ScalarFnVTable` when the natural kernel is columnar rather than row-oriented. pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// Options for this function, if any. Use [`EmptyOptions`](crate::scalar_fn::EmptyOptions) /// for none. diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index a75da4c44ca..3e520718b99 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -20,7 +20,24 @@ use crate::scalar_fn::OutputElement; /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; /// skip-invalid execution can omit invalid rows when /// [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) is present. -pub trait OutputSink: 'static + Sized { +/// +/// # Safety +/// +/// An implementation must uphold all of these requirements: +/// +/// - When [`row_count_matches`](Self::row_count_matches) returns `true`, every index in +/// `0..row_count` **must** identify one distinct row owned by this sink. +/// - A row must either be initialized before the callback or require a +/// [`WriteToken`](Self::WriteToken) that safe code cannot produce without initializing that exact +/// row. Evidence for an uninitialized row **must not** be safely forgeable, reusable, or +/// substitutable for another row. +/// - A present [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) **must** initialize +/// every row handed to it. +/// - [`finish`](Self::finish) **must** be sound once every visited callback returned its required +/// token and the skipped-row initializer, when present, ran successfully. +/// +/// The executor relies on these guarantees when it calls `finish`. +pub unsafe trait OutputSink: 'static + Sized { /// A loop-local view of all output rows. /// /// Borrowed once before execution so the sink's buffer descriptor and shape become loop @@ -84,6 +101,14 @@ pub trait OutputSink: 'static + Sized { } /// Proof that one uninitialized element row was initialized. +/// +/// The private field prevents safe construction without calling [`write`](Self::write): +/// +/// ```compile_fail,E0423 +/// use vortex_array::scalar_fn::InitializedElement; +/// +/// let _evidence = InitializedElement(()); +/// ``` #[must_use = "return this token from the row closure to prove that it initialized the output"] pub struct InitializedElement( /// Private so constructing initialization evidence requires an unsafe operation. @@ -98,6 +123,16 @@ impl InitializedElement { /// `row` must be the [`UninitElementSink`] row supplied to the current callback. The caller must /// return the token from that callback. Using another row or returning the token from another /// callback can cause undefined behavior. + /// + /// Safe code cannot construct initialization evidence: + /// + /// ```compile_fail,E0133 + /// use std::mem::MaybeUninit; + /// use vortex_array::scalar_fn::InitializedElement; + /// + /// let mut unrelated = MaybeUninit::::uninit(); + /// let _evidence = InitializedElement::write(&mut unrelated, 42); + /// ``` #[inline] pub unsafe fn write(row: &mut MaybeUninit, value: T) -> Self { row.write(value); @@ -121,7 +156,13 @@ pub struct UninitElementSink { row_count: usize, } -impl OutputSink for UninitElementSink { +// SAFETY: the row slice covers exactly the reserved spare-capacity range, so each accepted index +// names one distinct slot. `InitializedElement` cannot be constructed by safe code; its unsafe +// constructor writes the supplied slot and requires the caller to return that exact evidence. The +// skipped-row initializer writes `T::default()` into every slot before masked traversal. +unsafe impl OutputSink + for UninitElementSink +{ type Rows<'a> = &'a mut [MaybeUninit]; const SKIPPED_ROWS_INITIALIZER: Option fn(&mut Self::Rows<'a>)> = Some(|rows| { diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index d31fde89cf2..96a24f8f42d 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -24,90 +24,6 @@ use crate::scalar_fn::row::visitor::ExecuteRows; use crate::scalar_fn::row::visitor::ExecuteValidRows; use crate::scalar_fn::row::visitor::PlanRows; -/// Implement [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) for one concrete [`RowFn`]. -/// -/// Adoption is explicit so a function that needs custom coercion, simplification, reduction, or -/// formatting hooks can implement the vtable itself and delegate only execution to -/// [`execute_rows`](crate::scalar_fn::execute_rows). -#[macro_export] -macro_rules! impl_row_fn_vtable { - ($function:ty) => { - impl $crate::scalar_fn::ScalarFnVTable for $function { - type Options = <$function as $crate::scalar_fn::RowFn>::Options; - - fn id(&self) -> $crate::scalar_fn::ScalarFnId { - $crate::scalar_fn::RowFn::id(self) - } - - fn serialize( - &self, - options: &Self::Options, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult>> { - $crate::scalar_fn::RowFn::serialize(self, options) - } - - fn deserialize( - &self, - metadata: &[u8], - session: &$crate::scalar_fn::row_fn_macro_support::VortexSession, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult { - $crate::scalar_fn::RowFn::deserialize(self, metadata, session) - } - - fn arity(&self, _options: &Self::Options) -> $crate::scalar_fn::Arity { - $crate::scalar_fn::Arity::Exact( - <$function as $crate::scalar_fn::RowFn>::ARG_NAMES.len(), - ) - } - - fn child_name( - &self, - _options: &Self::Options, - child_index: usize, - ) -> $crate::scalar_fn::ChildName { - $crate::scalar_fn::ChildName::from( - <$function as $crate::scalar_fn::RowFn>::ARG_NAMES[child_index], - ) - } - - fn return_dtype( - &self, - options: &Self::Options, - args: &[$crate::dtype::DType], - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult<$crate::dtype::DType> { - $crate::scalar_fn::row_fn_return_dtype(self, options, args) - } - - fn execute( - &self, - options: &Self::Options, - args: &dyn $crate::scalar_fn::ExecutionArgs, - ctx: &mut $crate::ExecutionCtx, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult<$crate::ArrayRef> { - $crate::scalar_fn::execute_rows(self, options, args, ctx) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &$crate::expr::Expression, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult< - Option<$crate::expr::Expression>, - > { - $crate::expr::union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - <$function as $crate::scalar_fn::RowFn>::FALLIBLE - } - } - }; -} - /// Compute the return dtype for a [`RowFn`] without adopting its complete scalar-function vtable. pub fn row_fn_return_dtype( function: &F, @@ -122,7 +38,8 @@ pub fn row_fn_return_dtype( /// Execute a [`RowFn`] while preserving a caller-owned scalar-function vtable. /// /// Existing vtables delegate here when they need row execution but retain custom hooks for other -/// capabilities. A simple RowFn-only function can use [`impl_row_fn_vtable`] instead. +/// capabilities. A RowFn-only function implements the remaining vtable methods mechanically and +/// delegates its return-dtype and execution methods here. pub fn execute_rows( function: &F, options: &F::Options, diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 88854f24da0..514b76593b7 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -93,7 +93,7 @@ impl RowFn for SpatialContains { } } -vortex_array::impl_row_fn_vtable!(SpatialContains); +impl_row_fn_scalar_vtable!(SpatialContains); /// Per-batch state for the contains row kernel: the prepared form of whichever operand is /// constant for the batch. `None` marks an operand that varies by row. diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index 751e402f118..3b6b086b3a0 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -75,7 +75,7 @@ impl RowFn for SpatialDistance { } } -vortex_array::impl_row_fn_vtable!(SpatialDistance); +impl_row_fn_scalar_vtable!(SpatialDistance); #[cfg(test)] mod tests { diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index 6f04e86b14a..37d5a72ca61 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -84,7 +84,7 @@ impl RowFn for SpatialIntersects { } } -vortex_array::impl_row_fn_vtable!(SpatialIntersects); +impl_row_fn_scalar_vtable!(SpatialIntersects); /// Per-batch state for the intersects row kernel: the bounding rect of each operand that is /// constant for the batch. diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index bcdb15e51e6..3d02d9fbd56 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -3,6 +3,83 @@ //! Geometry scalar functions over the native geometry extension types. +/// Adopt the standard scalar-function behavior for a row function defined in this module. +macro_rules! impl_row_fn_scalar_vtable { + ($function:ty) => { + impl vortex_array::scalar_fn::ScalarFnVTable for $function { + type Options = <$function as vortex_array::scalar_fn::RowFn>::Options; + + fn id(&self) -> vortex_array::scalar_fn::ScalarFnId { + vortex_array::scalar_fn::RowFn::id(self) + } + + fn serialize( + &self, + options: &Self::Options, + ) -> vortex_error::VortexResult>> { + vortex_array::scalar_fn::RowFn::serialize(self, options) + } + + fn deserialize( + &self, + metadata: &[u8], + session: &vortex_session::VortexSession, + ) -> vortex_error::VortexResult { + vortex_array::scalar_fn::RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::Arity { + vortex_array::scalar_fn::Arity::Exact( + <$function as vortex_array::scalar_fn::RowFn>::ARG_NAMES.len(), + ) + } + + fn child_name( + &self, + _options: &Self::Options, + child_index: usize, + ) -> vortex_array::scalar_fn::ChildName { + vortex_array::scalar_fn::ChildName::from( + <$function as vortex_array::scalar_fn::RowFn>::ARG_NAMES[child_index], + ) + } + + fn return_dtype( + &self, + options: &Self::Options, + args: &[vortex_array::dtype::DType], + ) -> vortex_error::VortexResult { + vortex_array::scalar_fn::row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn vortex_array::scalar_fn::ExecutionArgs, + ctx: &mut vortex_array::ExecutionCtx, + ) -> vortex_error::VortexResult { + vortex_array::scalar_fn::execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &vortex_array::expr::Expression, + ) -> vortex_error::VortexResult> { + vortex_array::expr::union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + <$function as vortex_array::scalar_fn::RowFn>::FALLIBLE + } + } + }; +} + pub mod contains; pub mod distance; pub mod envelope; diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 91e4e837d64..3b08c6d46fc 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -140,7 +140,7 @@ impl RowFn for CosineSimilarity { } } -vortex_array::impl_row_fn_vtable!(CosineSimilarity); +impl_row_fn_scalar_vtable!(CosineSimilarity); impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index c0968d30474..394cd3b158a 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -125,7 +125,7 @@ impl RowFn for InnerProduct { } } -vortex_array::impl_row_fn_vtable!(InnerProduct); +impl_row_fn_scalar_vtable!(InnerProduct); impl ScalarFnArrayVTable for InnerProduct { fn serialize( diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 3bad1d6d1d9..d0a9b1f951b 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -122,7 +122,7 @@ pub(super) struct L2NormMetadata { input_dtype: Option, } -vortex_array::impl_row_fn_vtable!(L2Norm); +impl_row_fn_scalar_vtable!(L2Norm); impl ScalarFnArrayVTable for L2Norm { fn serialize( diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index da9b8950e7a..fd343f2c8fa 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -3,6 +3,83 @@ //! Scalar function expressions defined on tensor and tensor-like extension types. +/// Adopt the standard scalar-function behavior for a row function defined in this module. +macro_rules! impl_row_fn_scalar_vtable { + ($function:ty) => { + impl vortex_array::scalar_fn::ScalarFnVTable for $function { + type Options = <$function as vortex_array::scalar_fn::RowFn>::Options; + + fn id(&self) -> vortex_array::scalar_fn::ScalarFnId { + vortex_array::scalar_fn::RowFn::id(self) + } + + fn serialize( + &self, + options: &Self::Options, + ) -> vortex_error::VortexResult>> { + vortex_array::scalar_fn::RowFn::serialize(self, options) + } + + fn deserialize( + &self, + metadata: &[u8], + session: &vortex_session::VortexSession, + ) -> vortex_error::VortexResult { + vortex_array::scalar_fn::RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::Arity { + vortex_array::scalar_fn::Arity::Exact( + <$function as vortex_array::scalar_fn::RowFn>::ARG_NAMES.len(), + ) + } + + fn child_name( + &self, + _options: &Self::Options, + child_index: usize, + ) -> vortex_array::scalar_fn::ChildName { + vortex_array::scalar_fn::ChildName::from( + <$function as vortex_array::scalar_fn::RowFn>::ARG_NAMES[child_index], + ) + } + + fn return_dtype( + &self, + options: &Self::Options, + args: &[vortex_array::dtype::DType], + ) -> vortex_error::VortexResult { + vortex_array::scalar_fn::row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn vortex_array::scalar_fn::ExecutionArgs, + ctx: &mut vortex_array::ExecutionCtx, + ) -> vortex_error::VortexResult { + vortex_array::scalar_fn::execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &vortex_array::expr::Expression, + ) -> vortex_error::VortexResult> { + vortex_array::expr::union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + <$function as vortex_array::scalar_fn::RowFn>::FALLIBLE + } + } + }; +} + pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index 77248329ae5..6313b0f8d68 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -57,7 +57,7 @@ impl RowFn for L1Norm { } } -vortex_array::impl_row_fn_vtable!(L1Norm); +impl_row_fn_scalar_vtable!(L1Norm); fn l1_norm_row(row: &[T]) -> T { row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) From ee83291ff235c06c8f0fe4555ff51ef49bddb3e9 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:53:58 -0400 Subject: [PATCH 062/160] Generalize indexed RowFn tuple sources Signed-off-by: Connor Tsui --- .../src/scalar_fn/row/execute/owned.rs | 4 +- .../src/scalar_fn/row/types/element/tuple.rs | 159 ++++++++++++++---- vortex-compute/src/lane_kernels/source.rs | 5 +- 3 files changed, 135 insertions(+), 33 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/execute/owned.rs b/vortex-array/src/scalar_fn/row/execute/owned.rs index 97519408937..490d93d1040 100644 --- a/vortex-array/src/scalar_fn/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/row/execute/owned.rs @@ -83,7 +83,9 @@ where "a decoded row input does not address exactly {row_count} rows", ); - failure = Args::indexed_source(&varying) + // SAFETY: `varying_len_matches` proved every column addresses exactly `row_count` + // rows immediately above. + failure = unsafe { Args::indexed_source(varying, row_count) } .map_checked_into(output, |elements| apply(&prepared, elements)); } else { // A batch-constant input was collapsed to one row during decoding. This path reads that diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple.rs b/vortex-array/src/scalar_fn/row/types/element/tuple.rs index 680ae0b4e9c..04821d78363 100644 --- a/vortex-array/src/scalar_fn/row/types/element/tuple.rs +++ b/vortex-array/src/scalar_fn/row/types/element/tuple.rs @@ -15,7 +15,6 @@ use crate::arrays::Masked; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::masked::MaskedArraySlotsExt; use crate::dtype::DType; -use crate::dtype::NativePType; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::InputElement; @@ -205,11 +204,8 @@ pub trait ElementTuple: 'static + private::Sealed { /// An argument tuple that supports a validated dense indexed traversal. /// -/// This is separate from [`ElementTuple`] because many row elements have no contiguous source, and -/// stable Rust cannot provide a blanket fallback plus a more specific primitive implementation. -/// The trait is sealed so shared execution can rely on its unchecked-read contract. A tuple only -/// implements it when the source can be validated once and every lane can then be read -/// independently. +/// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's +/// unchecked varying-row access after batch execution validates every decoded column length once. pub trait IndexedElementTuple: ElementTuple { /// The source shared execution uses for a dense all-varying loop. /// @@ -218,27 +214,77 @@ pub trait IndexedElementTuple: ElementTuple { /// read contract of [`IndexedSource`]. type Source<'a>: IndexedSource>; - /// Borrow a source from columns already validated to vary within the batch. - fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a>; + /// Build a source from columns already validated to vary within the batch. + /// + /// # Safety + /// + /// Every column in `columns` **must** address exactly `row_count` rows. Violating this + /// requirement can make a safe lane kernel read outside a column's allocation. + unsafe fn indexed_source<'a>( + columns: Self::VaryingColumns<'a>, + row_count: usize, + ) -> Self::Source<'a>; } -/// An indexed native slice yielding the one-tuples expected by a unary row closure. -#[derive(Clone, Copy)] -pub struct UnaryTupleSource<'a, T>( - /// The native values read by the row loop. - &'a [T], -); +/// Indexed access to one varying element column. +pub struct ElementSource<'a, T: InputElement> { + column: T::Varying<'a>, +} + +impl<'a, T: InputElement> ElementSource<'a, T> { + fn new(column: T::Varying<'a>) -> Self { + Self { column } + } +} -impl IndexedSource for UnaryTupleSource<'_, T> { - type Item = (T,); +impl<'a, T: InputElement> IndexedSource for ElementSource<'a, T> { + type Item = T::Elem<'a>; + + fn len(&self) -> usize { + T::varying_len(&self.column) + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the source length is the number of rows addressable by `column`, and the caller + // guarantees that `index` is below that length. + unsafe { T::get_varying_unchecked(&self.column, index) } + } +} + +/// An indexed element source yielding the one-tuples expected by a unary row closure. +pub struct UnaryTupleSource(Source); + +impl IndexedSource for UnaryTupleSource { + type Item = (Source::Item,); fn len(&self) -> usize { self.0.len() } unsafe fn get_unchecked(&self, index: usize) -> Self::Item { - // SAFETY: the caller guarantees that `index` is in bounds. - (unsafe { *self.0.get_unchecked(index) },) + // SAFETY: forwarded from this method's contract. + (unsafe { self.0.get_unchecked(index) },) + } +} + +/// Indexed access to the varying columns of an element tuple. +pub struct ElementTupleSource<'a, Args: ElementTuple> { + columns: Args::VaryingColumns<'a>, + row_count: usize, +} + +impl<'a, Args: ElementTuple> IndexedSource for ElementTupleSource<'a, Args> { + type Item = Args::Elems<'a>; + + fn len(&self) -> usize { + self.row_count + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the caller guarantees that `index` is below `row_count`. Batch execution checks + // that every varying column addresses exactly `row_count` rows before constructing this + // source. + unsafe { Args::get_varying_unchecked(&self.columns, index) } } } @@ -400,22 +446,65 @@ element_tuple!(10; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9); element_tuple!(11; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10); element_tuple!(12; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11); -impl IndexedElementTuple for (T,) { - type Source<'a> = UnaryTupleSource<'a, T>; +impl IndexedElementTuple for () { + type Source<'a> = ElementTupleSource<'a, ()>; - fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a> { - UnaryTupleSource(columns.0) + unsafe fn indexed_source<'a>( + columns: Self::VaryingColumns<'a>, + row_count: usize, + ) -> Self::Source<'a> { + ElementTupleSource { columns, row_count } } } -impl IndexedElementTuple for (Left, Right) { - type Source<'a> = LaneZip<&'a [Left], &'a [Right]>; +impl IndexedElementTuple for (A,) { + type Source<'a> = UnaryTupleSource>; - fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a> { - LaneZip::new(columns.0, columns.1) + unsafe fn indexed_source<'a>( + columns: Self::VaryingColumns<'a>, + _row_count: usize, + ) -> Self::Source<'a> { + UnaryTupleSource(ElementSource::new(columns.0)) } } +impl IndexedElementTuple for (A, B) { + type Source<'a> = LaneZip, ElementSource<'a, B>>; + + unsafe fn indexed_source<'a>( + columns: Self::VaryingColumns<'a>, + _row_count: usize, + ) -> Self::Source<'a> { + LaneZip::new(ElementSource::new(columns.0), ElementSource::new(columns.1)) + } +} + +macro_rules! indexed_element_tuple { + ($($t:ident),+) => { + impl<$($t: InputElement),+> IndexedElementTuple for ($($t,)+) { + type Source<'a> = ElementTupleSource<'a, ($($t,)+)>; + + unsafe fn indexed_source<'a>( + columns: Self::VaryingColumns<'a>, + row_count: usize, + ) -> Self::Source<'a> { + ElementTupleSource { columns, row_count } + } + } + }; +} + +indexed_element_tuple!(A, B, C); +indexed_element_tuple!(A, B, C, D); +indexed_element_tuple!(A, B, C, D, E); +indexed_element_tuple!(A, B, C, D, E, F); +indexed_element_tuple!(A, B, C, D, E, F, G); +indexed_element_tuple!(A, B, C, D, E, F, G, H); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K, L); + #[cfg(test)] mod tests { use vortex_compute::lane_kernels::IndexedSource; @@ -423,7 +512,7 @@ mod tests { use vortex_error::vortex_bail; use vortex_mask::Mask; - use super::UnaryTupleSource; + use super::IndexedElementTuple; use super::batch_constant; use crate::IntoArray; use crate::arrays::ConstantArray; @@ -435,14 +524,26 @@ mod tests { use crate::validity::Validity; #[test] - fn test_unary_tuple_source_reads_one_tuple_per_row() { - let source = UnaryTupleSource(&[10, 20, 30]); + fn test_unary_element_source_reads_one_tuple_per_row() { + // SAFETY: the only column has exactly three rows. + let source = unsafe { <(i32,)>::indexed_source((&[10, 20, 30][..],), 3) }; assert_eq!(source.len(), 3); // SAFETY: index one is within the three-element source. assert_eq!(unsafe { source.get_unchecked(1) }, (20,)); } + #[test] + fn test_binary_element_source_zips_columns() { + // SAFETY: both columns have exactly three rows. + let source = + unsafe { <(i32, i64)>::indexed_source((&[10, 20, 30][..], &[100, 200, 300][..]), 3) }; + assert_eq!(source.len(), 3); + + // SAFETY: index one is within the three-element source. + assert_eq!(unsafe { source.get_unchecked(1) }, (20, 200)); + } + #[test] fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> { let child = ConstantArray::new(7_i64, 3).into_array(); diff --git a/vortex-compute/src/lane_kernels/source.rs b/vortex-compute/src/lane_kernels/source.rs index 9a486b4789b..5f295e76361 100644 --- a/vortex-compute/src/lane_kernels/source.rs +++ b/vortex-compute/src/lane_kernels/source.rs @@ -10,9 +10,8 @@ /// reads carry no inter-iteration data dependency — the autovectorizer treats each /// lane independently. pub trait IndexedSource { - /// The per-lane item type. Must be `Copy` so the kernels can pass it through - /// the closure by value without extra moves. - type Item: Copy; + /// The per-lane item type passed through the kernel by value. + type Item; /// Logical lane count. fn len(&self) -> usize; /// Returns true when there are no lanes. From 33f9c34adbcb1266024c4c49a807c0acf2d84b94 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:02:12 -0400 Subject: [PATCH 063/160] Automate RowFn benchmark comparisons Signed-off-by: Connor Tsui --- scripts/benchmark-rowfn.sh | 328 ++++++++++++++++++++++++++ scripts/rowfn_benchmark.py | 306 ++++++++++++++++++++++++ scripts/tests/test_rowfn_benchmark.py | 111 +++++++++ 3 files changed, 745 insertions(+) create mode 100755 scripts/benchmark-rowfn.sh create mode 100755 scripts/rowfn_benchmark.py create mode 100644 scripts/tests/test_rowfn_benchmark.py diff --git a/scripts/benchmark-rowfn.sh b/scripts/benchmark-rowfn.sh new file mode 100755 index 00000000000..20f705dd194 --- /dev/null +++ b/scripts/benchmark-rowfn.sh @@ -0,0 +1,328 @@ +#!/usr/bin/env bash + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +set -Eeu -o pipefail + +usage() { + cat >&2 <<'EOF' +Usage: benchmark-rowfn.sh [OPTIONS] + +Options: + --suite NAME Select a preset or benchmark label. Repeatable; defaults to full. + --filter PATTERN Pass a Divan benchmark filter. Repeatable. + --config NAME primary (1 CGU/fat LTO) or repository (16 CGUs/no LTO). + --target-root PATH New directory for separate baseline and candidate Cargo targets. + --codegen-units N Override the selected configuration. + --lto VALUE Override LTO with false, thin, or fat. + --rustflags FLAGS Override RUSTFLAGS; defaults to -C target-cpu=native. + --build-jobs N Jobs per concurrent revision build; defaults to 8 and cannot exceed 8. + --bench-cpu N Logical CPU used for every timed process; defaults to 4. + --warm-runs N Warm runs per revision; defaults to 2. + --measured-pairs N Alternating measured pairs; defaults to 7. + --sample-count N Divan sample count; defaults to 100. + --min-time SECONDS Divan minimum time; defaults to 0.25. + --max-time SECONDS Divan maximum time; defaults to 0.5. + --lock-file PATH Global timed-run lock; defaults to /tmp/vortex-rowfn-benchmark.lock. + --list-suites Print presets and benchmark labels, then exit. +EOF +} + +suite_catalog=( + "array-binary_ops|vortex-array|binary_ops|array,numeric,design-a-matrix,full" + "array-compare|vortex-array|compare|array,compare,full" + "array-row_fn_executor|vortex-array|row_fn_executor|array,framework,full" + "array-strict_validity|vortex-array|strict_validity|array,framework,full" + "array-like|vortex-array|like|array,full" + "array-take_filter|vortex-array|take_filter|array,full" + "array-varbinview_compact|vortex-array|varbinview_compact|array,full" + "tensor-l2_norm|vortex-tensor|l2_norm|tensor,full" + "tensor-inner_product|vortex-tensor|inner_product|tensor,full" + "tensor-cosine_similarity|vortex-tensor|cosine_similarity|tensor,full" + "tensor-normalized|vortex-tensor|normalized|tensor,full" + "spatial-binary_predicates|vortex-spatial|binary_predicates|spatial,full" + "spatial-distance|vortex-spatial|distance|spatial,full" + "spatial-envelope|vortex-spatial|envelope|spatial,full" + "spatial-predicate_bbox|vortex-spatial|predicate_bbox|spatial,full" +) + +requested_suites=() +filters=() +configuration=primary +target_root= +codegen_units_override= +lto_override= +rustflags_override= +build_jobs=8 +bench_cpu=4 +warm_runs=2 +measured_pairs=7 +sample_count=100 +min_time=0.25 +max_time=0.5 +lock_file=/tmp/vortex-rowfn-benchmark.lock + +while [[ $# -gt 0 ]]; do + case $1 in + --suite) requested_suites+=("$2"); shift 2 ;; + --filter) filters+=("$2"); shift 2 ;; + --config) configuration=$2; shift 2 ;; + --target-root) target_root=$2; shift 2 ;; + --codegen-units) codegen_units_override=$2; shift 2 ;; + --lto) lto_override=$2; shift 2 ;; + --rustflags) rustflags_override=$2; shift 2 ;; + --build-jobs) build_jobs=$2; shift 2 ;; + --bench-cpu) bench_cpu=$2; shift 2 ;; + --warm-runs) warm_runs=$2; shift 2 ;; + --measured-pairs) measured_pairs=$2; shift 2 ;; + --sample-count) sample_count=$2; shift 2 ;; + --min-time) min_time=$2; shift 2 ;; + --max-time) max_time=$2; shift 2 ;; + --lock-file) lock_file=$2; shift 2 ;; + --list-suites) + echo "Presets: full array framework numeric design-a-matrix compare tensor spatial" + printf '%s\n' "${suite_catalog[@]}" | cut -d '|' -f 1 + exit 0 + ;; + -h|--help) usage; exit 0 ;; + --*) echo "Unknown option: $1" >&2; usage; exit 1 ;; + *) break ;; + esac +done + +if [[ $# -ne 3 ]]; then + usage + exit 1 +fi +if [[ $(uname -m) != x86_64 ]]; then + echo "RowFn native performance decisions require an x86_64 host." >&2 + exit 1 +fi +if ((build_jobs < 1 || build_jobs > 8)); then + echo "--build-jobs must be between 1 and 8 so two builds cannot exceed 16 jobs." >&2 + exit 1 +fi +command -v flock >/dev/null || { echo "benchmark-rowfn.sh requires flock." >&2; exit 1; } + +baseline=$(realpath "$1") +candidate=$(realpath "$2") +output=$(realpath -m "$3") +if [[ -e $output ]]; then + echo "Output path already exists: $output" >&2 + exit 1 +fi + +case $configuration in + primary) codegen_units=1; lto=fat ;; + repository) codegen_units=16; lto=false ;; + *) echo "Unknown configuration: $configuration" >&2; exit 1 ;; +esac +codegen_units=${codegen_units_override:-$codegen_units} +lto=${lto_override:-$lto} +rustflags=${rustflags_override:--C target-cpu=native} + +if ((${#requested_suites[@]} == 0)); then + requested_suites=(full) +fi +selected_suites=() +declare -A selected_labels=() +for request in "${requested_suites[@]}"; do + matched=false + for entry in "${suite_catalog[@]}"; do + IFS='|' read -r label _ _ groups <<<"$entry" + if [[ $request == "$label" || ,$groups, == *,$request,* ]]; then + matched=true + if [[ -z ${selected_labels[$label]:-} ]]; then + selected_suites+=("$entry") + selected_labels[$label]=1 + fi + fi + done + if [[ $matched == false ]]; then + echo "Unknown suite or benchmark label: $request" >&2 + exit 1 + fi +done + +common_git_dir=$(git -C "$candidate" rev-parse --path-format=absolute --git-common-dir) +repository_root=$(dirname "$common_git_dir") +if [[ -z $target_root ]]; then + target_root="$repository_root/target/rowfn-benchmark/$(basename "$output")" +fi +target_root=$(realpath -m "$target_root") +if [[ -e $target_root ]]; then + echo "Target root already exists: $target_root" >&2 + exit 1 +fi + +mkdir -p "$output/build" "$output/warm" "$output/measured" "$target_root" +baseline_target="$target_root/baseline" +candidate_target="$target_root/candidate" +parser="$candidate/scripts/rowfn_benchmark.py" + +{ + echo "RowFn benchmark machine record" + echo "Date: $(date --iso-8601=seconds)" + echo "Host: $(hostname)" + echo "Kernel: $(uname -srvmo)" + echo "Benchmark CPU: $bench_cpu" + echo "Configuration: $configuration" + echo "Cargo profile: bench, $codegen_units codegen units, LTO $lto" + echo "RUSTFLAGS: $rustflags" + echo "Warm runs: $warm_runs" + echo "Measured pairs: $measured_pairs" + echo "Divan: TSC timer, $sample_count samples, min $min_time s, max $max_time s" + echo + rustc -vV + cargo -V + echo + lscpu + echo + rg -m1 '^microcode' /proc/cpuinfo || true + for path in \ + /sys/devices/system/cpu/cpu"$bench_cpu"/cpufreq/scaling_governor \ + /sys/devices/system/cpu/cpu"$bench_cpu"/cpufreq/energy_performance_preference \ + /sys/devices/system/cpu/cpufreq/boost; do + [[ -r $path ]] && echo "$path: $(<"$path")" + done +} >"$output/machine.txt" + +build_revision() { + local worktree=$1 + local target=$2 + local log=$3 + + ( + cd "$worktree" + export CARGO_TARGET_DIR=$target + export CARGO_PROFILE_BENCH_CODEGEN_UNITS=$codegen_units + export CARGO_PROFILE_BENCH_LTO=$lto + export RUSTFLAGS=$rustflags + for entry in "${selected_suites[@]}"; do + IFS='|' read -r _ package bench _ <<<"$entry" + cargo bench --no-run -j "$build_jobs" -p "$package" --bench "$bench" + done + ) >"$log" 2>&1 +} + +echo "Building baseline and candidate with $build_jobs jobs each." +build_revision "$baseline" "$baseline_target" "$output/build/baseline.txt" & +baseline_pid=$! +build_revision "$candidate" "$candidate_target" "$output/build/candidate.txt" & +candidate_pid=$! +baseline_status=0 +candidate_status=0 +wait "$baseline_pid" || baseline_status=$? +wait "$candidate_pid" || candidate_status=$? +if [[ $baseline_status -ne 0 || $candidate_status -ne 0 ]]; then + echo "Benchmark build failed; see $output/build/." >&2 + exit 1 +fi + +find_benchmark() { + local target=$1 + local name=$2 + local binary + + binary=$(find "$target/release/deps" -maxdepth 1 -type f -executable -name "$name-*" \ + -printf '%T@ %p\n' | sort -nr | head -n 1 | cut -d ' ' -f 2-) + [[ -n $binary ]] || { echo "Cannot find benchmark $name under $target." >&2; exit 1; } + echo "$binary" +} + +declare -A baseline_binaries=() +declare -A candidate_binaries=() +manifest_args=( + manifest + --output "$output/manifest.json" + --machine-record "$output/machine.txt" + --baseline-worktree "$baseline" + --candidate-worktree "$candidate" + --baseline-target "$baseline_target" + --candidate-target "$candidate_target" + --setting "configuration=$configuration" + --setting "codegen_units=$codegen_units" + --setting "lto=$lto" + --setting "rustflags=$rustflags" + --setting "bench_cpu=$bench_cpu" + --setting "warm_runs=$warm_runs" + --setting "measured_pairs=$measured_pairs" + --setting "sample_count=$sample_count" + --setting "min_time=$min_time" + --setting "max_time=$max_time" +) +for filter in "${filters[@]}"; do + manifest_args+=(--filter "$filter") +done +for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ bench _ <<<"$entry" + baseline_binaries[$label]=$(find_benchmark "$baseline_target" "$bench") + candidate_binaries[$label]=$(find_benchmark "$candidate_target" "$bench") + manifest_args+=( + --suite "$label" + --baseline-binary "$label=${baseline_binaries[$label]}" + --candidate-binary "$label=${candidate_binaries[$label]}" + ) +done +python3 "$parser" "${manifest_args[@]}" + +run_suite() { + local revision=$1 + local label=$2 + local destination=$3 + local binary + local command + + if [[ $revision == baseline ]]; then + binary=${baseline_binaries[$label]} + else + binary=${candidate_binaries[$label]} + fi + command=( + taskset -c "$bench_cpu" "$binary" + --bench --timer tsc --sample-count "$sample_count" + --min-time "$min_time" --max-time "$max_time" --color never + "${filters[@]}" + ) + echo "Running $label ($revision) -> $destination" + "${command[@]}" >"$destination" 2>&1 +} + +echo "Waiting for the global timed benchmark lock: $lock_file" +exec {benchmark_lock}>"$lock_file" +flock "$benchmark_lock" +if pgrep -x cargo >/dev/null || pgrep -x rustc >/dev/null; then + echo "Cargo or rustc is active after acquiring the benchmark lock; refusing to measure." >&2 + exit 1 +fi + +for ((round = 1; round <= warm_runs; round++)); do + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + if ((round % 2 == 1)); then + run_suite baseline "$label" "$output/warm/$label-baseline-$round.txt" + run_suite candidate "$label" "$output/warm/$label-candidate-$round.txt" + else + run_suite candidate "$label" "$output/warm/$label-candidate-$round.txt" + run_suite baseline "$label" "$output/warm/$label-baseline-$round.txt" + fi + done +done + +for ((pair = 1; pair <= measured_pairs; pair++)); do + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + if ((pair % 2 == 1)); then + run_suite baseline "$label" "$output/measured/$label-baseline-$pair.txt" + run_suite candidate "$label" "$output/measured/$label-candidate-$pair.txt" + else + run_suite candidate "$label" "$output/measured/$label-candidate-$pair.txt" + run_suite baseline "$label" "$output/measured/$label-baseline-$pair.txt" + fi + done +done + +python3 "$parser" summarize "$output" +echo "Raw results: $output" +echo "Summary: $output/summary.md" diff --git a/scripts/rowfn_benchmark.py b/scripts/rowfn_benchmark.py new file mode 100755 index 00000000000..3446744812c --- /dev/null +++ b/scripts/rowfn_benchmark.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Capture and summarize evidence from ``benchmark-rowfn.sh`` runs.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import re +import statistics +import subprocess +from collections.abc import Iterable +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from pathlib import Path + +RESULT_FILE = re.compile(r"^(?P.+)-(?Pbaseline|candidate)-(?P\d+)\.txt$") +TREE_ROW = re.compile(r"^(?P(?:│ | )*)(?:├─ |╰─ )(?P.*)$") +TIMING = re.compile(r"(?P\d+(?:\.\d+)?)\s*(?Pps|ns|µs|us|ms|s)\s*$") +UNIT_TO_NS = { + "ps": 0.001, + "ns": 1.0, + "µs": 1_000.0, + "us": 1_000.0, + "ms": 1_000_000.0, + "s": 1_000_000_000.0, +} + + +@dataclass(frozen=True) +class BenchmarkSummary: + suite: str + benchmark: str + pairs: int + baseline_median_ns: float + candidate_median_ns: float + median_ratio: float + minimum_ratio: float + maximum_ratio: float + ratio_mad: float + + +def run_git(worktree: Path, *args: str, binary: bool = False) -> str | bytes: + """Run one read-only Git command in ``worktree``.""" + + result = subprocess.run( + ["git", "-C", str(worktree), *args], + check=True, + capture_output=True, + text=not binary, + ) + return result.stdout if binary else result.stdout.strip() + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + + return digest.hexdigest() + + +def revision_record(worktree: Path, target: Path, binaries: Iterable[str]) -> dict[str, object]: + """Describe the exact revision, dirty patch, targets, and benchmark executables.""" + + status = str(run_git(worktree, "status", "--short")).splitlines() + diff = run_git(worktree, "diff", "--binary", "HEAD", binary=True) + assert isinstance(diff, bytes) + + untracked = run_git(worktree, "ls-files", "--others", "--exclude-standard", "-z", binary=True) + assert isinstance(untracked, bytes) + dirty_digest = hashlib.sha256(diff) + dirty_digest.update(untracked) + for relative_path in filter(None, untracked.decode().split("\0")): + path = worktree / relative_path + if path.is_file(): + dirty_digest.update(relative_path.encode()) + dirty_digest.update(bytes.fromhex(sha256_file(path))) + + executable_records: dict[str, object] = {} + for entry in binaries: + label, separator, raw_path = entry.partition("=") + if not separator: + raise ValueError(f"expected LABEL=PATH for benchmark binary, got {entry!r}") + path = Path(raw_path).resolve() + executable_records[label] = { + "path": str(path), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + + return { + "worktree": str(worktree.resolve()), + "head": run_git(worktree, "rev-parse", "HEAD"), + "changed_paths": status, + "tracked_diff_sha256": hashlib.sha256(diff).hexdigest(), + "dirty_state_sha256": dirty_digest.hexdigest(), + "target": str(target.resolve()), + "binaries": executable_records, + } + + +def write_manifest(args: argparse.Namespace) -> None: + settings = dict(setting.split("=", 1) for setting in args.setting) + manifest = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + "settings": settings, + "suites": args.suite, + "filters": args.filter, + "machine_record": str(Path(args.machine_record).resolve()), + "baseline": revision_record( + Path(args.baseline_worktree), + Path(args.baseline_target), + args.baseline_binary, + ), + "candidate": revision_record( + Path(args.candidate_worktree), + Path(args.candidate_target), + args.candidate_binary, + ), + } + output = Path(args.output) + output.write_text(f"{json.dumps(manifest, indent=2, sort_keys=True)}\n", encoding="utf-8") + + +def timing_ns(field: str) -> float: + match = TIMING.search(field.strip()) + if match is None: + raise ValueError(f"cannot parse Divan timing from {field!r}") + + return float(match.group("value")) * UNIT_TO_NS[match.group("unit")] + + +def parse_divan(path: Path) -> dict[str, float]: + """Return benchmark paths and median nanoseconds from one Divan table.""" + + parents: dict[int, str] = {} + timings: dict[str, float] = {} + for line in path.read_text(encoding="utf-8").splitlines(): + fields = re.split(r"\s+│\s+", line) + tree_match = TREE_ROW.match(fields[0]) + if tree_match is None: + continue + + depth = len(tree_match.group("prefix")) // 3 + body = tree_match.group("body").rstrip() + timing_match = TIMING.search(body) + name = body[: timing_match.start()].rstrip() if timing_match else body.strip() + parents = {level: parent for level, parent in parents.items() if level < depth} + + if timing_match is None: + parents[depth] = name + continue + if len(fields) < 3: + raise ValueError(f"timed Divan row has no median column in {path}: {line}") + + components = [parents[level] for level in sorted(parents) if level < depth] + benchmark = "/".join([*components, name]) + if benchmark in timings: + raise ValueError(f"duplicate benchmark {benchmark!r} in {path}") + timings[benchmark] = timing_ns(fields[2]) + + if not timings: + raise ValueError(f"no Divan benchmark timings found in {path}") + + return timings + + +def read_measurements(directory: Path) -> dict[tuple[str, str, int, str], float]: + measurements: dict[tuple[str, str, int, str], float] = {} + for path in sorted(directory.glob("*.txt")): + match = RESULT_FILE.match(path.name) + if match is None: + continue + suite = match.group("suite") + revision = match.group("revision") + pair = int(match.group("pair")) + for benchmark, median_ns in parse_divan(path).items(): + measurements[suite, revision, pair, benchmark] = median_ns + + if not measurements: + raise ValueError(f"no measured result files found in {directory}") + + return measurements + + +def summarize(measurements: dict[tuple[str, str, int, str], float]) -> list[BenchmarkSummary]: + groups = {(suite, pair, benchmark) for suite, _, pair, benchmark in measurements} + incomplete = [ + group + for group in groups + if (group[0], "baseline", group[1], group[2]) not in measurements + or (group[0], "candidate", group[1], group[2]) not in measurements + ] + if incomplete: + raise ValueError(f"unpaired benchmark measurements: {sorted(incomplete)!r}") + + by_benchmark: dict[tuple[str, str], list[tuple[float, float]]] = {} + for suite, pair, benchmark in sorted(groups): + baseline = measurements[suite, "baseline", pair, benchmark] + candidate = measurements[suite, "candidate", pair, benchmark] + by_benchmark.setdefault((suite, benchmark), []).append((baseline, candidate)) + + summaries = [] + for (suite, benchmark), pairs in sorted(by_benchmark.items()): + baseline_values = [baseline for baseline, _ in pairs] + candidate_values = [candidate for _, candidate in pairs] + ratios = [candidate / baseline for baseline, candidate in pairs] + median_ratio = statistics.median(ratios) + summaries.append( + BenchmarkSummary( + suite=suite, + benchmark=benchmark, + pairs=len(pairs), + baseline_median_ns=statistics.median(baseline_values), + candidate_median_ns=statistics.median(candidate_values), + median_ratio=median_ratio, + minimum_ratio=min(ratios), + maximum_ratio=max(ratios), + ratio_mad=statistics.median(abs(ratio - median_ratio) for ratio in ratios), + ) + ) + + return summaries + + +def format_ns(value: float) -> str: + for divisor, unit in ((1_000_000_000, "s"), (1_000_000, "ms"), (1_000, "µs")): + if value >= divisor: + return f"{value / divisor:.3f} {unit}" + + return f"{value:.3f} ns" + + +def write_summary(output_directory: Path, summaries: list[BenchmarkSummary]) -> None: + csv_path = output_directory / "ratios.csv" + with csv_path.open("w", encoding="utf-8", newline="") as file: + writer = csv.DictWriter(file, fieldnames=list(asdict(summaries[0]))) + writer.writeheader() + writer.writerows(asdict(summary) for summary in summaries) + + markdown = [ + "# RowFn benchmark comparison", + "", + "Ratios are paired candidate/baseline medians. Lower is faster.", + "", + "| Suite | Benchmark | Pairs | Baseline | Candidate | Ratio | Change | MAD |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for summary in sorted(summaries, key=lambda result: result.median_ratio, reverse=True): + change = (summary.median_ratio - 1.0) * 100.0 + markdown.append( + f"| {summary.suite} | `{summary.benchmark}` | {summary.pairs} " + f"| {format_ns(summary.baseline_median_ns)} " + f"| {format_ns(summary.candidate_median_ns)} " + f"| {summary.median_ratio:.6f} | {change:+.2f}% | {summary.ratio_mad:.6f} |" + ) + markdown.append("") + (output_directory / "summary.md").write_text("\n".join(markdown), encoding="utf-8") + + +def summarize_directory(args: argparse.Namespace) -> None: + output_directory = Path(args.output_directory) + summaries = summarize(read_measurements(output_directory / "measured")) + write_summary(output_directory, summaries) + + +def argument_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(required=True) + + manifest = subparsers.add_parser("manifest", help="capture revisions and executable hashes") + manifest.add_argument("--output", required=True) + manifest.add_argument("--machine-record", required=True) + manifest.add_argument("--baseline-worktree", required=True) + manifest.add_argument("--candidate-worktree", required=True) + manifest.add_argument("--baseline-target", required=True) + manifest.add_argument("--candidate-target", required=True) + manifest.add_argument("--setting", action="append", default=[]) + manifest.add_argument("--suite", action="append", default=[]) + manifest.add_argument("--filter", action="append", default=[]) + manifest.add_argument("--baseline-binary", action="append", default=[]) + manifest.add_argument("--candidate-binary", action="append", default=[]) + manifest.set_defaults(function=write_manifest) + + summary = subparsers.add_parser("summarize", help="write ratios.csv and summary.md") + summary.add_argument("output_directory") + summary.set_defaults(function=summarize_directory) + + return parser + + +def main() -> None: + args = argument_parser().parse_args() + args.function(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/tests/test_rowfn_benchmark.py b/scripts/tests/test_rowfn_benchmark.py new file mode 100644 index 00000000000..e224d6229e0 --- /dev/null +++ b/scripts/tests/test_rowfn_benchmark.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SCRIPT = REPO_ROOT / "scripts" / "rowfn_benchmark.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("rowfn_benchmark", SCRIPT) + assert spec is not None + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def write_divan(path: Path, rows: list[str]) -> None: + path.write_text( + "\n".join( + [ + "Timer precision: 20 ns", + "bench fastest │ slowest │ median │ mean │ samples │ iters", + *rows, + "", + ] + ), + encoding="utf-8", + ) + + +class RowFnBenchmarkTest(unittest.TestCase): + def setUp(self) -> None: + self.module = load_module() + self.temporary_directory = tempfile.TemporaryDirectory() + self.directory = Path(self.temporary_directory.name) + + def tearDown(self) -> None: + self.temporary_directory.cleanup() + + def test_parse_divan_preserves_nested_benchmark_names_and_converts_units(self) -> None: + output = self.directory / "result.txt" + write_divan( + output, + [ + "├─ non_nullable │ │ │ │ │", + "│ ├─ 2 17.18 µs │ 18 µs │ 17.33 µs │ 17.4 µs │ 100 │ 100", + "│ ╰─ 32 6.709 µs │ 8 µs │ 6.829 µs │ 7 µs │ 100 │ 100", + "╰─ nullable │ │ │ │ │", + " ╰─ 2 799.7 ns │ 1 µs │ 979.7 ns │ 986 ns │ 100 │ 100", + ], + ) + + self.assertEqual( + self.module.parse_divan(output), + { + "non_nullable/2": 17_330.0, + "non_nullable/32": 6_829.0, + "nullable/2": 979.7, + }, + ) + + def test_summarize_writes_paired_ratios_and_slowest_first_markdown(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + results = { + "numeric-baseline-1.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-1.txt": ["├─ add 12 ns │ 12 ns │ 12 ns │ 12 ns │ 100 │ 100"], + "numeric-baseline-2.txt": ["├─ add 20 ns │ 20 ns │ 20 ns │ 20 ns │ 100 │ 100"], + "numeric-candidate-2.txt": ["├─ add 18 ns │ 18 ns │ 18 ns │ 18 ns │ 100 │ 100"], + "numeric-baseline-3.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-3.txt": ["├─ add 11 ns │ 11 ns │ 11 ns │ 11 ns │ 100 │ 100"], + "numeric-baseline-4.txt": ["├─ mul 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-4.txt": ["├─ mul 9 ns │ 9 ns │ 9 ns │ 9 ns │ 100 │ 100"], + } + for name, rows in results.items(): + write_divan(measured / name, rows) + + summaries = self.module.summarize(self.module.read_measurements(measured)) + self.module.write_summary(self.directory, summaries) + + add = next(summary for summary in summaries if summary.benchmark == "add") + self.assertEqual(add.pairs, 3) + self.assertAlmostEqual(add.median_ratio, 1.1) + self.assertAlmostEqual(add.ratio_mad, 0.1) + + csv_output = (self.directory / "ratios.csv").read_text(encoding="utf-8") + markdown = (self.directory / "summary.md").read_text(encoding="utf-8") + self.assertIn("suite,benchmark,pairs", csv_output) + self.assertLess(markdown.index("`add`"), markdown.index("`mul`")) + + def test_summarize_rejects_unpaired_measurements(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + write_divan( + measured / "numeric-baseline-1.txt", + ["╰─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + ) + + with self.assertRaisesRegex(ValueError, "unpaired benchmark measurements"): + self.module.summarize(self.module.read_measurements(measured)) + + +if __name__ == "__main__": + unittest.main() From b1dcb838bf505edcc05b60e7528427432049e108 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:41:06 -0400 Subject: [PATCH 064/160] Handle changing benchmark inventories Signed-off-by: Connor Tsui --- scripts/rowfn_benchmark.py | 58 +++++++++++++++++++++++++-- scripts/tests/test_rowfn_benchmark.py | 23 +++++++++++ 2 files changed, 77 insertions(+), 4 deletions(-) diff --git a/scripts/rowfn_benchmark.py b/scripts/rowfn_benchmark.py index 3446744812c..b2f778ea3f7 100755 --- a/scripts/rowfn_benchmark.py +++ b/scripts/rowfn_benchmark.py @@ -192,7 +192,19 @@ def read_measurements(directory: Path) -> dict[tuple[str, str, int, str], float] def summarize(measurements: dict[tuple[str, str, int, str], float]) -> list[BenchmarkSummary]: - groups = {(suite, pair, benchmark) for suite, _, pair, benchmark in measurements} + inventories: dict[tuple[str, str], set[str]] = {} + for suite, revision, _, benchmark in measurements: + inventories.setdefault((suite, revision), set()).add(benchmark) + + suites = {suite for suite, _ in inventories} + comparable = { + (suite, benchmark) + for suite in suites + for benchmark in inventories.get((suite, "baseline"), set()) & inventories.get((suite, "candidate"), set()) + } + groups = { + (suite, pair, benchmark) for suite, _, pair, benchmark in measurements if (suite, benchmark) in comparable + } incomplete = [ group for group in groups @@ -201,6 +213,8 @@ def summarize(measurements: dict[tuple[str, str, int, str], float]) -> list[Benc ] if incomplete: raise ValueError(f"unpaired benchmark measurements: {sorted(incomplete)!r}") + if not groups: + raise ValueError("unpaired benchmark measurements: no comparable benchmarks") by_benchmark: dict[tuple[str, str], list[tuple[float, float]]] = {} for suite, pair, benchmark in sorted(groups): @@ -231,6 +245,25 @@ def summarize(measurements: dict[tuple[str, str, int, str], float]) -> list[Benc return summaries +def inventory_differences( + measurements: dict[tuple[str, str, int, str], float], +) -> list[tuple[str, str, str]]: + """Return benchmarks that exist in only one revision.""" + + inventories: dict[tuple[str, str], set[str]] = {} + for suite, revision, _, benchmark in measurements: + inventories.setdefault((suite, revision), set()).add(benchmark) + + differences = [] + for suite in sorted({suite for suite, _ in inventories}): + baseline = inventories.get((suite, "baseline"), set()) + candidate = inventories.get((suite, "candidate"), set()) + differences.extend((suite, "baseline only", benchmark) for benchmark in baseline - candidate) + differences.extend((suite, "candidate only", benchmark) for benchmark in candidate - baseline) + + return sorted(differences) + + def format_ns(value: float) -> str: for divisor, unit in ((1_000_000_000, "s"), (1_000_000, "ms"), (1_000, "µs")): if value >= divisor: @@ -239,7 +272,11 @@ def format_ns(value: float) -> str: return f"{value:.3f} ns" -def write_summary(output_directory: Path, summaries: list[BenchmarkSummary]) -> None: +def write_summary( + output_directory: Path, + summaries: list[BenchmarkSummary], + differences: Iterable[tuple[str, str, str]] = (), +) -> None: csv_path = output_directory / "ratios.csv" with csv_path.open("w", encoding="utf-8", newline="") as file: writer = csv.DictWriter(file, fieldnames=list(asdict(summaries[0]))) @@ -262,14 +299,27 @@ def write_summary(output_directory: Path, summaries: list[BenchmarkSummary]) -> f"| {format_ns(summary.candidate_median_ns)} " f"| {summary.median_ratio:.6f} | {change:+.2f}% | {summary.ratio_mad:.6f} |" ) + differences = list(differences) + if differences: + markdown.extend( + [ + "", + "## Unpaired benchmark inventory", + "", + "These benchmarks were recorded for only one revision and are excluded from ratios.", + "", + ] + ) + markdown.extend(f"- `{suite}/{benchmark}`: {revision}." for suite, revision, benchmark in differences) markdown.append("") (output_directory / "summary.md").write_text("\n".join(markdown), encoding="utf-8") def summarize_directory(args: argparse.Namespace) -> None: output_directory = Path(args.output_directory) - summaries = summarize(read_measurements(output_directory / "measured")) - write_summary(output_directory, summaries) + measurements = read_measurements(output_directory / "measured") + summaries = summarize(measurements) + write_summary(output_directory, summaries, inventory_differences(measurements)) def argument_parser() -> argparse.ArgumentParser: diff --git a/scripts/tests/test_rowfn_benchmark.py b/scripts/tests/test_rowfn_benchmark.py index e224d6229e0..308580f8d68 100644 --- a/scripts/tests/test_rowfn_benchmark.py +++ b/scripts/tests/test_rowfn_benchmark.py @@ -106,6 +106,29 @@ def test_summarize_rejects_unpaired_measurements(self) -> None: with self.assertRaisesRegex(ValueError, "unpaired benchmark measurements"): self.module.summarize(self.module.read_measurements(measured)) + def test_summarize_excludes_and_reports_revision_only_benchmarks(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + results = { + "numeric-baseline-1.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-1.txt": [ + "├─ add 11 ns │ 11 ns │ 11 ns │ 11 ns │ 100 │ 100", + "╰─ candidate 5 ns │ 5 ns │ 5 ns │ 5 ns │ 100 │ 100", + ], + } + for name, rows in results.items(): + write_divan(measured / name, rows) + + measurements = self.module.read_measurements(measured) + summaries = self.module.summarize(measurements) + differences = self.module.inventory_differences(measurements) + self.module.write_summary(self.directory, summaries, differences) + + self.assertEqual([summary.benchmark for summary in summaries], ["add"]) + self.assertEqual(differences, [("numeric", "candidate only", "candidate")]) + markdown = (self.directory / "summary.md").read_text(encoding="utf-8") + self.assertIn("`numeric/candidate`: candidate only.", markdown) + if __name__ == "__main__": unittest.main() From dfd1d4cd4869dc57b65e37ea1a379196cd365ee8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:12:13 -0400 Subject: [PATCH 065/160] Cover invalid encoded RowFn reductions Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/row/batch/tests.rs | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 06539114638..8f7d0dd4282 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -50,6 +50,9 @@ struct NullarySeven; #[derive(Clone)] struct OriginalInputReducer; +#[derive(Clone)] +struct InvalidEncodedReduction; + #[derive(Clone)] struct DeferredOriginalReducer; @@ -258,6 +261,37 @@ impl RowFn for OriginalInputReducer { } } +impl RowFn for InvalidEncodedReduction { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_encoded_reduction"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::Output( + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([false, true])).into_array(), + ))) + } +} + impl RowFn for DeferredOriginalReducer { type Options = EmptyOptions; @@ -438,6 +472,31 @@ fn test_reduce_encoded_defers_errors_behind_nulls() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::all_valid(Validity::AllValid)] +#[case::mixed(Validity::from_iter([true, false]))] +fn test_reduce_encoded_rejects_nulls_on_valid_rows(#[case] validity: Validity) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![10_i64, 20], validity).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&InvalidEncodedReduction, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an encoded reduction introduced a null on a valid row"), + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_encoded_reduction"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("encoded reduction produced nulls for valid rows"), + "the boundary error must identify invalid reduced output, got {error}", + ); + Ok(()) +} + #[test] fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { let input = ConstantArray::new(7_i64, 3).into_array(); From eb4e035a5cd4babeb0739c9d8f4db67c83f31acf Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:14:31 -0400 Subject: [PATCH 066/160] Validate encoded RowFn output validity Signed-off-by: Connor Tsui --- .../src/scalar_fn/row/batch/execution.rs | 28 +++++++++++++------ vortex-array/src/scalar_fn/row/batch/tests.rs | 25 ++++++++++++++--- 2 files changed, 40 insertions(+), 13 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index 83df651d4f4..c03956b3682 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -150,7 +150,7 @@ impl Batch { // broadcast and sees the original inputs before slicing or filtering changes them. if let Some(execution) = reduce(self.kernel_args(&self.inputs, self.row_count), ctx)? { match execution { - RowExecution::Output(values) => return self.finalize_reduced(values), + RowExecution::Output(values) => return self.finalize_reduced(values, ctx), RowExecution::DeferredError(error) => { return self.resolve_reduced_error(error, kernel, try_unfiltered, ctx); } @@ -347,15 +347,25 @@ impl Batch { } /// Reconcile an encoding-aware result and apply the batch's strict input validity. - fn finalize_reduced(&self, values: ArrayRef) -> VortexResult { - match self.validity.clone() { - Validity::NonNullable | Validity::AllValid => { - self.finalize_output(values, self.row_count) - } - Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + fn finalize_reduced(&self, values: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + validate_output(self.id, &self.result_dtype, self.row_count, &values)?; + + let input_valid = self.validity.execute_mask(self.row_count, ctx)?; + let output_valid = values.validity()?.execute_mask(self.row_count, ctx)?; + vortex_ensure!( + input_valid.bitand_not(&output_valid).all_false(), + "the {} encoded reduction produced nulls for valid rows", + self.id, + ); + + let values = match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => values, + Validity::Array(valid) => values.mask(valid)?, // Handled before the encoding-aware hook runs. - Validity::AllInvalid => Ok(self.all_null()), - } + Validity::AllInvalid => return Ok(self.all_null()), + }; + + cast_output_nullability(&self.result_dtype, values) } /// Resolve deferred evidence from the encoded path by executing only observable rows. diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 8f7d0dd4282..0c213491a39 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -262,7 +262,7 @@ impl RowFn for OriginalInputReducer { } impl RowFn for InvalidEncodedReduction { - type Options = EmptyOptions; + type Options = usize; const ARG_NAMES: &'static [&'static str] = &["value"]; @@ -282,12 +282,16 @@ impl RowFn for InvalidEncodedReduction { fn reduce_encoded( &self, - _options: &Self::Options, + null_index: &Self::Options, _args: &[ArrayRef], _ctx: &mut ExecutionCtx, ) -> VortexResult> { Ok(Some(RowExecution::Output( - PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([false, true])).into_array(), + PrimitiveArray::new( + vec![10_i64, 20], + Validity::from_iter((0..2).map(|index| index != *null_index)), + ) + .into_array(), ))) } } @@ -480,7 +484,7 @@ fn test_reduce_encoded_rejects_nulls_on_valid_rows(#[case] validity: Validity) - let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = array_session().create_execution_ctx(); - let error = match execute_rows(&InvalidEncodedReduction, &EmptyOptions, &args, &mut ctx) { + let error = match execute_rows(&InvalidEncodedReduction, &0, &args, &mut ctx) { Err(error) => error, Ok(_) => vortex_bail!("an encoded reduction introduced a null on a valid row"), }; @@ -497,6 +501,19 @@ fn test_reduce_encoded_rejects_nulls_on_valid_rows(#[case] validity: Validity) - Ok(()) } +#[test] +fn test_reduce_encoded_preserves_input_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&InvalidEncodedReduction, &1, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + #[test] fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { let input = ConstantArray::new(7_i64, 3).into_array(); From 90a97b36bc15b60464e255a3a314933a33edeb3e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:17:49 -0400 Subject: [PATCH 067/160] Clarify explicit RowFn adoption Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/row/row_fn.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index c84c07b62e2..ccdc9d6c70c 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -29,7 +29,9 @@ use crate::scalar_fn::ScalarFnId; /// [`row_fn_return_dtype`](crate::scalar_fn::row_fn_return_dtype), and delegates `execute` to /// [`execute_rows`](crate::scalar_fn::execute_rows). Explicit adoption leaves the function free to /// provide custom coercion, simplification, reduction, or formatting hooks. Implement only -/// `ScalarFnVTable` when the natural kernel is columnar rather than row-oriented. +/// `ScalarFnVTable` when the natural kernel is columnar rather than row-oriented. A serializable +/// adopter also delegates the vtable's `serialize` and `deserialize` methods to the matching +/// methods below; an internal or non-serializable adopter can keep the vtable defaults. pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// Options for this function, if any. Use [`EmptyOptions`](crate::scalar_fn::EmptyOptions) /// for none. From 878d22f234b54d3947bf1d4f25090b0203d7bb53 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:48:58 -0400 Subject: [PATCH 068/160] Tighten checked arithmetic contracts Signed-off-by: Connor Tsui --- .../src/scalar_fn/fns/binary/numeric/primitive.rs | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 42fe3fd3e03..f07ed8e9d5f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -23,7 +23,11 @@ pub(super) struct CheckedDiv; /// OR-reducible evidence that a row failed, with [`Default`] meaning success. pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} -impl Failure for T {} +impl Failure for bool {} +impl Failure for u8 {} +impl Failure for u16 {} +impl Failure for u32 {} +impl Failure for u64 {} /// One arithmetic operator at one width, split into its value and failure evidence. pub(super) trait CheckedPrimitiveOp: 'static + Sized { @@ -87,7 +91,10 @@ impl CheckedPrimitiveOp for CheckedDiv { } } -/// Per-width checked arithmetic. Every value method **must** be total over stored lane values. +/// Per-width arithmetic used to compute values and failure evidence. +/// +/// The add, subtract, and multiply value methods **must** be total over every stored lane value. +/// [`Self::div_value`] may assume that [`Self::div_error`] returned `false` for the same operands. pub(super) trait CheckedArithmetic: NativePType { /// How multiplication reports a failing row. /// @@ -100,7 +107,11 @@ pub(super) trait CheckedArithmetic: NativePType { fn sub_error(self, rhs: Self) -> bool; fn mul_value(self, rhs: Self) -> Self; fn mul_failure(self, rhs: Self) -> Self::MulFailure; + + /// Divide operands that [`Self::div_error`] accepted. fn div_value(self, rhs: Self) -> Self; + + /// Return whether [`Self::div_value`] would trap for these operands. fn div_error(self, rhs: Self) -> bool; } From 45aa5a35eb84a8109cdee4c37cf75f989dec3aca Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:23:06 -0400 Subject: [PATCH 069/160] Preserve containment with non-finite bounds Signed-off-by: Connor Tsui --- vortex-spatial/src/scalar_fn/contains.rs | 38 +++++++++++++++++++++--- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 514b76593b7..c0bee5ddcee 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -128,7 +128,7 @@ impl PreparedOperand { fn new(geometry: &Geometry) -> Self { Self { geometry: geometry.clone(), - bbox: geometry.bounding_rect(), + bbox: finite_bounding_rect(geometry), prepared: OnceCell::new(), } } @@ -140,6 +140,22 @@ impl PreparedOperand { } } +/// Returns a bounding rectangle only when ordered comparisons can conservatively reject a row. +/// +/// Geo permits non-finite coordinates. A rectangle containing NaN cannot prove non-containment, +/// because its ordered comparisons can return false even when the exact algorithm accepts the +/// geometry. +fn finite_bounding_rect(geometry: &Geometry) -> Option> { + let bbox = geometry.bounding_rect()?; + let min = bbox.min(); + let max = bbox.max(); + + [min.x, min.y, max.x, max.y] + .into_iter() + .all(f64::is_finite) + .then_some(bbox) +} + /// How geo's `a.contains(b)` computes its verdict for a pairing. enum ContainsRoute { /// `a.relate(b).is_contains()`. @@ -298,10 +314,9 @@ fn contains_row_prepared(operands: &ConstOperands, a: &Geometry, b: &Geomet .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), (Some(const_a), None) => const_a .bbox - .zip(b.bounding_rect()) + .zip(finite_bounding_rect(b)) .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), - (None, Some(const_b)) => a - .bounding_rect() + (None, Some(const_b)) => finite_bounding_rect(a) .zip(const_b.bbox) .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), }; @@ -428,6 +443,21 @@ mod tests { assert_contains(container, other, [expected; 3]) } + /// A non-finite bounding rectangle cannot reject a containment that the exact geometry + /// algorithm accepts. + #[test] + fn nan_bounding_rect_does_not_reject_containment() { + let container = multipoint(vec![(f64::NAN, f64::NAN), (1.0, 1.0)]); + let contained = point(1.0, 1.0); + let operands = ConstOperands { + a: Some(PreparedOperand::new(&container)), + b: Some(PreparedOperand::new(&contained)), + }; + + assert!(container.contains(&contained)); + assert!(contains_row_prepared(&operands, &container, &contained)); + } + /// Partially overlapping polygons contain each other in neither direction. #[test] fn overlapping_polygons_contain_neither_way() -> VortexResult<()> { From 54f8efd13655f65e0644352bcdd21443b7399306 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 22:55:49 -0400 Subject: [PATCH 070/160] Share the RowFn vtable adapter Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/mod.rs | 7 ++ vortex-array/src/scalar_fn/row/vtable.rs | 91 ++++++++++++++++++- vortex-spatial/src/scalar_fn/contains.rs | 2 +- vortex-spatial/src/scalar_fn/distance.rs | 2 +- vortex-spatial/src/scalar_fn/intersects.rs | 2 +- vortex-spatial/src/scalar_fn/mod.rs | 77 ---------------- .../src/scalar_fns/cosine_similarity.rs | 2 +- vortex-tensor/src/scalar_fns/inner_product.rs | 2 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 2 +- vortex-tensor/src/scalar_fns/mod.rs | 77 ---------------- vortex-tensor/src/scalar_fns/tests/row.rs | 2 +- 11 files changed, 103 insertions(+), 163 deletions(-) diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 5e73caefdfa..4d20e73c897 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -19,6 +19,13 @@ use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::literal::Literal; +/// Reexports used by [`impl_row_fn_vtable!`] without downstream transitive dependencies. +#[doc(hidden)] +pub mod row_fn_macro_support { + pub use vortex_error::VortexResult; + pub use vortex_session::VortexSession; +} + mod vtable; pub use vtable::*; diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index 96a24f8f42d..f0fdcb66186 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -24,6 +24,93 @@ use crate::scalar_fn::row::visitor::ExecuteRows; use crate::scalar_fn::row::visitor::ExecuteValidRows; use crate::scalar_fn::row::visitor::PlanRows; +/// Implement the standard [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) behavior for one +/// concrete [`RowFn`]. +/// +/// This opt-in is explicit so a function that needs custom coercion, simplification, reduction, +/// formatting, or validity hooks can implement the vtable itself and delegate only execution to +/// [`execute_rows`](crate::scalar_fn::execute_rows). The standard adapter delegates serialization +/// to [`RowFn`], derives arity and child names from [`RowFn::ARG_NAMES`], propagates child validity, +/// and reports the function as strict. +#[macro_export] +macro_rules! impl_row_fn_vtable { + ($function:ty) => { + impl $crate::scalar_fn::ScalarFnVTable for $function { + type Options = <$function as $crate::scalar_fn::RowFn>::Options; + + fn id(&self) -> $crate::scalar_fn::ScalarFnId { + $crate::scalar_fn::RowFn::id(self) + } + + fn serialize( + &self, + options: &Self::Options, + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult>> { + $crate::scalar_fn::RowFn::serialize(self, options) + } + + fn deserialize( + &self, + metadata: &[u8], + session: &$crate::scalar_fn::row_fn_macro_support::VortexSession, + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult { + $crate::scalar_fn::RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> $crate::scalar_fn::Arity { + $crate::scalar_fn::Arity::Exact( + <$function as $crate::scalar_fn::RowFn>::ARG_NAMES.len(), + ) + } + + fn child_name( + &self, + _options: &Self::Options, + child_index: usize, + ) -> $crate::scalar_fn::ChildName { + $crate::scalar_fn::ChildName::from( + <$function as $crate::scalar_fn::RowFn>::ARG_NAMES[child_index], + ) + } + + fn return_dtype( + &self, + options: &Self::Options, + args: &[$crate::dtype::DType], + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult<$crate::dtype::DType> { + $crate::scalar_fn::row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn $crate::scalar_fn::ExecutionArgs, + ctx: &mut $crate::ExecutionCtx, + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult<$crate::ArrayRef> { + $crate::scalar_fn::execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &$crate::expr::Expression, + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult< + Option<$crate::expr::Expression>, + > { + $crate::expr::union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + <$function as $crate::scalar_fn::RowFn>::FALLIBLE + } + } + }; +} + /// Compute the return dtype for a [`RowFn`] without adopting its complete scalar-function vtable. pub fn row_fn_return_dtype( function: &F, @@ -38,8 +125,8 @@ pub fn row_fn_return_dtype( /// Execute a [`RowFn`] while preserving a caller-owned scalar-function vtable. /// /// Existing vtables delegate here when they need row execution but retain custom hooks for other -/// capabilities. A RowFn-only function implements the remaining vtable methods mechanically and -/// delegates its return-dtype and execution methods here. +/// capabilities. A function that needs only the standard hooks can use [`impl_row_fn_vtable`] to +/// generate its complete vtable. pub fn execute_rows( function: &F, options: &F::Options, diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index c0bee5ddcee..52d01df4014 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -93,7 +93,7 @@ impl RowFn for SpatialContains { } } -impl_row_fn_scalar_vtable!(SpatialContains); +vortex_array::impl_row_fn_vtable!(SpatialContains); /// Per-batch state for the contains row kernel: the prepared form of whichever operand is /// constant for the batch. `None` marks an operand that varies by row. diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index 3b6b086b3a0..751e402f118 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -75,7 +75,7 @@ impl RowFn for SpatialDistance { } } -impl_row_fn_scalar_vtable!(SpatialDistance); +vortex_array::impl_row_fn_vtable!(SpatialDistance); #[cfg(test)] mod tests { diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index 37d5a72ca61..6f04e86b14a 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -84,7 +84,7 @@ impl RowFn for SpatialIntersects { } } -impl_row_fn_scalar_vtable!(SpatialIntersects); +vortex_array::impl_row_fn_vtable!(SpatialIntersects); /// Per-batch state for the intersects row kernel: the bounding rect of each operand that is /// constant for the batch. diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index 3d02d9fbd56..bcdb15e51e6 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -3,83 +3,6 @@ //! Geometry scalar functions over the native geometry extension types. -/// Adopt the standard scalar-function behavior for a row function defined in this module. -macro_rules! impl_row_fn_scalar_vtable { - ($function:ty) => { - impl vortex_array::scalar_fn::ScalarFnVTable for $function { - type Options = <$function as vortex_array::scalar_fn::RowFn>::Options; - - fn id(&self) -> vortex_array::scalar_fn::ScalarFnId { - vortex_array::scalar_fn::RowFn::id(self) - } - - fn serialize( - &self, - options: &Self::Options, - ) -> vortex_error::VortexResult>> { - vortex_array::scalar_fn::RowFn::serialize(self, options) - } - - fn deserialize( - &self, - metadata: &[u8], - session: &vortex_session::VortexSession, - ) -> vortex_error::VortexResult { - vortex_array::scalar_fn::RowFn::deserialize(self, metadata, session) - } - - fn arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::Arity { - vortex_array::scalar_fn::Arity::Exact( - <$function as vortex_array::scalar_fn::RowFn>::ARG_NAMES.len(), - ) - } - - fn child_name( - &self, - _options: &Self::Options, - child_index: usize, - ) -> vortex_array::scalar_fn::ChildName { - vortex_array::scalar_fn::ChildName::from( - <$function as vortex_array::scalar_fn::RowFn>::ARG_NAMES[child_index], - ) - } - - fn return_dtype( - &self, - options: &Self::Options, - args: &[vortex_array::dtype::DType], - ) -> vortex_error::VortexResult { - vortex_array::scalar_fn::row_fn_return_dtype(self, options, args) - } - - fn execute( - &self, - options: &Self::Options, - args: &dyn vortex_array::scalar_fn::ExecutionArgs, - ctx: &mut vortex_array::ExecutionCtx, - ) -> vortex_error::VortexResult { - vortex_array::scalar_fn::execute_rows(self, options, args, ctx) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &vortex_array::expr::Expression, - ) -> vortex_error::VortexResult> { - vortex_array::expr::union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - <$function as vortex_array::scalar_fn::RowFn>::FALLIBLE - } - } - }; -} - pub mod contains; pub mod distance; pub mod envelope; diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 3b08c6d46fc..91e4e837d64 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -140,7 +140,7 @@ impl RowFn for CosineSimilarity { } } -impl_row_fn_scalar_vtable!(CosineSimilarity); +vortex_array::impl_row_fn_vtable!(CosineSimilarity); impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 394cd3b158a..c0968d30474 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -125,7 +125,7 @@ impl RowFn for InnerProduct { } } -impl_row_fn_scalar_vtable!(InnerProduct); +vortex_array::impl_row_fn_vtable!(InnerProduct); impl ScalarFnArrayVTable for InnerProduct { fn serialize( diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index d0a9b1f951b..3bad1d6d1d9 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -122,7 +122,7 @@ pub(super) struct L2NormMetadata { input_dtype: Option, } -impl_row_fn_scalar_vtable!(L2Norm); +vortex_array::impl_row_fn_vtable!(L2Norm); impl ScalarFnArrayVTable for L2Norm { fn serialize( diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index fd343f2c8fa..da9b8950e7a 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -3,83 +3,6 @@ //! Scalar function expressions defined on tensor and tensor-like extension types. -/// Adopt the standard scalar-function behavior for a row function defined in this module. -macro_rules! impl_row_fn_scalar_vtable { - ($function:ty) => { - impl vortex_array::scalar_fn::ScalarFnVTable for $function { - type Options = <$function as vortex_array::scalar_fn::RowFn>::Options; - - fn id(&self) -> vortex_array::scalar_fn::ScalarFnId { - vortex_array::scalar_fn::RowFn::id(self) - } - - fn serialize( - &self, - options: &Self::Options, - ) -> vortex_error::VortexResult>> { - vortex_array::scalar_fn::RowFn::serialize(self, options) - } - - fn deserialize( - &self, - metadata: &[u8], - session: &vortex_session::VortexSession, - ) -> vortex_error::VortexResult { - vortex_array::scalar_fn::RowFn::deserialize(self, metadata, session) - } - - fn arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::Arity { - vortex_array::scalar_fn::Arity::Exact( - <$function as vortex_array::scalar_fn::RowFn>::ARG_NAMES.len(), - ) - } - - fn child_name( - &self, - _options: &Self::Options, - child_index: usize, - ) -> vortex_array::scalar_fn::ChildName { - vortex_array::scalar_fn::ChildName::from( - <$function as vortex_array::scalar_fn::RowFn>::ARG_NAMES[child_index], - ) - } - - fn return_dtype( - &self, - options: &Self::Options, - args: &[vortex_array::dtype::DType], - ) -> vortex_error::VortexResult { - vortex_array::scalar_fn::row_fn_return_dtype(self, options, args) - } - - fn execute( - &self, - options: &Self::Options, - args: &dyn vortex_array::scalar_fn::ExecutionArgs, - ctx: &mut vortex_array::ExecutionCtx, - ) -> vortex_error::VortexResult { - vortex_array::scalar_fn::execute_rows(self, options, args, ctx) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &vortex_array::expr::Expression, - ) -> vortex_error::VortexResult> { - vortex_array::expr::union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - <$function as vortex_array::scalar_fn::RowFn>::FALLIBLE - } - } - }; -} - pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index 6313b0f8d68..77248329ae5 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -57,7 +57,7 @@ impl RowFn for L1Norm { } } -impl_row_fn_scalar_vtable!(L1Norm); +vortex_array::impl_row_fn_vtable!(L1Norm); fn l1_norm_row(row: &[T]) -> T { row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) From e7c8977ee5c8a985e660648aecc79b15f4b26ccc Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 22:56:14 -0400 Subject: [PATCH 071/160] Document primitive comparison paths Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs | 3 +++ .../src/scalar_fn/fns/binary/compare/primitive/columnar.rs | 6 ++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 53ab5448585..35fb637d00b 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -95,6 +95,9 @@ impl RowFn for PrimitiveCompare { const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; fn id(&self) -> ScalarFnId { + // `PrimitiveCompare` is a private implementation detail of `Binary`: it is never registered + // or serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. ScalarFnVTable::id(&Binary) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs index ccd7ffb8719..6728437e6a8 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Fused comparison and bit-packing for wide x86 lanes. +//! Fused comparison and bit-packing for wide primitive lanes. +//! +//! Production uses this implementation only for measured x86 paths. Keeping it portable lets the +//! semantic tests exercise the RowFn and fused paths on every target. use vortex_buffer::BitBuffer; use vortex_error::VortexResult; @@ -85,7 +88,6 @@ fn compare_primitive_typed( Ok(BoolArray::try_new(bits, validity)?.into_array()) } -#[inline(always)] fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { match op { CompareOperator::Eq => lhs.is_eq(rhs), From b1ec2f37eb74c8bf7d6d4ae970ba76c6b6e864c8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 22:56:34 -0400 Subject: [PATCH 072/160] Validate tensor row storage at runtime Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/row.rs | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 68f7f44d1c3..877804535ee 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -19,6 +19,7 @@ use vortex_array::dtype::PType; use vortex_array::scalar_fn::InputElement; use vortex_buffer::Buffer; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; use crate::utils::extract_flat_elements; @@ -95,11 +96,28 @@ unsafe impl InputElement for TensorRow { let stride = flat.row_stride(); let elements = flat.into_buffer::(); - debug_assert!(if stride == 0 { - elements.len() == list_size + let expected_elements = if stride == 0 { + list_size } else { - stride == list_size && rows.checked_mul(stride) == Some(elements.len()) - }); + vortex_ensure_eq!( + stride, + list_size, + "varying tensor row stride must equal its width, got {stride}", + ); + let Some(expected_elements) = rows.checked_mul(stride) else { + vortex_bail!( + "tensor row storage length must fit usize, got {rows} rows of width {stride}", + ); + }; + + expected_elements + }; + vortex_ensure_eq!( + elements.len(), + expected_elements, + "tensor row storage must contain {expected_elements} elements, got {}", + elements.len(), + ); Ok(TensorRows { elements, From a50f375a6f54a8a341458c7dc1aa9cae90a5d953 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 22:59:02 -0400 Subject: [PATCH 073/160] Use the shared RowFn vtable in benchmarks Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 62 ++----------------------- vortex-array/benches/strict_validity.rs | 45 +----------------- 2 files changed, 4 insertions(+), 103 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index 77baf03ea83..bff6a29d2da 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -38,62 +38,6 @@ const ROWS: usize = 65_536; static SESSION: LazyLock = LazyLock::new(array_session); -/// Adopt the standard scalar-function behavior for a row function in this benchmark. -macro_rules! impl_row_fn_scalar_vtable { - ($function:ty) => { - impl ScalarFnVTable for $function { - type Options = <$function as RowFn>::Options; - - fn id(&self) -> ScalarFnId { - RowFn::id(self) - } - - fn arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::Arity { - vortex_array::scalar_fn::Arity::Exact(<$function as RowFn>::ARG_NAMES.len()) - } - - fn child_name( - &self, - _options: &Self::Options, - child_index: usize, - ) -> vortex_array::scalar_fn::ChildName { - vortex_array::scalar_fn::ChildName::from( - <$function as RowFn>::ARG_NAMES[child_index], - ) - } - - fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { - vortex_array::scalar_fn::row_fn_return_dtype(self, options, args) - } - - fn execute( - &self, - options: &Self::Options, - args: &dyn vortex_array::scalar_fn::ExecutionArgs, - ctx: &mut vortex_array::ExecutionCtx, - ) -> VortexResult { - vortex_array::scalar_fn::execute_rows(self, options, args, ctx) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &vortex_array::expr::Expression, - ) -> VortexResult> { - vortex_array::expr::union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - <$function as RowFn>::FALLIBLE - } - } - }; -} - fn main() { LazyLock::force(&SESSION); divan::main(); @@ -224,9 +168,9 @@ impl RowFn for RowSinkWrappingAdd { } } -impl_row_fn_scalar_vtable!(RowWrappingAdd); -impl_row_fn_scalar_vtable!(RowCheckedAdd); -impl_row_fn_scalar_vtable!(RowSinkWrappingAdd); +vortex_array::impl_row_fn_vtable!(RowWrappingAdd); +vortex_array::impl_row_fn_vtable!(RowCheckedAdd); +vortex_array::impl_row_fn_vtable!(RowSinkWrappingAdd); fn inputs() -> (ArrayRef, ArrayRef) { let lhs = (0..ROWS) diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs index e765889cbd4..e0a8b6a565e 100644 --- a/vortex-array/benches/strict_validity.rs +++ b/vortex-array/benches/strict_validity.rs @@ -106,50 +106,7 @@ impl RowFn for LazyDouble { } } -impl ScalarFnVTable for LazyDouble { - type Options = EmptyOptions; - - fn id(&self) -> ScalarFnId { - RowFn::id(self) - } - - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(Self::ARG_NAMES.len()) - } - - fn child_name(&self, _options: &Self::Options, child_index: usize) -> ChildName { - ChildName::from(Self::ARG_NAMES[child_index]) - } - - fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { - vortex_array::scalar_fn::row_fn_return_dtype(self, options, args) - } - - fn execute( - &self, - options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - vortex_array::scalar_fn::execute_rows(self, options, args, ctx) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - Self::FALLIBLE - } -} +vortex_array::impl_row_fn_vtable!(LazyDouble); /// The same function, applying validity the way the adapter used to: materialize a mask first. #[derive(Clone)] From 2d3cf47208180454f7a020a6564b92e02fa532a6 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:34:03 -0400 Subject: [PATCH 074/160] Avoid allocating numeric execution arguments Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/numeric/row.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index bc38dd0e5eb..f0efaca929d 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -23,11 +23,11 @@ use crate::dtype::NativePType; use crate::dtype::PType; use crate::match_each_native_ptype; use crate::scalar::NumericOperator; +use crate::scalar_fn::BorrowedExecutionArgs; use crate::scalar_fn::RowFn; use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::execute_rows; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::row::InitializedElement; @@ -39,7 +39,8 @@ pub(super) fn execute_numeric_primitive( op: NumericOperator, ctx: &mut ExecutionCtx, ) -> VortexResult { - let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + let inputs = [lhs.clone(), rhs.clone()]; + let args = BorrowedExecutionArgs::new(&inputs, lhs.len()); execute_rows(&NumericBinary, &op, &args, ctx) } From ee0c12b2f97564b9ac72948f93ab311375cf13b3 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:34:33 -0400 Subject: [PATCH 075/160] Polish the RowFn vtable adapter Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/row/row_fn.rs | 15 +++---- vortex-array/src/scalar_fn/row/vtable.rs | 54 +++++++++++++++++++++++- 2 files changed, 59 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index ccdc9d6c70c..811ea22c509 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -24,14 +24,13 @@ use crate::scalar_fn::ScalarFnId; /// sink types for each accepted dtype combination. /// /// A `RowFn` does not automatically implement -/// [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable). A function adopted by the expression -/// system implements that trait explicitly, delegates its `return_dtype` method to -/// [`row_fn_return_dtype`](crate::scalar_fn::row_fn_return_dtype), and delegates `execute` to -/// [`execute_rows`](crate::scalar_fn::execute_rows). Explicit adoption leaves the function free to -/// provide custom coercion, simplification, reduction, or formatting hooks. Implement only -/// `ScalarFnVTable` when the natural kernel is columnar rather than row-oriented. A serializable -/// adopter also delegates the vtable's `serialize` and `deserialize` methods to the matching -/// methods below; an internal or non-serializable adopter can keep the vtable defaults. +/// [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable). Invoke +/// [`impl_row_fn_vtable!`](crate::impl_row_fn_vtable) to adopt its standard scalar-function +/// behavior. A function that needs custom coercion, simplification, reduction, formatting, or +/// validity hooks implements `ScalarFnVTable` itself and can delegate its `return_dtype` and +/// `execute` methods to [`row_fn_return_dtype`](crate::scalar_fn::row_fn_return_dtype) and +/// [`execute_rows`](crate::scalar_fn::execute_rows). Implement only `ScalarFnVTable` when the +/// natural kernel is columnar rather than row-oriented. pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// Options for this function, if any. Use [`EmptyOptions`](crate::scalar_fn::EmptyOptions) /// for none. diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index f0fdcb66186..04fc4563446 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -45,7 +45,9 @@ macro_rules! impl_row_fn_vtable { fn serialize( &self, options: &Self::Options, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult>> { + ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult< + ::core::option::Option<::std::vec::Vec>, + > { $crate::scalar_fn::RowFn::serialize(self, options) } @@ -95,7 +97,7 @@ macro_rules! impl_row_fn_vtable { _options: &Self::Options, expression: &$crate::expr::Expression, ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult< - Option<$crate::expr::Expression>, + ::core::option::Option<$crate::expr::Expression>, > { $crate::expr::union_child_validities(expression) } @@ -207,3 +209,51 @@ fn prepare_batch( function.dispatch(options, arg_dtypes, PlanRows::::new(arg_dtypes, options)) }) } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_session::registry::CachedId; + + use crate::dtype::DType; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::RowFn; + use crate::scalar_fn::RowVisitor; + use crate::scalar_fn::ScalarFnId; + use crate::scalar_fn::ScalarFnVTable; + + struct Option; + struct Vec; + + #[derive(Clone)] + struct ShadowedPrelude; + + impl RowFn for ShadowedPrelude { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.shadowed_prelude"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + _visitor: V, + ) -> VortexResult { + vortex_bail!("compile-only RowFn must not execute") + } + } + + crate::impl_row_fn_vtable!(ShadowedPrelude); + + #[test] + fn adapter_macro_ignores_shadowed_prelude_types() { + _ = (Option, Vec); + _ = ScalarFnVTable::id(&ShadowedPrelude); + } +} From 882160b3ad6355e2aa27447336c343a5942be84c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:34:52 -0400 Subject: [PATCH 076/160] Enforce LaneZip length equality in release builds Signed-off-by: Connor Tsui --- vortex-compute/src/lane_kernels/source.rs | 25 +++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/vortex-compute/src/lane_kernels/source.rs b/vortex-compute/src/lane_kernels/source.rs index 5f295e76361..80fe101d96a 100644 --- a/vortex-compute/src/lane_kernels/source.rs +++ b/vortex-compute/src/lane_kernels/source.rs @@ -54,8 +54,9 @@ impl IndexedSource for &mut [T] { /// Pair of two [`IndexedSource`]s of equal length. Yields `(A::Item, B::Item)` per lane. /// -/// Use this to drive a binary kernel from two columns. Length equality is enforced -/// at construction. +/// Use this to drive a binary kernel from two columns. [`LaneZip::new`] enforces length equality +/// at construction, and [`IndexedSource::len`] checks it for callers of the public tuple +/// constructor. #[derive(Clone, Copy)] pub struct LaneZip(pub A, pub B); @@ -79,12 +80,28 @@ impl IndexedSource for LaneZip { type Item = (A::Item, B::Item); #[inline] fn len(&self) -> usize { - debug_assert_eq!(self.0.len(), self.1.len()); + assert_eq!( + self.0.len(), + self.1.len(), + "LaneZip operands must have the same length" + ); self.0.len() } #[inline] unsafe fn get_unchecked(&self, i: usize) -> (A::Item, B::Item) { - // SAFETY: caller guarantees i < self.len(); `new` enforces matching lengths. + // SAFETY: caller guarantees i < self.len(), which also verifies matching lengths. unsafe { (self.0.get_unchecked(i), self.1.get_unchecked(i)) } } } + +#[cfg(test)] +mod tests { + use super::IndexedSource; + use super::LaneZip; + + #[test] + #[should_panic(expected = "LaneZip operands must have the same length")] + fn direct_construction_checks_lengths() { + LaneZip(&[1_u8][..], &[2_u8, 3][..]).len(); + } +} From e2f4a2bbe8f7a4986809787fddbdd16c0b15b31e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:35:45 -0400 Subject: [PATCH 077/160] Repair RowFn benchmark revision comparisons Signed-off-by: Connor Tsui --- scripts/benchmark-rowfn.sh | 36 +++++++++++++++++++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/scripts/benchmark-rowfn.sh b/scripts/benchmark-rowfn.sh index 20f705dd194..4c6ae4152cc 100755 --- a/scripts/benchmark-rowfn.sh +++ b/scripts/benchmark-rowfn.sh @@ -5,6 +5,8 @@ set -Eeu -o pipefail +script_directory=$(dirname "$(realpath "${BASH_SOURCE[0]}")") + usage() { cat >&2 <<'EOF' Usage: benchmark-rowfn.sh [OPTIONS] @@ -145,6 +147,33 @@ for request in "${requested_suites[@]}"; do fi done +common_suites=() +skipped_suites=() +for entry in "${selected_suites[@]}"; do + IFS='|' read -r label package bench _ <<<"$entry" + baseline_source="$baseline/$package/benches/$bench.rs" + candidate_source="$candidate/$package/benches/$bench.rs" + + if [[ -f $baseline_source && -f $candidate_source ]]; then + common_suites+=("$entry") + elif [[ -f $baseline_source ]]; then + skipped_suites+=("$label (baseline only)") + elif [[ -f $candidate_source ]]; then + skipped_suites+=("$label (candidate only)") + else + skipped_suites+=("$label (missing from both revisions)") + fi +done +if ((${#common_suites[@]} == 0)); then + echo "No requested benchmark targets exist in both revisions; no comparison is possible." >&2 + printf 'Skipped: %s\n' "${skipped_suites[@]}" >&2 + exit 1 +fi +selected_suites=("${common_suites[@]}") +if ((${#skipped_suites[@]} != 0)); then + printf 'Skipping one-sided benchmark target: %s\n' "${skipped_suites[@]}" >&2 +fi + common_git_dir=$(git -C "$candidate" rev-parse --path-format=absolute --git-common-dir) repository_root=$(dirname "$common_git_dir") if [[ -z $target_root ]]; then @@ -159,7 +188,7 @@ fi mkdir -p "$output/build" "$output/warm" "$output/measured" "$target_root" baseline_target="$target_root/baseline" candidate_target="$target_root/candidate" -parser="$candidate/scripts/rowfn_benchmark.py" +parser="$script_directory/rowfn_benchmark.py" { echo "RowFn benchmark machine record" @@ -173,6 +202,11 @@ parser="$candidate/scripts/rowfn_benchmark.py" echo "Warm runs: $warm_runs" echo "Measured pairs: $measured_pairs" echo "Divan: TSC timer, $sample_count samples, min $min_time s, max $max_time s" + if ((${#skipped_suites[@]} == 0)); then + echo "Skipped one-sided benchmark targets: none" + else + printf 'Skipped one-sided benchmark target: %s\n' "${skipped_suites[@]}" + fi echo rustc -vV cargo -V From 2029bef090af488a1645ffa9e3195140b0be4f1e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:40:24 -0400 Subject: [PATCH 078/160] Explain signed multiplication overflow probes Signed-off-by: Connor Tsui --- .../scalar_fn/fns/binary/numeric/primitive.rs | 38 +++++++++++-------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index f07ed8e9d5f..7cbe88c7c8e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -206,6 +206,10 @@ macro_rules! impl_checked_signed { let kept = wide as $ty; let discarded = (wide >> <$ty>::BITS) as $ty; + // A product fits exactly when its discarded half is the sign extension of the kept + // half. XOR reduces that comparison to zero evidence for success and nonzero evidence + // for overflow without converting the wide product to a branch. + (discarded ^ (kept >> (<$ty>::BITS - 1))) as $failure }); }; @@ -294,23 +298,25 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; + /// Values around zero, signed extrema, and 32- and 64-bit boundaries where the discarded + /// multiplication half or its sign extension changes. const PROBES: &[i64] = &[ - 0, - 1, - -1, - 2, - -2, - 3, - i64::MIN, - i64::MIN + 1, - i64::MAX, - i64::MAX - 1, - 1 << 31, - 1 << 32, - 1 << 62, - -(1 << 62), - 0x7FFF_FFFF, - -0x8000_0000, + 0, // Additive identity. + 1, // Smallest positive value. + -1, // All sign bits set. + 2, // Small positive power of two. + -2, // Small negative power of two. + 3, // Small non-power of two. + i64::MIN, // Minimum signed value. + i64::MIN + 1, // Minimum signed value's neighbor. + i64::MAX, // Maximum signed value. + i64::MAX - 1, // Maximum signed value's neighbor. + 1 << 31, // First positive value outside i32. + 1 << 32, // First value with bit 32 set. + 1 << 62, // Largest positive power of two in i64. + -(1 << 62), // Negative counterpart of the largest power of two. + 0x7FFF_FFFF, // Maximum i32 represented as i64. + -0x8000_0000, // Minimum i32 represented as i64. ]; #[track_caller] From 8d3357a2e7805a63be7557d5a86b5126053479be Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:40:45 -0400 Subject: [PATCH 079/160] Polish primitive comparison execution Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/compare/mod.rs | 6 +++--- vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs | 5 +++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index 2ef3d7b424e..d5fdcbae831 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -4,9 +4,9 @@ //! Native comparison kernels. //! //! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every -//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from -//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise -//! comparator for nested types. There is no Arrow fallback. +//! comparison directly over Vortex canonical arrays: bit buffers for booleans, row or fused lane +//! kernels for primitives, lane kernels for decimals, binary views for strings and bytes, and a +//! row-wise comparator for nested types. There is no Arrow fallback. //! //! Floating point values compare with Vortex's total ordering (`NaN` is the largest value, //! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics. diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 35fb637d00b..4d4a5e8175f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -15,11 +15,11 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::PType; use crate::match_each_native_ptype; +use crate::scalar_fn::BorrowedExecutionArgs; use crate::scalar_fn::RowFn; use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::execute_rows; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; @@ -80,7 +80,8 @@ pub(super) fn compare_primitive_with_path( return columnar::compare_primitive(lhs, rhs, op, ctx); } - let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + let inputs = [lhs.clone(), rhs.clone()]; + let args = BorrowedExecutionArgs::new(&inputs, lhs.len()); execute_rows(&PrimitiveCompare, &op, &args, ctx) } From 409dd4f161751576982267faf64360786960dd21 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:53:42 -0400 Subject: [PATCH 080/160] fixup! Execute tensor L2 norm with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/benches/l2_norm.rs | 30 +++++++++++++++++++ vortex-tensor/src/scalar_fns/l2_norm.rs | 19 ++++++++++++ vortex-tensor/src/scalar_fns/tests/l2_norm.rs | 23 ++++++++++++++ vortex-tensor/src/utils.rs | 12 ++++++++ 4 files changed, 84 insertions(+) diff --git a/vortex-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index d96e4877af9..bf8832f2520 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -15,12 +15,18 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; @@ -56,6 +62,16 @@ fn vectors(width: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { let session = vortex_array::array_session(); bencher @@ -84,3 +100,17 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_l2_norm(bencher, input); } + +#[divan::bench(args = WIDTHS)] +fn constant(bencher: Bencher, width: usize) { + bench_l2_norm(bencher, constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let input = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_l2_norm(bencher, input); +} diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 3bad1d6d1d9..27deb825e0f 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -7,6 +7,7 @@ use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; +use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; @@ -20,6 +21,7 @@ use vortex_array::scalar_fn::RowExecution; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::scalar_fn::UninitElementSink; use vortex_array::serde::ArrayChildren; use vortex_error::VortexResult; @@ -52,6 +54,23 @@ use crate::utils::validate_tensor_float_input; #[derive(Clone, Debug, Default)] pub struct L2Norm; +impl L2Norm { + /// Creates a new [`TypedScalarFnInstance`] wrapping the L2 norm operation. + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) + } + + /// Constructs a [`ScalarFnArray`] that lazily computes the L2 norm over `child`. + /// + /// # Errors + /// + /// Returns an error if the array cannot be constructed, such as when the input dtype is + /// unsupported. + pub fn try_new_array(child: ArrayRef) -> VortexResult { + ScalarFnArray::try_new(Self::new().erased(), vec![child]) + } +} + impl RowFn for L2Norm { type Options = EmptyOptions; diff --git a/vortex-tensor/src/scalar_fns/tests/l2_norm.rs b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs index a9fda0326d8..930ef423d47 100644 --- a/vortex-tensor/src/scalar_fns/tests/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs @@ -32,6 +32,7 @@ use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::literal_vector_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. fn eval_l2_norm(input: ArrayRef) -> VortexResult> { @@ -42,6 +43,28 @@ fn eval_l2_norm(input: ArrayRef) -> VortexResult> { Ok(prim.as_slice::().to_vec()) } +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = L2Norm::new(); + let array = L2Norm::try_new_array(tensor_array(&[1], &[3.0])?)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + assert_close( + &eval_l2_norm(zero_width_vector_array::(3)?)?, + &[0.0, 0.0, 0.0], + ); + assert!(eval_l2_norm(vector_array(2, &[] as &[f64])?)?.is_empty()); + + let constant = Vector::constant_array::(&[], 3)?; + assert_close(&eval_l2_norm(constant)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + #[rstest] #[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] #[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 460dde82ea7..43086d6e036 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -411,6 +411,18 @@ pub mod test_helpers { Vector::try_new_vector_array(flat_fsl(elements, dim)) } + /// Builds `rows` zero-width vectors over an empty typed element buffer. + pub fn zero_width_vector_array(rows: usize) -> VortexResult { + let storage = FixedSizeListArray::new( + Buffer::::empty().into_array(), + 0, + Validity::NonNullable, + rows, + ) + .into_array(); + Vector::try_new_vector_array(storage) + } + /// Builds a [`FixedShapeTensor`] extension array whose storage is a [`ConstantArray`], /// representing a single query tensor broadcast to `len` rows. pub fn constant_tensor_array>( From 9354a2548bd0fb7254e00b74df83295dfcc76474 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 23:53:52 -0400 Subject: [PATCH 081/160] fixup! Execute tensor product functions with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/benches/cosine_similarity.rs | 27 ++++++++--- vortex-tensor/benches/inner_product.rs | 35 ++++++++++++++ .../src/scalar_fns/cosine_similarity.rs | 20 ++++++++ vortex-tensor/src/scalar_fns/inner_product.rs | 20 ++++++++ .../src/scalar_fns/tests/cosine_similarity.rs | 47 ++++++++++++++----- .../src/scalar_fns/tests/inner_product.rs | 29 ++++++++++++ 6 files changed, 158 insertions(+), 20 deletions(-) diff --git a/vortex-tensor/benches/cosine_similarity.rs b/vortex-tensor/benches/cosine_similarity.rs index fef94a0aa91..49551fbf701 100644 --- a/vortex-tensor/benches/cosine_similarity.rs +++ b/vortex-tensor/benches/cosine_similarity.rs @@ -21,6 +21,7 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; @@ -43,15 +44,11 @@ fn main() { } /// Total `f64` elements per operand, held constant across widths: the row count is -/// `ELEMENTS / width`. This budget is a quarter of the one the other tensor benches use, because -/// the constant arms recompute the broadcast vector's norm per row and cost roughly ten times the -/// column arms per element. It is what keeps every arm inside the 1 ms per-iteration limit from -/// `docs/developer-guide/benchmarking.md`, measured against CodSpeed's CPU simulation. +/// `ELEMENTS / width`. The smaller budget keeps the wider cosine kernels inside the 1 ms +/// per-iteration limit from `docs/developer-guide/benchmarking.md` under CodSpeed simulation. const ELEMENTS: usize = 2_048; -/// Widths chosen to separate the two costs, as in `l2_norm.rs`: the redundant norm pass is -/// `O(rows * width)`, one third of the closure's arithmetic, so wide tensors show the hoist -/// while a narrow one is dominated by per-row framework costs. +/// Widths that expose both fixed row-framework costs and the `O(width)` kernel work. const WIDTHS: &[usize] = &[2, 32, 256]; /// `ELEMENTS / width` vectors of `width` `f64` elements, non-nullable. `seed` offsets the values so @@ -108,6 +105,22 @@ fn column_x_constant(bencher: Bencher, width: usize) { bench_cosine(bencher, vectors(width, 0), constant_vector(width)); } +/// The lhs is a broadcast query vector, whose norm is the same in every row. +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_cosine(bencher, constant_vector(width), vectors(width, 31)); +} + +/// A nullable broadcast rhs exercises constant preparation and output validity together. +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_cosine(bencher, vectors(width, 0), rhs); +} + /// One query vector represented as an extension array over constant storage. fn extension_constant_vector(width: usize) -> ArrayRef { let ext_dtype = vectors(width, 0).dtype().as_extension().clone(); diff --git a/vortex-tensor/benches/inner_product.rs b/vortex-tensor/benches/inner_product.rs index c0918f87ba2..5c4adf1c7ec 100644 --- a/vortex-tensor/benches/inner_product.rs +++ b/vortex-tensor/benches/inner_product.rs @@ -15,12 +15,18 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; @@ -58,6 +64,16 @@ fn vectors(width: usize, seed: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_inner_product(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { let session = vortex_array::array_session(); bencher @@ -86,3 +102,22 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_inner_product(bencher, lhs, vectors(width, 31)); } + +#[divan::bench(args = WIDTHS)] +fn column_x_constant(bencher: Bencher, width: usize) { + bench_inner_product(bencher, vectors(width, 0), constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_inner_product(bencher, constant_vector(width), vectors(width, 31)); +} + +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_inner_product(bencher, vectors(width, 0), rhs); +} diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 91e4e837d64..9cccf05d17b 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -9,6 +9,7 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; @@ -22,6 +23,7 @@ use vortex_array::scalar_fn::RowExecution; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::scalar_fn::UninitElementSink; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; @@ -61,6 +63,24 @@ use crate::utils::l2_norm_row; #[derive(Clone, Debug, Default)] pub struct CosineSimilarity; +impl CosineSimilarity { + /// Creates a new [`TypedScalarFnInstance`] wrapping the cosine similarity operation. + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) + } + + /// Constructs a [`ScalarFnArray`] that lazily computes the cosine similarity between `lhs` and + /// `rhs`. + /// + /// # Errors + /// + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. + pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) + } +} + impl RowFn for CosineSimilarity { type Options = EmptyOptions; diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index c0968d30474..6f09e2b17f1 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -6,6 +6,7 @@ use num_traits::Float; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; @@ -20,6 +21,7 @@ use vortex_array::scalar_fn::RowExecution; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::scalar_fn::UninitElementSink; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::serde::ArrayChildren; @@ -47,6 +49,24 @@ use crate::utils::extract_normalized_children; #[derive(Clone, Debug, Default)] pub struct InnerProduct; +impl InnerProduct { + /// Creates a new [`TypedScalarFnInstance`] wrapping the inner product operation. + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) + } + + /// Constructs a [`ScalarFnArray`] that lazily computes the inner product between `lhs` and + /// `rhs`. + /// + /// # Errors + /// + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. + pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) + } +} + impl RowFn for InnerProduct { type Options = EmptyOptions; diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs index 60e75792109..0e2687690a9 100644 --- a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -29,6 +29,7 @@ use crate::utils::test_helpers::literal_vector_array; use crate::utils::test_helpers::normalized_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { @@ -39,6 +40,33 @@ fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult Ok(prim.as_slice::().to_vec()) } +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = CosineSimilarity::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = CosineSimilarity::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_cosine_similarity(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + /// Like [`eval_cosine_similarity`], but returns the executed array for exact array comparisons. fn eval_cosine_similarity_array( lhs: ArrayRef, @@ -410,8 +438,8 @@ fn both_constant_tensors() -> VortexResult<()> { #[test] fn constant_zero_norm_query() -> VortexResult<()> { - // A zero-norm constant query must produce `0.0` for every row via the zero-norm guard in - // `cosine_one_normalized` and `execute_both_normalized`. + // A zero-norm constant query must produce `0.0` through the prepared row kernel's + // zero-denominator guard. let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; let rhs = tensor_array( &[3], @@ -427,9 +455,7 @@ fn constant_zero_norm_query() -> VortexResult<()> { #[test] fn constant_self_similarity_nonunit() -> VortexResult<()> { - // A non-unit constant query compared to itself must produce `1.0`. This exercises the - // helper's division: after normalization, both sides must be exactly unit so the - // Normalized fast path's inner product yields 1. + // The prepared path hoists both norms and computes the same dot product for every row. let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); @@ -437,9 +463,7 @@ fn constant_self_similarity_nonunit() -> VortexResult<()> { } /// An extension array over constant storage (what [`Vector::constant_array`] builds) is a batch -/// constant like any other: the row layer sees through the wrapper, so `prepare` hoists its norm -/// exactly as it does for the literal shape. This used to be intercepted by a hand-written -/// `reduce_encoded` rewrite into `Normalized`, deleted in favor of the framework path. +/// constant like any other. The row layer sees through the wrapper, so `prepare` hoists its norm. #[test] fn vector_constant_matches_plain() -> VortexResult<()> { let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; @@ -465,11 +489,8 @@ fn vector_constant_matches_plain() -> VortexResult<()> { Ok(()) } -/// The literal-constant shape (a [`ConstantArray`] over a [`Vector`] extension scalar, what a -/// `lit(query)` expression produces) reaches the row loop, unlike an extension-wrapped constant, -/// which `reduce_encoded` rewrites into `Normalized`. There the prepared kernel hoists the query's -/// norm once per batch, and the result must be exactly the result of expanding the same query -/// into a full column, which hoists nothing. +/// Both literal and extension-wrapped constant storage reach the prepared row path. The probe +/// ensures that the literal query remains a batch constant instead of becoming a varying column. /// /// [`ConstantArray`]: vortex_array::arrays::ConstantArray #[test] diff --git a/vortex-tensor/src/scalar_fns/tests/inner_product.rs b/vortex-tensor/src/scalar_fns/tests/inner_product.rs index af7fbb7bc1a..9a5a60f5dda 100644 --- a/vortex-tensor/src/scalar_fns/tests/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/tests/inner_product.rs @@ -19,10 +19,12 @@ use vortex_error::VortexResult; use crate::encodings::normalized::Normalized; use crate::scalar_fns::inner_product::InnerProduct; use crate::tests::SESSION; +use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::normalized_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; /// Evaluates inner product between two tensor arrays and returns the result as `Vec`. fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { @@ -33,6 +35,33 @@ fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { Ok(prim.as_slice::().to_vec()) } +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = InnerProduct::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = InnerProduct::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_inner_product(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + /// Single-row inner product for various vector pairs. #[rstest] // Orthogonal: [1, 0] . [0, 1] = 0. From 10977c44078baac61f8207babcfc0321f7ac1fb4 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 01:20:46 -0400 Subject: [PATCH 082/160] fixup! Execute tensor L2 norm with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/l2_norm.rs | 52 ++++++++++++++++++++----- 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 27deb825e0f..f8d7dc617b7 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -6,6 +6,9 @@ use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; @@ -15,6 +18,7 @@ use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; use vortex_array::dtype::proto::dtype as pb; use vortex_array::match_each_float_ptype; +use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::InitializedElement; use vortex_array::scalar_fn::RowExecution; @@ -24,6 +28,7 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::scalar_fn::UninitElementSink; use vortex_array::serde::ArrayChildren; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; @@ -117,18 +122,45 @@ impl RowFn for L2Norm { _ctx: &mut ExecutionCtx, ) -> VortexResult> { let input = &args[0]; - if !input.is::() { - return Ok(None); + if input.is::() { + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let (_, norms) = extract_normalized_children(input); + vortex_ensure!( + norms.dtype().is_primitive(), + "normalized norms must be primitive, got {}", + norms.dtype(), + ); + vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); + return Ok(Some(RowExecution::Output(norms))); } + + let Some(constant) = input.as_opt::() else { + return Ok(None); + }; let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); - let (_, norms) = extract_normalized_children(input); - vortex_ensure!( - norms.dtype().is_primitive(), - "normalized norms must be primitive, got {}", - norms.dtype(), - ); - vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); - Ok(Some(RowExecution::Output(norms))) + let norm_dtype = + DType::Primitive(element_ptype, input.dtype().as_extension().nullability()); + let storage = constant.scalar().as_extension().to_storage_scalar(); + + let Some(elements) = storage.as_list().elements() else { + let output = ConstantArray::new(Scalar::null(norm_dtype), input.len()); + return Ok(Some(RowExecution::Output(output.into_array()))); + }; + + let norm = match_each_float_ptype!(element_ptype, |T| { + let values: Vec = elements + .iter() + .map(|element| { + element + .as_primitive() + .as_::() + .vortex_expect("tensor element must match its declared ptype") + }) + .collect(); + Scalar::try_new(norm_dtype, Some(l2_norm_row::(&values).into())) + })?; + let output = ConstantArray::new(norm, input.len()); + Ok(Some(RowExecution::Output(output.into_array()))) } } From 15fef920268d746c5322c37360e69ccefa4b074d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 01:25:05 -0400 Subject: [PATCH 083/160] fixup! Exercise both primitive comparison paths Signed-off-by: Connor Tsui --- .../scalar_fn/fns/binary/compare/primitive.rs | 14 +-- .../src/scalar_fn/fns/binary/compare/tests.rs | 91 +++++++++++++++++-- 2 files changed, 88 insertions(+), 17 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 4d4a5e8175f..a3f50f907fd 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -122,17 +122,17 @@ fn use_columnar_comparison( rhs: &ArrayRef, op: CompareOperator, ) -> VortexResult { - if matches!(op, CompareOperator::Eq | CompareOperator::NotEq) { - return Ok(false); - } - let ptype = PType::try_from(lhs.dtype())?; - Ok(match ptype { + Ok(match (ptype, op) { + // Equality bit-packs efficiently for every type supported by the columnar path. + (PType::I64 | PType::U64 | PType::F64, CompareOperator::Eq | CompareOperator::NotEq) => { + true + } // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. - PType::I64 | PType::F64 => true, + (PType::I64 | PType::F64, _) => true, // LLVM vectorizes varying u64 inputs, but not the mixed-constant RowFn loop. - PType::U64 => lhs.as_constant().is_some() || rhs.as_constant().is_some(), + (PType::U64, _) => lhs.as_constant().is_some() || rhs.as_constant().is_some(), _ => false, }) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index 06d8e20ba11..d92d368ab3f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -433,10 +433,15 @@ fn float_total_order() { } #[rstest] -#[case::row(PrimitiveComparisonPath::Row)] -#[case::columnar(PrimitiveComparisonPath::Columnar)] +#[case::row_eq(PrimitiveComparisonPath::Row, CompareOperator::Eq)] +#[case::row_not_eq(PrimitiveComparisonPath::Row, CompareOperator::NotEq)] +#[case::row_lt(PrimitiveComparisonPath::Row, CompareOperator::Lt)] +#[case::columnar_eq(PrimitiveComparisonPath::Columnar, CompareOperator::Eq)] +#[case::columnar_not_eq(PrimitiveComparisonPath::Columnar, CompareOperator::NotEq)] +#[case::columnar_lt(PrimitiveComparisonPath::Columnar, CompareOperator::Lt)] fn test_primitive_comparison_paths_preserve_semantics_and_encoding( #[case] path: PrimitiveComparisonPath, + #[case] op: CompareOperator, ) -> VortexResult<()> { let lhs = PrimitiveArray::new( vec![ @@ -474,14 +479,25 @@ fn test_primitive_comparison_paths_preserve_semantics_and_encoding( .into_array(); let mut ctx = array_session().create_execution_ctx(); - let actual = compare_primitive_with_path(&lhs, &rhs, CompareOperator::Lt, path, &mut ctx)?; - let expected = BoolArray::from_iter([ - Some(false), // - None, // - Some(true), // - Some(true), // - None, // - ]); + let actual = compare_primitive_with_path(&lhs, &rhs, op, path, &mut ctx)?; + let expected = match op { + CompareOperator::Eq => [ + Some(true), // Equal NaNs. + None, // Null on the left. + Some(false), // Distinct signed zeroes. + Some(false), // A finite value and NaN. + None, // Null on the right. + ], + CompareOperator::NotEq | CompareOperator::Lt => [ + Some(false), // Equal NaNs. + None, // Null on the left. + Some(true), // Distinct signed zeroes. + Some(true), // A finite value and NaN. + None, // Null on the right. + ], + _ => unreachable!(), + }; + let expected = BoolArray::from_iter(expected); assert_arrays_eq!(&actual, &expected, &mut ctx); assert_eq!(actual.dtype(), &DType::Bool(Nullability::Nullable)); @@ -496,6 +512,61 @@ fn test_primitive_comparison_paths_preserve_semantics_and_encoding( Ok(()) } +#[cfg(target_arch = "x86_64")] +#[rstest] +#[case::i64_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::i64_not_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::u64_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::u64_not_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::f64_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::f64_not_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +fn test_primitive_equality_auto_uses_columnar_for_supported_ptype( + #[case] lhs: ArrayRef, + #[case] rhs: ArrayRef, + #[case] op: CompareOperator, + #[case] expected: [bool; 3], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let actual = + compare_primitive_with_path(&lhs, &rhs, op, PrimitiveComparisonPath::Auto, &mut ctx)?; + + assert_eq!(actual.encoding_id(), Bool.id()); + assert_arrays_eq!(actual, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +} + #[rstest] #[case(Operator::Eq, [true, false, true, true])] #[case(Operator::Lt, [false, true, false, false])] From b05b965c5c189aeb33cb7e128c1cc7df72560dc7 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 01:25:08 -0400 Subject: [PATCH 084/160] fixup! Execute spatial predicates with RowFn Signed-off-by: Connor Tsui --- vortex-spatial/src/scalar_fn/contains.rs | 33 ++++++++++++------------ 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 52d01df4014..9376ceaeb10 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -105,38 +105,37 @@ struct ConstOperands { b: Option, } -/// One batch-constant operand: the geometry cloned out of its decoded column (the state must not -/// borrow from the columns), plus its [`PreparedGeometry`], built on the first row whose pairing -/// routes through relate. +/// One batch-constant operand: its bounding rectangle and the [`PreparedGeometry`] built on the +/// first row whose pairing routes through relate. /// /// The build is lazy because preparation (self-noding the topology graph plus an R*-tree over the /// edges) costs `O(edges log edges)` and pays off only on relate-routed pairings; a batch of /// point rows against a constant polygon never touches it, and preparing a large constant eagerly /// would charge such a batch for nothing. struct PreparedOperand { - /// The constant's decoded geometry, owned so [`prepared`](Self::prepared) can be `'static`. - geometry: Geometry, - /// The constant's bounding rectangle, folded once for conservative row rejection. bbox: Option>, - /// The lazily built prepared form of [`geometry`](Self::geometry). + /// The constant's prepared form, initialized only when a relate route needs it. prepared: OnceCell, f64>>, } impl PreparedOperand { fn new(geometry: &Geometry) -> Self { Self { - geometry: geometry.clone(), bbox: finite_bounding_rect(geometry), prepared: OnceCell::new(), } } - /// The prepared geometry, built on first use. - fn get(&self) -> &PreparedGeometry<'static, Geometry, f64> { + /// Return the prepared geometry, cloning the decoded constant only on first use. + /// + /// `geometry` **must** be the constant represented by this state. The row kernel maintains + /// that relationship by passing the operand from the same decoded constant column that + /// produced this [`PreparedOperand`]. + fn get(&self, geometry: &Geometry) -> &PreparedGeometry<'static, Geometry, f64> { self.prepared - .get_or_init(|| PreparedGeometry::from(self.geometry.clone())) + .get_or_init(|| PreparedGeometry::from(geometry.clone())) } } @@ -328,15 +327,15 @@ fn contains_row_prepared(operands: &ConstOperands, a: &Geometry, b: &Geomet match contains_route(a, b) { ContainsRoute::Direct => a.contains(b), ContainsRoute::ForwardRelate => match (&operands.a, &operands.b) { - (Some(const_a), Some(const_b)) => const_a.get().relate(const_b.get()).is_contains(), - (Some(const_a), None) => const_a.get().relate(b).is_contains(), - (None, Some(const_b)) => a.relate(const_b.get()).is_contains(), + (Some(const_a), Some(const_b)) => const_a.get(a).relate(const_b.get(b)).is_contains(), + (Some(const_a), None) => const_a.get(a).relate(b).is_contains(), + (None, Some(const_b)) => a.relate(const_b.get(b)).is_contains(), (None, None) => a.contains(b), }, ContainsRoute::ReversedRelate => match (&operands.a, &operands.b) { - (Some(const_a), Some(const_b)) => const_b.get().relate(const_a.get()).is_within(), - (Some(const_a), None) => b.relate(const_a.get()).is_within(), - (None, Some(const_b)) => const_b.get().relate(a).is_within(), + (Some(const_a), Some(const_b)) => const_b.get(b).relate(const_a.get(a)).is_within(), + (Some(const_a), None) => b.relate(const_a.get(a)).is_within(), + (None, Some(const_b)) => const_b.get(b).relate(a).is_within(), (None, None) => a.contains(b), }, } From 4c7c4b949dd6ada45f978205cfa6a2529e24689c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:10:27 -0400 Subject: [PATCH 085/160] Simplify the RowFn public API Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/mod.rs | 7 - .../src/scalar_fn/row/batch/execution.rs | 10 +- .../src/scalar_fn/row/batch/policy.rs | 10 +- vortex-array/src/scalar_fn/row/batch/tests.rs | 6 +- vortex-array/src/scalar_fn/row/execute/mod.rs | 6 +- .../src/scalar_fn/row/execute/owned.rs | 23 +- .../src/scalar_fn/row/execute/sink.rs | 30 +- vortex-array/src/scalar_fn/row/row_fn.rs | 10 +- .../src/scalar_fn/row/types/element/bool.rs | 18 +- .../src/scalar_fn/row/types/element/mod.rs | 34 +- .../scalar_fn/row/types/element/primitive.rs | 18 +- .../src/scalar_fn/row/types/element/tuple.rs | 576 ------------------ .../row/types/element/tuple/indexed.rs | 140 +++++ .../scalar_fn/row/types/element/tuple/mod.rs | 366 +++++++++++ .../row/types/element/tuple/tests.rs | 70 +++ .../src/scalar_fn/row/types/result.rs | 8 +- vortex-array/src/scalar_fn/row/types/sink.rs | 62 +- vortex-array/src/scalar_fn/row/vtable.rs | 179 +++--- 18 files changed, 768 insertions(+), 805 deletions(-) delete mode 100644 vortex-array/src/scalar_fn/row/types/element/tuple.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/tuple/indexed.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/tuple/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/tuple/tests.rs diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 4d20e73c897..5e73caefdfa 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -19,13 +19,6 @@ use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::literal::Literal; -/// Reexports used by [`impl_row_fn_vtable!`] without downstream transitive dependencies. -#[doc(hidden)] -pub mod row_fn_macro_support { - pub use vortex_error::VortexResult; - pub use vortex_session::VortexSession; -} - mod vtable; pub use vtable::*; diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs index c03956b3682..7f99def3866 100644 --- a/vortex-array/src/scalar_fn/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -19,6 +19,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::BoolArray; +use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; @@ -138,10 +139,11 @@ impl Batch { // Strictness: an all-null batch has no observable row work. Keep the literal-constant // check explicit alongside the conjoined validity invariant. if matches!(self.validity, Validity::AllInvalid) - || self - .inputs - .iter() - .any(|input| input.as_constant().is_some_and(|scalar| scalar.is_null())) + || self.inputs.iter().any(|input| { + input + .as_opt::() + .is_some_and(|constant| constant.scalar().is_null()) + }) { return Ok(self.all_null()); } diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs index d7e024ce763..d2f75f170a8 100644 --- a/vortex-array/src/scalar_fn/row/batch/policy.rs +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -85,11 +85,11 @@ mod tests { struct SparseFallibleElement; - // SAFETY: the varying view reports length zero, so no index satisfies the unchecked-read + // SAFETY: the per-row view reports length zero, so no index satisfies the unchecked-read // precondition. unsafe impl InputElement for SparseFallibleElement { type Column = (); - type Varying<'a> = (); + type View<'a> = (); type Elem<'a> = (); const DENSE_SAFE: bool = false; @@ -105,13 +105,13 @@ mod tests { fn get(_column: &Self::Column, _index: usize) -> Self::Elem<'_> {} - fn varying(_column: &Self::Column) -> Self::Varying<'_> {} + fn view(_column: &Self::Column) -> Self::View<'_> {} - fn varying_len(_column: &Self::Varying<'_>) -> usize { + fn view_len(_view: &Self::View<'_>) -> usize { 0 } - fn get_varying<'a>(_column: &Self::Varying<'a>, _index: usize) -> Self::Elem<'a> {} + fn get_from_view<'a>(_view: &Self::View<'a>, _index: usize) -> Self::Elem<'a> {} } #[test] diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs index 0c213491a39..d5bf5962d2a 100644 --- a/vortex-array/src/scalar_fn/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -667,11 +667,11 @@ fn assert_invalid_kernel_output(input: ArrayRef) -> VortexResult<()> { #[rstest] #[case::owned_constant(PreparedVisit::Owned, true)] -#[case::owned_varying(PreparedVisit::Owned, false)] +#[case::owned_per_row(PreparedVisit::Owned, false)] #[case::sink_constant(PreparedVisit::Sink, true)] -#[case::sink_varying(PreparedVisit::Sink, false)] +#[case::sink_per_row(PreparedVisit::Sink, false)] #[case::deferred_constant(PreparedVisit::Deferred, true)] -#[case::deferred_varying(PreparedVisit::Deferred, false)] +#[case::deferred_per_row(PreparedVisit::Deferred, false)] fn test_prepared_visits( #[case] visit: PreparedVisit, #[case] constant_rhs: bool, diff --git a/vortex-array/src/scalar_fn/row/execute/mod.rs b/vortex-array/src/scalar_fn/row/execute/mod.rs index 9c07363dd7c..c3b5c37733a 100644 --- a/vortex-array/src/scalar_fn/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/row/execute/mod.rs @@ -56,11 +56,11 @@ impl From for VortexResult { /// Ensure that every decoded input addresses the complete row loop. pub(super) fn ensure_decoded_lengths( columns: &Args::Columns, - varying: Option<&Args::VaryingColumns<'_>>, + views: Option<&Args::Views<'_>>, row_count: usize, ) -> VortexResult<()> { - let lengths_match = match varying { - Some(varying) => Args::varying_len_matches(varying, row_count), + let lengths_match = match views { + Some(views) => Args::view_lens_match(views, row_count), None => Args::decoded_lens_match(columns, row_count), }; vortex_ensure!( diff --git a/vortex-array/src/scalar_fn/row/execute/owned.rs b/vortex-array/src/scalar_fn/row/execute/owned.rs index 490d93d1040..3437b1c3c69 100644 --- a/vortex-array/src/scalar_fn/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/row/execute/owned.rs @@ -71,25 +71,26 @@ where { let output = &mut values.spare_capacity_mut()[..row_count]; - // When every input varies, the indexed source removes argument-shape dispatch from the hot - // loop and lets the lane kernel optimize the traversal as one operation. Keep the varying - // view and its length proof in this branch: hoisting them through the shared validation - // helper changed mixed-constant add, subtract, and multiply from 9.219, 9.229, and 18.94 us - // to 30.46, 31.11, and 37.73 us on a Ryzen 9 7950X with rustc 1.91.0 and LLVM 21.1.2. + // When every input stores one value per row, the indexed source removes argument-shape + // dispatch from the hot loop and lets the lane kernel optimize the traversal as one + // operation. Keep view construction and its length proof in this branch. Hoisting them + // through the shared validation helper changed mixed-constant add, subtract, and multiply + // from 9.219, 9.229, and 18.94 us to 30.46, 31.11, and 37.73 us on a Ryzen 9 7950X with + // rustc 1.91.0 and LLVM 21.1.2. // Restoring this placement recovered the fast code under the 16-CGU, no-LTO bench profile. - if let Some(varying) = Args::varying(&columns) { + if let Some(views) = Args::per_row_views(&columns) { vortex_ensure!( - Args::varying_len_matches(&varying, row_count), + Args::view_lens_match(&views, row_count), "a decoded row input does not address exactly {row_count} rows", ); - // SAFETY: `varying_len_matches` proved every column addresses exactly `row_count` - // rows immediately above. - failure = unsafe { Args::indexed_source(varying, row_count) } + // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows + // immediately above. + failure = unsafe { Args::indexed_source(views, row_count) } .map_checked_into(output, |elements| apply(&prepared, elements)); } else { // A batch-constant input was collapsed to one row during decoding. This path reads that - // row repeatedly while indexing only the inputs that vary. + // row repeatedly while indexing only the per-row inputs. vortex_ensure!( Args::decoded_lens_match(&columns, row_count), "a decoded row input does not address exactly {row_count} rows", diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs index cbf0b536b89..4473562c7b7 100644 --- a/vortex-array/src/scalar_fn/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -38,8 +38,8 @@ where let mut sink = >::with_capacity(row_count, sink_dtype)?; let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); - let varying = Args::varying(&columns); - ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; let mut accumulated = ApplyResult::Accumulated::default(); { @@ -51,13 +51,13 @@ where "the output sink does not address exactly {row_count} rows", ); - // The all-varying representation removes argument-shape dispatch from the hot loop. The + // The all-per-row representation removes argument-shape dispatch from the hot loop. The // mixed path instead reads collapsed batch constants at row zero. - if let Some(varying) = varying { + if let Some(views) = views { for index in 0..row_count { - // SAFETY: `ensure_decoded_lengths` proved every varying column has `row_count` - // rows before the loop. - let elements = unsafe { Args::get_varying_unchecked(&varying, index) }; + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before + // the loop. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; apply( &prepared, elements, @@ -96,7 +96,7 @@ where { // Decline before input decoding or sink allocation when this sink cannot initialize rows that // the mask skips. The capability and the operation are the same function pointer. - let Some(initialize_skipped_rows) = >::SKIPPED_ROWS_INITIALIZER + let Some(initialize_skipped_rows) = >::skipped_rows_initializer() else { return Ok(None); }; @@ -127,8 +127,8 @@ where "the output sink does not address exactly {row_count} rows", ); - let varying = Args::varying(&columns); - ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; // The loop writes only valid indices, but the sink still finishes a full-length output. // Initialize placeholders now; batch execution masks them before the result escapes. @@ -142,12 +142,12 @@ where return; } - let result = match &varying { - Some(varying) => apply( + let result = match &views { + Some(views) => apply( &prepared, - // SAFETY: `ensure_decoded_lengths` proved every varying column has - // `row_count` rows, and mask indices are below `row_count`. - unsafe { Args::get_varying_unchecked(varying, index) }, + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows, and + // mask indices are below `row_count`. + unsafe { Args::get_from_views_unchecked(views, index) }, >::row(&mut rows, index), ), None => apply( diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index 811ea22c509..e39a8195a4d 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -23,12 +23,10 @@ use crate::scalar_fn::ScalarFnId; /// Declare the argument names and use [`dispatch`](Self::dispatch) to choose concrete element and /// sink types for each accepted dtype combination. /// -/// A `RowFn` does not automatically implement -/// [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable). Invoke -/// [`impl_row_fn_vtable!`](crate::impl_row_fn_vtable) to adopt its standard scalar-function -/// behavior. A function that needs custom coercion, simplification, reduction, formatting, or -/// validity hooks implements `ScalarFnVTable` itself and can delegate its `return_dtype` and -/// `execute` methods to [`row_fn_return_dtype`](crate::scalar_fn::row_fn_return_dtype) and +/// Every `RowFn` receives the standard [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) +/// implementation. A function that needs custom scalar-function hooks instead implements +/// `ScalarFnVTable` on its public type and delegates row execution to a private `RowFn` kernel with +/// [`row_fn_return_dtype`](crate::scalar_fn::row_fn_return_dtype) and /// [`execute_rows`](crate::scalar_fn::execute_rows). Implement only `ScalarFnVTable` when the /// natural kernel is columnar rather than row-oriented. pub trait RowFn: 'static + Sized + Clone + Send + Sync { diff --git a/vortex-array/src/scalar_fn/row/types/element/bool.rs b/vortex-array/src/scalar_fn/row/types/element/bool.rs index 3b68dfb1e80..30625b9cc4c 100644 --- a/vortex-array/src/scalar_fn/row/types/element/bool.rs +++ b/vortex-array/src/scalar_fn/row/types/element/bool.rs @@ -15,10 +15,10 @@ use crate::scalar_fn::InputElement; use crate::scalar_fn::OutputElement; use crate::validity::Validity; -// SAFETY: the varying view is a bit buffer, and its reported length is the buffer length. +// SAFETY: the per-row view is a bit buffer, and its reported length is the buffer length. unsafe impl InputElement for bool { type Column = BitBuffer; - type Varying<'a> = &'a BitBuffer; + type View<'a> = &'a BitBuffer; type Elem<'a> = bool; // Every bit of the buffer is readable, valid or not. @@ -41,27 +41,27 @@ unsafe impl InputElement for bool { column.value(index) } - fn varying(column: &Self::Column) -> Self::Varying<'_> { + fn view(column: &Self::Column) -> Self::View<'_> { column } - fn varying_len(column: &Self::Varying<'_>) -> usize { - column.len() + fn view_len(view: &Self::View<'_>) -> usize { + view.len() } - fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> bool + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> bool where Self: 'a, { - column.value(index) + view.value(index) } - unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> bool + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> bool where Self: 'a, { // SAFETY: forwarded from this method's contract. - unsafe { column.value_unchecked(index) } + unsafe { view.value_unchecked(index) } } } diff --git a/vortex-array/src/scalar_fn/row/types/element/mod.rs b/vortex-array/src/scalar_fn/row/types/element/mod.rs index 304e52b7d00..9ed297805b4 100644 --- a/vortex-array/src/scalar_fn/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -26,20 +26,20 @@ pub use tuple::batch_constant; /// /// # Safety /// -/// For every view returned by [`varying`](Self::varying), every index below -/// [`varying_len`](Self::varying_len) **must** satisfy the safety contract of -/// [`get_varying_unchecked`](Self::get_varying_unchecked). Shared execution relies on this proof to -/// perform unchecked reads after one pre-loop length check. +/// For every view returned by [`view`](Self::view), every index below +/// [`view_len`](Self::view_len) **must** satisfy the safety contract of +/// [`get_from_view_unchecked`](Self::get_from_view_unchecked). Shared execution relies on this +/// proof to perform unchecked reads after one pre-loop length check. pub unsafe trait InputElement: 'static { /// The decoded column representation supporting `O(1)` row access. type Column; - /// The view of a varying decoded column read by the hot row loop. + /// The view of a per-row decoded column read by the hot row loop. /// /// This may borrow a cheaper representation than [`Column`](Self::Column). Primitive elements, /// for example, expose a slice so its pointer and length are loop invariants rather than /// re-reading a [`Buffer`](vortex_buffer::Buffer) descriptor for every row. - type Varying<'a>; + type View<'a>; /// The borrowed element value handed to a row closure. type Elem<'a>; @@ -48,8 +48,8 @@ pub unsafe trait InputElement: 'static { /// /// Arrays only guarantee payloads for valid rows. This is `false` when a null row's stored /// offset or pointer may not address anything, and `true` only when [`decode`](Self::decode), - /// [`get`](Self::get), [`varying`](Self::varying), [`varying_len`](Self::varying_len), and - /// [`get_varying`](Self::get_varying) remain safe and correct for null rows. + /// [`get`](Self::get), [`view`](Self::view), [`view_len`](Self::view_len), and + /// [`get_from_view`](Self::get_from_view) remain safe and correct for null rows. /// /// Dense execution requires this of every argument; otherwise the row layer executes only /// valid rows. @@ -105,16 +105,16 @@ pub unsafe trait InputElement: 'static { /// /// Called once before the hot loop. Constants do not use this view because the tuple adapter /// keeps their one-row decoded representation separate. - fn varying(column: &Self::Column) -> Self::Varying<'_>; + fn view(column: &Self::Column) -> Self::View<'_>; - /// Number of rows addressable through a [`Varying`](Self::Varying) view. + /// Number of rows addressable through a [`View`](Self::View). /// /// Every index below this length must be valid for - /// [`get_varying_unchecked`](Self::get_varying_unchecked). - fn varying_len(column: &Self::Varying<'_>) -> usize; + /// [`get_from_view_unchecked`](Self::get_from_view_unchecked). + fn view_len(view: &Self::View<'_>) -> usize; - /// Read one row from a [`Varying`](Self::Varying) view. - fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> Self::Elem<'a> + /// Read one row from a [`View`](Self::View). + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> where Self: 'a; @@ -122,12 +122,12 @@ pub unsafe trait InputElement: 'static { /// /// # Safety /// - /// `index` must be less than [`varying_len`](Self::varying_len) for `column`. - unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> Self::Elem<'a> + /// `index` must be less than [`view_len`](Self::view_len) for `view`. + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> where Self: 'a, { - Self::get_varying(column, index) + Self::get_from_view(view, index) } } diff --git a/vortex-array/src/scalar_fn/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/row/types/element/primitive.rs index 9e80d937ef5..d48ba6494df 100644 --- a/vortex-array/src/scalar_fn/row/types/element/primitive.rs +++ b/vortex-array/src/scalar_fn/row/types/element/primitive.rs @@ -17,10 +17,10 @@ use crate::scalar_fn::InputElement; use crate::scalar_fn::OutputElement; use crate::validity::Validity; -// SAFETY: the varying view is a native slice, and its reported length is the slice length. +// SAFETY: the per-row view is a native slice, and its reported length is the slice length. unsafe impl InputElement for T { type Column = Buffer; - type Varying<'a> = &'a [T]; + type View<'a> = &'a [T]; type Elem<'a> = T; // Every lane of the buffer holds a `T`, valid or not. @@ -48,27 +48,27 @@ unsafe impl InputElement for T { column[index] } - fn varying(column: &Self::Column) -> Self::Varying<'_> { + fn view(column: &Self::Column) -> Self::View<'_> { column.as_slice() } - fn varying_len(column: &Self::Varying<'_>) -> usize { - column.len() + fn view_len(view: &Self::View<'_>) -> usize { + view.len() } - fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> T + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> T where Self: 'a, { - column[index] + view[index] } - unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> T + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> T where Self: 'a, { // SAFETY: forwarded from this method's contract. - unsafe { *column.get_unchecked(index) } + unsafe { *view.get_unchecked(index) } } } diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple.rs b/vortex-array/src/scalar_fn/row/types/element/tuple.rs deleted file mode 100644 index 04821d78363..00000000000 --- a/vortex-array/src/scalar_fn/row/types/element/tuple.rs +++ /dev/null @@ -1,576 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Argument lists built from [`InputElement`]s, and the per-argument decode behind them. - -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure_eq; - -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::arrays::Extension; -use crate::arrays::Masked; -use crate::arrays::extension::ExtensionArrayExt; -use crate::arrays::masked::MaskedArraySlotsExt; -use crate::dtype::DType; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::InputElement; - -mod private { - pub trait Sealed {} -} - -/// One decoded input, collapsed to a single row when it is constant for the batch. -pub struct ArgColumn( - /// The decoded column, classified by whether it varies within the batch. - ArgColumnKind, -); - -enum ArgColumnKind { - Varying(T::Column), - Constant(T::Column), -} - -impl ArgColumn { - fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - // An empty input has no row 0 to slice, and its row loop runs zero times either way. - if let Some(constant) = batch_constant(&array) - && !array.is_empty() - { - return Ok(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?))); - } - - Ok(Self(ArgColumnKind::Varying(T::decode(array, ctx)?))) - } - - fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { - // Batch execution short-circuits null constants before selecting a strategy, so a - // constant reaching this path is non-null and can use the ordinary decode. - if let Some(constant) = batch_constant(&array) - && !array.is_empty() - { - return Ok(Some(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?)))); - } - - Ok(T::decode_null_tolerant(array, ctx)? - .map(ArgColumnKind::Varying) - .map(Self)) - } - - fn get(&self, index: usize) -> T::Elem<'_> { - match &self.0 { - ArgColumnKind::Varying(column) => T::get(column, index), - ArgColumnKind::Constant(column) => T::get(column, 0), - } - } - - fn varying(&self) -> Option<&T::Column> { - match &self.0 { - ArgColumnKind::Varying(column) => Some(column), - ArgColumnKind::Constant(_) => None, - } - } - - fn addresses_rows(&self, row_count: usize) -> bool { - // A constant is always read at index zero, so it addresses any batch length. - match &self.0 { - ArgColumnKind::Varying(column) => T::varying_len(&T::varying(column)) == row_count, - ArgColumnKind::Constant(_) => true, - } - } - - fn constant(&self) -> Option> { - match &self.0 { - ArgColumnKind::Varying(_) => None, - ArgColumnKind::Constant(column) => Some(T::get(column, 0)), - } - } -} - -/// Return the batch-constant array, looking through masked and extension wrappers. -/// -/// Batch execution owns mask validity, so a masked constant may expose its constant child here. An -/// extension over constant storage remains wrapped to preserve its extension dtype. -pub fn batch_constant(array: &ArrayRef) -> Option { - if array.as_constant().is_some() { - return Some(array.clone()); - } - - if let Some(masked) = array.as_opt::() { - return Some(masked.child().clone()).filter(|child| child.as_constant().is_some()); - } - - array - .as_opt::() - .is_some_and(|ext| ext.storage_array().as_constant().is_some()) - .then(|| array.clone()) -} - -/// Typed argument tuples for arities zero through twelve. -/// -/// This trait is sealed; add a new row representation by implementing [`InputElement`] and placing -/// it in one of the supplied tuples. -pub trait ElementTuple: 'static + private::Sealed { - /// The decoded column representations. - type Columns; - - /// Direct references to decoded columns when every argument varies within the batch. - type VaryingColumns<'a>; - - /// The borrowed row of element values. - type Elems<'a>; - - /// The batch-constant element values: [`Elems`](Self::Elems) with every argument wrapped in - /// `Option`. - /// - /// `Some` marks an argument whose operand is constant for the batch and carries the element - /// every row reads; `None` marks one that varies by row. This is what - /// [`RowVisitor`](crate::scalar_fn::RowVisitor) hands to a visit's prepare closure, so a kernel - /// can hoist work that depends only on a constant argument out of the row loop. - type ConstElems<'a>; - - /// The number of arguments. - const ARITY: usize; - - /// Whether every argument is [`InputElement::DENSE_SAFE`]. - const DENSE_SAFE: bool; - - /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. - const DECODE_FALLIBLE: bool; - - /// Validate the input dtypes, including that `dtypes` has exactly `ARITY` entries. - /// - /// The expression layer checks the count against [`Arity`](crate::scalar_fn::Arity) before it - /// builds a call, but this is also the entry point of the public - /// [`return_dtype`](crate::scalar_fn::ScalarFnVTable::return_dtype), so the count is enforced - /// here rather than assumed. - fn validate(dtypes: &[DType]) -> VortexResult<()>; - - /// Decode every input column once. Called once per batch. - fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; - - /// Decode every input column once while tolerating null rows. - /// - /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid - /// strategy calls this once per batch. - fn decode_null_tolerant( - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; - - /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. - fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; - - /// Borrow every decoded column directly, or `None` when any argument is batch-constant. - /// - /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple - /// gives the optimizer ordinary contiguous column access without a per-row constant check. - fn varying(columns: &Self::Columns) -> Option>; - - /// Whether every varying column contains exactly `row_count` rows. - fn varying_len_matches(columns: &Self::VaryingColumns<'_>, row_count: usize) -> bool; - - /// Whether every argument that varies within the batch contains exactly `row_count` rows. - /// - /// The same guarantee as [`varying_len_matches`](Self::varying_len_matches), for the mixed case - /// [`varying`](Self::varying) declines: a batch-constant argument is exempt because it was - /// collapsed to one row, while every argument beside it still has to address the whole batch. - fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; - - /// Read one row from columns already known to vary within the batch. - fn get_varying<'a>(columns: &Self::VaryingColumns<'a>, index: usize) -> Self::Elems<'a>; - - /// Read one row from varying columns without checking bounds. - /// - /// # Safety - /// - /// `index` must be in bounds for every column. - unsafe fn get_varying_unchecked<'a>( - columns: &Self::VaryingColumns<'a>, - index: usize, - ) -> Self::Elems<'a>; - - /// Read the batch-constant elements out of the decoded columns. Called once per batch. - fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; -} - -/// An argument tuple that supports a validated dense indexed traversal. -/// -/// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's -/// unchecked varying-row access after batch execution validates every decoded column length once. -pub trait IndexedElementTuple: ElementTuple { - /// The source shared execution uses for a dense all-varying loop. - /// - /// Its length must be the common varying-column length. For every valid index it must preserve - /// row order, return the same value as [`ElementTuple::get_varying`], and uphold the unchecked - /// read contract of [`IndexedSource`]. - type Source<'a>: IndexedSource>; - - /// Build a source from columns already validated to vary within the batch. - /// - /// # Safety - /// - /// Every column in `columns` **must** address exactly `row_count` rows. Violating this - /// requirement can make a safe lane kernel read outside a column's allocation. - unsafe fn indexed_source<'a>( - columns: Self::VaryingColumns<'a>, - row_count: usize, - ) -> Self::Source<'a>; -} - -/// Indexed access to one varying element column. -pub struct ElementSource<'a, T: InputElement> { - column: T::Varying<'a>, -} - -impl<'a, T: InputElement> ElementSource<'a, T> { - fn new(column: T::Varying<'a>) -> Self { - Self { column } - } -} - -impl<'a, T: InputElement> IndexedSource for ElementSource<'a, T> { - type Item = T::Elem<'a>; - - fn len(&self) -> usize { - T::varying_len(&self.column) - } - - unsafe fn get_unchecked(&self, index: usize) -> Self::Item { - // SAFETY: the source length is the number of rows addressable by `column`, and the caller - // guarantees that `index` is below that length. - unsafe { T::get_varying_unchecked(&self.column, index) } - } -} - -/// An indexed element source yielding the one-tuples expected by a unary row closure. -pub struct UnaryTupleSource(Source); - -impl IndexedSource for UnaryTupleSource { - type Item = (Source::Item,); - - fn len(&self) -> usize { - self.0.len() - } - - unsafe fn get_unchecked(&self, index: usize) -> Self::Item { - // SAFETY: forwarded from this method's contract. - (unsafe { self.0.get_unchecked(index) },) - } -} - -/// Indexed access to the varying columns of an element tuple. -pub struct ElementTupleSource<'a, Args: ElementTuple> { - columns: Args::VaryingColumns<'a>, - row_count: usize, -} - -impl<'a, Args: ElementTuple> IndexedSource for ElementTupleSource<'a, Args> { - type Item = Args::Elems<'a>; - - fn len(&self) -> usize { - self.row_count - } - - unsafe fn get_unchecked(&self, index: usize) -> Self::Item { - // SAFETY: the caller guarantees that `index` is below `row_count`. Batch execution checks - // that every varying column addresses exactly `row_count` rows before constructing this - // source. - unsafe { Args::get_varying_unchecked(&self.columns, index) } - } -} - -impl private::Sealed for () {} - -impl ElementTuple for () { - type Columns = (); - type VaryingColumns<'a> = (); - type Elems<'a> = (); - type ConstElems<'a> = (); - - const ARITY: usize = 0; - const DENSE_SAFE: bool = true; - const DECODE_FALLIBLE: bool = false; - - fn validate(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure_eq!( - dtypes.len(), - 0, - "expected 0 argument dtypes, got {}", - dtypes.len(), - ); - Ok(()) - } - - fn decode(_args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(()) - } - - fn decode_null_tolerant( - _args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Ok(Some(())) - } - - fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} - - fn varying(_columns: &Self::Columns) -> Option> { - Some(()) - } - - fn varying_len_matches(_columns: &Self::VaryingColumns<'_>, _row_count: usize) -> bool { - true - } - - fn decoded_lens_match(_columns: &Self::Columns, _row_count: usize) -> bool { - true - } - - fn get_varying<'a>(_columns: &Self::VaryingColumns<'a>, _index: usize) -> Self::Elems<'a> {} - - unsafe fn get_varying_unchecked<'a>( - _columns: &Self::VaryingColumns<'a>, - _index: usize, - ) -> Self::Elems<'a> { - } - - fn constants(_columns: &Self::Columns) -> Self::ConstElems<'_> {} -} - -macro_rules! element_tuple { - ($arity:literal; $($t:ident : $idx:tt),+) => { - impl<$($t: InputElement),+> private::Sealed for ($($t,)+) {} - - impl<$($t: InputElement),+> ElementTuple for ($($t,)+) { - type Columns = ($(ArgColumn<$t>,)+); - type VaryingColumns<'a> = ($($t::Varying<'a>,)+); - type Elems<'a> = ($($t::Elem<'a>,)+); - type ConstElems<'a> = ($(Option<$t::Elem<'a>>,)+); - - const ARITY: usize = $arity; - const DENSE_SAFE: bool = $($t::DENSE_SAFE &&)+ true; - const DECODE_FALLIBLE: bool = $($t::DECODE_FALLIBLE ||)+ false; - - fn validate(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure_eq!( - dtypes.len(), - $arity, - "expected {} argument dtypes, got {}", - $arity, - dtypes.len(), - ); - - $($t::validate(&dtypes[$idx])?;)+ - Ok(()) - } - - fn decode( - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - Ok(($(ArgColumn::<$t>::decode(args.get($idx)?, ctx)?,)+)) - } - - fn decode_null_tolerant( - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Ok(Some(( - $(match ArgColumn::<$t>::decode_null_tolerant(args.get($idx)?, ctx)? { - Some(column) => column, - None => return Ok(None), - },)+ - ))) - } - - fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_> { - ($(columns.$idx.get(index),)+) - } - - fn varying(columns: &Self::Columns) -> Option> { - Some(($($t::varying(columns.$idx.varying()?),)+)) - } - - fn varying_len_matches( - columns: &Self::VaryingColumns<'_>, - row_count: usize, - ) -> bool { - $($t::varying_len(&columns.$idx) == row_count &&)+ true - } - - fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { - $(columns.$idx.addresses_rows(row_count) &&)+ true - } - - fn get_varying<'a>( - columns: &Self::VaryingColumns<'a>, - index: usize, - ) -> Self::Elems<'a> { - ($($t::get_varying(&columns.$idx, index),)+) - } - - unsafe fn get_varying_unchecked<'a>( - columns: &Self::VaryingColumns<'a>, - index: usize, - ) -> Self::Elems<'a> { - // SAFETY: forwarded from this method's contract. - ($(unsafe { $t::get_varying_unchecked(&columns.$idx, index) },)+) - } - - fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { - ($(columns.$idx.constant(),)+) - } - } - }; -} - -element_tuple!(1; A:0); -element_tuple!(2; A:0, B:1); -element_tuple!(3; A:0, B:1, C:2); -element_tuple!(4; A:0, B:1, C:2, D:3); -element_tuple!(5; A:0, B:1, C:2, D:3, E:4); -element_tuple!(6; A:0, B:1, C:2, D:3, E:4, F:5); -element_tuple!(7; A:0, B:1, C:2, D:3, E:4, F:5, G:6); -element_tuple!(8; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7); -element_tuple!(9; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8); -element_tuple!(10; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9); -element_tuple!(11; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10); -element_tuple!(12; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11); - -impl IndexedElementTuple for () { - type Source<'a> = ElementTupleSource<'a, ()>; - - unsafe fn indexed_source<'a>( - columns: Self::VaryingColumns<'a>, - row_count: usize, - ) -> Self::Source<'a> { - ElementTupleSource { columns, row_count } - } -} - -impl IndexedElementTuple for (A,) { - type Source<'a> = UnaryTupleSource>; - - unsafe fn indexed_source<'a>( - columns: Self::VaryingColumns<'a>, - _row_count: usize, - ) -> Self::Source<'a> { - UnaryTupleSource(ElementSource::new(columns.0)) - } -} - -impl IndexedElementTuple for (A, B) { - type Source<'a> = LaneZip, ElementSource<'a, B>>; - - unsafe fn indexed_source<'a>( - columns: Self::VaryingColumns<'a>, - _row_count: usize, - ) -> Self::Source<'a> { - LaneZip::new(ElementSource::new(columns.0), ElementSource::new(columns.1)) - } -} - -macro_rules! indexed_element_tuple { - ($($t:ident),+) => { - impl<$($t: InputElement),+> IndexedElementTuple for ($($t,)+) { - type Source<'a> = ElementTupleSource<'a, ($($t,)+)>; - - unsafe fn indexed_source<'a>( - columns: Self::VaryingColumns<'a>, - row_count: usize, - ) -> Self::Source<'a> { - ElementTupleSource { columns, row_count } - } - } - }; -} - -indexed_element_tuple!(A, B, C); -indexed_element_tuple!(A, B, C, D); -indexed_element_tuple!(A, B, C, D, E); -indexed_element_tuple!(A, B, C, D, E, F); -indexed_element_tuple!(A, B, C, D, E, F, G); -indexed_element_tuple!(A, B, C, D, E, F, G, H); -indexed_element_tuple!(A, B, C, D, E, F, G, H, I); -indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J); -indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K); -indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K, L); - -#[cfg(test)] -mod tests { - use vortex_compute::lane_kernels::IndexedSource; - use vortex_error::VortexResult; - use vortex_error::vortex_bail; - use vortex_mask::Mask; - - use super::IndexedElementTuple; - use super::batch_constant; - use crate::IntoArray; - use crate::arrays::ConstantArray; - use crate::arrays::ExtensionArray; - use crate::arrays::MaskedArray; - use crate::dtype::Nullability; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::validity::Validity; - - #[test] - fn test_unary_element_source_reads_one_tuple_per_row() { - // SAFETY: the only column has exactly three rows. - let source = unsafe { <(i32,)>::indexed_source((&[10, 20, 30][..],), 3) }; - assert_eq!(source.len(), 3); - - // SAFETY: index one is within the three-element source. - assert_eq!(unsafe { source.get_unchecked(1) }, (20,)); - } - - #[test] - fn test_binary_element_source_zips_columns() { - // SAFETY: both columns have exactly three rows. - let source = - unsafe { <(i32, i64)>::indexed_source((&[10, 20, 30][..], &[100, 200, 300][..]), 3) }; - assert_eq!(source.len(), 3); - - // SAFETY: index one is within the three-element source. - assert_eq!(unsafe { source.get_unchecked(1) }, (20, 200)); - } - - #[test] - fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> { - let child = ConstantArray::new(7_i64, 3).into_array(); - let masked = - MaskedArray::try_new(child, Validity::from_iter([true, false, true]))?.into_array(); - let filtered = masked.filter(Mask::from_iter([true, true, false]))?; - - let Some(constant) = batch_constant(&filtered) else { - vortex_bail!("filtered masked constant must remain batch-constant"); - }; - - assert!(constant.as_constant().is_some()); - Ok(()) - } - - #[test] - fn test_batch_constant_preserves_filtered_extension() -> VortexResult<()> { - let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); - let extension = - ExtensionArray::new(ext_dtype, ConstantArray::new(7_i64, 3).into_array()).into_array(); - let filtered = extension.filter(Mask::from_iter([true, false, true]))?; - - let Some(constant) = batch_constant(&filtered) else { - vortex_bail!("filtered extension storage must remain batch-constant"); - }; - - assert_eq!(constant.dtype(), extension.dtype()); - Ok(()) - } -} diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/row/types/element/tuple/indexed.rs new file mode 100644 index 00000000000..1c1c02adb73 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/tuple/indexed.rs @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_compute::lane_kernels::IndexedSource; +use vortex_compute::lane_kernels::LaneZip; + +use super::ElementTuple; +use crate::scalar_fn::InputElement; + +/// An argument tuple that supports a validated dense indexed traversal. +/// +/// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's +/// unchecked view access after batch execution validates every decoded column length once. +pub trait IndexedElementTuple: ElementTuple { + /// The source shared execution uses for a dense all-per-row loop. + /// + /// Its length must be the common view length. For every valid index it must preserve row order, + /// return the same value as [`ElementTuple::get_from_views`], and uphold the unchecked read + /// contract of [`IndexedSource`]. + type Source<'a>: IndexedSource>; + + /// Build a source from views already validated to cover the complete batch. + /// + /// # Safety + /// + /// Every view in `views` **must** address exactly `row_count` rows. Violating this requirement + /// can make a safe lane kernel read outside a column's allocation. + unsafe fn indexed_source<'a>(views: Self::Views<'a>, row_count: usize) -> Self::Source<'a>; +} + +/// Indexed access to one element view. +pub struct ElementSource<'a, T: InputElement> { + view: T::View<'a>, +} + +impl<'a, T: InputElement> ElementSource<'a, T> { + fn new(view: T::View<'a>) -> Self { + Self { view } + } +} + +impl<'a, T: InputElement> IndexedSource for ElementSource<'a, T> { + type Item = T::Elem<'a>; + + fn len(&self) -> usize { + T::view_len(&self.view) + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the source length is the number of rows addressable by `view`, and the caller + // guarantees that `index` is below that length. + unsafe { T::get_from_view_unchecked(&self.view, index) } + } +} + +/// An indexed element source yielding the one-tuples expected by a unary row closure. +pub struct UnaryTupleSource(Source); + +impl IndexedSource for UnaryTupleSource { + type Item = (Source::Item,); + + fn len(&self) -> usize { + self.0.len() + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: forwarded from this method's contract. + (unsafe { self.0.get_unchecked(index) },) + } +} + +/// Indexed access to the views of an element tuple. +pub struct ElementTupleSource<'a, Args: ElementTuple> { + views: Args::Views<'a>, + row_count: usize, +} + +impl<'a, Args: ElementTuple> IndexedSource for ElementTupleSource<'a, Args> { + type Item = Args::Elems<'a>; + + fn len(&self) -> usize { + self.row_count + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the caller guarantees that `index` is below `row_count`. Batch execution checks + // that every view addresses exactly `row_count` rows before constructing this source. + unsafe { Args::get_from_views_unchecked(&self.views, index) } + } +} + +impl IndexedElementTuple for () { + type Source<'a> = ElementTupleSource<'a, ()>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, row_count: usize) -> Self::Source<'a> { + ElementTupleSource { views, row_count } + } +} + +impl IndexedElementTuple for (A,) { + type Source<'a> = UnaryTupleSource>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, _row_count: usize) -> Self::Source<'a> { + UnaryTupleSource(ElementSource::new(views.0)) + } +} + +impl IndexedElementTuple for (A, B) { + type Source<'a> = LaneZip, ElementSource<'a, B>>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, _row_count: usize) -> Self::Source<'a> { + LaneZip::new(ElementSource::new(views.0), ElementSource::new(views.1)) + } +} + +macro_rules! indexed_element_tuple { + ($($t:ident),+) => { + impl<$($t: InputElement),+> IndexedElementTuple for ($($t,)+) { + type Source<'a> = ElementTupleSource<'a, ($($t,)+)>; + + unsafe fn indexed_source<'a>( + views: Self::Views<'a>, + row_count: usize, + ) -> Self::Source<'a> { + ElementTupleSource { views, row_count } + } + } + }; +} + +indexed_element_tuple!(A, B, C); +indexed_element_tuple!(A, B, C, D); +indexed_element_tuple!(A, B, C, D, E); +indexed_element_tuple!(A, B, C, D, E, F); +indexed_element_tuple!(A, B, C, D, E, F, G); +indexed_element_tuple!(A, B, C, D, E, F, G, H); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K, L); diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/row/types/element/tuple/mod.rs new file mode 100644 index 00000000000..3c44c8aae05 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/tuple/mod.rs @@ -0,0 +1,366 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Argument lists built from [`InputElement`]s, and the per-argument decode behind them. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Constant; +use crate::arrays::Extension; +use crate::arrays::Masked; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::InputElement; + +mod indexed; +pub use indexed::IndexedElementTuple; + +/// One decoded input, collapsed to a single row when it is constant for the batch. +pub struct ArgColumn( + /// The decoded column, classified by whether it stores one value per row. + ArgColumnKind, +); + +enum ArgColumnKind { + PerRow(T::Column), + Constant(T::Column), +} + +impl ArgColumn { + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // An empty input has no row 0 to slice, and its row loop runs zero times either way. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?))); + } + + Ok(Self(ArgColumnKind::PerRow(T::decode(array, ctx)?))) + } + + fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + // Batch execution short-circuits null constants before selecting a strategy, so a + // constant reaching this path is non-null and can use the ordinary decode. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Some(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?)))); + } + + Ok(T::decode_null_tolerant(array, ctx)? + .map(ArgColumnKind::PerRow) + .map(Self)) + } + + fn get(&self, index: usize) -> T::Elem<'_> { + match &self.0 { + ArgColumnKind::PerRow(column) => T::get(column, index), + ArgColumnKind::Constant(column) => T::get(column, 0), + } + } + + fn per_row_column(&self) -> Option<&T::Column> { + match &self.0 { + ArgColumnKind::PerRow(column) => Some(column), + ArgColumnKind::Constant(_) => None, + } + } + + fn addresses_rows(&self, row_count: usize) -> bool { + // A constant is always read at index zero, so it addresses any batch length. + match &self.0 { + ArgColumnKind::PerRow(column) => T::view_len(&T::view(column)) == row_count, + ArgColumnKind::Constant(_) => true, + } + } + + fn constant(&self) -> Option> { + match &self.0 { + ArgColumnKind::PerRow(_) => None, + ArgColumnKind::Constant(column) => Some(T::get(column, 0)), + } + } +} + +/// Return the batch-constant array, looking through masked and extension wrappers. +/// +/// Batch execution owns mask validity, so a masked constant may expose its constant child here. An +/// extension over constant storage remains wrapped to preserve its extension dtype. +pub fn batch_constant(array: &ArrayRef) -> Option { + if array.is::() { + return Some(array.clone()); + } + + if let Some(masked) = array.as_opt::() { + return Some(masked.child().clone()).filter(|child| child.is::()); + } + + array + .as_opt::() + .is_some_and(|ext| ext.storage_array().is::()) + .then(|| array.clone()) +} + +/// Typed argument tuples for arities zero through twelve. +/// +/// This trait is sealed; add a new row representation by implementing [`InputElement`] and placing +/// it in one of the supplied tuples. +pub trait ElementTuple: 'static + private::Sealed { + /// The decoded column representations. + type Columns; + + /// Borrowed views of decoded columns when every argument stores one value per row. + type Views<'a>; + + /// The borrowed row of element values. + type Elems<'a>; + + /// The batch-constant element values: [`Elems`](Self::Elems) with every argument wrapped in + /// `Option`. + /// + /// `Some` marks an argument whose operand is constant for the batch and carries the element + /// every row reads; `None` marks one that stores a separate value per row. This is what + /// [`RowVisitor`](crate::scalar_fn::RowVisitor) hands to a visit's prepare closure, so a kernel + /// can hoist work that depends only on a constant argument out of the row loop. + type ConstElems<'a>; + + /// The number of arguments. + const ARITY: usize; + + /// Whether every argument is [`InputElement::DENSE_SAFE`]. + const DENSE_SAFE: bool; + + /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. + const DECODE_FALLIBLE: bool; + + /// Validate the input dtypes, including that `dtypes` has exactly `ARITY` entries. + /// + /// The expression layer checks the count against [`Arity`](crate::scalar_fn::Arity) before it + /// builds a call, but this is also the entry point of the public + /// [`return_dtype`](crate::scalar_fn::ScalarFnVTable::return_dtype), so the count is enforced + /// here rather than assumed. + fn validate(dtypes: &[DType]) -> VortexResult<()>; + + /// Decode every input column once. Called once per batch. + fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Decode every input column once while tolerating null rows. + /// + /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid + /// strategy calls this once per batch. + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; + + /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; + + /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// + /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple + /// gives the optimizer ordinary contiguous column access without a per-row constant check. + fn per_row_views(columns: &Self::Columns) -> Option>; + + /// Whether every view contains exactly `row_count` rows. + fn view_lens_match(views: &Self::Views<'_>, row_count: usize) -> bool; + + /// Whether every per-row argument contains exactly `row_count` rows. + /// + /// This provides the same guarantee as [`view_lens_match`](Self::view_lens_match) when + /// [`per_row_views`](Self::per_row_views) declines a mixed per-row and batch-constant tuple. A + /// batch constant is exempt because it was collapsed to one row. + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; + + /// Read one row from borrowed views. + fn get_from_views<'a>(views: &Self::Views<'a>, index: usize) -> Self::Elems<'a>; + + /// Read one row from borrowed views without checking bounds. + /// + /// # Safety + /// + /// `index` must be in bounds for every column. + unsafe fn get_from_views_unchecked<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a>; + + /// Read the batch-constant elements out of the decoded columns. Called once per batch. + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; +} + +impl private::Sealed for () {} + +impl ElementTuple for () { + type Columns = (); + type Views<'a> = (); + type Elems<'a> = (); + type ConstElems<'a> = (); + + const ARITY: usize = 0; + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + 0, + "expected 0 argument dtypes, got {}", + dtypes.len(), + ); + Ok(()) + } + + fn decode(_args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(()) + } + + fn decode_null_tolerant( + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(())) + } + + fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} + + fn per_row_views(_columns: &Self::Columns) -> Option> { + Some(()) + } + + fn view_lens_match(_views: &Self::Views<'_>, _row_count: usize) -> bool { + true + } + + fn decoded_lens_match(_columns: &Self::Columns, _row_count: usize) -> bool { + true + } + + fn get_from_views<'a>(_views: &Self::Views<'a>, _index: usize) -> Self::Elems<'a> {} + + unsafe fn get_from_views_unchecked<'a>( + _views: &Self::Views<'a>, + _index: usize, + ) -> Self::Elems<'a> { + } + + fn constants(_columns: &Self::Columns) -> Self::ConstElems<'_> {} +} + +macro_rules! element_tuple { + ($arity:literal; $($t:ident : $idx:tt),+) => { + impl<$($t: InputElement),+> private::Sealed for ($($t,)+) {} + + impl<$($t: InputElement),+> ElementTuple for ($($t,)+) { + type Columns = ($(ArgColumn<$t>,)+); + type Views<'a> = ($($t::View<'a>,)+); + type Elems<'a> = ($($t::Elem<'a>,)+); + type ConstElems<'a> = ($(Option<$t::Elem<'a>>,)+); + + const ARITY: usize = $arity; + const DENSE_SAFE: bool = $($t::DENSE_SAFE &&)+ true; + const DECODE_FALLIBLE: bool = $($t::DECODE_FALLIBLE ||)+ false; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + $arity, + "expected {} argument dtypes, got {}", + $arity, + dtypes.len(), + ); + + $($t::validate(&dtypes[$idx])?;)+ + Ok(()) + } + + fn decode( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(($(ArgColumn::<$t>::decode(args.get($idx)?, ctx)?,)+)) + } + + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(( + $(match ArgColumn::<$t>::decode_null_tolerant(args.get($idx)?, ctx)? { + Some(column) => column, + None => return Ok(None), + },)+ + ))) + } + + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_> { + ($(columns.$idx.get(index),)+) + } + + fn per_row_views(columns: &Self::Columns) -> Option> { + Some(($($t::view(columns.$idx.per_row_column()?),)+)) + } + + fn view_lens_match( + views: &Self::Views<'_>, + row_count: usize, + ) -> bool { + $($t::view_len(&views.$idx) == row_count &&)+ true + } + + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { + $(columns.$idx.addresses_rows(row_count) &&)+ true + } + + fn get_from_views<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a> { + ($($t::get_from_view(&views.$idx, index),)+) + } + + unsafe fn get_from_views_unchecked<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a> { + // SAFETY: forwarded from this method's contract. + ($(unsafe { $t::get_from_view_unchecked(&views.$idx, index) },)+) + } + + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { + ($(columns.$idx.constant(),)+) + } + } + }; +} + +element_tuple!(1; A:0); +element_tuple!(2; A:0, B:1); +element_tuple!(3; A:0, B:1, C:2); +element_tuple!(4; A:0, B:1, C:2, D:3); +element_tuple!(5; A:0, B:1, C:2, D:3, E:4); +element_tuple!(6; A:0, B:1, C:2, D:3, E:4, F:5); +element_tuple!(7; A:0, B:1, C:2, D:3, E:4, F:5, G:6); +element_tuple!(8; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7); +element_tuple!(9; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8); +element_tuple!(10; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9); +element_tuple!(11; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10); +element_tuple!(12; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11); + +mod private { + pub trait Sealed {} +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/row/types/element/tuple/tests.rs new file mode 100644 index 00000000000..1a8efc3ecff --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/tuple/tests.rs @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_compute::lane_kernels::IndexedSource; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; + +use super::IndexedElementTuple; +use super::batch_constant; +use crate::IntoArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; +use crate::arrays::MaskedArray; +use crate::dtype::Nullability; +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; +use crate::validity::Validity; + +#[test] +fn test_unary_element_source_reads_one_tuple_per_row() { + // SAFETY: the only view has exactly three rows. + let source = unsafe { <(i32,)>::indexed_source((&[10, 20, 30][..],), 3) }; + assert_eq!(source.len(), 3); + + // SAFETY: index one is within the three-element source. + assert_eq!(unsafe { source.get_unchecked(1) }, (20,)); +} + +#[test] +fn test_binary_element_source_zips_views() { + // SAFETY: both views have exactly three rows. + let source = + unsafe { <(i32, i64)>::indexed_source((&[10, 20, 30][..], &[100, 200, 300][..]), 3) }; + assert_eq!(source.len(), 3); + + // SAFETY: index one is within the three-element source. + assert_eq!(unsafe { source.get_unchecked(1) }, (20, 200)); +} + +#[test] +fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> { + let child = ConstantArray::new(7_i64, 3).into_array(); + let masked = + MaskedArray::try_new(child, Validity::from_iter([true, false, true]))?.into_array(); + let filtered = masked.filter(Mask::from_iter([true, true, false]))?; + + let Some(constant) = batch_constant(&filtered) else { + vortex_bail!("filtered masked constant must remain batch-constant"); + }; + + assert!(constant.is::()); + Ok(()) +} + +#[test] +fn test_batch_constant_preserves_filtered_extension() -> VortexResult<()> { + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); + let extension = + ExtensionArray::new(ext_dtype, ConstantArray::new(7_i64, 3).into_array()).into_array(); + let filtered = extension.filter(Mask::from_iter([true, false, true]))?; + + let Some(constant) = batch_constant(&filtered) else { + vortex_bail!("filtered extension storage must remain batch-constant"); + }; + + assert_eq!(constant.dtype(), extension.dtype()); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/types/result.rs b/vortex-array/src/scalar_fn/row/types/result.rs index cda44802c40..cee24d50185 100644 --- a/vortex-array/src/scalar_fn/row/types/result.rs +++ b/vortex-array/src/scalar_fn/row/types/result.rs @@ -7,10 +7,6 @@ use vortex_error::VortexResult; use super::InitializedElement; -mod private { - pub trait Sealed {} -} - /// The result of writing one row: success or an immediate error. /// /// This trait is sealed; row functions choose one of its supplied implementations. @@ -79,3 +75,7 @@ impl SinkResult for VortexResult { self.map(|_| ()) } } + +mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs index 3e520718b99..421dfa93965 100644 --- a/vortex-array/src/scalar_fn/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -19,7 +19,7 @@ use crate::scalar_fn::OutputElement; /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; /// skip-invalid execution can omit invalid rows when -/// [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) is present. +/// [`skipped_rows_initializer`](Self::skipped_rows_initializer) returns an initializer. /// /// # Safety /// @@ -31,8 +31,12 @@ use crate::scalar_fn::OutputElement; /// [`WriteToken`](Self::WriteToken) that safe code cannot produce without initializing that exact /// row. Evidence for an uninitialized row **must not** be safely forgeable, reusable, or /// substitutable for another row. -/// - A present [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) **must** initialize -/// every row handed to it. +/// - An initializer returned by +/// [`skipped_rows_initializer`](Self::skipped_rows_initializer) **must** initialize every row +/// handed to it. +/// - `Self` and every borrowed [`Rows`](Self::Rows) view **must** remain safe to drop if decoding, +/// preparation, skipped-row initialization, or a row callback returns an error or unwinds. The +/// executor can abandon a sink after any prefix of rows. /// - [`finish`](Self::finish) **must** be sound once every visited callback returned its required /// token and the skipped-row initializer, when present, ran successfully. /// @@ -46,14 +50,6 @@ pub unsafe trait OutputSink: 'static + Sized { where Self: 'a; - /// The operation that initializes every output position before skip-invalid execution. - /// - /// `None` declines skip-invalid execution before input decoding or sink allocation. A present - /// initializer **must** leave a legal arbitrary value in every row. Encoding support as the - /// initializer's presence prevents a separate capability flag from disagreeing with a no-op - /// method. - const SKIPPED_ROWS_INITIALIZER: Option fn(&mut Self::Rows<'a>)> = None; - /// The place a row closure writes one row through, borrowed from the sink. type Row<'a> where @@ -67,6 +63,16 @@ pub unsafe trait OutputSink: 'static + Sized { /// unsafe when Rust cannot tie the token to the supplied row handle. type WriteToken: 'static; + /// The operation that initializes every output position before skip-invalid execution. + /// + /// `None` declines skip-invalid execution before input decoding or sink allocation. A present + /// initializer **must** leave a legal arbitrary value in every row. Encoding support as the + /// initializer's presence prevents a separate capability flag from disagreeing with a no-op + /// method. + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + None + } + /// The dtype of the column this sink builds, given the function options and input dtypes. /// /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the @@ -96,7 +102,8 @@ pub unsafe trait OutputSink: 'static + Sized { /// /// The executor must have completed every row callback successfully, and each callback must /// have returned this sink's [`WriteToken`](Self::WriteToken). When skipped rows are allowed, - /// [`SKIPPED_ROWS_INITIALIZER`](Self::SKIPPED_ROWS_INITIALIZER) must have run before traversal. + /// the initializer returned by + /// [`skipped_rows_initializer`](Self::skipped_rows_initializer) must have run before traversal. unsafe fn finish(self) -> VortexResult; } @@ -120,19 +127,9 @@ impl InitializedElement { /// /// # Safety /// - /// `row` must be the [`UninitElementSink`] row supplied to the current callback. The caller must - /// return the token from that callback. Using another row or returning the token from another - /// callback can cause undefined behavior. - /// - /// Safe code cannot construct initialization evidence: - /// - /// ```compile_fail,E0133 - /// use std::mem::MaybeUninit; - /// use vortex_array::scalar_fn::InitializedElement; - /// - /// let mut unrelated = MaybeUninit::::uninit(); - /// let _evidence = InitializedElement::write(&mut unrelated, 42); - /// ``` + /// `row` must be the [`UninitElementSink`] row supplied to the current callback. The caller + /// must return the token from that callback. Using another row or returning the token from + /// another callback can cause undefined behavior. #[inline] pub unsafe fn write(row: &mut MaybeUninit, value: T) -> Self { row.write(value); @@ -164,16 +161,17 @@ unsafe impl OutputSink for UninitElementSink { type Rows<'a> = &'a mut [MaybeUninit]; - - const SKIPPED_ROWS_INITIALIZER: Option fn(&mut Self::Rows<'a>)> = Some(|rows| { - for row in rows.iter_mut() { - row.write(T::default()); - } - }); - type Row<'a> = &'a mut MaybeUninit; type WriteToken = InitializedElement; + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + for row in rows.iter_mut() { + row.write(T::default()); + } + }) + } + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(T::element_dtype()) } diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index 04fc4563446..159d0bb5e2c 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The explicit [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) adapter for [`RowFn`]. +//! The [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) adapter for [`RowFn`]. //! //! The [`visitor`](super::visitor) module validates and executes the concrete row signature //! selected by dispatch. This module connects those visits to batch execution and exposes the @@ -9,13 +9,20 @@ use vortex_error::VortexResult; use vortex_mask::Mask; +use vortex_session::VortexSession; use super::row_fn::RowFn; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; +use crate::expr::Expression; +use crate::expr::union_child_validities; +use crate::scalar_fn::Arity; use crate::scalar_fn::BorrowedExecutionArgs; +use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::row::batch::Batch; use crate::scalar_fn::row::batch::KernelArgs; use crate::scalar_fn::row::batch::finalize_kernel_output; @@ -24,96 +31,60 @@ use crate::scalar_fn::row::visitor::ExecuteRows; use crate::scalar_fn::row::visitor::ExecuteValidRows; use crate::scalar_fn::row::visitor::PlanRows; -/// Implement the standard [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) behavior for one -/// concrete [`RowFn`]. -/// -/// This opt-in is explicit so a function that needs custom coercion, simplification, reduction, -/// formatting, or validity hooks can implement the vtable itself and delegate only execution to -/// [`execute_rows`](crate::scalar_fn::execute_rows). The standard adapter delegates serialization -/// to [`RowFn`], derives arity and child names from [`RowFn::ARG_NAMES`], propagates child validity, -/// and reports the function as strict. -#[macro_export] -macro_rules! impl_row_fn_vtable { - ($function:ty) => { - impl $crate::scalar_fn::ScalarFnVTable for $function { - type Options = <$function as $crate::scalar_fn::RowFn>::Options; - - fn id(&self) -> $crate::scalar_fn::ScalarFnId { - $crate::scalar_fn::RowFn::id(self) - } - - fn serialize( - &self, - options: &Self::Options, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult< - ::core::option::Option<::std::vec::Vec>, - > { - $crate::scalar_fn::RowFn::serialize(self, options) - } - - fn deserialize( - &self, - metadata: &[u8], - session: &$crate::scalar_fn::row_fn_macro_support::VortexSession, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult { - $crate::scalar_fn::RowFn::deserialize(self, metadata, session) - } - - fn arity(&self, _options: &Self::Options) -> $crate::scalar_fn::Arity { - $crate::scalar_fn::Arity::Exact( - <$function as $crate::scalar_fn::RowFn>::ARG_NAMES.len(), - ) - } - - fn child_name( - &self, - _options: &Self::Options, - child_index: usize, - ) -> $crate::scalar_fn::ChildName { - $crate::scalar_fn::ChildName::from( - <$function as $crate::scalar_fn::RowFn>::ARG_NAMES[child_index], - ) - } - - fn return_dtype( - &self, - options: &Self::Options, - args: &[$crate::dtype::DType], - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult<$crate::dtype::DType> { - $crate::scalar_fn::row_fn_return_dtype(self, options, args) - } - - fn execute( - &self, - options: &Self::Options, - args: &dyn $crate::scalar_fn::ExecutionArgs, - ctx: &mut $crate::ExecutionCtx, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult<$crate::ArrayRef> { - $crate::scalar_fn::execute_rows(self, options, args, ctx) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &$crate::expr::Expression, - ) -> $crate::scalar_fn::row_fn_macro_support::VortexResult< - ::core::option::Option<$crate::expr::Expression>, - > { - $crate::expr::union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - <$function as $crate::scalar_fn::RowFn>::FALLIBLE - } - } - }; +impl ScalarFnVTable for F { + type Options = F::Options; + + fn id(&self) -> ScalarFnId { + RowFn::id(self) + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + RowFn::serialize(self, options) + } + + fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { + RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(F::ARG_NAMES.len()) + } + + fn child_name(&self, _options: &Self::Options, child_index: usize) -> ChildName { + ChildName::from(F::ARG_NAMES[child_index]) + } + + fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { + row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + F::FALLIBLE + } } -/// Compute the return dtype for a [`RowFn`] without adopting its complete scalar-function vtable. +/// Compute the return dtype of a [`RowFn`] kernel without invoking its blanket vtable. pub fn row_fn_return_dtype( function: &F, options: &F::Options, @@ -124,11 +95,11 @@ pub fn row_fn_return_dtype( Ok(plan.result_dtype(args)) } -/// Execute a [`RowFn`] while preserving a caller-owned scalar-function vtable. +/// Execute a [`RowFn`] without using its blanket [`ScalarFnVTable`] implementation. /// -/// Existing vtables delegate here when they need row execution but retain custom hooks for other -/// capabilities. A function that needs only the standard hooks can use [`impl_row_fn_vtable`] to -/// generate its complete vtable. +/// A type cannot implement both [`RowFn`] and [`ScalarFnVTable`] because every `RowFn` receives the +/// standard vtable automatically. Existing vtables can keep their custom hooks on one type and +/// delegate row execution to a private `RowFn` kernel through this function. pub fn execute_rows( function: &F, options: &F::Options, @@ -217,25 +188,23 @@ mod tests { use vortex_session::registry::CachedId; use crate::dtype::DType; + use crate::scalar_fn::Arity; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::RowFn; use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; - struct Option; - struct Vec; - #[derive(Clone)] - struct ShadowedPrelude; + struct TestRowFn; - impl RowFn for ShadowedPrelude { + impl RowFn for TestRowFn { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &[]; fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("test.shadowed_prelude"); + static ID: CachedId = CachedId::new("test.row_fn_vtable"); *ID } @@ -249,11 +218,13 @@ mod tests { } } - crate::impl_row_fn_vtable!(ShadowedPrelude); - #[test] - fn adapter_macro_ignores_shadowed_prelude_types() { - _ = (Option, Vec); - _ = ScalarFnVTable::id(&ShadowedPrelude); + fn test_row_fn_implements_standard_vtable() { + let function = TestRowFn; + let options = EmptyOptions; + + assert_eq!(ScalarFnVTable::arity(&function, &options), Arity::Exact(0)); + assert!(ScalarFnVTable::is_strict(&function, &options)); + assert!(!ScalarFnVTable::is_fallible(&function, &options)); } } From 83037fe98ed0e51a1c8c2a34fa666c272c16430f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:08:16 -0400 Subject: [PATCH 086/160] Avoid cloning primitive comparison constants Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs | 5 +++-- .../src/scalar_fn/fns/binary/compare/primitive/operand.rs | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index a3f50f907fd..6ca461622e1 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::arrays::Constant; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::PType; @@ -131,8 +132,8 @@ fn use_columnar_comparison( // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. (PType::I64 | PType::F64, _) => true, - // LLVM vectorizes varying u64 inputs, but not the mixed-constant RowFn loop. - (PType::U64, _) => lhs.as_constant().is_some() || rhs.as_constant().is_some(), + // LLVM vectorizes per-row u64 inputs, but not the mixed-constant RowFn loop. + (PType::U64, _) => lhs.is::() || rhs.is::(), _ => false, }) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs index 55d81153b1f..1563b8e68e6 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs @@ -15,7 +15,7 @@ use crate::validity::Validity; /// A materialized primitive column, a non-null constant, or an all-null constant. pub(super) enum PrimitiveOperand { - /// A varying primitive column and its validity. + /// A per-row primitive column and its validity. Array { /// The materialized values. values: Buffer, From da911cdf8922c61550f5ce478df387ebf31e56ee Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:08:55 -0400 Subject: [PATCH 087/160] Update tensor L2 for the RowFn API Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/row.rs | 23 ++++++++++------------- vortex-tensor/src/scalar_fns/tests/row.rs | 2 -- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 877804535ee..72cbc636898 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -60,7 +60,7 @@ pub struct TensorRows { // unchecked access use the same stride and row width. unsafe impl InputElement for TensorRow { type Column = TensorRows; - type Varying<'a> = &'a TensorRows; + type View<'a> = &'a TensorRows; type Elem<'a> = &'a [T]; // Tensor storage is a fully materialized non-nullable primitive buffer, so the elements behind @@ -102,7 +102,7 @@ unsafe impl InputElement for TensorRow { vortex_ensure_eq!( stride, list_size, - "varying tensor row stride must equal its width, got {stride}", + "per-row tensor stride must equal its width, got {stride}", ); let Some(expected_elements) = rows.checked_mul(stride) else { vortex_bail!( @@ -132,34 +132,31 @@ unsafe impl InputElement for TensorRow { &column.elements.as_slice()[start..start + column.list_size] } - fn varying(column: &Self::Column) -> Self::Varying<'_> { + fn view(column: &Self::Column) -> Self::View<'_> { column } - fn varying_len(column: &Self::Varying<'_>) -> usize { - column.rows + fn view_len(view: &Self::View<'_>) -> usize { + view.rows } - fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] where Self: 'a, { - Self::get(column, index) + Self::get(view, index) } - unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] where Self: 'a, { - let start = index * column.stride; + let start = index * view.stride; // SAFETY: decode established one complete stored row for stride 0, or `rows` contiguous // `list_size`-element rows otherwise. The caller guarantees `index < rows`. unsafe { - std::slice::from_raw_parts( - column.elements.as_slice().as_ptr().add(start), - column.list_size, - ) + std::slice::from_raw_parts(view.elements.as_slice().as_ptr().add(start), view.list_size) } } } diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index 77248329ae5..b88a22838aa 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -57,8 +57,6 @@ impl RowFn for L1Norm { } } -vortex_array::impl_row_fn_vtable!(L1Norm); - fn l1_norm_row(row: &[T]) -> T { row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) } From de475458ac0ccab047b1748d3a71aba903f9fad2 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:00 -0400 Subject: [PATCH 088/160] Update tensor products for the RowFn API Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 2 -- vortex-tensor/src/scalar_fns/inner_product.rs | 2 -- vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs | 2 +- 3 files changed, 1 insertion(+), 5 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 9cccf05d17b..dba407ff3bc 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -160,8 +160,6 @@ impl RowFn for CosineSimilarity { } } -vortex_array::impl_row_fn_vtable!(CosineSimilarity); - impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 6f09e2b17f1..10edf5d040a 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -145,8 +145,6 @@ impl RowFn for InnerProduct { } } -vortex_array::impl_row_fn_vtable!(InnerProduct); - impl ScalarFnArrayVTable for InnerProduct { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs index 0e2687690a9..ecdbe23f835 100644 --- a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -490,7 +490,7 @@ fn vector_constant_matches_plain() -> VortexResult<()> { } /// Both literal and extension-wrapped constant storage reach the prepared row path. The probe -/// ensures that the literal query remains a batch constant instead of becoming a varying column. +/// ensures that the literal query remains a batch constant instead of becoming a per-row column. /// /// [`ConstantArray`]: vortex_array::arrays::ConstantArray #[test] From f56bf932be882652078472808807e97c358012b8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:43 -0400 Subject: [PATCH 089/160] Update spatial distance for RowFn views Signed-off-by: Connor Tsui --- vortex-spatial/src/scalar_fn/distance.rs | 2 -- vortex-spatial/src/scalar_fn/row.rs | 26 +++++++++++------------- 2 files changed, 12 insertions(+), 16 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index 751e402f118..6f6b7e7bd93 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -75,8 +75,6 @@ impl RowFn for SpatialDistance { } } -vortex_array::impl_row_fn_vtable!(SpatialDistance); - #[cfg(test)] mod tests { use vortex_array::ArrayRef; diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index 52d3818f73f..22237c2a57c 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -24,11 +24,11 @@ use crate::extension::is_native_geometry; /// column is *some* native geometry. pub struct GeometryRow; -// SAFETY: the varying view is the decoded geometry slice, and its reported length is that slice's -// length. +// SAFETY: [`view`](InputElement::view) returns the decoded geometry slice and +// [`view_len`](InputElement::view_len) reports that slice's exact length. unsafe impl InputElement for GeometryRow { type Column = Vec>; - type Varying<'a> = &'a [Geometry]; + type View<'a> = &'a [Geometry]; type Elem<'a> = &'a Geometry; // A geometry row is decoded from its coordinate storage, which behind a null row holds arbitrary @@ -53,30 +53,28 @@ unsafe impl InputElement for GeometryRow { &column[index] } - fn varying(column: &Self::Column) -> Self::Varying<'_> { + fn view(column: &Self::Column) -> Self::View<'_> { column.as_slice() } - fn varying_len(column: &Self::Varying<'_>) -> usize { - column.len() + fn view_len(view: &Self::View<'_>) -> usize { + view.len() } - fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a Geometry + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry where Self: 'a, { - &column[index] + &view[index] } - unsafe fn get_varying_unchecked<'a>( - column: &Self::Varying<'a>, - index: usize, - ) -> &'a Geometry + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry where Self: 'a, { - // SAFETY: forwarded from this method's contract. - unsafe { column.get_unchecked(index) } + // SAFETY: The caller established that `index` is below the slice length returned by + // `view_len` for this exact view. + unsafe { view.get_unchecked(index) } } /// Null rows decode to a placeholder geometry that the branch-and-skip row loop never reads. From bae08f6c47a08f95edf237e935a71b03e0146b86 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:45 -0400 Subject: [PATCH 090/160] Use the blanket RowFn vtable for spatial predicates Signed-off-by: Connor Tsui --- vortex-spatial/src/scalar_fn/contains.rs | 2 -- vortex-spatial/src/scalar_fn/intersects.rs | 2 -- 2 files changed, 4 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 9376ceaeb10..10fdc961237 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -93,8 +93,6 @@ impl RowFn for SpatialContains { } } -vortex_array::impl_row_fn_vtable!(SpatialContains); - /// Per-batch state for the contains row kernel: the prepared form of whichever operand is /// constant for the batch. `None` marks an operand that varies by row. struct ConstOperands { diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index 6f04e86b14a..6edc9f1bc66 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -84,8 +84,6 @@ impl RowFn for SpatialIntersects { } } -vortex_array::impl_row_fn_vtable!(SpatialIntersects); - /// Per-batch state for the intersects row kernel: the bounding rect of each operand that is /// constant for the batch. /// From 1d9fba694ae9c8b29c8479d8fe0dc4f95a25eac0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:47 -0400 Subject: [PATCH 091/160] Use the blanket RowFn vtable in benchmarks Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 7 +------ vortex-array/benches/strict_validity.rs | 2 -- 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index bff6a29d2da..df9bbb5aa8c 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -24,7 +24,6 @@ use vortex_array::scalar_fn::OutputSink; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -168,10 +167,6 @@ impl RowFn for RowSinkWrappingAdd { } } -vortex_array::impl_row_fn_vtable!(RowWrappingAdd); -vortex_array::impl_row_fn_vtable!(RowCheckedAdd); -vortex_array::impl_row_fn_vtable!(RowSinkWrappingAdd); - fn inputs() -> (ArrayRef, ArrayRef) { let lhs = (0..ROWS) .map(|index| index as i64) @@ -208,7 +203,7 @@ fn nullable_inputs() -> (ArrayRef, ArrayRef) { fn bench_row_fn(bencher: Bencher, function: F, make_inputs: fn() -> (ArrayRef, ArrayRef)) where - F: RowFn + ScalarFnVTable, + F: RowFn, { bencher .with_inputs(make_inputs) diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs index e0a8b6a565e..8064c8582e8 100644 --- a/vortex-array/benches/strict_validity.rs +++ b/vortex-array/benches/strict_validity.rs @@ -106,8 +106,6 @@ impl RowFn for LazyDouble { } } -vortex_array::impl_row_fn_vtable!(LazyDouble); - /// The same function, applying validity the way the adapter used to: materialize a mask first. #[derive(Clone)] struct EagerDouble; From 7b3c40451be1f082b737efb7eb488fa6213e82c9 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:22:56 -0400 Subject: [PATCH 092/160] Finish the RowFn umbrella cleanup Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/l2_norm.rs | 2 -- vortex-tensor/src/scalar_fns/row.rs | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index f8d7dc617b7..2a4fb0368cf 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -173,8 +173,6 @@ pub(super) struct L2NormMetadata { input_dtype: Option, } -vortex_array::impl_row_fn_vtable!(L2Norm); - impl ScalarFnArrayVTable for L2Norm { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 72cbc636898..cf56a6ce3d3 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -163,7 +163,7 @@ unsafe impl InputElement for TensorRow { /// Test-only probe recording which operands the last `prepare` step saw as batch-constant, so a /// test can assert its inputs took the stride-0 decode path rather than merely producing the right -/// values through the varying path. +/// values through the per-row path. #[cfg(test)] pub(crate) mod probe { use std::cell::Cell; From f7abcba6336dc4cc2666487be81386a366d05212 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 14:06:01 -0400 Subject: [PATCH 093/160] Define experimental RowFn contracts Signed-off-by: Connor Tsui --- vortex-array/Cargo.toml | 8 +- vortex-array/src/scalar_fn/mod.rs | 12 + vortex-array/src/scalar_fn/unstable/mod.rs | 9 + .../src/scalar_fn/unstable/row/mod.rs | 37 ++ .../src/scalar_fn/unstable/row/row_fn.rs | 81 ++++ .../unstable/row/types/element/bool.rs | 77 ++++ .../unstable/row/types/element/input.rs | 117 ++++++ .../unstable/row/types/element/mod.rs | 22 ++ .../unstable/row/types/element/output.rs | 20 + .../unstable/row/types/element/primitive.rs | 83 ++++ .../row/types/element/tuple/element_tuple.rs | 359 ++++++++++++++++++ .../row/types/element/tuple/indexed.rs | 140 +++++++ .../unstable/row/types/element/tuple/mod.rs | 13 + .../unstable/row/types/element/tuple/tests.rs | 47 +++ .../src/scalar_fn/unstable/row/types/mod.rs | 22 ++ .../scalar_fn/unstable/row/types/result.rs | 70 ++++ .../src/scalar_fn/unstable/row/types/sink.rs | 229 +++++++++++ .../scalar_fn/unstable/row/visitor/check.rs | 122 ++++++ .../src/scalar_fn/unstable/row/visitor/mod.rs | 14 + .../scalar_fn/unstable/row/visitor/plan.rs | 183 +++++++++ .../unstable/row/visitor/row_visitor.rs | 165 ++++++++ .../src/scalar_fn/unstable/row/vtable.rs | 192 ++++++++++ vortex/Cargo.toml | 2 + 23 files changed, 2023 insertions(+), 1 deletion(-) create mode 100644 vortex-array/src/scalar_fn/unstable/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/row_fn.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/input.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/output.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/result.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/sink.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/check.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/vtable.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 7c076bbb7b3..2af2eacf238 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -79,6 +79,8 @@ cudarc = ["dep:cudarc"] table-display = ["dep:tabled"] _test-harness = ["dep:goldenfile", "dep:rstest", "dep:rstest_reuse"] serde = ["dep:serde", "vortex-buffer/serde", "vortex-mask/serde"] +# Exposes experimental row-function APIs without compatibility guarantees. +unstable_row_fns = [] [dev-dependencies] divan = { workspace = true } @@ -90,7 +92,11 @@ rstest = { workspace = true } serde_json = { workspace = true } serde_test = { workspace = true } test-with = { workspace = true } -vortex-array = { path = ".", features = ["_test-harness", "table-display"] } +vortex-array = { path = ".", features = [ + "_test-harness", + "table-display", + "unstable_row_fns", +] } [[bench]] name = "aggregate_max" diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 003dbc2ca48..cbbb08e708f 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -6,6 +6,12 @@ //! This module contains the [`ScalarFnVTable`] trait and all built-in scalar function //! implementations. Expressions ([`crate::expr::Expression`]) reference scalar functions //! at each node. +//! +//! Use `unstable::row::RowFn` for strict functions whose natural kernel computes one row at a +//! time. It derives decoding, constant handling, null propagation, output construction, and +//! validity. This experimental API requires the `unstable_row_fns` feature and has no compatibility +//! guarantees. Implement [`ScalarFnVTable`] directly when the natural kernel is columnar, aliases +//! an input, or may produce null from otherwise valid inputs. use vortex_session::registry::Id; @@ -35,6 +41,12 @@ pub use options::*; mod signature; pub use signature::*; +#[cfg(feature = "unstable_row_fns")] +pub mod unstable; +#[cfg(not(feature = "unstable_row_fns"))] +#[allow(dead_code, unused_imports)] +pub(crate) mod unstable; + pub mod fns; pub mod internal; pub mod session; diff --git a/vortex-array/src/scalar_fn/unstable/mod.rs b/vortex-array/src/scalar_fn/unstable/mod.rs new file mode 100644 index 00000000000..6849c0bc661 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/mod.rs @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Experimental scalar-function APIs without compatibility guarantees. +//! +//! These APIs can change or disappear without a deprecation period. External users must enable +//! the corresponding `unstable_*` Cargo feature before importing them. + +pub mod row; diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs new file mode 100644 index 00000000000..1a420e7490d --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions computed one row at a time. +//! +//! This module is experimental and has no compatibility guarantees. External users must enable +//! the `unstable_row_fns` Cargo feature before importing it. +//! +//! Start with [`RowFn`] to define an operation and [`RowVisitor`] to select one of its output +//! capabilities. [`InputElement`] and [`ElementTuple`] describe how columns become row values. +//! [`OutputElement`] covers independent owned values, while [`OutputSink`] supports outputs that +//! need row handles or shared batch state. [`SinkResult`] describes how a sink-writing closure +//! reports errors. +//! +//! The internal executor owns decoding, batch constants, null propagation, allocation, and +//! validity. A visitor's prepare closure may derive shared state from constant operands once per +//! batch. + +mod row_fn; +pub use row_fn::RowFn; + +mod types; +pub use types::ElementTuple; +pub use types::IndexedElementTuple; +pub use types::InitializedElement; +pub use types::InputElement; +pub use types::OutputElement; +pub use types::OutputSink; +pub use types::SinkResult; +pub use types::UninitElementSink; + +mod visitor; +pub use visitor::RowVisitor; + +mod vtable; +pub use vtable::execute_rows; +pub use vtable::row_fn_return_dtype; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs new file mode 100644 index 00000000000..947c12a3d53 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions computed one row at a time. + +use std::fmt::Debug; +use std::fmt::Display; +use std::hash::Hash; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use super::visitor::RowVisitor; +use crate::dtype::DType; +use crate::scalar_fn::ScalarFnId; + +/// A scalar function computed one row at a time. +/// +/// Declare the argument names and use [`dispatch`](Self::dispatch) to choose concrete element and +/// sink types for each accepted dtype combination. +/// +/// Every `RowFn` receives the standard [`ScalarFnVTable`] implementation. A function that needs +/// custom scalar-function hooks instead implements `ScalarFnVTable` on its public type and +/// delegates row execution to a private `RowFn` kernel with [`row_fn_return_dtype`] and +/// [`execute_rows`]. Implement only `ScalarFnVTable` when the natural kernel is columnar. +/// +/// [`ScalarFnVTable`]: crate::scalar_fn::ScalarFnVTable +/// [`execute_rows`]: crate::scalar_fn::unstable::row::execute_rows +/// [`row_fn_return_dtype`]: crate::scalar_fn::unstable::row::row_fn_return_dtype +pub trait RowFn: 'static + Sized + Clone + Send + Sync { + /// Options for this function, or [`EmptyOptions`](crate::scalar_fn::EmptyOptions) for none. + type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash; + + /// The arguments in display order. Its length is the function's exact arity. + const ARG_NAMES: &'static [&'static str]; + + /// Whether any legal dispatch can raise a semantic error. + /// + /// The framework checks this at compile time for every fallible dispatched element or result. + /// A conservative `true` is allowed when only some dtype choices are fallible. + /// + /// [`OutputSink`](crate::scalar_fn::unstable::row::OutputSink) lifecycle methods only report + /// incidental execution failures. Semantic sink errors must come from the row callback. + /// + /// Semantic errors are defined by + /// [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible). + const FALLIBLE: bool = false; + + /// Returns the ID of the scalar function. + fn id(&self) -> ScalarFnId; + + /// Serialize this function's options, or return `None` when the function is not serializable. + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + _ = options; + Ok(None) + } + + /// Restore options written by [`serialize`](Self::serialize). + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + vortex_bail!("Expression {} is not deserializable", self.id()) + } + + /// Choose element types for these input dtypes and visit the framework with them. + /// + /// Plan time and run time both call this method, so the choice **must** be a pure function of + /// `options` and `args`. The framework rejects a change to the derived nullable execution + /// policy before row execution. It cannot compare the remaining types, preparation values, or + /// closure behavior, so those must also remain stable. Cross-argument dtype validation belongs + /// here. + fn dispatch>( + &self, + options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult; +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs new file mode 100644 index 00000000000..5aedd7fca4d --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::validity::Validity; + +// SAFETY: the per-row view is a bit buffer, and its reported length is the buffer length. +unsafe impl InputElement for bool { + type Column = BitBuffer; + type View<'a> = &'a BitBuffer; + type Elem<'a> = bool; + + // Every bit of the buffer is readable, valid or not. + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + matches!(dtype, DType::Bool(_)), + "expected a Bool column, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_bit_buffer()) + } + + fn get(column: &Self::Column, index: usize) -> bool { + column.value(index) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.len() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> bool + where + Self: 'a, + { + view.value(index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> bool + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { view.value_unchecked(index) } + } +} + +impl OutputElement for bool { + fn element_dtype() -> DType { + DType::Bool(Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + // `From>` uses the bulk bit-packing path. + BoolArray::new(BitBuffer::from(values), Validity::NonNullable).into_array() + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs new file mode 100644 index 00000000000..a605f1d8ba4 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; + +/// An element type that can be read row-wise out of an input column. +/// +/// # Safety +/// +/// For every view returned by [`view`](Self::view), every index below +/// [`view_len`](Self::view_len) **must** satisfy the safety contract of +/// [`get_from_view_unchecked`](Self::get_from_view_unchecked). Shared execution relies on this +/// proof to perform unchecked reads after one pre-loop length check. +pub unsafe trait InputElement: 'static { + /// The decoded column representation supporting `O(1)` row access. + type Column; + + /// The view of a per-row decoded column read by the hot row loop. + /// + /// This may borrow a cheaper representation than [`Column`](Self::Column). Primitive elements, + /// for example, expose a slice so its pointer and length are loop invariants rather than + /// re-reading a [`Buffer`](vortex_buffer::Buffer) descriptor for every row. + type View<'a>; + + /// The borrowed element value handed to a row closure. + type Elem<'a>; + + /// Whether every dense decode and access path tolerates rows that are null in the input. + /// + /// Arrays only guarantee payloads for valid rows. This is `false` when a null row's stored + /// offset or pointer may not address anything, and `true` only when [`decode`](Self::decode), + /// [`get`](Self::get), [`view`](Self::view), [`view_len`](Self::view_len), and + /// [`get_from_view`](Self::get_from_view) remain safe and correct for null rows. + /// + /// Dense execution requires this of every argument; otherwise the row layer executes only + /// valid rows. + /// + /// Dense execution can pass unspecified values from null rows. The closure must be total over + /// every stored value: it cannot panic or cause side effects beyond its declared output. + const DENSE_SAFE: bool = false; + + /// Whether [`decode`](Self::decode) can fail on _legal_ input data. + /// + /// This excludes infrastructural failures such as IO or allocation. Set it when legal input may + /// contain a value that the decoder rejects. + const DECODE_FALLIBLE: bool = true; + + /// Validate that `dtype` is an acceptable input column dtype for this element type. + fn validate(dtype: &DType) -> VortexResult<()>; + + /// Decode `array` into its column representation. Called once per batch. + /// + /// Hoist dtype checks, downcasts, and other batch-invariant work into this method. + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Decode `array` _without_ assuming every row is valid, or `Ok(None)` when this element + /// cannot for this particular array. + /// + /// An element with [`DENSE_SAFE`](Self::DENSE_SAFE) set **should not** override this: its + /// ordinary decode already tolerates null payloads, so the default is already correct and an + /// override just restates it. Overriding is for an element that is _not_ dense-safe but can + /// still write an arbitrary placeholder into null slots; the caller guarantees + /// [`get`](Self::get) is never called for such a row. The skip-invalid strategy uses this + /// representation to avoid filtering the input. + /// + /// Return `Ok(None)` rather than an error when an array has no null-tolerant decode; the + /// batch execution falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if Self::DENSE_SAFE { + Self::decode(array, ctx).map(Some) + } else { + Ok(None) + } + } + + /// Read the element at `index`, the one function called once per row. + /// + /// This must not repeat work that is constant across the batch; do that work in + /// [`decode`](Self::decode). + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>; + + /// Borrow the representation used when this argument varies within the batch. + /// + /// Called once before the hot loop. Constants do not use this view because the tuple adapter + /// keeps their one-row decoded representation separate. + fn view(column: &Self::Column) -> Self::View<'_>; + + /// Number of rows addressable through a [`View`](Self::View). + /// + /// Every index below this length must be valid for + /// [`get_from_view_unchecked`](Self::get_from_view_unchecked). + fn view_len(view: &Self::View<'_>) -> usize; + + /// Read one row from a [`View`](Self::View). + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a; + + /// Read one row without checking that `index` is in bounds. + /// + /// # Safety + /// + /// `index` must be less than [`view_len`](Self::view_len) for `view`. + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a, + { + Self::get_from_view(view, index) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs new file mode 100644 index 00000000000..37120798908 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The element types a row function can read and produce. +//! +//! [`InputElement::Elem`] may borrow from its decoded column. [`OutputElement`] is returned by an +//! owned row computation; runtime-shaped output uses an +//! [`OutputSink`](crate::scalar_fn::unstable::row::OutputSink). + +mod bool; + +mod input; +pub use input::InputElement; + +mod output; +pub use output::OutputElement; + +mod primitive; + +mod tuple; +pub use tuple::ElementTuple; +pub use tuple::IndexedElementTuple; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs new file mode 100644 index 00000000000..218a0ee5c6f --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs @@ -0,0 +1,20 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::ArrayRef; +use crate::dtype::DType; + +/// An owned row value that can be built into an all-valid column. +pub trait OutputElement: 'static + Sized { + /// The dtype of columns built from this element type. **Must** be non-nullable: nullability is + /// derived from the inputs by batch execution. + /// + /// Because this method takes no arguments, the dtype must be a property of the Rust type. Use + /// an [`OutputSink`] when the output dtype depends on function options or input dtypes. + /// + /// [`OutputSink`]: crate::scalar_fn::unstable::row::OutputSink + fn element_dtype() -> DType; + + /// Build a column from one value per row. Called once per batch. + fn build(values: Vec) -> ArrayRef; +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs new file mode 100644 index 00000000000..fb15030cc64 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::validity::Validity; + +// SAFETY: the per-row view is a native slice, and its reported length is the slice length. +unsafe impl InputElement for T { + type Column = Buffer; + type View<'a> = &'a [T]; + type Elem<'a> = T; + + // Every lane of the buffer holds a `T`, valid or not. + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + let expected = T::PTYPE; + let DType::Primitive(ptype, _) = dtype else { + vortex_bail!("expected a {expected} column, got {dtype}"); + }; + vortex_ensure_eq!( + *ptype, + expected, + "expected a {expected} column, got {dtype}" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_buffer::()) + } + + fn get(column: &Self::Column, index: usize) -> T { + column[index] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column.as_slice() + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.len() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> T + where + Self: 'a, + { + view[index] + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> T + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { *view.get_unchecked(index) } + } +} + +impl OutputElement for T { + fn element_dtype() -> DType { + DType::Primitive(T::PTYPE, Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + PrimitiveArray::new(values, Validity::NonNullable).into_array() + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs new file mode 100644 index 00000000000..f07ce3e8e19 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -0,0 +1,359 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Constant; +use crate::arrays::Extension; +use crate::arrays::Masked; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; + +/// One decoded input, collapsed to a single row when it is constant for the batch. +pub struct ArgColumn( + /// The decoded column, classified by whether it stores one value per row. + ArgColumnKind, +); + +enum ArgColumnKind { + PerRow(T::Column), + Constant(T::Column), +} + +impl ArgColumn { + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // An empty input has no row 0 to slice, and its row loop runs zero times either way. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?))); + } + + Ok(Self(ArgColumnKind::PerRow(T::decode(array, ctx)?))) + } + + fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + // Batch execution short-circuits null constants before selecting a strategy, so a + // constant reaching this path is non-null and can use the ordinary decode. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Some(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?)))); + } + + Ok(T::decode_null_tolerant(array, ctx)? + .map(ArgColumnKind::PerRow) + .map(Self)) + } + + fn get(&self, index: usize) -> T::Elem<'_> { + match &self.0 { + ArgColumnKind::PerRow(column) => T::get(column, index), + ArgColumnKind::Constant(column) => T::get(column, 0), + } + } + + fn per_row_column(&self) -> Option<&T::Column> { + match &self.0 { + ArgColumnKind::PerRow(column) => Some(column), + ArgColumnKind::Constant(_) => None, + } + } + + fn addresses_rows(&self, row_count: usize) -> bool { + // A constant is always read at index zero, so it addresses any batch length. + match &self.0 { + ArgColumnKind::PerRow(column) => T::view_len(&T::view(column)) == row_count, + ArgColumnKind::Constant(_) => true, + } + } + + fn constant(&self) -> Option> { + match &self.0 { + ArgColumnKind::PerRow(_) => None, + ArgColumnKind::Constant(column) => Some(T::get(column, 0)), + } + } +} + +/// Return the batch-constant array, looking through masked and extension wrappers. +/// +/// Batch execution owns mask validity, so a masked constant may expose its constant child here. An +/// extension over constant storage remains wrapped to preserve its extension dtype. +pub fn batch_constant(array: &ArrayRef) -> Option { + if array.is::() { + return Some(array.clone()); + } + + if let Some(masked) = array.as_opt::() { + return Some(masked.child().clone()).filter(|child| child.is::()); + } + + array + .as_opt::() + .is_some_and(|ext| ext.storage_array().is::()) + .then(|| array.clone()) +} + +/// Typed argument tuples for arities zero through twelve. +/// +/// This trait is sealed; add a new row representation by implementing [`InputElement`] and placing +/// it in one of the supplied tuples. +pub trait ElementTuple: 'static + private::Sealed { + /// The decoded column representations. + type Columns; + + /// Borrowed views of decoded columns when every argument stores one value per row. + type Views<'a>; + + /// The borrowed row of element values. + type Elems<'a>; + + /// The batch-constant element values: [`Elems`](Self::Elems) with every argument wrapped in + /// `Option`. + /// + /// `Some` carries the value of a batch-constant argument; `None` marks a per-row argument. A + /// [`RowVisitor`] passes these values to its prepare closure so constant work can leave the row + /// loop. + /// + /// [`RowVisitor`]: crate::scalar_fn::unstable::row::RowVisitor + type ConstElems<'a>; + + /// The number of arguments. + const ARITY: usize; + + /// Whether every argument is [`InputElement::DENSE_SAFE`]. + const DENSE_SAFE: bool; + + /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. + const DECODE_FALLIBLE: bool; + + /// Validate the input dtypes, including that `dtypes` has exactly `ARITY` entries. + /// + /// The expression layer checks arity when it builds a call, but callers can invoke + /// [`ScalarFnVTable::return_dtype`] directly. This boundary therefore checks it again. + /// + /// [`ScalarFnVTable::return_dtype`]: crate::scalar_fn::ScalarFnVTable::return_dtype + fn validate(dtypes: &[DType]) -> VortexResult<()>; + + /// Decode every input column once. Called once per batch. + fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Decode every input column once while tolerating null rows. + /// + /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid + /// strategy calls this once per batch. + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; + + /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; + + /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// + /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple + /// gives the optimizer ordinary contiguous column access without a per-row constant check. + fn per_row_views(columns: &Self::Columns) -> Option>; + + /// Whether every view contains exactly `row_count` rows. + fn view_lens_match(views: &Self::Views<'_>, row_count: usize) -> bool; + + /// Whether every per-row argument contains exactly `row_count` rows. + /// + /// This provides the same guarantee as [`view_lens_match`](Self::view_lens_match) when + /// [`per_row_views`](Self::per_row_views) declines a mixed per-row and batch-constant tuple. A + /// batch constant is exempt because it was collapsed to one row. + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; + + /// Read one row from borrowed views. + fn get_from_views<'a>(views: &Self::Views<'a>, index: usize) -> Self::Elems<'a>; + + /// Read one row from borrowed views without checking bounds. + /// + /// # Safety + /// + /// `index` must be in bounds for every column. + unsafe fn get_from_views_unchecked<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a>; + + /// Read the batch-constant elements out of the decoded columns. Called once per batch. + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; +} + +impl private::Sealed for () {} + +impl ElementTuple for () { + type Columns = (); + type Views<'a> = (); + type Elems<'a> = (); + type ConstElems<'a> = (); + + const ARITY: usize = 0; + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + 0, + "expected 0 argument dtypes, got {}", + dtypes.len(), + ); + Ok(()) + } + + fn decode(_args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(()) + } + + fn decode_null_tolerant( + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(())) + } + + fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} + + fn per_row_views(_columns: &Self::Columns) -> Option> { + Some(()) + } + + fn view_lens_match(_views: &Self::Views<'_>, _row_count: usize) -> bool { + true + } + + fn decoded_lens_match(_columns: &Self::Columns, _row_count: usize) -> bool { + true + } + + fn get_from_views<'a>(_views: &Self::Views<'a>, _index: usize) -> Self::Elems<'a> {} + + unsafe fn get_from_views_unchecked<'a>( + _views: &Self::Views<'a>, + _index: usize, + ) -> Self::Elems<'a> { + } + + fn constants(_columns: &Self::Columns) -> Self::ConstElems<'_> {} +} + +macro_rules! element_tuple { + ($arity:literal; $($t:ident : $idx:tt),+) => { + impl<$($t: InputElement),+> private::Sealed for ($($t,)+) {} + + impl<$($t: InputElement),+> ElementTuple for ($($t,)+) { + type Columns = ($(ArgColumn<$t>,)+); + type Views<'a> = ($($t::View<'a>,)+); + type Elems<'a> = ($($t::Elem<'a>,)+); + type ConstElems<'a> = ($(Option<$t::Elem<'a>>,)+); + + const ARITY: usize = $arity; + const DENSE_SAFE: bool = $($t::DENSE_SAFE &&)+ true; + const DECODE_FALLIBLE: bool = $($t::DECODE_FALLIBLE ||)+ false; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + $arity, + "expected {} argument dtypes, got {}", + $arity, + dtypes.len(), + ); + + $($t::validate(&dtypes[$idx])?;)+ + Ok(()) + } + + fn decode( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(($(ArgColumn::<$t>::decode(args.get($idx)?, ctx)?,)+)) + } + + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(( + $(match ArgColumn::<$t>::decode_null_tolerant(args.get($idx)?, ctx)? { + Some(column) => column, + None => return Ok(None), + },)+ + ))) + } + + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_> { + ($(columns.$idx.get(index),)+) + } + + fn per_row_views(columns: &Self::Columns) -> Option> { + Some(($($t::view(columns.$idx.per_row_column()?),)+)) + } + + fn view_lens_match( + views: &Self::Views<'_>, + row_count: usize, + ) -> bool { + $($t::view_len(&views.$idx) == row_count &&)+ true + } + + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { + $(columns.$idx.addresses_rows(row_count) &&)+ true + } + + fn get_from_views<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a> { + ($($t::get_from_view(&views.$idx, index),)+) + } + + unsafe fn get_from_views_unchecked<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a> { + // SAFETY: forwarded from this method's contract. + ($(unsafe { $t::get_from_view_unchecked(&views.$idx, index) },)+) + } + + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { + ($(columns.$idx.constant(),)+) + } + } + }; +} + +element_tuple!(1; A:0); +element_tuple!(2; A:0, B:1); +element_tuple!(3; A:0, B:1, C:2); +element_tuple!(4; A:0, B:1, C:2, D:3); +element_tuple!(5; A:0, B:1, C:2, D:3, E:4); +element_tuple!(6; A:0, B:1, C:2, D:3, E:4, F:5); +element_tuple!(7; A:0, B:1, C:2, D:3, E:4, F:5, G:6); +element_tuple!(8; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7); +element_tuple!(9; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8); +element_tuple!(10; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9); +element_tuple!(11; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10); +element_tuple!(12; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11); + +mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs new file mode 100644 index 00000000000..5f1ddbd1d90 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_compute::lane_kernels::IndexedSource; +use vortex_compute::lane_kernels::LaneZip; + +use super::ElementTuple; +use crate::scalar_fn::unstable::row::InputElement; + +/// An argument tuple that supports a validated dense indexed traversal. +/// +/// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's +/// unchecked view access after batch execution validates every decoded column length once. +pub trait IndexedElementTuple: ElementTuple { + /// The source shared execution uses for a dense all-per-row loop. + /// + /// Its length must be the common view length. For every valid index it must preserve row order, + /// return the same value as [`ElementTuple::get_from_views`], and uphold the unchecked read + /// contract of [`IndexedSource`]. + type Source<'a>: IndexedSource>; + + /// Build a source from views already validated to cover the complete batch. + /// + /// # Safety + /// + /// Every view in `views` **must** address exactly `row_count` rows. Violating this requirement + /// can make a safe lane kernel read outside a column's allocation. + unsafe fn indexed_source<'a>(views: Self::Views<'a>, row_count: usize) -> Self::Source<'a>; +} + +/// Indexed access to one element view. +pub struct ElementSource<'a, T: InputElement> { + view: T::View<'a>, +} + +impl<'a, T: InputElement> ElementSource<'a, T> { + fn new(view: T::View<'a>) -> Self { + Self { view } + } +} + +impl<'a, T: InputElement> IndexedSource for ElementSource<'a, T> { + type Item = T::Elem<'a>; + + fn len(&self) -> usize { + T::view_len(&self.view) + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the source length is the number of rows addressable by `view`, and the caller + // guarantees that `index` is below that length. + unsafe { T::get_from_view_unchecked(&self.view, index) } + } +} + +/// An indexed element source yielding the one-tuples expected by a unary row closure. +pub struct UnaryTupleSource(Source); + +impl IndexedSource for UnaryTupleSource { + type Item = (Source::Item,); + + fn len(&self) -> usize { + self.0.len() + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: forwarded from this method's contract. + (unsafe { self.0.get_unchecked(index) },) + } +} + +/// Indexed access to the views of an element tuple. +pub struct ElementTupleSource<'a, Args: ElementTuple> { + views: Args::Views<'a>, + row_count: usize, +} + +impl<'a, Args: ElementTuple> IndexedSource for ElementTupleSource<'a, Args> { + type Item = Args::Elems<'a>; + + fn len(&self) -> usize { + self.row_count + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the caller guarantees that `index` is below `row_count`. Batch execution checks + // that every view addresses exactly `row_count` rows before constructing this source. + unsafe { Args::get_from_views_unchecked(&self.views, index) } + } +} + +impl IndexedElementTuple for () { + type Source<'a> = ElementTupleSource<'a, ()>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, row_count: usize) -> Self::Source<'a> { + ElementTupleSource { views, row_count } + } +} + +impl IndexedElementTuple for (A,) { + type Source<'a> = UnaryTupleSource>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, _row_count: usize) -> Self::Source<'a> { + UnaryTupleSource(ElementSource::new(views.0)) + } +} + +impl IndexedElementTuple for (A, B) { + type Source<'a> = LaneZip, ElementSource<'a, B>>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, _row_count: usize) -> Self::Source<'a> { + LaneZip::new(ElementSource::new(views.0), ElementSource::new(views.1)) + } +} + +macro_rules! indexed_element_tuple { + ($($t:ident),+) => { + impl<$($t: InputElement),+> IndexedElementTuple for ($($t,)+) { + type Source<'a> = ElementTupleSource<'a, ($($t,)+)>; + + unsafe fn indexed_source<'a>( + views: Self::Views<'a>, + row_count: usize, + ) -> Self::Source<'a> { + ElementTupleSource { views, row_count } + } + } + }; +} + +indexed_element_tuple!(A, B, C); +indexed_element_tuple!(A, B, C, D); +indexed_element_tuple!(A, B, C, D, E); +indexed_element_tuple!(A, B, C, D, E, F); +indexed_element_tuple!(A, B, C, D, E, F, G); +indexed_element_tuple!(A, B, C, D, E, F, G, H); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K, L); diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs new file mode 100644 index 00000000000..e1431af901c --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Argument lists built from [`InputElement`](super::InputElement)s and their per-argument decode. + +mod element_tuple; +pub use element_tuple::ElementTuple; + +mod indexed; +pub use indexed::IndexedElementTuple; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs new file mode 100644 index 00000000000..2b4ca800a91 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; + +use super::element_tuple::batch_constant; +use crate::IntoArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; +use crate::arrays::MaskedArray; +use crate::dtype::Nullability; +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; +use crate::validity::Validity; + +#[test] +fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> { + let child = ConstantArray::new(7_i64, 3).into_array(); + let masked = + MaskedArray::try_new(child, Validity::from_iter([true, false, true]))?.into_array(); + let filtered = masked.filter(Mask::from_iter([true, true, false]))?; + + let Some(constant) = batch_constant(&filtered) else { + vortex_bail!("filtered masked constant must remain batch-constant"); + }; + + assert!(constant.is::()); + Ok(()) +} + +#[test] +fn test_batch_constant_preserves_filtered_extension() -> VortexResult<()> { + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); + let extension = + ExtensionArray::new(ext_dtype, ConstantArray::new(7_i64, 3).into_array()).into_array(); + let filtered = extension.filter(Mask::from_iter([true, false, true]))?; + + let Some(constant) = batch_constant(&filtered) else { + vortex_bail!("filtered extension storage must remain batch-constant"); + }; + + assert_eq!(constant.dtype(), extension.dtype()); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs new file mode 100644 index 00000000000..ce119f32915 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Input decoding and output construction for row functions. +//! +//! [`element`] defines the Rust values decoded from input columns and built into simple output +//! columns. [`sink`] handles outputs that need row handles or batch-wide state. [`result`] defines +//! the immediate and deferred outcomes returned by sink-writing row closures. + +mod element; +pub use element::ElementTuple; +pub use element::IndexedElementTuple; +pub use element::InputElement; +pub use element::OutputElement; + +mod result; +pub use result::SinkResult; + +mod sink; +pub use sink::InitializedElement; +pub use sink::OutputSink; +pub use sink::UninitElementSink; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/result.rs b/vortex-array/src/scalar_fn/unstable/row/types/result.rs new file mode 100644 index 00000000000..cba30750ccd --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/result.rs @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What a sink-writing row closure may return. + +use vortex_error::VortexResult; + +use super::InitializedElement; + +/// The result of writing one row: success or an immediate error. +/// +/// This trait is sealed; row functions choose one of its supplied implementations. +pub trait SinkResult: 'static + private::Sealed { + /// The [`OutputSink::WriteToken`](super::OutputSink::WriteToken) carried by a success. + type WriteToken: 'static; + + /// Whether this return type can carry an error. + const FALLIBLE: bool; + + /// Convert this row's outcome into immediate success or failure. + fn into_result(self) -> VortexResult<()>; +} + +impl private::Sealed for () {} + +impl SinkResult for () { + type WriteToken = (); + const FALLIBLE: bool = false; + + fn into_result(self) -> VortexResult<()> { + Ok(()) + } +} + +impl private::Sealed for InitializedElement {} + +impl SinkResult for InitializedElement { + type WriteToken = InitializedElement; + const FALLIBLE: bool = false; + + fn into_result(self) -> VortexResult<()> { + Ok(()) + } +} + +impl private::Sealed for VortexResult<()> {} + +impl SinkResult for VortexResult<()> { + type WriteToken = (); + const FALLIBLE: bool = true; + + fn into_result(self) -> VortexResult<()> { + self + } +} + +impl private::Sealed for VortexResult {} + +impl SinkResult for VortexResult { + type WriteToken = InitializedElement; + const FALLIBLE: bool = true; + + fn into_result(self) -> VortexResult<()> { + self.map(|_| ()) + } +} + +mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs new file mode 100644 index 00000000000..77e103cf154 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -0,0 +1,229 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The column builders a row function can write its output into. + +use std::mem::MaybeUninit; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::unstable::row::OutputElement; + +/// A column allocated once per batch that a row closure writes into, one row at a time. +/// +/// A sink may use the function's `Options` and input dtypes to build a runtime-shaped output or own +/// shared batch state. The executor passes each row slot into an [`Fn`] closure, keeping mutable +/// state out of its capture. +/// +/// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; +/// skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an +/// initializer. +/// +/// # Errors +/// +/// Errors from [`sink_dtype`], [`with_capacity`], and [`finish`] are limited to incidental +/// execution failures such as allocation or array construction. A semantic error that depends on +/// the function's input values **must** be returned by the row callback through a fallible +/// [`SinkResult`]. Returning it from a sink lifecycle method hides it from [`RowFn::FALLIBLE`] and +/// can make optimizations such as dictionary push-down change the function's behavior. +/// +/// # Safety +/// +/// An implementation must uphold all of these requirements: +/// +/// - When [`row_count_matches`] returns `true`, every index in `0..row_count` **must** identify one +/// distinct row owned by this sink. +/// - A row must either be initialized before the callback or require a +/// [`WriteToken`] that safe code cannot produce without initializing that exact row. Evidence for +/// an uninitialized row **must not** be safely forgeable, reusable, or substitutable. +/// - An initializer returned by [`skipped_rows_initializer`] **must** initialize every row. +/// - `Self` and every borrowed [`Rows`] view **must** remain safe to drop if decoding, +/// preparation, skipped-row initialization, or a row callback returns an error or unwinds. The +/// executor can abandon a sink after any prefix of rows. +/// - [`finish`] **must** be sound once every visited callback returned its required token and the +/// skipped-row initializer, when present, ran successfully. +/// +/// The executor relies on these guarantees when it calls `finish`. +/// +/// [`Rows`]: Self::Rows +/// [`WriteToken`]: Self::WriteToken +/// [`finish`]: Self::finish +/// [`row_count_matches`]: Self::row_count_matches +/// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE +/// [`SinkResult`]: crate::scalar_fn::unstable::row::SinkResult +/// [`sink_dtype`]: Self::sink_dtype +/// [`skipped_rows_initializer`]: Self::skipped_rows_initializer +/// [`with_capacity`]: Self::with_capacity +pub unsafe trait OutputSink: 'static + Sized { + /// A loop-local view of all output rows. + /// + /// Borrowed once before execution so the sink's buffer descriptor and shape become loop + /// invariants rather than being re-read through `&mut Self` for every row. + type Rows<'a> + where + Self: 'a; + + /// The place a row closure writes one row through, borrowed from the sink. + type Row<'a> + where + Self: 'a; + + /// Proof that a successful row closure left its row handle initialized. + /// + /// Use `()` for initialized row handles. A sink exposing uninitialized storage uses a distinct + /// token returned after initialization. A sink that uses this token to justify unsafe code + /// **must** prevent safe construction that does not establish the invariant. Make construction + /// unsafe when Rust cannot tie the token to the supplied row handle. + type WriteToken: 'static; + + /// The operation that initializes every output position before skip-invalid execution. + /// + /// `Some(initializer)` enables skip-invalid execution and supplies the operation that prepares + /// output storage before callbacks run. The initializer **must** make every row safe to finish. + /// Callbacks overwrite valid rows, and batch execution masks skipped rows. + /// + /// `None` makes the executor fall back to filtering the inputs. + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + None + } + + /// The dtype of the column this sink builds, given the function options and input dtypes. + /// + /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the + /// result, and masks the null rows. + fn sink_dtype(options: &Options, args: &[DType]) -> VortexResult; + + /// Allocate a sink for `rows` rows of `dtype`, which is this sink's own + /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; + + /// Borrow all output rows for the hot loop. + fn rows(&mut self) -> Self::Rows<'_>; + + /// Whether every index in `0..row_count` is addressable through + /// [`row_unchecked`](Self::row_unchecked). + /// + /// Called once before the hot loop. Besides validating the sink contract, this gives the + /// optimizer the output bounds it needs to remove the bounds check hidden in each row accessor. + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool; + + /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. + /// + /// # Safety + /// + /// [`row_count_matches`](Self::row_count_matches) must have returned `true` for `rows` and the + /// same `row_count`, and `index` must be less than that `row_count`. + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; + + /// Finish into the built column, whose dtype **must** be this sink's + /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + /// + /// # Safety + /// + /// The executor must have completed every row callback successfully, and each callback must + /// have returned this sink's [`WriteToken`](Self::WriteToken). When skipped rows are allowed, + /// the initializer returned by + /// [`skipped_rows_initializer`](Self::skipped_rows_initializer) must have run before traversal. + unsafe fn finish(self) -> VortexResult; +} + +/// Proof that one uninitialized element row was initialized. +/// +/// The private field prevents safe construction without calling [`write`](Self::write): +/// +/// ```compile_fail,E0423 +/// use vortex_array::scalar_fn::unstable::row::InitializedElement; +/// +/// let _evidence = InitializedElement(()); +/// ``` +#[must_use = "return this token from the row closure to prove that it initialized the output"] +pub struct InitializedElement( + /// Private so constructing initialization evidence requires an unsafe operation. + (), +); + +impl InitializedElement { + /// Write `value` into an uninitialized row and return its proof token. + /// + /// # Safety + /// + /// `row` must be the [`UninitElementSink`] row supplied to the current callback. The caller + /// must return the token from that callback. Using another row or returning the token from + /// another callback can cause undefined behavior. + #[inline] + pub unsafe fn write(row: &mut MaybeUninit, value: T) -> Self { + row.write(value); + + Self(()) + } +} + +/// An element sink that leaves dense output uninitialized before the row loop. +/// +/// The row closure must return the [`InitializedElement`] from [`InitializedElement::write`] on +/// success. The token is zero-sized, so the proof adds no runtime row state. +/// +/// Skip-invalid execution initializes placeholders before omitting rows. Errors and unwinds are +/// safe because `values` keeps length zero until `finish`; `T: Copy` means initialized +/// spare-capacity elements require no destruction. +pub struct UninitElementSink { + /// Spare storage written in increasing row order. + values: Vec, + + /// The number of slots exposed to the row loop and initialized before finishing. + row_count: usize, +} + +// SAFETY: the row slice covers exactly the reserved spare-capacity range, so each accepted index +// names one distinct slot. `InitializedElement` cannot be constructed by safe code; its unsafe +// constructor writes the supplied slot and requires the caller to return that exact evidence. The +// skipped-row initializer writes `T::default()` into every slot before masked traversal. +unsafe impl OutputSink + for UninitElementSink +{ + type Rows<'a> = &'a mut [MaybeUninit]; + type Row<'a> = &'a mut MaybeUninit; + type WriteToken = InitializedElement; + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + for row in rows.iter_mut() { + row.write(T::default()); + } + }) + } + + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(T::element_dtype()) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self { + values: Vec::with_capacity(rows), + row_count: rows, + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.values.spare_capacity_mut()[..self.row_count] + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(mut self) -> VortexResult { + // SAFETY: the caller guarantees every slot in `0..row_count` was initialized, and + // `with_capacity` reserved every slot in that range. + unsafe { self.values.set_len(self.row_count) }; + + Ok(T::build(self.values)) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs new file mode 100644 index 00000000000..66ef90dd7b3 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Contract checks shared by planning and execution visits. +//! +//! Const assertions reject invalid generic visits during compilation. The validators compare a +//! selected visit with the input dtypes during planning and return its output dtype. + +use std::mem::needs_drop; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::dtype::DType; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::SinkResult; + +/// Assert the no-drop contract that makes partially initialized output safe to abandon on unwind. +pub(in crate::scalar_fn::unstable::row) const fn assert_owned_output_needs_no_drop() { + assert!( + !needs_drop::(), + "owned row outputs must not require drop glue" + ); +} + +/// Assert that the input arity and decode fallibility match the function-wide declarations. +const fn assert_input_visit_contract() { + assert!( + Args::ARITY == F::ARG_NAMES.len(), + "the visited argument tuple must have the arity declared by RowFn::ARG_NAMES", + ); + // Dictionary pushdown treats an infallible function as safe to evaluate over values no code + // references, so every dispatch must fit the function-wide declaration. + assert!( + !Args::DECODE_FALLIBLE || F::FALLIBLE, + "RowFn::FALLIBLE must be true when input decoding can fail", + ); +} + +/// Assert the input contract and that owned output values do not require drop glue. +pub(super) const fn assert_owned_visit_contract() +where + Function: RowFn, + Args: IndexedElementTuple, + Out: OutputElement, +{ + assert_input_visit_contract::(); + assert_owned_output_needs_no_drop::(); +} + +/// Assert that a sink visit obeys the input, fallibility, and deferred-error contracts. +pub(super) const fn assert_sink_visit_contract() +where + Function: RowFn, + Args: ElementTuple, + ApplyResult: SinkResult, +{ + assert_input_visit_contract::(); + assert!( + !ApplyResult::FALLIBLE || Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result can fail", + ); +} + +/// Assert the owned-output contract, fallibility declaration, and failure-evidence width bound. +pub(super) const fn assert_deferred_visit_contract() +where + Function: RowFn, + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, +{ + assert_owned_visit_contract::(); + assert!( + Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result defers failure evidence", + ); + assert!( + size_of::() <= size_of::(), + "failure evidence must be no wider than the value, or it bounds the vector width", + ); +} + +/// Validate the input dtypes and return the non-nullable dtype built by `Out`. +pub(super) fn validate_owned_visit( + dtypes: &[DType], +) -> VortexResult { + Args::validate(dtypes)?; + + let dtype = Out::element_dtype(); + vortex_ensure!( + !dtype.is_nullable(), + "row output elements must declare a non-nullable dtype, got {dtype}", + ); + + Ok(dtype) +} + +/// Validate the input dtypes and return the non-nullable dtype built by `Sink`. +pub(super) fn validate_sink_visit( + options: &Options, + dtypes: &[DType], +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, +{ + Args::validate(dtypes)?; + + let dtype = Sink::sink_dtype(options, dtypes)?; + vortex_ensure!( + !dtype.is_nullable(), + "row output sinks must declare a non-nullable dtype, got {dtype}", + ); + + Ok(dtype) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs new file mode 100644 index 00000000000..24d6e1f5b8d --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visits that plan or execute the concrete row signature selected by [`RowFn::dispatch`]. +//! +//! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch + +mod check; + +mod plan; +pub(super) use plan::PlanRows; + +mod row_visitor; +pub use row_visitor::RowVisitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs new file mode 100644 index 00000000000..31776d02d95 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The visitor that validates a concrete dispatch and plans its nullable execution. + +use std::marker::PhantomData; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; +use super::row_visitor::private; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::SinkResult; + +/// The plan-time visit that validates dtypes and derives the nullable execution policy. +pub(in crate::scalar_fn::unstable::row) struct PlanRows<'a, F: RowFn> { + /// The input dtypes for this plan. + dtypes: &'a [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'a F::Options, + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'a, F: RowFn> PlanRows<'a, F> { + pub(in crate::scalar_fn::unstable::row) fn new( + dtypes: &'a [DType], + options: &'a F::Options, + ) -> Self { + Self { + dtypes, + options, + function: PhantomData, + } + } +} + +impl private::Sealed for PlanRows<'_, F> {} + +impl RowVisitor for PlanRows<'_, F> { + type VisitResult = BatchPlan; + + fn visit_prepared( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + Ok(BatchPlan { + output_dtype: validate_owned_visit::(self.dtypes)?, + policy: RowPolicy::for_owned_output::(), + }) + } + + fn visit_prepared_into( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + Ok(BatchPlan { + output_dtype: validate_sink_visit::(self.options, self.dtypes)?, + policy: RowPolicy::for_sink::(), + }) + } + + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + _finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + Ok(BatchPlan { + output_dtype: validate_owned_visit::(self.dtypes)?, + policy: RowPolicy::for_deferred_output::(), + }) + } +} + +/// The execution policy and output dtype selected by a planning visit. +pub(in crate::scalar_fn::unstable::row) struct BatchPlan { + /// The non-nullable dtype built by the selected output capability. + pub(in crate::scalar_fn::unstable::row) output_dtype: DType, + + /// How this concrete dispatch executes nullable rows. + // TODO(connor)[RowFn]: The execution backend tracked by #9130 consumes this field. + #[allow(dead_code)] + pub(in crate::scalar_fn::unstable::row) policy: RowPolicy, +} + +impl BatchPlan { + /// Return the output dtype widened with strict input nullability. + pub(in crate::scalar_fn::unstable::row) fn result_dtype(self, args: &[DType]) -> DType { + let Self { + output_dtype, + policy: _, + } = self; + let nullability = + output_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); + + output_dtype.with_nullability(nullability) + } +} + +/// The nullable execution policy derived from one concrete dispatch. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(in crate::scalar_fn::unstable::row) enum RowPolicy { + /// Evaluate all rows and mask the result. + Dense, + + /// Evaluate all rows, retrying only valid rows if a deferred error is raised. + DenseWithRetry, + + /// Execute only valid rows, trying skip-invalid execution before filtering. + ValidOnly, +} + +impl RowPolicy { + /// The policy for an infallible owned output. + pub(in crate::scalar_fn::unstable::row) const fn for_owned_output() -> Self + { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + Self::Dense + } else { + Self::ValidOnly + } + } + + /// The policy for an owned output carrying batch-deferred failure evidence. + pub(in crate::scalar_fn::unstable::row) const fn for_deferred_output() + -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + Self::DenseWithRetry + } else { + Self::ValidOnly + } + } + + /// The policy for a sink-writing output. + pub(in crate::scalar_fn::unstable::row) const fn for_sink< + Args: ElementTuple, + ApplyResult: SinkResult, + >() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { + Self::Dense + } else { + Self::ValidOnly + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs new file mode 100644 index 00000000000..684c445ea8e --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::SinkResult; + +/// A planning or execution visit at concrete input and output types. +/// +/// Only the framework implements this trait. The `visit_prepared*` methods derive shared state +/// from constant arguments before visiting any rows. +pub trait RowVisitor: private::Sealed + Sized { + /// The framework result of visiting one concrete row signature. + /// + /// This is a batch plan or execution result, not the per-row `Out` returned by + /// [`RowVisitor::visit`] and [`RowVisitor::visit_deferred`]. + type VisitResult; + + /// Visit an infallible row computation that returns one independent output value. + /// + /// `apply` must be total over every stored element value: it must not panic or have side + /// effects. Dense execution can pass unspecified values from null rows. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of [`RowFn::ARG_NAMES`]. + /// - [`RowFn::FALLIBLE`] **must** be `true` when decoding `Args` can fail. + /// - `Out` **must not** require drop glue. + /// + /// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES + /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE + fn visit( + self, + apply: impl Fn(Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + self.visit_prepared::(|_| (), move |&(), args| apply(args)) + } + + /// The prepared form of [`visit`](Self::visit), with the same prerequisites. + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement; + + /// Visit a row computation that writes through a sink-provided row handle. + /// + /// `apply` must be total over every stored input value: it must not panic or cause side effects + /// other than writing the supplied row handle. Dense execution can pass unspecified values + /// from null rows. + /// + /// On success, `apply` must return the write token produced by writing the `Sink::Row` supplied + /// to that same invocation. It must not return evidence produced for another row, sink, or + /// unrelated local cell. Violating this requirement can make the unsafe + /// [`OutputSink::finish`] precondition false. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of [`RowFn::ARG_NAMES`]. + /// - [`RowFn::FALLIBLE`] **must** be `true` when decoding `Args` or computing the result can + /// fail. + /// + /// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES + /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE + fn visit_into( + self, + apply: impl Fn(Args::Elems<'_>, >::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + self.visit_prepared_into::( + |_| (), + move |&(), args, row| apply(args, row), + ) + } + + /// The prepared form of [`visit_into`](Self::visit_into), with the same prerequisites. + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>; + + /// Visit a row computation that returns an owned output and deferred failure evidence. + /// + /// `apply` must be total over every stored element value: it must not panic or have side + /// effects. Dense execution can pass unspecified values from null rows. + /// + /// `Fail` is OR-reduced across rows and handed to `finish_failure`. The value from + /// [`Default::default`] **must** mean success, including for an empty batch. The compiler + /// cannot check this semantic requirement. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of [`RowFn::ARG_NAMES`]. + /// - [`RowFn::FALLIBLE`] **must** be `true`. + /// - `Out` **must not** require drop glue. + /// - `Out` **must** be at least as wide as `Fail` so failure tracking does not reduce the + /// vector width. + /// + /// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES + /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE + fn visit_deferred( + self, + apply: impl Fn(Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + self.visit_prepared_deferred::( + |_| (), + move |&(), args| apply(args), + finish_failure, + ) + } + + /// The prepared form of [`visit_deferred`](Self::visit_deferred), with the same prerequisites. + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign; +} + +pub(super) mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs new file mode 100644 index 00000000000..d4d77562f65 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`ScalarFnVTable`] adapter for [`RowFn`]. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::row_fn::RowFn; +use super::visitor::PlanRows; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::expr::Expression; +use crate::expr::union_child_validities; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; + +impl ScalarFnVTable for F { + type Options = F::Options; + + fn id(&self) -> ScalarFnId { + RowFn::id(self) + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + RowFn::serialize(self, options) + } + + fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { + RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(F::ARG_NAMES.len()) + } + + fn child_name(&self, _options: &Self::Options, child_index: usize) -> ChildName { + ChildName::from(F::ARG_NAMES[child_index]) + } + + fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { + row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + F::FALLIBLE + } +} + +/// Compute the return dtype of a [`RowFn`] kernel without invoking its blanket vtable. +pub fn row_fn_return_dtype( + function: &F, + options: &F::Options, + args: &[DType], +) -> VortexResult { + ensure_arity(function, args.len())?; + + let plan = function.dispatch(options, args, PlanRows::::new(args, options))?; + + Ok(plan.result_dtype(args)) +} + +/// Execute a [`RowFn`] without using its blanket [`ScalarFnVTable`] implementation. +/// +/// A type cannot implement both [`RowFn`] and [`ScalarFnVTable`] because every `RowFn` receives the +/// standard vtable automatically. Existing vtables can keep their custom hooks on one type and +/// delegate row execution to a private `RowFn` kernel through this function. +pub fn execute_rows( + function: &F, + _options: &F::Options, + args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, +) -> VortexResult { + ensure_arity(function, args.num_inputs())?; + + // TODO(connor)[RowFn]: Replace this temporary error with the execution backend in #9129. + vortex_bail!( + "Row function {} does not yet have an execution backend", + RowFn::id(function) + ) +} + +/// Validate the number of arguments before calling user-defined dispatch code. +fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { + let expected = F::ARG_NAMES.len(); + vortex_ensure_eq!( + actual, + expected, + "row function {} must receive exactly {expected} input values, got {actual}", + RowFn::id(function), + ); + + Ok(()) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexError; + use vortex_error::VortexResult; + use vortex_session::registry::CachedId; + + use super::execute_rows; + use super::row_fn_return_dtype; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::dtype::DType; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnId; + use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::RowFn; + use crate::scalar_fn::unstable::row::RowVisitor; + + #[derive(Clone)] + struct IndexingRowFn; + + impl RowFn for IndexingRowFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.indexing_row_fn"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + _ = &args[0]; + + visitor.visit::<(i64,), i64>(|(value,)| value) + } + } + + #[test] + fn test_return_dtype_rejects_wrong_arity_before_dispatch() { + let error = row_fn_return_dtype(&IndexingRowFn, &EmptyOptions, &[]) + .expect_err("wrong arity must fail before dispatch"); + + assert_arity_error(error); + } + + #[test] + fn test_execute_rejects_wrong_arity_before_dispatch() { + let args = VecExecutionArgs::new(vec![], 0); + let mut ctx = array_session().create_execution_ctx(); + let error = execute_rows(&IndexingRowFn, &EmptyOptions, &args, &mut ctx) + .expect_err("wrong arity must fail before dispatch"); + + assert_arity_error(error); + } + + #[track_caller] + fn assert_arity_error(error: VortexError) { + assert!( + error + .to_string() + .contains("must receive exactly 1 input values, got 0"), + "unexpected error: {error}", + ); + } +} diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 6a2a840a500..57392ed627d 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -90,6 +90,8 @@ tokio = [ zstd = ["dep:vortex-zstd", "vortex-file?/zstd"] pretty = ["vortex-array/table-display"] serde = ["vortex-array/serde", "vortex-buffer/serde", "vortex-mask/serde"] +# Exposes experimental row-function APIs without compatibility guarantees. +unstable_row_fns = ["vortex-array/unstable_row_fns"] # This feature enabled unstable encodings for which we don't guarantee stability. unstable_encodings = [ "dep:vortex-tensor", From 4552c188bdf8c2c61515bc36d99fae1624b3ed55 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 17:01:01 -0400 Subject: [PATCH 094/160] Clarify RowFn retry preparation Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/unstable/row/mod.rs | 3 ++- .../src/scalar_fn/unstable/row/types/element/input.rs | 6 ++++-- .../unstable/row/types/element/tuple/element_tuple.rs | 7 +++++-- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index 1a420e7490d..f9a425ce6c3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -14,7 +14,8 @@ //! //! The internal executor owns decoding, batch constants, null propagation, allocation, and //! validity. A visitor's prepare closure may derive shared state from constant operands once per -//! batch. +//! row-kernel invocation. A dense deferred-error retry invokes the kernel again over filtered valid +//! rows. mod row_fn; pub use row_fn::RowFn; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index a605f1d8ba4..33a5991977a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -52,9 +52,11 @@ pub unsafe trait InputElement: 'static { /// Validate that `dtype` is an acceptable input column dtype for this element type. fn validate(dtype: &DType) -> VortexResult<()>; - /// Decode `array` into its column representation. Called once per batch. + /// Decode `array` into its column representation. /// - /// Hoist dtype checks, downcasts, and other batch-invariant work into this method. + /// The executor calls this once per row-kernel invocation. A dense deferred-error retry starts + /// another invocation over filtered valid rows. Hoist dtype checks, downcasts, and other + /// invocation-invariant work into this method. fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; /// Decode `array` _without_ assuming every row is valid, or `Ok(None)` when this element diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index f07ce3e8e19..b480e22b0c7 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -148,7 +148,9 @@ pub trait ElementTuple: 'static + private::Sealed { /// [`ScalarFnVTable::return_dtype`]: crate::scalar_fn::ScalarFnVTable::return_dtype fn validate(dtypes: &[DType]) -> VortexResult<()>; - /// Decode every input column once. Called once per batch. + /// Decode every input column once for one row-kernel invocation. + /// + /// A dense deferred-error retry starts another invocation over filtered valid rows. fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; /// Decode every input column once while tolerating null rows. @@ -192,7 +194,8 @@ pub trait ElementTuple: 'static + private::Sealed { index: usize, ) -> Self::Elems<'a>; - /// Read the batch-constant elements out of the decoded columns. Called once per batch. + /// Read the batch-constant elements out of the decoded columns once for one row-kernel + /// invocation. fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; } From ca222dcbcfedb7c2bbb13f06e112b308be7dc0ad Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 19:20:45 -0400 Subject: [PATCH 095/160] Tighten experimental RowFn documentation Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/mod.rs | 18 ++++---- .../src/scalar_fn/unstable/row/row_fn.rs | 35 ++++++--------- .../unstable/row/types/element/input.rs | 26 +++-------- .../row/types/element/tuple/element_tuple.rs | 7 +-- .../unstable/row/types/element/tuple/mod.rs | 5 ++- .../scalar_fn/unstable/row/types/result.rs | 5 ++- .../src/scalar_fn/unstable/row/types/sink.rs | 26 +++++------ .../scalar_fn/unstable/row/visitor/plan.rs | 5 ++- .../unstable/row/visitor/row_visitor.rs | 44 +++++-------------- .../src/scalar_fn/unstable/row/vtable.rs | 6 ++- 10 files changed, 68 insertions(+), 109 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index f9a425ce6c3..bcb3a008488 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -1,21 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Scalar functions computed one row at a time. +//! Experimental support for strict scalar functions computed one row at a time. //! //! This module is experimental and has no compatibility guarantees. External users must enable //! the `unstable_row_fns` Cargo feature before importing it. //! -//! Start with [`RowFn`] to define an operation and [`RowVisitor`] to select one of its output -//! capabilities. [`InputElement`] and [`ElementTuple`] describe how columns become row values. -//! [`OutputElement`] covers independent owned values, while [`OutputSink`] supports outputs that -//! need row handles or shared batch state. [`SinkResult`] describes how a sink-writing closure -//! reports errors. +//! A [`RowFn`] describes the typed operation while the framework owns columnar concerns such as +//! decoding, constant handling, null propagation, allocation, and validity. Its +//! [`RowFn::dispatch`] implementation uses a [`RowVisitor`] to select an [`ElementTuple`] and +//! either an [`OutputElement`] or [`OutputSink`] for each supported dtype combination. //! -//! The internal executor owns decoding, batch constants, null propagation, allocation, and -//! validity. A visitor's prepare closure may derive shared state from constant operands once per -//! row-kernel invocation. A dense deferred-error retry invokes the kernel again over filtered valid -//! rows. +//! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits +//! reduce compact failure evidence in that loop and retry only valid rows when null payloads may +//! have caused the failure. mod row_fn; pub use row_fn::RowFn; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 947c12a3d53..5264d0841ae 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Scalar functions computed one row at a time. +//! The [`RowFn`] contract for scalar functions whose natural kernel computes one row at a time. +//! +//! Implementations declare their arity and fallibility, then use [`RowFn::dispatch`] to select the +//! typed row signature for each supported dtype combination. Optional methods provide +//! serialization without putting persistence plumbing in the row kernel. use std::fmt::Debug; use std::fmt::Display; @@ -17,13 +21,9 @@ use crate::scalar_fn::ScalarFnId; /// A scalar function computed one row at a time. /// -/// Declare the argument names and use [`dispatch`](Self::dispatch) to choose concrete element and -/// sink types for each accepted dtype combination. -/// -/// Every `RowFn` receives the standard [`ScalarFnVTable`] implementation. A function that needs -/// custom scalar-function hooks instead implements `ScalarFnVTable` on its public type and -/// delegates row execution to a private `RowFn` kernel with [`row_fn_return_dtype`] and -/// [`execute_rows`]. Implement only `ScalarFnVTable` when the natural kernel is columnar. +/// Declare argument names and use [`dispatch`](Self::dispatch) to select element and output types. +/// Every implementation receives the standard [`ScalarFnVTable`]. A public type that needs custom +/// vtable hooks can delegate its row kernel through [`row_fn_return_dtype`] and [`execute_rows`]. /// /// [`ScalarFnVTable`]: crate::scalar_fn::ScalarFnVTable /// [`execute_rows`]: crate::scalar_fn::unstable::row::execute_rows @@ -35,16 +35,10 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// The arguments in display order. Its length is the function's exact arity. const ARG_NAMES: &'static [&'static str]; - /// Whether any legal dispatch can raise a semantic error. - /// - /// The framework checks this at compile time for every fallible dispatched element or result. - /// A conservative `true` is allowed when only some dtype choices are fallible. - /// - /// [`OutputSink`](crate::scalar_fn::unstable::row::OutputSink) lifecycle methods only report - /// incidental execution failures. Semantic sink errors must come from the row callback. + /// Whether any dispatch can raise a semantic error. /// - /// Semantic errors are defined by - /// [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible). + /// The framework checks dispatched element and result types. A conservative `true` is allowed. + /// Sink lifecycle errors are incidental; semantic sink errors come from the row callback. const FALLIBLE: bool = false; /// Returns the ID of the scalar function. @@ -67,11 +61,8 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// Choose element types for these input dtypes and visit the framework with them. /// - /// Plan time and run time both call this method, so the choice **must** be a pure function of - /// `options` and `args`. The framework rejects a change to the derived nullable execution - /// policy before row execution. It cannot compare the remaining types, preparation values, or - /// closure behavior, so those must also remain stable. Cross-argument dtype validation belongs - /// here. + /// Planning and execution both call this method, so its result **must** depend only on + /// `options` and `args`. Cross-argument dtype validation belongs here. fn dispatch>( &self, options: &Self::Options, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 33a5991977a..2c8ac416914 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -31,16 +31,9 @@ pub unsafe trait InputElement: 'static { /// Whether every dense decode and access path tolerates rows that are null in the input. /// - /// Arrays only guarantee payloads for valid rows. This is `false` when a null row's stored - /// offset or pointer may not address anything, and `true` only when [`decode`](Self::decode), - /// [`get`](Self::get), [`view`](Self::view), [`view_len`](Self::view_len), and - /// [`get_from_view`](Self::get_from_view) remain safe and correct for null rows. - /// - /// Dense execution requires this of every argument; otherwise the row layer executes only - /// valid rows. - /// - /// Dense execution can pass unspecified values from null rows. The closure must be total over - /// every stored value: it cannot panic or cause side effects beyond its declared output. + /// Arrays guarantee payloads only for valid rows. Set this to `true` only when every decode and + /// access method remains safe for null rows. Dense execution may pass unspecified values from + /// null rows to the row closure. const DENSE_SAFE: bool = false; /// Whether [`decode`](Self::decode) can fail on _legal_ input data. @@ -62,12 +55,8 @@ pub unsafe trait InputElement: 'static { /// Decode `array` _without_ assuming every row is valid, or `Ok(None)` when this element /// cannot for this particular array. /// - /// An element with [`DENSE_SAFE`](Self::DENSE_SAFE) set **should not** override this: its - /// ordinary decode already tolerates null payloads, so the default is already correct and an - /// override just restates it. Overriding is for an element that is _not_ dense-safe but can - /// still write an arbitrary placeholder into null slots; the caller guarantees - /// [`get`](Self::get) is never called for such a row. The skip-invalid strategy uses this - /// representation to avoid filtering the input. + /// Override this for a non-dense-safe representation that can still place safe placeholders in + /// null slots. The skip-invalid executor never reads those slots. /// /// Return `Ok(None)` rather than an error when an array has no null-tolerant decode; the /// batch execution falls back to the filter strategy. @@ -82,10 +71,7 @@ pub unsafe trait InputElement: 'static { } } - /// Read the element at `index`, the one function called once per row. - /// - /// This must not repeat work that is constant across the batch; do that work in - /// [`decode`](Self::decode). + /// Read one row without repeating batch-constant work from [`decode`](Self::decode). fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>; /// Borrow the representation used when this argument varies within the batch. diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index b480e22b0c7..00bc2d84c5b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -140,12 +140,7 @@ pub trait ElementTuple: 'static + private::Sealed { /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. const DECODE_FALLIBLE: bool; - /// Validate the input dtypes, including that `dtypes` has exactly `ARITY` entries. - /// - /// The expression layer checks arity when it builds a call, but callers can invoke - /// [`ScalarFnVTable::return_dtype`] directly. This boundary therefore checks it again. - /// - /// [`ScalarFnVTable::return_dtype`]: crate::scalar_fn::ScalarFnVTable::return_dtype + /// Validate the input dtypes and exact arity. fn validate(dtypes: &[DType]) -> VortexResult<()>; /// Decode every input column once for one row-kernel invocation. diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index e1431af901c..a2c143704a0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Argument lists built from [`InputElement`](super::InputElement)s and their per-argument decode. +//! Combines [`InputElement`](super::InputElement)s into typed row argument lists. +//! +//! [`ElementTuple`] owns decoding, constant classification, and row access for supported arities. +//! [`IndexedElementTuple`] adds the validated indexed source used by vectorizable dense loops. mod element_tuple; pub use element_tuple::ElementTuple; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/result.rs b/vortex-array/src/scalar_fn/unstable/row/types/result.rs index cba30750ccd..a3439f07ffe 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/result.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/result.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! What a sink-writing row closure may return. +//! Return types for sink-writing row closures. +//! +//! [`SinkResult`] lets the executor handle initialized sinks and sinks that require an +//! [`InitializedElement`] token, with either infallible or immediate-error callbacks. use vortex_error::VortexResult; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index 77e103cf154..cbe5acbdebe 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The column builders a row function can write its output into. +//! Output builders for row kernels that cannot return independent owned values. +//! +//! [`OutputSink`] allocates batch-wide state and lends one row handle to each callback. +//! [`UninitElementSink`] is the fixed-width implementation used when avoiding output +//! initialization matters. use std::mem::MaybeUninit; @@ -13,9 +17,8 @@ use crate::scalar_fn::unstable::row::OutputElement; /// A column allocated once per batch that a row closure writes into, one row at a time. /// -/// A sink may use the function's `Options` and input dtypes to build a runtime-shaped output or own -/// shared batch state. The executor passes each row slot into an [`Fn`] closure, keeping mutable -/// state out of its capture. +/// A sink may use function options and input dtypes to build a runtime-shaped output or own shared +/// batch state. The executor passes each row slot into an [`Fn`] closure. /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; /// skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an @@ -23,11 +26,9 @@ use crate::scalar_fn::unstable::row::OutputElement; /// /// # Errors /// -/// Errors from [`sink_dtype`], [`with_capacity`], and [`finish`] are limited to incidental -/// execution failures such as allocation or array construction. A semantic error that depends on -/// the function's input values **must** be returned by the row callback through a fallible -/// [`SinkResult`]. Returning it from a sink lifecycle method hides it from [`RowFn::FALLIBLE`] and -/// can make optimizations such as dictionary push-down change the function's behavior. +/// Lifecycle methods report only incidental failures such as allocation. A semantic error that +/// depends on input values **must** come from the row callback through a fallible [`SinkResult`], or +/// [`RowFn::FALLIBLE`] cannot protect optimizations such as dictionary push-down. /// /// # Safety /// @@ -53,9 +54,7 @@ use crate::scalar_fn::unstable::row::OutputElement; /// [`row_count_matches`]: Self::row_count_matches /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE /// [`SinkResult`]: crate::scalar_fn::unstable::row::SinkResult -/// [`sink_dtype`]: Self::sink_dtype /// [`skipped_rows_initializer`]: Self::skipped_rows_initializer -/// [`with_capacity`]: Self::with_capacity pub unsafe trait OutputSink: 'static + Sized { /// A loop-local view of all output rows. /// @@ -80,9 +79,8 @@ pub unsafe trait OutputSink: 'static + Sized { /// The operation that initializes every output position before skip-invalid execution. /// - /// `Some(initializer)` enables skip-invalid execution and supplies the operation that prepares - /// output storage before callbacks run. The initializer **must** make every row safe to finish. - /// Callbacks overwrite valid rows, and batch execution masks skipped rows. + /// `Some` enables skip-invalid execution. The initializer **must** make every row safe to + /// finish; callbacks overwrite valid rows and batch execution masks skipped rows. /// /// `None` makes the executor fall back to filtering the inputs. fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index 31776d02d95..43fe1d4c94e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The visitor that validates a concrete dispatch and plans its nullable execution. +//! Plans the concrete signature selected by [`RowFn::dispatch`]. +//! +//! [`PlanRows`] validates input and output dtypes, then records the output dtype and null-handling +//! policy that execution must reproduce. use std::marker::PhantomData; use std::ops::BitOrAssign; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 684c445ea8e..b8ad1bccfa2 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -14,7 +14,11 @@ use crate::scalar_fn::unstable::row::SinkResult; /// A planning or execution visit at concrete input and output types. /// /// Only the framework implements this trait. The `visit_prepared*` methods derive shared state -/// from constant arguments before visiting any rows. +/// from constant arguments before visiting any rows. Every visit verifies that the argument tuple +/// matches [`RowFn::ARG_NAMES`] and that fallible decoding agrees with [`RowFn::FALLIBLE`]. +/// +/// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES +/// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE pub trait RowVisitor: private::Sealed + Sized { /// The framework result of visiting one concrete row signature. /// @@ -27,16 +31,7 @@ pub trait RowVisitor: private::Sealed + Sized { /// `apply` must be total over every stored element value: it must not panic or have side /// effects. Dense execution can pass unspecified values from null rows. /// - /// # Prerequisites - /// - /// The framework checks these at compile time: - /// - /// - `Args::ARITY` **must** equal the length of [`RowFn::ARG_NAMES`]. - /// - [`RowFn::FALLIBLE`] **must** be `true` when decoding `Args` can fail. - /// - `Out` **must not** require drop glue. - /// - /// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES - /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE + /// The framework also verifies that `Out` does not require drop glue. fn visit( self, apply: impl Fn(Args::Elems<'_>) -> Out, @@ -69,16 +64,8 @@ pub trait RowVisitor: private::Sealed + Sized { /// unrelated local cell. Violating this requirement can make the unsafe /// [`OutputSink::finish`] precondition false. /// - /// # Prerequisites - /// - /// The framework checks these at compile time: - /// - /// - `Args::ARITY` **must** equal the length of [`RowFn::ARG_NAMES`]. - /// - [`RowFn::FALLIBLE`] **must** be `true` when decoding `Args` or computing the result can - /// fail. - /// - /// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES - /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE + /// A fallible `ApplyResult` requires + /// [`RowFn::FALLIBLE`](crate::scalar_fn::unstable::row::RowFn::FALLIBLE) to be `true`. fn visit_into( self, apply: impl Fn(Args::Elems<'_>, >::Row<'_>) -> ApplyResult, @@ -118,18 +105,9 @@ pub trait RowVisitor: private::Sealed + Sized { /// [`Default::default`] **must** mean success, including for an empty batch. The compiler /// cannot check this semantic requirement. /// - /// # Prerequisites - /// - /// The framework checks these at compile time: - /// - /// - `Args::ARITY` **must** equal the length of [`RowFn::ARG_NAMES`]. - /// - [`RowFn::FALLIBLE`] **must** be `true`. - /// - `Out` **must not** require drop glue. - /// - `Out` **must** be at least as wide as `Fail` so failure tracking does not reduce the - /// vector width. - /// - /// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES - /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE + /// [`RowFn::FALLIBLE`](crate::scalar_fn::unstable::row::RowFn::FALLIBLE) **must** be `true`, + /// `Out` must not require drop glue, and `Fail` must be no wider than `Out` so failure tracking + /// does not reduce the vector width. The framework checks each requirement. fn visit_deferred( self, apply: impl Fn(Args::Elems<'_>) -> (Out, Fail), diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index d4d77562f65..c774d3dea07 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The [`ScalarFnVTable`] adapter for [`RowFn`]. +//! Adapts [`RowFn`] implementations to the scalar-function interface. +//! +//! The blanket [`ScalarFnVTable`] implementation supplies common arity, validity, fallibility, and +//! execution behavior. [`row_fn_return_dtype`] and [`execute_rows`] expose the same planning and +//! execution paths to public vtables that delegate to a private row kernel. use vortex_error::VortexResult; use vortex_error::vortex_bail; From 1e8455d36c5d4f6a4ae2b87c96cb87d3368d5591 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:14:12 -0400 Subject: [PATCH 096/160] Address RowFn API review feedback Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/row_fn.rs | 5 +++-- .../unstable/row/types/element/input.rs | 4 ++-- .../src/scalar_fn/unstable/row/types/sink.rs | 17 ++++++++--------- .../src/scalar_fn/unstable/row/visitor/check.rs | 2 +- .../src/scalar_fn/unstable/row/vtable.rs | 2 ++ 5 files changed, 16 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 5264d0841ae..a47314b6c4f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -37,9 +37,10 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// Whether any dispatch can raise a semantic error. /// + /// See [`ScalarFnVTable::is_fallible`] for a more detailed explanation of semantic errors. + /// /// The framework checks dispatched element and result types. A conservative `true` is allowed. - /// Sink lifecycle errors are incidental; semantic sink errors come from the row callback. - const FALLIBLE: bool = false; + const FALLIBLE: bool; /// Returns the ID of the scalar function. fn id(&self) -> ScalarFnId; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 2c8ac416914..5302af15f74 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -34,13 +34,13 @@ pub unsafe trait InputElement: 'static { /// Arrays guarantee payloads only for valid rows. Set this to `true` only when every decode and /// access method remains safe for null rows. Dense execution may pass unspecified values from /// null rows to the row closure. - const DENSE_SAFE: bool = false; + const DENSE_SAFE: bool; /// Whether [`decode`](Self::decode) can fail on _legal_ input data. /// /// This excludes infrastructural failures such as IO or allocation. Set it when legal input may /// contain a value that the decoder rejects. - const DECODE_FALLIBLE: bool = true; + const DECODE_FALLIBLE: bool; /// Validate that `dtype` is an acceptable input column dtype for this element type. fn validate(dtype: &DType) -> VortexResult<()>; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index cbe5acbdebe..b935e45ffb6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -27,8 +27,8 @@ use crate::scalar_fn::unstable::row::OutputElement; /// # Errors /// /// Lifecycle methods report only incidental failures such as allocation. A semantic error that -/// depends on input values **must** come from the row callback through a fallible [`SinkResult`], or -/// [`RowFn::FALLIBLE`] cannot protect optimizations such as dictionary push-down. +/// depends on input values **must** come from the row callback through a fallible [`SinkResult`], +/// or [`RowFn::FALLIBLE`] cannot protect optimizations such as dictionary push-down. /// /// # Safety /// @@ -91,11 +91,10 @@ pub unsafe trait OutputSink: 'static + Sized { /// /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the /// result, and masks the null rows. - fn sink_dtype(options: &Options, args: &[DType]) -> VortexResult; + fn output_dtype(options: &Options, args: &[DType]) -> VortexResult; - /// Allocate a sink for `rows` rows of `dtype`, which is this sink's own - /// [`sink_dtype`](Self::sink_dtype). Called once per batch. - fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; + /// Allocate a sink for `rows` rows. + fn with_capacity(rows: usize) -> VortexResult; /// Borrow all output rows for the hot loop. fn rows(&mut self) -> Self::Rows<'_>; @@ -116,7 +115,7 @@ pub unsafe trait OutputSink: 'static + Sized { unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; /// Finish into the built column, whose dtype **must** be this sink's - /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + /// [`output_dtype`](Self::output_dtype). Called once per batch. /// /// # Safety /// @@ -193,11 +192,11 @@ unsafe impl OutputSink }) } - fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(T::element_dtype()) } - fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(rows: usize) -> VortexResult { Ok(Self { values: Vec::with_capacity(rows), row_count: rows, diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs index 66ef90dd7b3..a95f4ea749c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -112,7 +112,7 @@ where { Args::validate(dtypes)?; - let dtype = Sink::sink_dtype(options, dtypes)?; + let dtype = Sink::output_dtype(options, dtypes)?; vortex_ensure!( !dtype.is_nullable(), "row output sinks must declare a non-nullable dtype, got {dtype}", diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index c774d3dea07..b24ea3e66e6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -149,6 +149,8 @@ mod tests { const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.indexing_row_fn"); *ID From 209767dde6ed8440849841341c84fa23b7b4feae Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:09:43 -0400 Subject: [PATCH 097/160] Expose sink row counts Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/types/sink.rs | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index b935e45ffb6..e7623de6559 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -34,8 +34,7 @@ use crate::scalar_fn::unstable::row::OutputElement; /// /// An implementation must uphold all of these requirements: /// -/// - When [`row_count_matches`] returns `true`, every index in `0..row_count` **must** identify one -/// distinct row owned by this sink. +/// - Every index in `0..row_count(rows)` **must** identify one distinct row owned by this sink. /// - A row must either be initialized before the callback or require a /// [`WriteToken`] that safe code cannot produce without initializing that exact row. Evidence for /// an uninitialized row **must not** be safely forgeable, reusable, or substitutable. @@ -51,7 +50,7 @@ use crate::scalar_fn::unstable::row::OutputElement; /// [`Rows`]: Self::Rows /// [`WriteToken`]: Self::WriteToken /// [`finish`]: Self::finish -/// [`row_count_matches`]: Self::row_count_matches +/// [`row_count`]: Self::row_count /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE /// [`SinkResult`]: crate::scalar_fn::unstable::row::SinkResult /// [`skipped_rows_initializer`]: Self::skipped_rows_initializer @@ -99,19 +98,14 @@ pub unsafe trait OutputSink: 'static + Sized { /// Borrow all output rows for the hot loop. fn rows(&mut self) -> Self::Rows<'_>; - /// Whether every index in `0..row_count` is addressable through - /// [`row_unchecked`](Self::row_unchecked). - /// - /// Called once before the hot loop. Besides validating the sink contract, this gives the - /// optimizer the output bounds it needs to remove the bounds check hidden in each row accessor. - fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool; + /// The number of rows addressable through [`row_unchecked`](Self::row_unchecked). + fn row_count(rows: &Self::Rows<'_>) -> usize; /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. /// /// # Safety /// - /// [`row_count_matches`](Self::row_count_matches) must have returned `true` for `rows` and the - /// same `row_count`, and `index` must be less than that `row_count`. + /// `index` must be less than [`row_count`](Self::row_count) for `rows`. unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; /// Finish into the built column, whose dtype **must** be this sink's @@ -207,8 +201,8 @@ unsafe impl OutputSink &mut self.values.spare_capacity_mut()[..self.row_count] } - fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { - rows.len() == row_count + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { From e35d968a271776765ee705ae65db15b3e49ede85 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:21:15 -0400 Subject: [PATCH 098/160] Polish experimental RowFn internals Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/row_fn.rs | 3 +- .../unstable/row/types/element/input.rs | 5 ++++ .../unstable/row/types/element/output.rs | 4 +++ .../row/types/element/tuple/element_tuple.rs | 5 ++++ .../row/types/element/tuple/indexed.rs | 5 ++++ .../scalar_fn/unstable/row/visitor/check.rs | 2 +- .../scalar_fn/unstable/row/visitor/plan.rs | 28 +++++++------------ .../unstable/row/visitor/row_visitor.rs | 7 +++++ 8 files changed, 39 insertions(+), 20 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index a47314b6c4f..8a982c4fb37 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -37,7 +37,8 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// Whether any dispatch can raise a semantic error. /// - /// See [`ScalarFnVTable::is_fallible`] for a more detailed explanation of semantic errors. + /// See [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible) for a + /// more detailed explanation of semantic errors. /// /// The framework checks dispatched element and result types. A conservative `true` is allowed. const FALLIBLE: bool; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 5302af15f74..4b80c956981 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -1,6 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Typed decoding and row access for one input column. +//! +//! [`InputElement`] separates invocation-wide decoding from the checked and unchecked access paths +//! used by row kernels. + use vortex_error::VortexResult; use crate::ArrayRef; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs index 218a0ee5c6f..d1c75b7054a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs @@ -1,6 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Owned scalar values that can be collected into all-valid output columns. +//! +//! [`OutputElement`] describes fixed-dtype values returned independently by each row invocation. + use crate::ArrayRef; use crate::dtype::DType; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index 00bc2d84c5b..9fd5fed6ba4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -1,6 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Decoding and row access for tuples of input element types. +//! +//! [`ElementTuple`] combines per-column [`InputElement`] implementations, preserves batch +//! constants outside the hot loop, and supports row functions with up to twelve arguments. + use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs index 5f1ddbd1d90..d612d874935 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -1,6 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Validated indexed access to tuples of decoded input columns. +//! +//! [`IndexedElementTuple`] adapts row arguments to the lane-kernel interface after batch execution +//! proves that every input covers the requested row range. + use vortex_compute::lane_kernels::IndexedSource; use vortex_compute::lane_kernels::LaneZip; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs index a95f4ea749c..501ef896df1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -21,7 +21,7 @@ use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::SinkResult; /// Assert the no-drop contract that makes partially initialized output safe to abandon on unwind. -pub(in crate::scalar_fn::unstable::row) const fn assert_owned_output_needs_no_drop() { +pub(crate) const fn assert_owned_output_needs_no_drop() { assert!( !needs_drop::(), "owned row outputs must not require drop glue" diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index 43fe1d4c94e..e84e5e24e5d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -28,7 +28,7 @@ use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::SinkResult; /// The plan-time visit that validates dtypes and derives the nullable execution policy. -pub(in crate::scalar_fn::unstable::row) struct PlanRows<'a, F: RowFn> { +pub(crate) struct PlanRows<'a, F: RowFn> { /// The input dtypes for this plan. dtypes: &'a [DType], @@ -40,10 +40,7 @@ pub(in crate::scalar_fn::unstable::row) struct PlanRows<'a, F: RowFn> { } impl<'a, F: RowFn> PlanRows<'a, F> { - pub(in crate::scalar_fn::unstable::row) fn new( - dtypes: &'a [DType], - options: &'a F::Options, - ) -> Self { + pub(crate) fn new(dtypes: &'a [DType], options: &'a F::Options) -> Self { Self { dtypes, options, @@ -114,19 +111,19 @@ impl RowVisitor for PlanRows<'_, F> { } /// The execution policy and output dtype selected by a planning visit. -pub(in crate::scalar_fn::unstable::row) struct BatchPlan { +pub(crate) struct BatchPlan { /// The non-nullable dtype built by the selected output capability. - pub(in crate::scalar_fn::unstable::row) output_dtype: DType, + pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. // TODO(connor)[RowFn]: The execution backend tracked by #9130 consumes this field. #[allow(dead_code)] - pub(in crate::scalar_fn::unstable::row) policy: RowPolicy, + pub(crate) policy: RowPolicy, } impl BatchPlan { /// Return the output dtype widened with strict input nullability. - pub(in crate::scalar_fn::unstable::row) fn result_dtype(self, args: &[DType]) -> DType { + pub(crate) fn result_dtype(self, args: &[DType]) -> DType { let Self { output_dtype, policy: _, @@ -140,7 +137,7 @@ impl BatchPlan { /// The nullable execution policy derived from one concrete dispatch. #[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(in crate::scalar_fn::unstable::row) enum RowPolicy { +pub(crate) enum RowPolicy { /// Evaluate all rows and mask the result. Dense, @@ -153,8 +150,7 @@ pub(in crate::scalar_fn::unstable::row) enum RowPolicy { impl RowPolicy { /// The policy for an infallible owned output. - pub(in crate::scalar_fn::unstable::row) const fn for_owned_output() -> Self - { + pub(crate) const fn for_owned_output() -> Self { if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { Self::Dense } else { @@ -163,8 +159,7 @@ impl RowPolicy { } /// The policy for an owned output carrying batch-deferred failure evidence. - pub(in crate::scalar_fn::unstable::row) const fn for_deferred_output() - -> Self { + pub(crate) const fn for_deferred_output() -> Self { if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { Self::DenseWithRetry } else { @@ -173,10 +168,7 @@ impl RowPolicy { } /// The policy for a sink-writing output. - pub(in crate::scalar_fn::unstable::row) const fn for_sink< - Args: ElementTuple, - ApplyResult: SinkResult, - >() -> Self { + pub(crate) const fn for_sink() -> Self { if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { Self::Dense } else { diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index b8ad1bccfa2..14b393dc252 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -1,6 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! The typed dispatch interface implemented by [`RowFn`] planning and execution. +//! +//! [`RowVisitor`] lets a function select its concrete input and output capabilities without +//! exposing framework-specific planning or execution state. +//! +//! [`RowFn`]: crate::scalar_fn::unstable::row::RowFn + use std::ops::BitOrAssign; use vortex_error::VortexResult; From 35f7453de0e8d83d88f43db315c4109616fd956a Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:21:36 -0400 Subject: [PATCH 099/160] Make null-tolerant decoding opt in Signed-off-by: Connor Tsui --- .../unstable/row/types/element/bool.rs | 4 + .../unstable/row/types/element/input.rs | 15 ++- .../unstable/row/types/element/primitive.rs | 4 + .../row/types/element/tuple/element_tuple.rs | 31 +++++++ .../unstable/row/types/element/tuple/tests.rs | 93 +++++++++++++++++++ 5 files changed, 144 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs index 5aedd7fca4d..28105c893cf 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs @@ -37,6 +37,10 @@ unsafe impl InputElement for bool { Ok(array.execute::(ctx)?.into_bit_buffer()) } + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + fn get(column: &Self::Column, index: usize) -> bool { column.value(index) } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 4b80c956981..32da6d5081f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -57,8 +57,17 @@ pub unsafe trait InputElement: 'static { /// invocation-invariant work into this method. fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; - /// Decode `array` _without_ assuming every row is valid, or `Ok(None)` when this element - /// cannot for this particular array. + /// Whether [`decode_null_tolerant`](Self::decode_null_tolerant) can decode this array. + /// + /// The conservative default declines. An implementation whose ordinary decode is safe and + /// infallible over null payloads can return `true`. Other implementations can inspect `array` + /// and opt in only for supported representations. + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(false) + } + + /// Decode `array` _without_ assuming every row is valid, or return `Ok(None)` when this element + /// cannot decode this particular array. /// /// Override this for a non-dense-safe representation that can still place safe placeholders in /// null slots. The skip-invalid executor never reads those slots. @@ -69,7 +78,7 @@ pub unsafe trait InputElement: 'static { array: ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if Self::DENSE_SAFE { + if Self::can_decode_null_tolerant(&array)? { Self::decode(array, ctx).map(Some) } else { Ok(None) diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs index fb15030cc64..9bc5db0c8e4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs @@ -44,6 +44,10 @@ unsafe impl InputElement for T { Ok(array.execute::(ctx)?.into_buffer::()) } + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + fn get(column: &Self::Column, index: usize) -> T { column[index] } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index 9fd5fed6ba4..d1b0127c6ec 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -63,6 +63,16 @@ impl ArgColumn { .map(Self)) } + fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { + // Batch execution short-circuits null constants before selecting this path, so a + // non-empty constant can always use the ordinary decode. + if batch_constant(array).is_some() && !array.is_empty() { + return Ok(true); + } + + T::can_decode_null_tolerant(array) + } + fn get(&self, index: usize) -> T::Elem<'_> { match &self.0 { ArgColumnKind::PerRow(column) => T::get(column, index), @@ -153,6 +163,12 @@ pub trait ElementTuple: 'static + private::Sealed { /// A dense deferred-error retry starts another invocation over filtered valid rows. fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; + /// Whether every input can be decoded without assuming that all rows are valid. + /// + /// The tuple checks this before decoding any column, so a decline does not discard work from + /// earlier arguments. + fn can_decode_null_tolerant(args: &dyn ExecutionArgs) -> VortexResult; + /// Decode every input column once while tolerating null rows. /// /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid @@ -225,6 +241,10 @@ impl ElementTuple for () { Ok(()) } + fn can_decode_null_tolerant(_args: &dyn ExecutionArgs) -> VortexResult { + Ok(true) + } + fn decode_null_tolerant( _args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx, @@ -291,10 +311,21 @@ macro_rules! element_tuple { Ok(($(ArgColumn::<$t>::decode(args.get($idx)?, ctx)?,)+)) } + fn can_decode_null_tolerant(args: &dyn ExecutionArgs) -> VortexResult { + Ok($({ + let array = args.get($idx)?; + ArgColumn::<$t>::can_decode_null_tolerant(&array)? + } &&)+ true) + } + fn decode_null_tolerant( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if !Self::can_decode_null_tolerant(args)? { + return Ok(None); + } + Ok(Some(( $(match ArgColumn::<$t>::decode_null_tolerant(args.get($idx)?, ctx)? { Some(column) => column, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index 2b4ca800a91..044ad15fd4b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -1,21 +1,114 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; +use super::ElementTuple; use super::element_tuple::batch_constant; +use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; use crate::dtype::Nullability; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; use crate::validity::Validity; +static DECODE_CALLS: AtomicUsize = AtomicUsize::new(0); + +macro_rules! i64_test_element { + ($element:ident, $decode_fallible:literal $(, $can_decode:item)?) => { + struct $element; + + // SAFETY: the view and unchecked access delegate to the `i64` implementation. + unsafe impl InputElement for $element { + type Column = Buffer; + type View<'a> = &'a [i64]; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = $decode_fallible; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + DECODE_CALLS.fetch_add(1, Ordering::Relaxed); + ::decode(array, ctx) + } + + $($can_decode)? + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + ::view(column) + } + + fn view_len(view: &Self::View<'_>) -> usize { + ::view_len(view) + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_from_view(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { ::get_from_view_unchecked(view, index) } + } + } + }; +} + +i64_test_element!( + DecodeProbe, + false, + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } +); +i64_test_element!(DenseFallible, true); + +#[test] +fn test_null_tolerant_decline_precedes_decoding() -> VortexResult<()> { + DECODE_CALLS.store(0, Ordering::Relaxed); + let first = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let second = PrimitiveArray::from_iter([3_i64, 4]).into_array(); + let args = VecExecutionArgs::new(vec![first, second], 2); + let mut ctx = array_session().create_execution_ctx(); + + let columns = <(DecodeProbe, DenseFallible)>::decode_null_tolerant(&args, &mut ctx)?; + + assert!(columns.is_none()); + assert_eq!(DECODE_CALLS.load(Ordering::Relaxed), 0); + Ok(()) +} + #[test] fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> { let child = ConstantArray::new(7_i64, 3).into_array(); From 56ca4a8052464bffd260b5cfe04071438f0b2774 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 14:19:50 -0400 Subject: [PATCH 100/160] Polish RowFn visitor API Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/mod.rs | 10 +- .../unstable/row/types/element/input.rs | 17 +- .../unstable/row/types/element/mod.rs | 4 +- .../row/types/element/tuple/element_tuple.rs | 33 ++- .../scalar_fn/unstable/row/types/result.rs | 2 +- .../src/scalar_fn/unstable/row/types/sink.rs | 25 +-- .../scalar_fn/unstable/row/visitor/check.rs | 10 +- .../src/scalar_fn/unstable/row/visitor/mod.rs | 2 +- .../scalar_fn/unstable/row/visitor/plan.rs | 33 ++- .../unstable/row/visitor/row_visitor.rs | 192 ++++++++++++++++-- .../src/scalar_fn/unstable/row/vtable.rs | 10 +- 11 files changed, 232 insertions(+), 106 deletions(-) diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index cbbb08e708f..ac369de427d 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -7,11 +7,11 @@ //! implementations. Expressions ([`crate::expr::Expression`]) reference scalar functions //! at each node. //! -//! Use `unstable::row::RowFn` for strict functions whose natural kernel computes one row at a -//! time. It derives decoding, constant handling, null propagation, output construction, and -//! validity. This experimental API requires the `unstable_row_fns` feature and has no compatibility -//! guarantees. Implement [`ScalarFnVTable`] directly when the natural kernel is columnar, aliases -//! an input, or may produce null from otherwise valid inputs. +//! Strict functions with row-at-a-time kernels can implement [`unstable::row::RowFn`]. It handles +//! decoding, constants, null propagation, output construction, and validity. This API requires the +//! `unstable_row_fns` feature and has no compatibility guarantees. Implement [`ScalarFnVTable`] +//! directly for columnar kernels and functions that alias an input or can produce null from valid +//! inputs. use vortex_session::registry::Id; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 32da6d5081f..2764840da7a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -24,9 +24,9 @@ pub unsafe trait InputElement: 'static { /// The decoded column representation supporting `O(1)` row access. type Column; - /// The view of a per-row decoded column read by the hot row loop. + /// The row-loop view of a decoded column. /// - /// This may borrow a cheaper representation than [`Column`](Self::Column). Primitive elements, + /// This can borrow a cheaper representation than [`Column`](Self::Column). Primitive elements, /// for example, expose a slice so its pointer and length are loop invariants rather than /// re-reading a [`Buffer`](vortex_buffer::Buffer) descriptor for every row. type View<'a>; @@ -37,14 +37,13 @@ pub unsafe trait InputElement: 'static { /// Whether every dense decode and access path tolerates rows that are null in the input. /// /// Arrays guarantee payloads only for valid rows. Set this to `true` only when every decode and - /// access method remains safe for null rows. Dense execution may pass unspecified values from + /// access method remains safe for null rows. Dense execution can pass unspecified values from /// null rows to the row closure. const DENSE_SAFE: bool; /// Whether [`decode`](Self::decode) can fail on _legal_ input data. /// - /// This excludes infrastructural failures such as IO or allocation. Set it when legal input may - /// contain a value that the decoder rejects. + /// This excludes infrastructural failures such as IO or allocation. const DECODE_FALLIBLE: bool; /// Validate that `dtype` is an acceptable input column dtype for this element type. @@ -52,9 +51,8 @@ pub unsafe trait InputElement: 'static { /// Decode `array` into its column representation. /// - /// The executor calls this once per row-kernel invocation. A dense deferred-error retry starts - /// another invocation over filtered valid rows. Hoist dtype checks, downcasts, and other - /// invocation-invariant work into this method. + /// Called once per row-kernel invocation, including deferred-error retries. Hoist dtype checks, + /// downcasts, and other invocation-invariant work into this method. fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; /// Whether [`decode_null_tolerant`](Self::decode_null_tolerant) can decode this array. @@ -71,9 +69,6 @@ pub unsafe trait InputElement: 'static { /// /// Override this for a non-dense-safe representation that can still place safe placeholders in /// null slots. The skip-invalid executor never reads those slots. - /// - /// Return `Ok(None)` rather than an error when an array has no null-tolerant decode; the - /// batch execution falls back to the filter strategy. fn decode_null_tolerant( array: ArrayRef, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 37120798908..51d66594332 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -3,8 +3,8 @@ //! The element types a row function can read and produce. //! -//! [`InputElement::Elem`] may borrow from its decoded column. [`OutputElement`] is returned by an -//! owned row computation; runtime-shaped output uses an +//! [`InputElement::Elem`] can borrow from its decoded column. Owned row computations return an +//! [`OutputElement`]. Runtime-shaped outputs use an //! [`OutputSink`](crate::scalar_fn::unstable::row::OutputSink). mod bool; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index d1b0127c6ec..cd9c5eb5d12 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -105,7 +105,7 @@ impl ArgColumn { /// Return the batch-constant array, looking through masked and extension wrappers. /// -/// Batch execution owns mask validity, so a masked constant may expose its constant child here. An +/// Batch execution owns mask validity, so a masked constant can expose its constant child here. An /// extension over constant storage remains wrapped to preserve its extension dtype. pub fn batch_constant(array: &ArrayRef) -> Option { if array.is::() { @@ -124,7 +124,7 @@ pub fn batch_constant(array: &ArrayRef) -> Option { /// Typed argument tuples for arities zero through twelve. /// -/// This trait is sealed; add a new row representation by implementing [`InputElement`] and placing +/// This trait is sealed. Add a new row representation by implementing [`InputElement`] and placing /// it in one of the supplied tuples. pub trait ElementTuple: 'static + private::Sealed { /// The decoded column representations. @@ -136,10 +136,9 @@ pub trait ElementTuple: 'static + private::Sealed { /// The borrowed row of element values. type Elems<'a>; - /// The batch-constant element values: [`Elems`](Self::Elems) with every argument wrapped in - /// `Option`. + /// The batch-constant element values. /// - /// `Some` carries the value of a batch-constant argument; `None` marks a per-row argument. A + /// `Some` carries the value of a batch-constant argument. `None` marks a per-row argument. A /// [`RowVisitor`] passes these values to its prepare closure so constant work can leave the row /// loop. /// @@ -375,18 +374,18 @@ macro_rules! element_tuple { }; } -element_tuple!(1; A:0); -element_tuple!(2; A:0, B:1); -element_tuple!(3; A:0, B:1, C:2); -element_tuple!(4; A:0, B:1, C:2, D:3); -element_tuple!(5; A:0, B:1, C:2, D:3, E:4); -element_tuple!(6; A:0, B:1, C:2, D:3, E:4, F:5); -element_tuple!(7; A:0, B:1, C:2, D:3, E:4, F:5, G:6); -element_tuple!(8; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7); -element_tuple!(9; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8); -element_tuple!(10; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9); -element_tuple!(11; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10); -element_tuple!(12; A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7, I:8, J:9, K:10, L:11); +element_tuple!(1; A: 0); +element_tuple!(2; A: 0, B: 1); +element_tuple!(3; A: 0, B: 1, C: 2); +element_tuple!(4; A: 0, B: 1, C: 2, D: 3); +element_tuple!(5; A: 0, B: 1, C: 2, D: 3, E: 4); +element_tuple!(6; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5); +element_tuple!(7; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6); +element_tuple!(8; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7); +element_tuple!(9; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8); +element_tuple!(10; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9); +element_tuple!(11; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10); +element_tuple!(12; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10, L: 11); mod private { pub trait Sealed {} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/result.rs b/vortex-array/src/scalar_fn/unstable/row/types/result.rs index a3439f07ffe..6163687932d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/result.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/result.rs @@ -12,7 +12,7 @@ use super::InitializedElement; /// The result of writing one row: success or an immediate error. /// -/// This trait is sealed; row functions choose one of its supplied implementations. +/// This trait is sealed. Row functions choose one of its supplied implementations. pub trait SinkResult: 'static + private::Sealed { /// The [`OutputSink::WriteToken`](super::OutputSink::WriteToken) carried by a success. type WriteToken: 'static; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index e7623de6559..6cc3ce06d30 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -17,11 +17,11 @@ use crate::scalar_fn::unstable::row::OutputElement; /// A column allocated once per batch that a row closure writes into, one row at a time. /// -/// A sink may use function options and input dtypes to build a runtime-shaped output or own shared +/// A sink can use function options and input dtypes to build a runtime-shaped output or own shared /// batch state. The executor passes each row slot into an [`Fn`] closure. /// -/// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once; -/// skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an +/// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once. +/// Skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an /// initializer. /// /// # Errors @@ -45,8 +45,6 @@ use crate::scalar_fn::unstable::row::OutputElement; /// - [`finish`] **must** be sound once every visited callback returned its required token and the /// skipped-row initializer, when present, ran successfully. /// -/// The executor relies on these guarantees when it calls `finish`. -/// /// [`Rows`]: Self::Rows /// [`WriteToken`]: Self::WriteToken /// [`finish`]: Self::finish @@ -70,16 +68,15 @@ pub unsafe trait OutputSink: 'static + Sized { /// Proof that a successful row closure left its row handle initialized. /// - /// Use `()` for initialized row handles. A sink exposing uninitialized storage uses a distinct - /// token returned after initialization. A sink that uses this token to justify unsafe code - /// **must** prevent safe construction that does not establish the invariant. Make construction - /// unsafe when Rust cannot tie the token to the supplied row handle. + /// Use `()` for initialized row handles. A sink exposing uninitialized storage uses a token + /// returned after initialization. If a sink uses the token to justify unsafe code, safe code + /// **must not** be able to construct one without establishing the invariant. type WriteToken: 'static; /// The operation that initializes every output position before skip-invalid execution. /// /// `Some` enables skip-invalid execution. The initializer **must** make every row safe to - /// finish; callbacks overwrite valid rows and batch execution masks skipped rows. + /// finish. Callbacks overwrite valid rows, and batch execution masks skipped rows. /// /// `None` makes the executor fall back to filtering the inputs. fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { @@ -157,8 +154,8 @@ impl InitializedElement { /// success. The token is zero-sized, so the proof adds no runtime row state. /// /// Skip-invalid execution initializes placeholders before omitting rows. Errors and unwinds are -/// safe because `values` keeps length zero until `finish`; `T: Copy` means initialized -/// spare-capacity elements require no destruction. +/// safe because `values` keeps length zero until `finish`. The `T: Copy` bound means that +/// initialized spare-capacity elements require no destruction. pub struct UninitElementSink { /// Spare storage written in increasing row order. values: Vec, @@ -168,8 +165,8 @@ pub struct UninitElementSink { } // SAFETY: the row slice covers exactly the reserved spare-capacity range, so each accepted index -// names one distinct slot. `InitializedElement` cannot be constructed by safe code; its unsafe -// constructor writes the supplied slot and requires the caller to return that exact evidence. The +// names one distinct slot. Safe code cannot construct `InitializedElement`. Its unsafe constructor +// writes the supplied slot and requires the caller to return that exact evidence. The // skipped-row initializer writes `T::default()` into every slot before masked traversal. unsafe impl OutputSink for UninitElementSink diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs index 501ef896df1..12d5bd7c7ea 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -28,21 +28,19 @@ pub(crate) const fn assert_owned_output_needs_no_drop() { ); } -/// Assert that the input arity and decode fallibility match the function-wide declarations. const fn assert_input_visit_contract() { assert!( Args::ARITY == F::ARG_NAMES.len(), "the visited argument tuple must have the arity declared by RowFn::ARG_NAMES", ); - // Dictionary pushdown treats an infallible function as safe to evaluate over values no code - // references, so every dispatch must fit the function-wide declaration. + // Dictionary push-down can evaluate values that no input row references. Every dispatch must + // therefore match the function-wide fallibility declaration. assert!( !Args::DECODE_FALLIBLE || F::FALLIBLE, "RowFn::FALLIBLE must be true when input decoding can fail", ); } -/// Assert the input contract and that owned output values do not require drop glue. pub(super) const fn assert_owned_visit_contract() where Function: RowFn, @@ -53,7 +51,6 @@ where assert_owned_output_needs_no_drop::(); } -/// Assert that a sink visit obeys the input, fallibility, and deferred-error contracts. pub(super) const fn assert_sink_visit_contract() where Function: RowFn, @@ -67,7 +64,6 @@ where ); } -/// Assert the owned-output contract, fallibility declaration, and failure-evidence width bound. pub(super) const fn assert_deferred_visit_contract() where Function: RowFn, @@ -86,7 +82,6 @@ where ); } -/// Validate the input dtypes and return the non-nullable dtype built by `Out`. pub(super) fn validate_owned_visit( dtypes: &[DType], ) -> VortexResult { @@ -101,7 +96,6 @@ pub(super) fn validate_owned_visit( Ok(dtype) } -/// Validate the input dtypes and return the non-nullable dtype built by `Sink`. pub(super) fn validate_sink_visit( options: &Options, dtypes: &[DType], diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index 24d6e1f5b8d..57da5f4691b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -8,7 +8,7 @@ mod check; mod plan; -pub(super) use plan::PlanRows; +pub(super) use plan::BatchPlanner; mod row_visitor; pub use row_visitor::RowVisitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index e84e5e24e5d..c522376d491 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -3,8 +3,8 @@ //! Plans the concrete signature selected by [`RowFn::dispatch`]. //! -//! [`PlanRows`] validates input and output dtypes, then records the output dtype and null-handling -//! policy that execution must reproduce. +//! [`BatchPlanner`] validates input and output dtypes, then records the output dtype and +//! null-handling policy that execution must reproduce. use std::marker::PhantomData; use std::ops::BitOrAssign; @@ -27,19 +27,17 @@ use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::SinkResult; -/// The plan-time visit that validates dtypes and derives the nullable execution policy. -pub(crate) struct PlanRows<'a, F: RowFn> { - /// The input dtypes for this plan. +/// A planning visitor that validates dtypes and selects the nullable execution policy. +pub(crate) struct BatchPlanner<'a, F: RowFn> { dtypes: &'a [DType], - /// The function options used to derive a sink's runtime dtype. options: &'a F::Options, - /// The visited function, carried only so the dispatch check can name its contract. + /// Ties the planner to the function used by its compile-time contract checks. function: PhantomData, } -impl<'a, F: RowFn> PlanRows<'a, F> { +impl<'a, F: RowFn> BatchPlanner<'a, F> { pub(crate) fn new(dtypes: &'a [DType], options: &'a F::Options) -> Self { Self { dtypes, @@ -49,9 +47,9 @@ impl<'a, F: RowFn> PlanRows<'a, F> { } } -impl private::Sealed for PlanRows<'_, F> {} +impl private::Sealed for BatchPlanner<'_, F> {} -impl RowVisitor for PlanRows<'_, F> { +impl RowVisitor for BatchPlanner<'_, F> { type VisitResult = BatchPlan; fn visit_prepared( @@ -116,7 +114,8 @@ pub(crate) struct BatchPlan { pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. - // TODO(connor)[RowFn]: The execution backend tracked by #9130 consumes this field. + // TODO(connor)[RowFn]: Remove this allowance when the execution backend from #9130 consumes + // this policy. #[allow(dead_code)] pub(crate) policy: RowPolicy, } @@ -124,14 +123,10 @@ pub(crate) struct BatchPlan { impl BatchPlan { /// Return the output dtype widened with strict input nullability. pub(crate) fn result_dtype(self, args: &[DType]) -> DType { - let Self { - output_dtype, - policy: _, - } = self; - let nullability = - output_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); - - output_dtype.with_nullability(nullability) + let nullability = self.output_dtype.nullability() + | Nullability::from(args.iter().any(DType::is_nullable)); + + self.output_dtype.with_nullability(nullability) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 14b393dc252..0e7abaa4322 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -29,16 +29,35 @@ use crate::scalar_fn::unstable::row::SinkResult; pub trait RowVisitor: private::Sealed + Sized { /// The framework result of visiting one concrete row signature. /// - /// This is a batch plan or execution result, not the per-row `Out` returned by - /// [`RowVisitor::visit`] and [`RowVisitor::visit_deferred`]. + /// This is a batch plan or execution result, not a per-row output. type VisitResult; - /// Visit an infallible row computation that returns one independent output value. + /// Visit an infallible row computation that returns one output value per row. /// - /// `apply` must be total over every stored element value: it must not panic or have side - /// effects. Dense execution can pass unspecified values from null rows. + /// `apply` must not panic or have side effects. Dense execution can pass unspecified values + /// from null rows. + /// + /// The framework verifies that `Out` does not require drop glue. + /// + /// # Examples + /// + /// Apply infallible wrapping arithmetic. /// - /// The framework also verifies that `Out` does not require drop glue. + /// ```ignore + /// visitor.visit::<(i64, i64), i64>(|(lhs, rhs)| lhs.wrapping_add(rhs)) + /// ``` + /// + /// Dispatch an equality helper over its primitive element type. + /// + /// ```ignore + /// fn visit_equal(visitor: V) -> VortexResult + /// where + /// T: NativePType, + /// V: RowVisitor, + /// { + /// visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)) + /// } + /// ``` fn visit( self, apply: impl Fn(Args::Elems<'_>) -> Out, @@ -51,6 +70,27 @@ pub trait RowVisitor: private::Sealed + Sized { } /// The prepared form of [`visit`](Self::visit), with the same prerequisites. + /// + /// # Examples + /// + /// Test whether each string occurs in its allowed-values list. The prepare closure builds one + /// lookup table for a batch-constant list. The row closure scans a varying list directly. + /// + /// ```ignore + /// visitor.visit_prepared::< + /// (StringRow, StringListRow), + /// bool, + /// Option, + /// >( + /// |(_value, allowed_values)| allowed_values.map(PreparedAllowedValues::new), + /// |prepared_allowed_values, (value, allowed_values)| { + /// match prepared_allowed_values { + /// Some(allowed_values) => allowed_values.contains(value), + /// None => allowed_values.iter().any(|allowed| allowed == value), + /// } + /// }, + /// ) + /// ``` fn visit_prepared( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, @@ -60,19 +100,41 @@ pub trait RowVisitor: private::Sealed + Sized { Args: IndexedElementTuple, Out: OutputElement; - /// Visit a row computation that writes through a sink-provided row handle. + /// Visit a row computation that writes through a row handle from an output sink. /// - /// `apply` must be total over every stored input value: it must not panic or cause side effects - /// other than writing the supplied row handle. Dense execution can pass unspecified values - /// from null rows. + /// `apply` must not panic or have side effects except for writes to the supplied row handle. + /// Dense execution can pass unspecified values from null rows. /// - /// On success, `apply` must return the write token produced by writing the `Sink::Row` supplied - /// to that same invocation. It must not return evidence produced for another row, sink, or - /// unrelated local cell. Violating this requirement can make the unsafe - /// [`OutputSink::finish`] precondition false. + /// On success, `apply` must return the write token for the supplied row handle. A token from + /// another row, sink, or local cell can violate the safety contract of [`OutputSink::finish`]. /// /// A fallible `ApplyResult` requires /// [`RowFn::FALLIBLE`](crate::scalar_fn::unstable::row::RowFn::FALLIBLE) to be `true`. + /// + /// # Examples + /// + /// Checked integer division reports errors immediately and writes successful rows into + /// uninitialized output. The cold, non-inlined helper keeps error construction out of the row + /// callback. + /// + /// ```ignore + /// #[cold] + /// #[inline(never)] + /// fn integer_division_error() -> VortexError { + /// vortex_err!(InvalidArgument: "integer division by zero or overflow") + /// } + /// + /// visitor.visit_into::<(i64, i64), UninitElementSink, _>( + /// |(lhs, rhs), output| { + /// let Some(value) = lhs.checked_div(rhs) else { + /// return Err(integer_division_error()); + /// }; + /// + /// // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + /// Ok(unsafe { InitializedElement::write(output, value) }) + /// }, + /// ) + /// ``` fn visit_into( self, apply: impl Fn(Args::Elems<'_>, >::Row<'_>) -> ApplyResult, @@ -89,6 +151,32 @@ pub trait RowVisitor: private::Sealed + Sized { } /// The prepared form of [`visit_into`](Self::visit_into), with the same prerequisites. + /// + /// # Examples + /// + /// Compute the cosine similarity of each vector pair: their dot product divided by their + /// magnitudes. The prepare closure computes each batch-constant vector's magnitude once. + /// + /// ```ignore + /// visitor.visit_prepared_into::< + /// (TensorRow, TensorRow), + /// UninitElementSink, + /// ConstantVectorMagnitudes, + /// InitializedElement, + /// >( + /// |(lhs, rhs)| ConstantVectorMagnitudes { + /// lhs: lhs.map(vector_magnitude), + /// rhs: rhs.map(vector_magnitude), + /// }, + /// |constant_magnitudes, (lhs, rhs), output| { + /// let similarity = + /// cosine_similarity_with_constant_magnitudes(constant_magnitudes, lhs, rhs); + /// + /// // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + /// unsafe { InitializedElement::write(output, similarity) } + /// }, + /// ) + /// ``` fn visit_prepared_into( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, @@ -103,18 +191,44 @@ pub trait RowVisitor: private::Sealed + Sized { Sink: OutputSink, ApplyResult: SinkResult>::WriteToken>; - /// Visit a row computation that returns an owned output and deferred failure evidence. + /// Visit a row computation that returns an owned output value and deferred failure evidence. /// - /// `apply` must be total over every stored element value: it must not panic or have side - /// effects. Dense execution can pass unspecified values from null rows. + /// `apply` must not panic or have side effects. Dense execution can pass unspecified values + /// from null rows. /// - /// `Fail` is OR-reduced across rows and handed to `finish_failure`. The value from + /// The executor OR-reduces `Fail` across rows and passes the result to `finish_failure`. /// [`Default::default`] **must** mean success, including for an empty batch. The compiler - /// cannot check this semantic requirement. + /// cannot check this requirement. + /// + /// [`RowFn::FALLIBLE`](crate::scalar_fn::unstable::row::RowFn::FALLIBLE) **must** be `true`. + /// `Out` must not require drop glue. `Fail` must be no wider than `Out`, or failure tracking + /// reduces the vector width. The framework checks these requirements. + /// + /// # Examples + /// + /// Checked addition returns a wrapping value and compact overflow flag without branching. The + /// executor reduces the flags after the loop. The cold, non-inlined helper keeps error + /// construction out of the row loop. /// - /// [`RowFn::FALLIBLE`](crate::scalar_fn::unstable::row::RowFn::FALLIBLE) **must** be `true`, - /// `Out` must not require drop glue, and `Fail` must be no wider than `Out` so failure tracking - /// does not reduce the vector width. The framework checks each requirement. + /// ```ignore + /// #[cold] + /// #[inline(never)] + /// fn integer_addition_error() -> VortexError { + /// vortex_err!(InvalidArgument: "integer overflow in checked add") + /// } + /// + /// visitor.visit_deferred::<(i64, i64), i64, bool>( + /// // `overflowing_add` returns `(i64, bool)`. + /// |(lhs, rhs)| lhs.overflowing_add(rhs), + /// |overflowed| { + /// if overflowed { + /// return Err(integer_addition_error()); + /// } + /// + /// Ok(()) + /// }, + /// ) + /// ``` fn visit_deferred( self, apply: impl Fn(Args::Elems<'_>) -> (Out, Fail), @@ -133,6 +247,40 @@ pub trait RowVisitor: private::Sealed + Sized { } /// The prepared form of [`visit_deferred`](Self::visit_deferred), with the same prerequisites. + /// + /// # Examples + /// + /// Rescale each unscaled decimal by multiplying it by `10^scale`. The prepare closure computes + /// the multiplier once for a batch-constant scale. Each row returns the rescaled value and an + /// overflow flag, which the executor reduces after the loop. + /// + /// ```ignore + /// #[cold] + /// #[inline(never)] + /// fn decimal_rescaling_overflow() -> VortexError { + /// vortex_err!(InvalidArgument: "decimal rescaling overflowed") + /// } + /// + /// visitor.visit_prepared_deferred::< + /// (i64, DecimalScale), + /// i64, + /// Option, + /// bool, + /// >( + /// |(_value, scale)| scale.map(PreparedDecimalScale::new), + /// |prepared_scale, (value, scale)| match prepared_scale { + /// Some(scale) => scale.apply_checked(value), + /// None => PreparedDecimalScale::new(scale).apply_checked(value), + /// }, + /// |overflowed| { + /// if overflowed { + /// return Err(decimal_rescaling_overflow()); + /// } + /// + /// Ok(()) + /// }, + /// ) + /// ``` fn visit_prepared_deferred( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index b24ea3e66e6..b6c14585c48 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -13,7 +13,7 @@ use vortex_error::vortex_ensure_eq; use vortex_session::VortexSession; use super::row_fn::RowFn; -use super::visitor::PlanRows; +use super::visitor::BatchPlanner; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; @@ -86,7 +86,7 @@ pub fn row_fn_return_dtype( ) -> VortexResult { ensure_arity(function, args.len())?; - let plan = function.dispatch(options, args, PlanRows::::new(args, options))?; + let plan = function.dispatch(options, args, BatchPlanner::::new(args, options))?; Ok(plan.result_dtype(args)) } @@ -117,7 +117,7 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { vortex_ensure_eq!( actual, expected, - "row function {} must receive exactly {expected} input values, got {actual}", + "row function {} requires arity {expected}, got {actual}", RowFn::id(function), ); @@ -189,9 +189,7 @@ mod tests { #[track_caller] fn assert_arity_error(error: VortexError) { assert!( - error - .to_string() - .contains("must receive exactly 1 input values, got 0"), + error.to_string().contains("requires arity 1, got 0"), "unexpected error: {error}", ); } From 06a86d9519edfa61f0d46ef7163ba3a22cb14e34 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 14:48:08 -0400 Subject: [PATCH 101/160] Clarify RowFn length validation docs Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/mod.rs | 2 +- .../unstable/row/types/element/tuple/element_tuple.rs | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index ac369de427d..6be34ce1f34 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -7,7 +7,7 @@ //! implementations. Expressions ([`crate::expr::Expression`]) reference scalar functions //! at each node. //! -//! Strict functions with row-at-a-time kernels can implement [`unstable::row::RowFn`]. It handles +//! Strict functions with row-at-a-time kernels can implement `unstable::row::RowFn`. It handles //! decoding, constants, null propagation, output construction, and validity. This API requires the //! `unstable_row_fns` feature and has no compatibility guarantees. Implement [`ScalarFnVTable`] //! directly for columnar kernels and functions that alias an input or can produce null from valid diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index cd9c5eb5d12..b0a3e709696 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -187,13 +187,17 @@ pub trait ElementTuple: 'static + private::Sealed { fn per_row_views(columns: &Self::Columns) -> Option>; /// Whether every view contains exactly `row_count` rows. + /// + /// The executor calls this once before the all-per-row hot loop. A successful check gives LLVM + /// a dominating equality between the loop bound and every source length, which lets it optimize + /// the tuple access as one fixed-length traversal. fn view_lens_match(views: &Self::Views<'_>, row_count: usize) -> bool; /// Whether every per-row argument contains exactly `row_count` rows. /// - /// This provides the same guarantee as [`view_lens_match`](Self::view_lens_match) when - /// [`per_row_views`](Self::per_row_views) declines a mixed per-row and batch-constant tuple. A - /// batch constant is exempt because it was collapsed to one row. + /// This is the mixed-shape equivalent of [`view_lens_match`](Self::view_lens_match) when + /// [`per_row_views`](Self::per_row_views) declines. It runs once before the hot loop for the + /// same LLVM optimization. A batch constant is exempt because decoding collapsed it to one row. fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; /// Read one row from borrowed views. From a061da6fac851c378f9a3af69cb16db1912d2a7b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 14:06:47 -0400 Subject: [PATCH 102/160] Implement RowFn batch execution Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 93 +++ .../scalar_fn/unstable/row/batch/execution.rs | 481 ++++++++++++++ .../src/scalar_fn/unstable/row/batch/mod.rs | 25 + .../src/scalar_fn/unstable/row/batch/tests.rs | 587 ++++++++++++++++++ .../src/scalar_fn/unstable/row/execute/mod.rs | 19 + .../scalar_fn/unstable/row/execute/outcome.rs | 43 ++ .../scalar_fn/unstable/row/execute/owned.rs | 124 ++++ .../scalar_fn/unstable/row/execute/sink.rs | 262 ++++++++ .../src/scalar_fn/unstable/row/mod.rs | 5 + .../unstable/row/types/element/mod.rs | 1 + .../unstable/row/types/element/tuple/mod.rs | 1 + .../unstable/row/types/element/tuple/tests.rs | 2 +- .../src/scalar_fn/unstable/row/types/mod.rs | 1 + .../scalar_fn/unstable/row/visitor/execute.rs | 263 ++++++++ .../src/scalar_fn/unstable/row/visitor/mod.rs | 7 + .../scalar_fn/unstable/row/visitor/plan.rs | 5 +- .../src/scalar_fn/unstable/row/vtable.rs | 149 ++++- 17 files changed, 2054 insertions(+), 14 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/args.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execution.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/tests.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/owned.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/sink.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs new file mode 100644 index 00000000000..d5ffe042613 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A borrowed execution view passed to one row-kernel invocation. + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::visitor::RowPolicy; + +/// A borrowed [`ExecutionArgs`] view with the planning metadata selected for its row kernel. +/// +/// `arrays` may be filtered or sliced, while `dtypes` and `output_dtype` always describe the +/// original planned batch. Keeping them together prevents an execution path from pairing an input +/// view with unrelated planning metadata. +#[derive(Clone, Copy)] +pub(in crate::scalar_fn::unstable::row) struct BorrowedExecutionArgs<'a> { + /// The input arrays for this kernel invocation. + arrays: &'a [ArrayRef], + + /// The number of rows in this kernel invocation. + row_count: usize, + + /// The original input dtypes used to select the row implementation. + dtypes: &'a [DType], + + /// The non-nullable dtype built by the selected output capability. + output_dtype: &'a DType, + + /// The nullable execution policy selected during planning. + policy: RowPolicy, +} + +impl<'a> BorrowedExecutionArgs<'a> { + /// Pair one input view with the planning metadata selected for its batch. + pub(in crate::scalar_fn::unstable::row) fn new( + arrays: &'a [ArrayRef], + row_count: usize, + dtypes: &'a [DType], + output_dtype: &'a DType, + policy: RowPolicy, + ) -> Self { + Self { + arrays, + row_count, + dtypes, + output_dtype, + policy, + } + } + + /// Return the concrete arrays used by this row-kernel invocation. + pub(in crate::scalar_fn::unstable::row) fn arrays(&self) -> &'a [ArrayRef] { + self.arrays + } + + /// Return the original input dtypes used to select the row implementation. + pub(in crate::scalar_fn::unstable::row) fn dtypes(&self) -> &'a [DType] { + self.dtypes + } + + /// Return the non-nullable dtype built by the selected output capability. + pub(in crate::scalar_fn::unstable::row) fn output_dtype(&self) -> &'a DType { + self.output_dtype + } + + /// Return the nullable execution policy selected during planning. + pub(in crate::scalar_fn::unstable::row) fn policy(&self) -> RowPolicy { + self.policy + } +} + +impl ExecutionArgs for BorrowedExecutionArgs<'_> { + fn get(&self, index: usize) -> VortexResult { + self.arrays.get(index).cloned().ok_or_else(|| { + vortex_err!( + "row-function input index must be less than {}, got {index}", + self.arrays.len(), + ) + }) + } + + fn num_inputs(&self) -> usize { + self.arrays.len() + } + + fn row_count(&self) -> usize { + self.row_count + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs new file mode 100644 index 00000000000..671b8f4ecb0 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -0,0 +1,481 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Null propagation, constant folding, and strategy execution for one columnar batch. + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::BatchPlan; +use super::RowPolicy; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::types::batch_constant; +use crate::validity::Validity; + +/// How far [`Batch::resolve_validity`] got before a mixed-mask strategy became necessary. +enum ResolvedMask { + /// The all-valid or all-null batch was answered without a mixed-mask strategy. + Decided(ArrayRef), + + /// A mask with both set and unset bits, which a strategy must now execute. + Mixed(Mask), +} + +/// One batch of inputs and the metadata needed before its row kernel runs. +pub struct Batch { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The number of rows in the original execution scope. + row_count: usize, + + /// The input columns, collected once: constant folding inspects them and the filter strategy + /// filters them. + inputs: SmallVec<[ArrayRef; 4]>, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: SmallVec<[DType; 4]>, + + /// The conjoined input validity, so a row of the output is valid iff it is valid in every + /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. + validity: Validity, + + /// The dtype the function declares for these inputs, which the kernel's output is reconciled + /// against. Already widened to nullable if any input is nullable. + result_dtype: DType, + + /// The non-nullable dtype the dispatched output capability builds, computed while planning. + output_dtype: DType, + + /// How the concrete dispatch executes nullable rows. + policy: RowPolicy, +} + +impl Batch { + /// Collect the inputs and derive their dtype, validity, and execution policy. + /// + /// **Not** for a nullary function: with no inputs there is no validity to propagate and no + /// per-row work to fold, and the all-constant check below would vacuously pass. + pub fn new( + id: ScalarFnId, + args: &dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let row_count = args.row_count(); + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>()?; + + for (index, input) in inputs.iter().enumerate() { + vortex_ensure_eq!( + input.len(), + row_count, + "the {id} input {index} must have {row_count} rows, got {}", + input.len(), + ); + } + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let result_dtype = plan.result_dtype(&arg_dtypes); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + row_count, + inputs, + arg_dtypes, + validity, + result_dtype, + output_dtype: plan.output_dtype, + policy: plan.policy, + }) + } + + /// Add null propagation, constant folding, and strategy selection around `kernel`. + /// + /// The kernel may ignore input validity. It receives valid-only rows when required, and its + /// output **must** match the planned dtype up to nullability. `try_unfiltered` receives the + /// originals plus a mixed validity mask; `Ok(None)` selects filter-and-scatter. + pub fn execute( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.validity, Validity::AllInvalid) + || self.inputs.iter().any(|input| { + input + .as_opt::() + .is_some_and(|constant| constant.scalar().is_null()) + }) + { + return Ok(self.all_null()); + } + + // All inputs constant, and their conjoined validity proves every row non-null. This sees + // through extension and masked wrappers just like argument decoding does. + if self.row_count > 0 + && self.validity.definitely_no_nulls() + && self + .inputs + .iter() + .all(|input| batch_constant(input).is_some()) + { + return self.broadcast_one_row(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, false, ctx), + RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), + RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx), + } + } + + /// Evaluate a single row of all-constant inputs and broadcast its value. + /// + /// Reconciling the row's dtype before reading the scalar keeps this path on the same + /// kernel/declaration agreement check as the dense and filter paths, rather than letting `cast` + /// paper over a disagreement. + fn broadcast_one_row( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one_row: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.slice(0..1)) + .collect::>()?; + + let result = VortexResult::from(kernel(self.execution_args(&one_row, 1), ctx)?)?; + let result = self.validate_kernel_output(result, 1, ctx)?; + let result = self.finalize_output(result, 1)?; + let scalar = result.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.row_count).into_array()) + } + + /// Run the kernel over every row, including the rows behind nulls, then mask its result. + /// + /// The arguments reach the kernel untouched, so the inputs keep their original encoding, and + /// the conjoined validity is handed to `mask` as an array rather than materialized into a + /// [`Mask`] first. + fn execute_dense( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + retry_deferred_error: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let values = match kernel(self.execution_args(&self.inputs, self.row_count), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Unlike `resolve_validity`, all-true preserves the deferred error and all-false + // suppresses evidence that came entirely from null rows. An empty loop cannot + // produce deferred evidence, so the ambiguous empty mask cannot reach this arm. + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + // Deferred retry receives only the dense kernel. Filter first so the retry cannot + // evaluate null rows again. + return self.filter_and_scatter(kernel, &valid, ctx); + } + RowExecution::DeferredError(error) => return Err(error), + }; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.row_count) + } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // Handled by the guard above, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } + + /// Materialize validity and answer all-valid or all-null batches before selecting a mixed-mask + /// strategy. + fn resolve_validity( + &self, + kernel: &impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Check all-true before all-false: an empty mask is both, and must not be treated as + // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). + if valid.all_true() { + let values = VortexResult::from(kernel( + self.execution_args(&self.inputs, self.row_count), + ctx, + )?)?; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = self.finalize_output(values, self.row_count)?; + + return Ok(ResolvedMask::Decided(values)); + } + + if valid.all_false() { + return Ok(ResolvedMask::Decided(self.all_null())); + } + + Ok(ResolvedMask::Mixed(valid)) + } + + /// Resolve validity, try unfiltered execution, then fall back to filtering. + fn execute_valid_only( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedMask::Decided(result) => return Ok(result), + ResolvedMask::Mixed(valid) => valid, + }; + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Try execution against the original inputs, then mask a returned full-length result. + fn try_execute_unfiltered( + &self, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(execution) = try_unfiltered( + self.execution_args(&self.inputs, self.row_count), + valid, + ctx, + )? + else { + return Ok(None); + }; + let values = VortexResult::from(execution)?; + let values = self.validate_kernel_output(values, valid.len(), ctx)?; + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + self.finalize_output(values.mask(mask)?, valid.len()) + .map(Some) + } + + /// The filter strategy for a mixed mask: filter every input down to the rows set in `valid`, + /// run the kernel over those, and scatter its results back into a null-padded output. + fn filter_and_scatter( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(valid.clone())) + .collect::>()?; + + let values = VortexResult::from(kernel( + self.execution_args(&filtered, valid.true_count()), + ctx, + )?)?; + let values = self.validate_kernel_output(values, valid.true_count(), ctx)?; + + self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) + } + + /// An all-null result of the function's declared return dtype. + fn all_null(&self) -> ArrayRef { + ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() + } + + /// Pair an input view with this batch's planning metadata. + fn execution_args<'b>( + &'b self, + arrays: &'b [ArrayRef], + row_count: usize, + ) -> BorrowedExecutionArgs<'b> { + BorrowedExecutionArgs::new( + arrays, + row_count, + &self.arg_dtypes, + &self.output_dtype, + self.policy, + ) + } + + /// Finalize an output against this batch's expected length and declared return dtype. + fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { + reconcile_output(self.id, &self.result_dtype, expected_len, values) + } + + /// Validate the output from a row kernel before batch validity is attached. + fn validate_kernel_output( + &self, + values: ArrayRef, + expected_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx) + } + + /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set + /// bits, producing an array of length `valid.len()` that is null at every unset position. + fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult { + vortex_ensure_eq!( + values.len(), + valid.true_count(), + "the {} kernel produced {} rows for {} filtered rows", + self.id, + values.len(), + valid.true_count(), + ); + + let AllOr::Some(slices) = valid.slices() else { + // The caller handles the all-true and all-false masks. + vortex_bail!("scatter_valid requires a mixed mask"); + }; + + // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index + // 0, and any in-bounds index would do since they are masked out below (values is non-empty + // here). + let mut indices = vec![0u64; valid.len()]; + let mut rank = 0u64; + for &(start, end) in slices { + for index in &mut indices[start..end] { + *index = rank; + rank += 1; + } + } + let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array(); + + let scattered = values.take(indices)?; + + // A nullable gathered array cannot be wrapped because a `Masked` child must be all valid. + // The general masking pass unions its nulls with the batch validity instead. + if scattered.dtype().is_nullable() { + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + return scattered.mask(mask); + } + + // The gathered values are all valid, so attaching validity is sufficient. + Ok(MaskedArray::try_new( + scattered, + Validity::from_mask(valid.clone(), Nullability::Nullable), + )? + .into_array()) + } +} + +/// Validate the output produced directly by a row kernel. +/// +/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring +/// nullability, and every produced row **must** be valid. Batch execution owns strict null +/// propagation and attaches input-derived validity only after this boundary. +pub fn finalize_kernel_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + vortex_ensure!( + values.all_valid(ctx)?, + "the {id} row kernel produced nulls for valid rows", + ); + + cast_output_nullability(result_dtype, values) +} + +/// Reconcile an output with the function's declared shape and nullability. +fn reconcile_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + + cast_output_nullability(result_dtype, values) +} + +/// Validate an output's shape and logical dtype without executing a nullability cast. +fn validate_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: &ArrayRef, +) -> VortexResult<()> { + vortex_ensure_eq!( + values.len(), + expected_len, + "the {id} kernel produced {} rows for {expected_len} input rows", + values.len(), + ); + vortex_ensure!( + values.dtype().eq_ignore_nullability(result_dtype), + "the {id} kernel produced {} but the function declares {result_dtype}", + values.dtype(), + ); + + Ok(()) +} + +/// Cast only the output nullability after its shape, dtype, and validity are accepted. +fn cast_output_nullability(result_dtype: &DType, values: ArrayRef) -> VortexResult { + if values.dtype() == result_dtype { + Ok(values) + } else { + values.cast(result_dtype.clone()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs new file mode 100644 index 00000000000..4bb238564ba --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batch execution around a non-null row kernel. +//! +//! A row kernel handles typed values for one row. This module adds the columnar concerns around it: +//! planning the output and null strategy, preserving batch constants and encodings, propagating +//! strict validity, selecting an execution strategy, and validating the finished output. +//! +//! [`BatchPlan`] carries the nullable execution strategy selected by a concrete dispatch. [`Batch`] +//! applies that strategy, and [`BorrowedExecutionArgs`] pairs each kernel invocation with its +//! planning metadata. + +mod args; +pub(super) use args::BorrowedExecutionArgs; + +mod execution; +pub(super) use execution::Batch; +pub(super) use execution::finalize_kernel_output; + +pub(super) use super::visitor::BatchPlan; +pub(super) use super::visitor::RowPolicy; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs new file mode 100644 index 00000000000..6919028715e --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -0,0 +1,587 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use rstest::rstest; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::super::execute::RowExecution; +use super::Batch; +use super::BatchPlan; +use super::RowPolicy; +use super::finalize_kernel_output; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; +use crate::scalar_fn::unstable::row::row_fn_return_dtype; +use crate::validity::Validity; + +#[derive(Clone)] +struct RetryConstantAdd; + +#[derive(Clone)] +struct NullarySeven; + +#[derive(Clone)] +struct AddThree; + +#[derive(Clone)] +struct Identity; + +#[derive(Clone)] +struct SinkOptions; + +struct OptionsCheckingSink; + +#[derive(Clone)] +struct InvalidKernelOutput; + +/// Deliberately violates [`OutputElement::build`] to test validation at the public boundary. +struct NullProducingI64(i64); + +#[derive(Clone)] +struct PreparedAdd { + visit: PreparedVisit, + prepares: Arc, +} + +#[derive(Clone, Copy)] +enum PreparedVisit { + Owned, + Sink, + Deferred, +} + +// SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or +// `finish` through the executor. The row-initialization requirements are therefore vacuous. +unsafe impl OutputSink for OptionsCheckingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn sink_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { + if !enabled { + vortex_bail!(InvalidArgument: "the test sink is disabled"); + } + + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + vortex_bail!("the planning-only test sink must not be allocated") + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { + true + } + + unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} + + unsafe fn finish(self) -> VortexResult { + vortex_bail!("the planning-only test sink must not finish") + } +} + +impl OutputElement for NullProducingI64 { + fn element_dtype() -> DType { + DType::from(i64::PTYPE) + } + + fn build(values: Vec) -> ArrayRef { + let values: Vec<_> = values.into_iter().map(|value| value.0).collect(); + let validity = Validity::from_iter((0..values.len()).map(|index| index != 0)); + + PrimitiveArray::new(values, validity).into_array() + } +} + +struct I64Sink(BufferMut); + +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +impl RowFn for NullarySeven { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.nullary_seven"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), I64Sink, _>(|(), output| { + *output = 7; + }) + } +} + +impl RowFn for AddThree { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["first", "second", "third"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.add_three"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64, i64), i64>(|(first, second, third)| first + second + third) + } +} + +impl RowFn for RetryConstantAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.retry_constant_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(u8, u8), u8, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "checked add overflowed")); + } + + Ok(()) + }, + ) + } +} + +impl RowFn for Identity { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.identity"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } +} + +impl RowFn for SinkOptions { + type Options = bool; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.sink_options"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), OptionsCheckingSink, _>(|(), ()| ()) + } +} + +impl RowFn for InvalidKernelOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_kernel_output"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), NullProducingI64>(|(value,)| NullProducingI64(value)) + } +} + +impl RowFn for PreparedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.prepared_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + let prepares = Arc::clone(&self.prepares); + let prepare = move |(_lhs, rhs): (Option, Option)| { + prepares.fetch_add(1, Ordering::Relaxed); + rhs + }; + + match self.visit { + PreparedVisit::Owned => visitor + .visit_prepared::<(i64, i64), i64, _>(prepare, |constant_rhs, (lhs, rhs)| { + lhs.wrapping_add(constant_rhs.unwrap_or(rhs)) + }), + PreparedVisit::Sink => visitor.visit_prepared_into::<(i64, i64), I64Sink, _, ()>( + prepare, + |constant_rhs, (lhs, rhs), output| { + *output = lhs.wrapping_add(constant_rhs.unwrap_or(rhs)); + }, + ), + PreparedVisit::Deferred => visitor.visit_prepared_deferred::<(i64, i64), i64, _, bool>( + prepare, + |constant_rhs, (lhs, rhs)| lhs.overflowing_add(constant_rhs.unwrap_or(rhs)), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "prepared add overflowed")); + } + + Ok(()) + }, + ), + } + } +} + +#[test] +fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.row_batch"); + + let input = PrimitiveArray::new(vec![1i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let result = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::Dense, + }) + }); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let result = execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![1, u8::MAX], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1_u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::new(vec![2_u8, 0], Validity::from_iter([true, false])); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} + +#[test] +fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&Identity, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid([true, true])] +#[case::all_invalid([false, false])] +fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.resolve_validity"); + + let validity = Validity::Array(BoolArray::from_iter(validity).into_array()); + let input = PrimitiveArray::new(vec![4_i64, 5], validity).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_filters_and_scatters() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.filter_and_scatter"); + + let input = PrimitiveArray::new( + vec![10_i64, 20, 30, 40], + Validity::from_iter([true, false, true, false]), + ) + .into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 4); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.finalize_kernel_output"); + + let values = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let result_dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); + let mut ctx = array_session().create_execution_ctx(); + + let actual = finalize_kernel_output(*ID, &result_dtype, 2, values.clone(), &mut ctx)?; + let expected = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + assert_eq!(actual.dtype(), &result_dtype); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + assert!(finalize_kernel_output(*ID, &result_dtype, 3, values, &mut ctx).is_err()); + + let bools = BoolArray::from_iter([true, false]).into_array(); + assert!(finalize_kernel_output(*ID, &result_dtype, 2, bools, &mut ctx).is_err()); + Ok(()) +} + +#[test] +fn test_sink_dtype_receives_function_options() -> VortexResult<()> { + assert_eq!( + row_fn_return_dtype(&SinkOptions, &true, &[])?, + DType::from(i64::PTYPE) + ); + assert!(row_fn_return_dtype(&SinkOptions, &false, &[]).is_err()); + Ok(()) +} + +#[test] +fn test_nonnullable_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + + assert_invalid_kernel_output(input) +} + +#[test] +fn test_all_valid_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + + assert_invalid_kernel_output(input) +} + +#[track_caller] +fn assert_invalid_kernel_output(input: ArrayRef) -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + let execution = execute_rows(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); + let error = match execution { + Err(error) => error, + Ok(output) => match output.execute::(&mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an invalid row kernel output passed boundary validation"), + }, + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_kernel_output"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("row kernel produced nulls for valid rows"), + "the boundary error must identify invalid row output, got {error}", + ); + Ok(()) +} + +#[rstest] +#[case::owned_constant(PreparedVisit::Owned, true)] +#[case::owned_per_row(PreparedVisit::Owned, false)] +#[case::sink_constant(PreparedVisit::Sink, true)] +#[case::sink_per_row(PreparedVisit::Sink, false)] +#[case::deferred_constant(PreparedVisit::Deferred, true)] +#[case::deferred_per_row(PreparedVisit::Deferred, false)] +fn test_prepared_visits( + #[case] visit: PreparedVisit, + #[case] constant_rhs: bool, +) -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let rhs = if constant_rhs { + ConstantArray::new(3_i64, 2).into_array() + } else { + PrimitiveArray::from_iter([3_i64, 4]).into_array() + }; + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let prepares = Arc::new(AtomicUsize::new(0)); + let function = PreparedAdd { + visit, + prepares: Arc::clone(&prepares), + }; + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; + let expected = if constant_rhs { + PrimitiveArray::from_iter([4_i64, 5]).into_array() + } else { + PrimitiveArray::from_iter([4_i64, 6]).into_array() + }; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(prepares.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[test] +fn test_nullary_row_function_broadcasts() -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([7i64, 7, 7]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_owned_execution_traverses_three_per_row_inputs() -> VortexResult<()> { + let args = VecExecutionArgs::new( + vec![ + PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(), + PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(), + PrimitiveArray::from_iter([100_i64, 200, 300]).into_array(), + ], + 3, + ); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&AddThree, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([111_i64, 222, 333]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs new file mode 100644 index 00000000000..73722549f41 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -0,0 +1,19 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Row-loop execution for owned outputs and output sinks. +//! +//! [`owned`] stores one independent value per row and reduces compact failure evidence. [`sink`] +//! drives output builders whose row handles may share batch state. Both return [`RowExecution`], +//! which distinguishes a completed array from a deferred error that batch validity may suppress. + +mod owned; +pub(super) use owned::execute_owned; +pub(super) use owned::execute_owned_infallible; + +mod outcome; +pub use outcome::RowExecution; + +mod sink; +pub(super) use sink::execute_sink; +pub(super) use sink::execute_sink_valid_rows; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs new file mode 100644 index 00000000000..fc013e7a317 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The result of a completed row loop before batch-level null handling. +//! +//! [`RowExecution`] preserves deferred failure evidence until batch execution can determine whether +//! the failing payload belonged to a valid row. + +use vortex_error::VortexError; +use vortex_error::VortexResult; + +use crate::ArrayRef; + +/// The outcome of a row loop before batch execution decides whether an error is observable. +/// +/// Together with the surrounding [`VortexResult`], this represents three outcomes: +/// +/// - `Err(error)` is a non-retryable execution or immediate row error. +/// - [`Output`](Self::Output) is a successful row loop. +/// - [`DeferredError`](Self::DeferredError) is failure evidence from a completed row loop. +/// +/// A dense loop can evaluate null payloads, so its deferred error is not always observable. Batch +/// execution can retry only valid rows to discard errors caused by null payloads. A plain +/// `VortexResult` cannot distinguish these errors from failures that a retry cannot fix. +/// +/// Once execution is known to contain only valid rows, converting this outcome into a +/// `VortexResult` turns [`DeferredError`](Self::DeferredError) into an ordinary error. +pub enum RowExecution { + /// The successfully built, full-length output column. + Output(ArrayRef), + + /// An error constructed from failure evidence reduced across a completed row loop. + DeferredError(VortexError), +} + +impl From for VortexResult { + fn from(execution: RowExecution) -> Self { + match execution { + RowExecution::Output(output) => Ok(output), + RowExecution::DeferredError(error) => Err(error), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs new file mode 100644 index 00000000000..3ac0c553891 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution that stores one owned output value per row. + +use std::ops::BitOrAssign; + +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; + +/// Zero-sized evidence used to erase failure reduction from infallible owned visits. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + +/// Decode every input column once, then store one infallible owned output per row. +pub fn execute_owned_infallible( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, +{ + execute_owned::( + args, + ctx, + prepare, + move |prepared, args| (apply(prepared, args), NoFailure), + |_| Ok(()), + ) +} + +/// Decode every input column once, then store owned row outputs and reduce deferred failures. +pub fn execute_owned( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, +{ + const { assert_owned_output_needs_no_drop::() }; + + // Keep the vector length at zero until every row succeeds. An unwind then abandons partially + // initialized spare capacity without treating it as initialized output. The no-drop assertion + // above proves that no initialized value requires its destructor to run. + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let failure; + + { + let output = &mut values.spare_capacity_mut()[..row_count]; + + // When every input stores one value per row, the indexed source removes argument-shape + // dispatch from the hot loop and lets the lane kernel optimize the traversal as one + // operation. Keep view construction and its length proof in this branch. Hoisting them + // through the shared validation helper changed mixed-constant add, subtract, and multiply + // from 9.219, 9.229, and 18.94 us to 30.46, 31.11, and 37.73 us on a Ryzen 9 7950X with + // rustc 1.91.0 and LLVM 21.1.2. + // Restoring this placement recovered the fast code under the 16-CGU, no-LTO bench profile. + if let Some(views) = Args::per_row_views(&columns) { + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows + // immediately above. + failure = unsafe { Args::indexed_source(views, row_count) } + .map_checked_into(output, |elements| apply(&prepared, elements)); + } else { + // A batch-constant input was collapsed to one row during decoding. This path reads that + // row repeatedly while indexing only the per-row inputs. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // Keep the output-slot iterator as the loop bound. `row_count` is address-taken by the + // validation error formatting above. With rustc 1.97.1 and LLVM 22.1.6 under 16 CGUs + // without LTO, indexing `output` by a `0..row_count` range retains an early-exit bounds + // check and prevents vectorization of mixed constant and per-row arithmetic. Recheck + // the optimized IR and mixed-constant benchmarks before restoring that range loop. + let mut accumulated = Fail::default(); + for (index, slot) in output.iter_mut().enumerate() { + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + slot.write(value); + accumulated |= row_failure; + } + failure = accumulated; + } + } + + // SAFETY: normal completion of either loop initializes every slot in `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + // Failure evidence is reduced inside the loop so its richer error construction stays cold. + // Preserve that provenance so batch execution may retry over only valid rows. + match finish_failure(failure) { + Ok(()) => Ok(RowExecution::Output(Out::build(values))), + Err(error) => Ok(RowExecution::DeferredError(error)), + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs new file mode 100644 index 00000000000..fce891cf613 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -0,0 +1,262 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution that writes through an output sink. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::SinkResult; + +/// Ensure that every decoded input addresses the complete row loop. +fn ensure_decoded_lengths( + columns: &Args::Columns, + views: Option<&Args::Views<'_>>, + row_count: usize, +) -> VortexResult<()> { + let lengths_match = match views { + Some(views) => Args::view_lens_match(views, row_count), + None => Args::decoded_lens_match(columns, row_count), + }; + vortex_ensure!( + lengths_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + Ok(()) +} + +/// Decode every input column once, allocate the sink once, then write one row at a time. +/// +/// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state +/// does not need to be captured by the closure. +pub fn execute_sink( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + let row_count = args.row_count(); + let mut sink = >::with_capacity(row_count, sink_dtype)?; + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + + { + // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This + // scope releases the borrow before `finish_sink` consumes the sink. + let mut rows = >::rows(&mut sink); + vortex_ensure!( + >::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + // The all-per-row representation removes argument-shape dispatch from the hot loop. The + // mixed path instead reads collapsed batch constants at row zero. + if let Some(views) = views { + for index in 0..row_count { + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before + // the loop. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + // SAFETY: `row_count_matches` proved the sink addresses every loop index. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + apply(&prepared, elements, output).into_result()?; + } + } else { + for index in 0..row_count { + // SAFETY: `row_count_matches` proved the sink addresses every loop index. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + apply(&prepared, Args::get(&columns, index), output).into_result()?; + } + } + } + + finish_sink::(sink) +} + +/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +pub fn execute_sink_valid_rows( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + valid: &Mask, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult> +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + // Decline before input decoding or sink allocation when this sink cannot initialize rows that + // the mask skips. The capability and the operation are the same function pointer. + let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + else { + return Ok(None); + }; + + // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An + // element representation may decline when it cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + let prepared = prepare(Args::constants(&columns)); + let row_count = args.row_count(); + let mut sink = >::with_capacity(row_count, sink_dtype)?; + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid) = valid.bit_buffer() else { + vortex_bail!("execute_sink_valid_rows requires a mixed mask"); + }; + vortex_ensure!( + valid.len() == row_count, + "the validity mask does not address exactly {row_count} rows", + ); + + { + let mut rows = >::rows(&mut sink); + vortex_ensure!( + >::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + + // The loop writes only valid indices, but the sink still finishes a full-length output. + // Initialize placeholders now; batch execution masks them before the result escapes. + initialize_skipped_rows(&mut rows); + + // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first + // immediate error, turn later callbacks into no-ops, and return before finishing the sink. + let mut error = None; + valid.for_each_set_index(|index| { + if error.is_some() { + return; + } + + // SAFETY: `row_count_matches` proved that the sink addresses every mask index, which + // is below the mask's validated `row_count`. + let output = unsafe { >::row_unchecked(&mut rows, index) }; + let result = match &views { + Some(views) => { + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows, and + // mask indices are below `row_count`. + let elements = unsafe { Args::get_from_views_unchecked(views, index) }; + apply(&prepared, elements, output) + } + None => apply(&prepared, Args::get(&columns, index), output), + }; + if let Err(err) = result.into_result() { + error = Some(err); + } + }); + + if let Some(error) = error { + return Err(error); + } + } + + finish_sink::(sink).map(Some) +} + +fn finish_sink(sink: S) -> VortexResult +where + S: OutputSink, +{ + // SAFETY: callers reach this helper only after every completed callback returned the sink's + // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. + // The sink contract defines how that evidence establishes initialization of its row storage. + unsafe { >::finish(sink) }.map(RowExecution::Output) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_mask::Mask; + + use super::execute_sink_valid_rows; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + use crate::dtype::NativePType; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::OutputSink; + use crate::validity::Validity; + + struct NonSkippingSink; + + // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or + // `finish` through the executor. The row-initialization requirements are therefore vacuous. + unsafe impl OutputSink for NonSkippingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + Err(vortex_err!( + "a non-skipping sink must decline before allocation" + )) + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { + true + } + + unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> { + } + + unsafe fn finish(self) -> VortexResult { + Err(vortex_err!("a non-skipping sink must not finish")) + } + } + + #[test] + fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([true, false]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( + &args, + &DType::from(i64::PTYPE), + &valid, + &mut ctx, + |_| (), + |_, _, _| (), + )?; + + assert!(execution.is_none()); + Ok(()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index bcb3a008488..e62f16639bc 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -15,6 +15,11 @@ //! reduce compact failure evidence in that loop and retry only valid rows when null payloads may //! have caused the failure. +mod execute; +pub use execute::RowExecution; + +mod batch; + mod row_fn; pub use row_fn::RowFn; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 51d66594332..7a8a8e92e41 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -20,3 +20,4 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; +pub use tuple::batch_constant; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index a2c143704a0..69b5cf686f6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -8,6 +8,7 @@ mod element_tuple; pub use element_tuple::ElementTuple; +pub use element_tuple::batch_constant; mod indexed; pub use indexed::IndexedElementTuple; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index 044ad15fd4b..560e1e48f4f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_bail; use vortex_mask::Mask; use super::ElementTuple; -use super::element_tuple::batch_constant; +use super::batch_constant; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs index ce119f32915..e47f195410c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -12,6 +12,7 @@ pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; +pub(super) use element::batch_constant; mod result; pub use result::SinkResult; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs new file mode 100644 index 00000000000..f5a93ed65d9 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -0,0 +1,263 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visitors that execute dense and skip-invalid row loops. +//! +//! Each method verifies that execution selected the nullable policy derived during planning before +//! handing its typed closures to the matching loop. Valid-row execution can decline without +//! running a loop. Batch execution then filters the inputs and retries the dense loop. + +use std::marker::PhantomData; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; + +use super::RowPolicy; +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::row_visitor::private; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::execute::execute_owned; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_sink; +use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; + +/// The run-time visit that decodes every column once and runs the selected row loop. +pub struct ExecuteRows<'args, 'ctx, F> { + /// The inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'args, 'ctx, F> ExecuteRows<'args, 'ctx, F> { + pub fn new( + args: &'args dyn ExecutionArgs, + output_dtype: &'args DType, + policy: RowPolicy, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + output_dtype, + policy, + ctx, + function: PhantomData, + } + } +} + +impl private::Sealed for ExecuteRows<'_, '_, F> {} + +impl RowVisitor for ExecuteRows<'_, '_, F> { + type VisitResult = RowExecution; + + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + ensure_policy(self.policy, RowPolicy::for_owned_output::())?; + + execute_owned_infallible::(self.args, self.ctx, prepare, apply) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_policy(self.policy, RowPolicy::for_sink::())?; + + execute_sink::( + self.args, + self.output_dtype, + self.ctx, + prepare, + apply, + ) + } + + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + ensure_policy(self.policy, RowPolicy::for_deferred_output::())?; + + execute_owned::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} + +/// The run-time visit that tries skip-invalid execution over the original input columns. +/// +/// Only output sinks have a contract for skipped output positions. Owned visits therefore decline +/// so batch execution can use its filter-and-scatter fallback. +pub struct ExecuteValidRows<'args, 'ctx, F> { + /// The original inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The conjoined validity, materialized by batch execution and guaranteed mixed. + valid: &'args Mask, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'args, 'ctx, F> ExecuteValidRows<'args, 'ctx, F> { + pub fn new( + args: &'args dyn ExecutionArgs, + output_dtype: &'args DType, + policy: RowPolicy, + valid: &'args Mask, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + output_dtype, + policy, + valid, + ctx, + function: PhantomData, + } + } +} + +impl private::Sealed for ExecuteValidRows<'_, '_, F> {} + +impl RowVisitor for ExecuteValidRows<'_, '_, F> { + type VisitResult = Option; + + fn visit_prepared( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + ensure_policy(self.policy, RowPolicy::for_owned_output::())?; + + // Owned execution has no sink that can initialize skipped output positions. Decline so + // batch execution filters the inputs and retries with the dense visitor. + Ok(None) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_policy(self.policy, RowPolicy::for_sink::())?; + + execute_sink_valid_rows::( + self.args, + self.output_dtype, + self.valid, + self.ctx, + prepare, + apply, + ) + } + + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + _finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + ensure_policy(self.policy, RowPolicy::for_deferred_output::())?; + + // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. + Ok(None) + } +} + +/// Validate that execution selected the nullable policy used to build the batch plan. +fn ensure_policy(planned: RowPolicy, actual: RowPolicy) -> VortexResult<()> { + vortex_ensure_eq!( + actual, + planned, + "row dispatch must select the planned nullable execution policy: planned {planned:?}, got {actual:?}", + ); + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index 57da5f4691b..c7f9baf6a62 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -6,9 +6,16 @@ //! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch mod check; +pub(super) use check::assert_owned_output_needs_no_drop; + +mod execute; +pub(super) use execute::ExecuteRows; +pub(super) use execute::ExecuteValidRows; mod plan; +pub(super) use plan::BatchPlan; pub(super) use plan::BatchPlanner; +pub(super) use plan::RowPolicy; mod row_visitor; pub use row_visitor::RowVisitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index c522376d491..214d81c29f4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -114,15 +114,12 @@ pub(crate) struct BatchPlan { pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. - // TODO(connor)[RowFn]: Remove this allowance when the execution backend from #9130 consumes - // this policy. - #[allow(dead_code)] pub(crate) policy: RowPolicy, } impl BatchPlan { /// Return the output dtype widened with strict input nullability. - pub(crate) fn result_dtype(self, args: &[DType]) -> DType { + pub(crate) fn result_dtype(&self, args: &[DType]) -> DType { let nullability = self.output_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index b6c14585c48..55e2a197c10 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -4,12 +4,13 @@ //! Adapts [`RowFn`] implementations to the scalar-function interface. //! //! The blanket [`ScalarFnVTable`] implementation supplies common arity, validity, fallibility, and -//! execution behavior. [`row_fn_return_dtype`] and [`execute_rows`] expose the same planning and -//! execution paths to public vtables that delegate to a private row kernel. +//! execution behavior. The visitor layer validates and executes the concrete signature selected by +//! dispatch. [`row_fn_return_dtype`] and [`execute_rows`] expose the same paths to public vtables +//! that delegate to a private row kernel. use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; use vortex_session::VortexSession; use super::row_fn::RowFn; @@ -24,6 +25,12 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::unstable::row::batch::Batch; +use crate::scalar_fn::unstable::row::batch::BorrowedExecutionArgs; +use crate::scalar_fn::unstable::row::batch::finalize_kernel_output; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::visitor::ExecuteRows; +use crate::scalar_fn::unstable::row::visitor::ExecuteValidRows; impl ScalarFnVTable for F { type Options = F::Options; @@ -98,16 +105,36 @@ pub fn row_fn_return_dtype( /// delegate row execution to a private `RowFn` kernel through this function. pub fn execute_rows( function: &F, - _options: &F::Options, + options: &F::Options, args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { ensure_arity(function, args.num_inputs())?; - // TODO(connor)[RowFn]: Replace this temporary error with the execution backend in #9129. - vortex_bail!( - "Row function {} does not yet have an execution backend", - RowFn::id(function) + // Nullary functions have no input validity to propagate, so they skip batch execution. + if args.num_inputs() == 0 { + let plan = function.dispatch(options, &[], BatchPlanner::::new(&[], options))?; + let result_dtype = plan.result_dtype(&[]); + let nullary_args = + BorrowedExecutionArgs::new(&[], args.row_count(), &[], &plan.output_dtype, plan.policy); + + let execution = execute_row_kernel(function, options, nullary_args, ctx)?; + let values = VortexResult::from(execution)?; + + return finalize_kernel_output( + RowFn::id(function), + &result_dtype, + args.row_count(), + values, + ctx, + ); + } + + let batch = prepare_batch(function, options, args)?; + batch.execute( + |args, ctx| execute_row_kernel(function, options, args, ctx), + |args, valid, ctx| try_execute_rows_unfiltered(function, options, args, valid, ctx), + ctx, ) } @@ -124,26 +151,82 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } +/// Execute the row loop selected by dispatch. +fn execute_row_kernel( + function: &F, + options: &F::Options, + args: BorrowedExecutionArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + function.dispatch( + options, + args.dtypes(), + ExecuteRows::::new(&args, args.output_dtype(), args.policy(), ctx), + ) +} + +/// Try execution against the original inputs, returning `None` when batch execution must filter. +fn try_execute_rows_unfiltered( + function: &F, + options: &F::Options, + args: BorrowedExecutionArgs<'_>, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + function.dispatch( + options, + args.dtypes(), + ExecuteValidRows::::new(&args, args.output_dtype(), args.policy(), valid, ctx), + ) +} + +/// Prepare the batch inputs and execution plan for `function`. +fn prepare_batch( + function: &F, + options: &F::Options, + args: &dyn ExecutionArgs, +) -> VortexResult { + Batch::new(RowFn::id(function), args, |arg_dtypes| { + function.dispatch( + options, + arg_dtypes, + BatchPlanner::::new(arg_dtypes, options), + ) + }) +} + #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use super::execute_rows; use super::row_fn_return_dtype; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::RowVisitor; + use crate::validity::Validity; #[derive(Clone)] struct IndexingRowFn; + #[derive(Clone)] + struct ChangingDispatchRowFn { + dispatches: Arc, + } + impl RowFn for IndexingRowFn { type Options = EmptyOptions; @@ -168,6 +251,31 @@ mod tests { } } + impl RowFn for ChangingDispatchRowFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.changing_dispatch_row_fn"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { + visitor.visit::<(i64,), i64>(|(value,)| value) + } else { + visitor.visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())) + } + } + } + #[test] fn test_return_dtype_rejects_wrong_arity_before_dispatch() { let error = row_fn_return_dtype(&IndexingRowFn, &EmptyOptions, &[]) @@ -186,6 +294,29 @@ mod tests { assert_arity_error(error); } + #[test] + fn test_execute_rejects_dispatch_that_changes_after_planning() { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&function, &EmptyOptions, &args, &mut ctx) + .expect_err("dispatch must not change after planning"); + let message = error.to_string(); + + assert!( + message.contains("row dispatch must select the planned nullable execution policy"), + "unexpected error: {error}", + ); + assert!( + message.contains("planned Dense, got DenseWithRetry"), + "unexpected error: {error}", + ); + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( From edab9a7f5730a57e1deb23e59d0136e894c4e3e8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 17:01:48 -0400 Subject: [PATCH 103/160] Clarify RowFn kernel invocation docs Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/unstable/row/execute/owned.rs | 4 ++-- vortex-array/src/scalar_fn/unstable/row/execute/sink.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index 3ac0c553891..57a0b3bf67a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -24,7 +24,7 @@ impl BitOrAssign for NoFailure { fn bitor_assign(&mut self, _rhs: Self) {} } -/// Decode every input column once, then store one infallible owned output per row. +/// Decode every input column for one kernel invocation, then store one infallible output per row. pub fn execute_owned_infallible( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, @@ -44,7 +44,7 @@ where ) } -/// Decode every input column once, then store owned row outputs and reduce deferred failures. +/// Decode every input column for one kernel invocation, then store outputs and reduce failures. pub fn execute_owned( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index fce891cf613..606a6fa262f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -35,7 +35,7 @@ fn ensure_decoded_lengths( Ok(()) } -/// Decode every input column once, allocate the sink once, then write one row at a time. +/// Decode every input column and allocate one sink for one kernel invocation. /// /// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state /// does not need to be captured by the closure. From 893a1cfd949de1d690174c51eb9fbb2ce4ac3e30 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 19:23:37 -0400 Subject: [PATCH 104/160] Harden RowFn execution edge cases Signed-off-by: Connor Tsui --- vortex-array/src/arrays/masked/tests.rs | 31 +++++ vortex-array/src/arrays/masked/vtable/mod.rs | 5 + .../src/scalar_fn/unstable/row/batch/args.rs | 5 +- .../scalar_fn/unstable/row/batch/execution.rs | 6 +- .../src/scalar_fn/unstable/row/batch/tests.rs | 23 +--- .../scalar_fn/unstable/row/execute/owned.rs | 6 +- .../scalar_fn/unstable/row/execute/sink.rs | 129 ++++++++++++++++-- .../scalar_fn/unstable/row/visitor/execute.rs | 110 +++++++++++---- .../src/scalar_fn/unstable/row/vtable.rs | 59 +++++++- 9 files changed, 311 insertions(+), 63 deletions(-) diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index 92ec4eb474f..55af652725b 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -4,17 +4,22 @@ use rstest::rstest; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use super::*; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::Array; use crate::array_session; +use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; +use crate::dtype::NativePType; use crate::dtype::Nullability; +use crate::scalar::Scalar; use crate::validity::Validity; #[rstest] @@ -55,6 +60,32 @@ fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> { Ok(()) } +#[test] +fn test_try_from_parts_rejects_null_child() -> VortexResult<()> { + let child = PrimitiveArray::from_iter([1_i64]).into_array(); + let masked = MaskedArray::try_new(child, Validity::AllValid)?; + let mut parts = match masked.try_into_parts() { + Ok(parts) => parts, + Err(_) => vortex_bail!("the uniquely owned masked array must expose its parts"), + }; + let dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); + parts.slots[MaskedSlots::CHILD] = + Some(ConstantArray::new(Scalar::null(dtype), parts.len).into_array()); + + let error = match Array::::try_from_parts(parts) { + Err(error) => error, + Ok(_) => vortex_bail!("rebuilding must reject a child containing nulls"), + }; + + assert!( + error + .to_string() + .contains("MaskedArray children must not have nulls"), + "unexpected error: {error}", + ); + Ok(()) +} + #[test] fn test_masked_child_with_validity() { // When validity has nulls, masked_child should apply inverted mask. diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index c7e32a7ee3c..c0608388135 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -73,6 +73,7 @@ impl VTable for Masked { *ID } + #[allow(clippy::disallowed_methods)] fn validate( &self, _data: &MaskedData, @@ -92,6 +93,10 @@ impl VTable for Masked { child.dtype().as_nullable() == *dtype, "MaskedArray dtype does not match child and validity" ); + vortex_ensure!( + child.all_valid(&mut legacy_session().create_execution_ctx())?, + "MaskedArray children must not have nulls", + ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index d5ffe042613..e3759b0dcf1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A borrowed execution view passed to one row-kernel invocation. +//! Execution arguments paired with the metadata selected during planning. +//! +//! [`BorrowedExecutionArgs`] can point at original, sliced, or filtered arrays while retaining the +//! dtypes, output dtype, and null policy of the original batch plan. use vortex_error::VortexResult; use vortex_error::vortex_err; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs index 671b8f4ecb0..0f880bf606e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Null propagation, constant folding, and strategy execution for one columnar batch. +//! Applies columnar semantics around one typed row kernel invocation. +//! +//! [`Batch`] owns strict null propagation, constant broadcasting, execution strategy selection, and +//! output validation. The row kernel therefore handles only decoded values and its selected output +//! capability. use smallvec::SmallVec; use vortex_error::VortexResult; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 6919028715e..35d948a8833 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -479,22 +479,13 @@ fn test_sink_dtype_receives_function_options() -> VortexResult<()> { Ok(()) } -#[test] -fn test_nonnullable_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { - let input = PrimitiveArray::from_iter([1_i64, 2]).into_array(); - - assert_invalid_kernel_output(input) -} - -#[test] -fn test_all_valid_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { - let input = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); - - assert_invalid_kernel_output(input) -} - -#[track_caller] -fn assert_invalid_kernel_output(input: ArrayRef) -> VortexResult<()> { +#[rstest] +#[case::nonnullable(Validity::NonNullable)] +#[case::all_valid(Validity::AllValid)] +fn test_kernel_output_rejects_nulls_at_function_boundary( + #[case] validity: Validity, +) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = array_session().create_execution_ctx(); let execution = execute_rows(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index 57a0b3bf67a..11fff70417a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Execution that stores one owned output value per row. +//! Executes row kernels that return one independent owned value per row. +//! +//! [`execute_owned`] decodes inputs once, prepares constant state, writes into spare vector +//! capacity, and reduces compact failure evidence without putting error construction in the hot +//! loop. [`execute_owned_infallible`] removes that failure path for infallible kernels. use std::ops::BitOrAssign; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index 606a6fa262f..c32491abedb 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Execution that writes through an output sink. +//! Executes row kernels that write through an [`OutputSink`]. +//! +//! Dense execution visits every row. Skip-invalid execution can instead initialize omitted output +//! positions and visit only the set bits of a mixed validity mask, falling back when either the +//! input representation or sink lacks that capability. use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -17,7 +21,6 @@ use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::SinkResult; -/// Ensure that every decoded input addresses the complete row loop. fn ensure_decoded_lengths( columns: &Args::Columns, views: Option<&Args::Views<'_>>, @@ -133,17 +136,17 @@ where { let mut rows = >::rows(&mut sink); - vortex_ensure!( - >::row_count_matches(&rows, row_count), - "the output sink does not address exactly {row_count} rows", - ); let views = Args::per_row_views(&columns); ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; - // The loop writes only valid indices, but the sink still finishes a full-length output. - // Initialize placeholders now; batch execution masks them before the result escapes. + // Initialize every slot before skipping rows. Recheck addressability afterward because the + // initializer mutably borrows the row representation. initialize_skipped_rows(&mut rows); + vortex_ensure!( + >::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows after initializing skipped rows", + ); // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first // immediate error, turn later callbacks into no-ops, and return before finishing the sink. @@ -191,24 +194,31 @@ where #[cfg(test)] mod tests { use vortex_error::VortexResult; + use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_mask::Mask; + use super::RowExecution; use super::execute_sink_valid_rows; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::NativePType; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::InitializedElement; use crate::scalar_fn::unstable::row::OutputSink; + use crate::scalar_fn::unstable::row::UninitElementSink; use crate::validity::Validity; struct NonSkippingSink; + struct ShrinkingSink(Vec); + // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or // `finish` through the executor. The row-initialization requirements are therefore vacuous. unsafe impl OutputSink for NonSkippingSink { @@ -240,6 +250,45 @@ mod tests { } } + // SAFETY: the initializer deliberately shrinks the row collection to exercise the executor's + // post-initialization length check. If execution incorrectly continues, safe indexing in + // `row_unchecked` panics instead of accessing invalid memory. + unsafe impl OutputSink for ShrinkingSink { + type Rows<'a> = &'a mut Vec; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + rows.pop(); + }) + } + + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(vec![0; rows])) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.0 + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::from_iter(self.0).into_array()) + } + } + #[test] fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); @@ -259,4 +308,68 @@ mod tests { assert!(execution.is_none()); Ok(()) } + + #[test] + fn test_skip_invalid_sink_initializes_and_writes_addressed_rows() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let valid = Mask::from_iter([true, false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::< + (i64,), + (), + UninitElementSink, + InitializedElement, + EmptyOptions, + >( + &args, + &DType::from(i64::PTYPE), + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + // SAFETY: `output` is the row supplied to this callback. + unsafe { InitializedElement::write(output, value * 2) } + }, + )?; + let Some(RowExecution::Output(actual)) = execution else { + vortex_bail!("the skip-invalid sink must produce an output"); + }; + let expected = PrimitiveArray::from_iter([20_i64, 0, 60]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) + } + + #[test] + fn test_skip_invalid_sink_rechecks_rows_after_initialization() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>( + &args, + &DType::from(i64::PTYPE), + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + *output = value; + }, + ); + + let error = match result { + Err(error) => error, + Ok(_) => vortex_bail!("the sink must reject rows changed by its initializer"), + }; + assert!( + error + .to_string() + .contains("after initializing skipped rows"), + "unexpected error: {error}", + ); + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index f5a93ed65d9..010dc079599 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -3,11 +3,10 @@ //! Visitors that execute dense and skip-invalid row loops. //! -//! Each method verifies that execution selected the nullable policy derived during planning before -//! handing its typed closures to the matching loop. Valid-row execution can decline without -//! running a loop. Batch execution then filters the inputs and retries the dense loop. +//! Each visit revalidates its concrete signature and checks that its output dtype and null policy +//! match the plan before entering a row loop. [`ExecuteValidRows`] can decline unsupported +//! skip-invalid execution so the batch layer filters the inputs and retries with [`ExecuteRows`]. -use std::marker::PhantomData; use std::ops::BitOrAssign; use vortex_error::VortexResult; @@ -19,6 +18,8 @@ use super::RowVisitor; use super::check::assert_deferred_visit_contract; use super::check::assert_owned_visit_contract; use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; use super::row_visitor::private; use crate::ExecutionCtx; use crate::dtype::DType; @@ -36,10 +37,16 @@ use crate::scalar_fn::unstable::row::execute::execute_sink; use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; /// The run-time visit that decodes every column once and runs the selected row loop. -pub struct ExecuteRows<'args, 'ctx, F> { +pub struct ExecuteRows<'args, 'ctx, F: RowFn> { /// The inputs for this kernel invocation. args: &'args dyn ExecutionArgs, + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + /// The output dtype computed by the planning visit. output_dtype: &'args DType, @@ -48,29 +55,29 @@ pub struct ExecuteRows<'args, 'ctx, F> { /// The execution context used to decode the input columns. ctx: &'ctx mut ExecutionCtx, - - /// The visited function, carried only so the dispatch check can name its contract. - function: PhantomData, } -impl<'args, 'ctx, F> ExecuteRows<'args, 'ctx, F> { +impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { pub fn new( args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, output_dtype: &'args DType, policy: RowPolicy, ctx: &'ctx mut ExecutionCtx, ) -> Self { Self { args, + dtypes, + options, output_dtype, policy, ctx, - function: PhantomData, } } } -impl private::Sealed for ExecuteRows<'_, '_, F> {} +impl private::Sealed for ExecuteRows<'_, '_, F> {} impl RowVisitor for ExecuteRows<'_, '_, F> { type VisitResult = RowExecution; @@ -85,7 +92,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { Out: OutputElement, { const { assert_owned_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_owned_output::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; execute_owned_infallible::(self.args, self.ctx, prepare, apply) } @@ -105,7 +117,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { ApplyResult: SinkResult>::WriteToken>, { const { assert_sink_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_sink::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; execute_sink::( self.args, @@ -128,7 +145,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { Fail: Copy + Default + BitOrAssign, { const { assert_deferred_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_deferred_output::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; execute_owned::( self.args, @@ -144,10 +166,16 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { /// /// Only output sinks have a contract for skipped output positions. Owned visits therefore decline /// so batch execution can use its filter-and-scatter fallback. -pub struct ExecuteValidRows<'args, 'ctx, F> { +pub struct ExecuteValidRows<'args, 'ctx, F: RowFn> { /// The original inputs for this kernel invocation. args: &'args dyn ExecutionArgs, + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + /// The output dtype computed by the planning visit. output_dtype: &'args DType, @@ -159,14 +187,13 @@ pub struct ExecuteValidRows<'args, 'ctx, F> { /// The execution context used to decode the input columns. ctx: &'ctx mut ExecutionCtx, - - /// The visited function, carried only so the dispatch check can name its contract. - function: PhantomData, } -impl<'args, 'ctx, F> ExecuteValidRows<'args, 'ctx, F> { +impl<'args, 'ctx, F: RowFn> ExecuteValidRows<'args, 'ctx, F> { pub fn new( args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, output_dtype: &'args DType, policy: RowPolicy, valid: &'args Mask, @@ -174,16 +201,17 @@ impl<'args, 'ctx, F> ExecuteValidRows<'args, 'ctx, F> { ) -> Self { Self { args, + dtypes, + options, output_dtype, policy, valid, ctx, - function: PhantomData, } } } -impl private::Sealed for ExecuteValidRows<'_, '_, F> {} +impl private::Sealed for ExecuteValidRows<'_, '_, F> {} impl RowVisitor for ExecuteValidRows<'_, '_, F> { type VisitResult = Option; @@ -198,7 +226,12 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { Out: OutputElement, { const { assert_owned_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_owned_output::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; // Owned execution has no sink that can initialize skipped output positions. Decline so // batch execution filters the inputs and retries with the dense visitor. @@ -220,7 +253,12 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { ApplyResult: SinkResult>::WriteToken>, { const { assert_sink_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_sink::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; execute_sink_valid_rows::( self.args, @@ -244,19 +282,33 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { Fail: Copy + Default + BitOrAssign, { const { assert_deferred_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_deferred_output::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. Ok(None) } } -/// Validate that execution selected the nullable policy used to build the batch plan. -fn ensure_policy(planned: RowPolicy, actual: RowPolicy) -> VortexResult<()> { +fn ensure_plan( + planned_output: &DType, + planned_policy: RowPolicy, + actual_output: DType, + actual_policy: RowPolicy, +) -> VortexResult<()> { + vortex_ensure_eq!( + actual_policy, + planned_policy, + "row dispatch must select the planned nullable execution policy: planned {planned_policy:?}, got {actual_policy:?}", + ); vortex_ensure_eq!( - actual, - planned, - "row dispatch must select the planned nullable execution policy: planned {planned:?}, got {actual:?}", + actual_output, + *planned_output, + "row dispatch must select the planned output dtype: planned {planned_output}, got {actual_output}", ); Ok(()) diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 55e2a197c10..9675219d547 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -138,7 +138,6 @@ pub fn execute_rows( ) } -/// Validate the number of arguments before calling user-defined dispatch code. fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { let expected = F::ARG_NAMES.len(); vortex_ensure_eq!( @@ -151,7 +150,6 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } -/// Execute the row loop selected by dispatch. fn execute_row_kernel( function: &F, options: &F::Options, @@ -161,11 +159,17 @@ fn execute_row_kernel( function.dispatch( options, args.dtypes(), - ExecuteRows::::new(&args, args.output_dtype(), args.policy(), ctx), + ExecuteRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + ctx, + ), ) } -/// Try execution against the original inputs, returning `None` when batch execution must filter. fn try_execute_rows_unfiltered( function: &F, options: &F::Options, @@ -176,11 +180,18 @@ fn try_execute_rows_unfiltered( function.dispatch( options, args.dtypes(), - ExecuteValidRows::::new(&args, args.output_dtype(), args.policy(), valid, ctx), + ExecuteValidRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + valid, + ctx, + ), ) } -/// Prepare the batch inputs and execution plan for `function`. fn prepare_batch( function: &F, options: &F::Options, @@ -225,6 +236,13 @@ mod tests { #[derive(Clone)] struct ChangingDispatchRowFn { dispatches: Arc, + change: DispatchChange, + } + + #[derive(Clone, Copy)] + enum DispatchChange { + Policy, + Element, } impl RowFn for IndexingRowFn { @@ -271,7 +289,11 @@ mod tests { if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { visitor.visit::<(i64,), i64>(|(value,)| value) } else { - visitor.visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())) + match self.change { + DispatchChange::Policy => visitor + .visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())), + DispatchChange::Element => visitor.visit::<(u64,), u64>(|(value,)| value), + } } } } @@ -298,6 +320,7 @@ mod tests { fn test_execute_rejects_dispatch_that_changes_after_planning() { let function = ChangingDispatchRowFn { dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Policy, }; let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); let args = VecExecutionArgs::new(vec![input], 2); @@ -317,6 +340,28 @@ mod tests { ); } + #[test] + fn test_execute_revalidates_element_types_after_planning() -> VortexResult<()> { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Element, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&function, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_error::vortex_bail!("dispatch must preserve its planned element types"), + }; + + assert!( + error.to_string().contains("expected a u64 column"), + "unexpected error: {error}", + ); + Ok(()) + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( From 76c5988809de645c9896ba71b946871a616575f2 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 19:25:28 -0400 Subject: [PATCH 105/160] Tighten RowFn executor documentation Signed-off-by: Connor Tsui --- .../scalar_fn/unstable/row/batch/execution.rs | 24 +++++-------------- 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs index 0f880bf606e..93aad88810e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -121,11 +121,10 @@ impl Batch { }) } - /// Add null propagation, constant folding, and strategy selection around `kernel`. + /// Apply constant folding and null handling around `kernel`. /// - /// The kernel may ignore input validity. It receives valid-only rows when required, and its - /// output **must** match the planned dtype up to nullability. `try_unfiltered` receives the - /// originals plus a mixed validity mask; `Ok(None)` selects filter-and-scatter. + /// For a mixed validity mask, `try_unfiltered` may avoid filtering; `Ok(None)` selects + /// filter-and-scatter. Every kernel result is checked against the planned shape and dtype. pub fn execute( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -167,11 +166,7 @@ impl Batch { } } - /// Evaluate a single row of all-constant inputs and broadcast its value. - /// - /// Reconciling the row's dtype before reading the scalar keeps this path on the same - /// kernel/declaration agreement check as the dense and filter paths, rather than letting `cast` - /// paper over a disagreement. + /// Evaluate one row of constant inputs and broadcast the validated result. fn broadcast_one_row( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -191,11 +186,7 @@ impl Batch { Ok(ConstantArray::new(scalar, self.row_count).into_array()) } - /// Run the kernel over every row, including the rows behind nulls, then mask its result. - /// - /// The arguments reach the kernel untouched, so the inputs keep their original encoding, and - /// the conjoined validity is handed to `mask` as an array rather than materialized into a - /// [`Mask`] first. + /// Run every stored payload, then attach the input validity without materializing its mask. fn execute_dense( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -314,8 +305,7 @@ impl Batch { .map(Some) } - /// The filter strategy for a mixed mask: filter every input down to the rows set in `valid`, - /// run the kernel over those, and scatter its results back into a null-padded output. + /// Filter to valid rows, run the kernel, then scatter into a null-padded output. fn filter_and_scatter( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -337,7 +327,6 @@ impl Batch { self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) } - /// An all-null result of the function's declared return dtype. fn all_null(&self) -> ArrayRef { ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() } @@ -357,7 +346,6 @@ impl Batch { ) } - /// Finalize an output against this batch's expected length and declared return dtype. fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { reconcile_output(self.id, &self.result_dtype, expected_len, values) } From 83355fff89f43bb7e57243b50a499ecf9a2bb750 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:18:13 -0400 Subject: [PATCH 106/160] Adapt RowFn framework to explicit contracts Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 15 ++++++++++----- .../src/scalar_fn/unstable/row/execute/sink.rs | 18 ++++++------------ .../scalar_fn/unstable/row/visitor/execute.rs | 13 ++----------- 3 files changed, 18 insertions(+), 28 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 35d948a8833..f801fa35e4c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -82,7 +82,7 @@ unsafe impl OutputSink for OptionsCheckingSink { type Row<'a> = (); type WriteToken = (); - fn sink_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { + fn output_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { if !enabled { vortex_bail!(InvalidArgument: "the test sink is disabled"); } @@ -90,7 +90,7 @@ unsafe impl OutputSink for OptionsCheckingSink { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(_rows: usize) -> VortexResult { vortex_bail!("the planning-only test sink must not be allocated") } @@ -129,11 +129,11 @@ unsafe impl OutputSink for I64Sink { type Row<'a> = &'a mut i64; type WriteToken = (); - fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(rows: usize) -> VortexResult { Ok(Self(BufferMut::zeroed(rows))) } @@ -159,6 +159,7 @@ impl RowFn for NullarySeven { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &[]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.nullary_seven"); @@ -181,6 +182,7 @@ impl RowFn for AddThree { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["first", "second", "third"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.add_three"); @@ -231,6 +233,7 @@ impl RowFn for Identity { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.identity"); @@ -251,6 +254,7 @@ impl RowFn for SinkOptions { type Options = bool; const ARG_NAMES: &'static [&'static str] = &[]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.sink_options"); @@ -271,6 +275,7 @@ impl RowFn for InvalidKernelOutput { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.invalid_kernel_output"); @@ -470,7 +475,7 @@ fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { } #[test] -fn test_sink_dtype_receives_function_options() -> VortexResult<()> { +fn test_output_dtype_receives_function_options() -> VortexResult<()> { assert_eq!( row_fn_return_dtype(&SinkOptions, &true, &[])?, DType::from(i64::PTYPE) diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index c32491abedb..2c6ebcebd4a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -15,7 +15,6 @@ use vortex_mask::Mask; use super::RowExecution; use crate::ExecutionCtx; -use crate::dtype::DType; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::OutputSink; @@ -44,7 +43,6 @@ fn ensure_decoded_lengths( /// does not need to be captured by the closure. pub fn execute_sink( args: &dyn ExecutionArgs, - sink_dtype: &DType, ctx: &mut ExecutionCtx, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, @@ -55,7 +53,7 @@ where ApplyResult: SinkResult>::WriteToken>, { let row_count = args.row_count(); - let mut sink = >::with_capacity(row_count, sink_dtype)?; + let mut sink = >::with_capacity(row_count)?; let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); let views = Args::per_row_views(&columns); @@ -98,7 +96,6 @@ where /// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. pub fn execute_sink_valid_rows( args: &dyn ExecutionArgs, - sink_dtype: &DType, valid: &Mask, ctx: &mut ExecutionCtx, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, @@ -123,7 +120,7 @@ where }; let prepared = prepare(Args::constants(&columns)); let row_count = args.row_count(); - let mut sink = >::with_capacity(row_count, sink_dtype)?; + let mut sink = >::with_capacity(row_count)?; // Batch execution resolves all-valid and all-null inputs before selecting this path. let AllOr::Some(valid) = valid.bit_buffer() else { @@ -226,11 +223,11 @@ mod tests { type Row<'a> = (); type WriteToken = (); - fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(_rows: usize) -> VortexResult { Err(vortex_err!( "a non-skipping sink must decline before allocation" )) @@ -264,11 +261,11 @@ mod tests { }) } - fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(rows: usize) -> VortexResult { Ok(Self(vec![0; rows])) } @@ -298,7 +295,6 @@ mod tests { let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( &args, - &DType::from(i64::PTYPE), &valid, &mut ctx, |_| (), @@ -324,7 +320,6 @@ mod tests { EmptyOptions, >( &args, - &DType::from(i64::PTYPE), &valid, &mut ctx, |_| (), @@ -351,7 +346,6 @@ mod tests { let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>( &args, - &DType::from(i64::PTYPE), &valid, &mut ctx, |_| (), diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index 010dc079599..24d2cd1cf6c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -125,11 +125,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { )?; execute_sink::( - self.args, - self.output_dtype, - self.ctx, - prepare, - apply, + self.args, self.ctx, prepare, apply, ) } @@ -261,12 +257,7 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { )?; execute_sink_valid_rows::( - self.args, - self.output_dtype, - self.valid, - self.ctx, - prepare, - apply, + self.args, self.valid, self.ctx, prepare, apply, ) } From 05fa94b3ebf5307ad699c4f383769858ef506feb Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:23:09 -0400 Subject: [PATCH 107/160] Use sink row counts directly Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 8 ++++---- .../src/scalar_fn/unstable/row/execute/sink.rs | 18 +++++++++--------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index f801fa35e4c..b3ce797bc79 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -96,8 +96,8 @@ unsafe impl OutputSink for OptionsCheckingSink { fn rows(&mut self) -> Self::Rows<'_> {} - fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { - true + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 } unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} @@ -141,8 +141,8 @@ unsafe impl OutputSink for I64Sink { self.0.as_mut_slice() } - fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { - rows.len() == row_count + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index 2c6ebcebd4a..45b65520d3a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -64,7 +64,7 @@ where // scope releases the borrow before `finish_sink` consumes the sink. let mut rows = >::rows(&mut sink); vortex_ensure!( - >::row_count_matches(&rows, row_count), + >::row_count(&rows) == row_count, "the output sink does not address exactly {row_count} rows", ); @@ -75,14 +75,14 @@ where // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before // the loop. let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; - // SAFETY: `row_count_matches` proved the sink addresses every loop index. + // SAFETY: `row_count` proved the sink addresses every loop index. let output = unsafe { >::row_unchecked(&mut rows, index) }; apply(&prepared, elements, output).into_result()?; } } else { for index in 0..row_count { - // SAFETY: `row_count_matches` proved the sink addresses every loop index. + // SAFETY: `row_count` proved the sink addresses every loop index. let output = unsafe { >::row_unchecked(&mut rows, index) }; apply(&prepared, Args::get(&columns, index), output).into_result()?; @@ -141,7 +141,7 @@ where // initializer mutably borrows the row representation. initialize_skipped_rows(&mut rows); vortex_ensure!( - >::row_count_matches(&rows, row_count), + >::row_count(&rows) == row_count, "the output sink does not address exactly {row_count} rows after initializing skipped rows", ); @@ -153,7 +153,7 @@ where return; } - // SAFETY: `row_count_matches` proved that the sink addresses every mask index, which + // SAFETY: `row_count` proved that the sink addresses every mask index, which // is below the mask's validated `row_count`. let output = unsafe { >::row_unchecked(&mut rows, index) }; let result = match &views { @@ -235,8 +235,8 @@ mod tests { fn rows(&mut self) -> Self::Rows<'_> {} - fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { - true + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 } unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> { @@ -273,8 +273,8 @@ mod tests { &mut self.0 } - fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { - rows.len() == row_count + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { From d5499fb8725fb1174e3cc59a954ae269bdf1834b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:23:14 -0400 Subject: [PATCH 108/160] Simplify RowFn internal visibility Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index e3759b0dcf1..d1f5b67fbd0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -20,7 +20,7 @@ use crate::scalar_fn::unstable::row::visitor::RowPolicy; /// original planned batch. Keeping them together prevents an execution path from pairing an input /// view with unrelated planning metadata. #[derive(Clone, Copy)] -pub(in crate::scalar_fn::unstable::row) struct BorrowedExecutionArgs<'a> { +pub(crate) struct BorrowedExecutionArgs<'a> { /// The input arrays for this kernel invocation. arrays: &'a [ArrayRef], @@ -39,7 +39,7 @@ pub(in crate::scalar_fn::unstable::row) struct BorrowedExecutionArgs<'a> { impl<'a> BorrowedExecutionArgs<'a> { /// Pair one input view with the planning metadata selected for its batch. - pub(in crate::scalar_fn::unstable::row) fn new( + pub(crate) fn new( arrays: &'a [ArrayRef], row_count: usize, dtypes: &'a [DType], @@ -56,22 +56,22 @@ impl<'a> BorrowedExecutionArgs<'a> { } /// Return the concrete arrays used by this row-kernel invocation. - pub(in crate::scalar_fn::unstable::row) fn arrays(&self) -> &'a [ArrayRef] { + pub(crate) fn arrays(&self) -> &'a [ArrayRef] { self.arrays } /// Return the original input dtypes used to select the row implementation. - pub(in crate::scalar_fn::unstable::row) fn dtypes(&self) -> &'a [DType] { + pub(crate) fn dtypes(&self) -> &'a [DType] { self.dtypes } /// Return the non-nullable dtype built by the selected output capability. - pub(in crate::scalar_fn::unstable::row) fn output_dtype(&self) -> &'a DType { + pub(crate) fn output_dtype(&self) -> &'a DType { self.output_dtype } /// Return the nullable execution policy selected during planning. - pub(in crate::scalar_fn::unstable::row) fn policy(&self) -> RowPolicy { + pub(crate) fn policy(&self) -> RowPolicy { self.policy } } From 11ce2c424e7dc020086c91a5efb3a3f667a5c1dc Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:24:19 -0400 Subject: [PATCH 109/160] Restore decoded-length regression coverage Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index b3ce797bc79..5e01b3da94b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -6,6 +6,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use rstest::rstest; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -18,6 +19,7 @@ use super::BatchPlan; use super::RowPolicy; use super::finalize_kernel_output; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -31,6 +33,7 @@ use crate::dtype::Nullability; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; @@ -48,6 +51,12 @@ struct NullarySeven; #[derive(Clone)] struct AddThree; +#[derive(Clone)] +struct AddShort; + +/// An element whose decode drops the last row, standing in for an invalid element implementation. +struct ShortDecodeI64; + #[derive(Clone)] struct Identity; @@ -75,6 +84,54 @@ enum PreparedVisit { Deferred, } +// SAFETY: the view and unchecked access delegate to the `i64` implementation. The implementation +// deliberately returns a short column so the executor's pre-loop length guard can be tested. +unsafe impl InputElement for ShortDecodeI64 { + type Column = Buffer; + type View<'a> = &'a [i64]; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let column = ::decode(array, ctx)?; + + Ok(column.slice(0..column.len().saturating_sub(1))) + } + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + ::view(column) + } + + fn view_len(view: &Self::View<'_>) -> usize { + ::view_len(view) + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_from_view(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { ::get_from_view_unchecked(view, index) } + } +} + // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or // `finish` through the executor. The row-initialization requirements are therefore vacuous. unsafe impl OutputSink for OptionsCheckingSink { @@ -199,6 +256,29 @@ impl RowFn for AddThree { } } +impl RowFn for AddShort { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.add_short"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(ShortDecodeI64, i64), I64Sink, _>(|(lhs, rhs), output| { + *output = lhs + rhs; + }) + } +} + impl RowFn for RetryConstantAdd { type Options = EmptyOptions; @@ -358,6 +438,27 @@ fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { Ok(()) } +#[test] +fn test_short_decode_beside_constant_is_rejected() -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter(0..64_i64).into_array(); + let rhs = ConstantArray::new(10_i64, 64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 64); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short decoded column passed the pre-loop length check"), + }; + + assert!( + error + .to_string() + .contains("does not address exactly 64 rows"), + "unexpected error: {error}", + ); + Ok(()) +} + #[test] fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> { let lhs = From 219e00aa359fb58c72f4e4523fb50f888684aa6c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 12:09:53 -0400 Subject: [PATCH 110/160] Remove unused RowFn batch array access Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/unstable/row/batch/args.rs | 5 ----- vortex-array/src/scalar_fn/unstable/row/batch/tests.rs | 5 +++-- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index d1f5b67fbd0..781d15711be 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -55,11 +55,6 @@ impl<'a> BorrowedExecutionArgs<'a> { } } - /// Return the concrete arrays used by this row-kernel invocation. - pub(crate) fn arrays(&self) -> &'a [ArrayRef] { - self.arrays - } - /// Return the original input dtypes used to select the row implementation. pub(crate) fn dtypes(&self) -> &'a [DType] { self.dtypes diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 5e01b3da94b..6c92fd52c40 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -31,6 +31,7 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::InputElement; @@ -518,7 +519,7 @@ fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResul let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; @@ -546,7 +547,7 @@ fn test_valid_only_filters_and_scatters() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; From 8be707fd0bf16924bb87db1da32264b6dfe517de Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 21:51:32 -0400 Subject: [PATCH 111/160] Execute primitive numeric operators with RowFn Route primitive numeric kernels through RowFn while retaining Binary as the public scalar function identity. Signed-off-by: Connor Tsui --- vortex-array/benches/binary_ops.rs | 8 + .../typed_view/primitive/numeric_operator.rs | 2 +- .../scalar_fn/fns/binary/numeric/checked.rs | 88 +---- .../src/scalar_fn/fns/binary/numeric/mod.rs | 12 +- .../scalar_fn/fns/binary/numeric/primitive.rs | 355 ++++-------------- .../src/scalar_fn/fns/binary/numeric/row.rs | 136 +++++++ .../src/scalar_fn/fns/binary/numeric/tests.rs | 9 +- 7 files changed, 241 insertions(+), 369 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/numeric/row.rs diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 6a07d03f50b..3bd466da0b1 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -170,6 +170,14 @@ fn div_i64_nonnull(bencher: Bencher) { bench_primitive(bencher, lhs, rhs, Operator::Div); } +#[divan::bench] +fn div_i64_nullable(bencher: Bencher) { + let lhs = primitive_nullable(1_000_000, 7).into_array(); + let rhs = primitive_nullable(17, 5).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Div); +} + #[divan::bench] fn sub_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index 6b389a0554b..d7bcc8d0b4c 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_err; use crate::scalar_fn::fns::operators::Operator; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. pub enum NumericOperator { /// Binary element-wise addition of two arrays or of two scalars. diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index afb709abe1c..054846b7ef7 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Checked-lane execution for numeric kernels, driven by the shared +//! Checked-lane execution for the decimal kernels, driven by the shared //! `vortex-compute` lane kernels. - -use std::ops::BitOrAssign; +//! +//! The primitive widths do not come through here: they are computed one row at a time by +//! [`row`](super::row), which writes a value for every row and reduces failure evidence without +//! scanning the finished output. use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -13,33 +15,18 @@ use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_mask::AllOr; use vortex_mask::Mask; -/// Evidence that a lane failed, anything other than [`Default`] meaning failure. -/// -/// `bool` is the ordinary choice; [`map_checked_into`] explains why the wider members exist and -/// asserts the width bound that membership here does **not** imply. -/// -/// [`map_checked_into`]: IndexedSourceExt::map_checked_into -pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - -impl Failure for bool {} -impl Failure for u8 {} -impl Failure for u16 {} -impl Failure for u32 {} -impl Failure for u64 {} - /// Apply the fallible `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it returns `None` on a _valid_ lane. +/// `Err(first_failing_valid_lane)` only when it returns `None` on a valid lane. /// /// `apply` also runs on invalid lanes, whose failures are masked out and whose values are /// unspecified, so it must be total: no panics or side effects on any stored lane value. /// -/// This drives the one-pass early-exit kernels, which abort at the end of the enclosing 64-lane -/// chunk. Use it when per-lane failure handling is cheap relative to the operation, as in integer -/// division. Prefer [`checked_apply_lanes`] when the failure check itself vectorizes. +/// This drives the one-pass early-exit kernels: failures abort at the end of the enclosing +/// 64-lane chunk. It suits an operation whose per-lane failure handling is cheap relative to the +/// operation itself, which is what the decimal kernels and their per-lane casts are. /// -/// `#[inline]`: this must inline into the caller that builds the closure, so that a captured -/// constant operand flattens into a register rather than living behind a pointer the loop reloads -/// on every lane, which blocks vectorization. +/// Keep this wrapper inlineable so captured constants can become loop invariants in the caller. +/// The lane kernels retain their own inlining decisions. #[inline] pub(super) fn checked_lanes( source: S, @@ -61,7 +48,6 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, @@ -73,55 +59,3 @@ where Ok(values.freeze()) } - -/// Apply the split value/failure `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it flags a _valid_ lane. -/// -/// The hot pass writes every value unconditionally and OR-reduces one piece of evidence, leaving -/// the loop free of per-lane selects and per-chunk exit branches. Only if a lane flagged does a -/// cold second pass re-run `apply` through the early-exit kernels, which drop null-lane failures -/// and attribute the first valid one. -/// -/// Like [`checked_lanes`], `apply` runs on invalid lanes and must be total. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -pub(super) fn checked_apply_lanes( - source: S, - valid_rows: &Mask, - mut apply: Apply, -) -> Result, usize> -where - S: IndexedSource + Copy, - T: Copy + Default, - Fail: Failure, - Apply: FnMut(S::Item) -> (T, Fail), -{ - let len = source.len(); - debug_assert_eq!(len, valid_rows.len()); - - let valid_bits = match valid_rows.bit_buffer() { - AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), - AllOr::Some(valid_bits) => Some(valid_bits), - }; - - let mut values = BufferMut::::with_capacity(len); - - let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &mut apply) != Fail::default() { - let mut checked = |item: S::Item| { - let (value, failure) = apply(item); - (failure == Fail::default()).then_some(value) - }; - match valid_bits { - None => source.try_map_into(out, &mut checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?, - } - } - - // SAFETY: the kernels initialize every lane in `out`. - unsafe { values.set_len(len) }; - - Ok(values.freeze()) -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 6dc0de0fbea..c7ae86b93c9 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,16 +4,19 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! +//! The primitive widths are computed by a [`RowFn`](crate::scalar_fn::RowFn), which owns null +//! handling, constants, and validity for them; see [`row`]. Decimal keeps its own columnar +//! implementation in [`decimal`]. +//! //! [`Binary`]: super::Binary mod checked; mod decimal; mod primitive; -#[cfg(test)] -mod tests; +mod row; use decimal::execute_numeric_decimal; -use primitive::execute_numeric_primitive; +use row::execute_numeric_primitive; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -81,3 +84,6 @@ fn build_empty_result( Ok(Canonical::empty(&result_dtype).into_array()) } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 8fd53d15216..42fe3fd3e03 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -1,73 +1,48 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::checked::Failure; -use super::checked::checked_apply_lanes; -use super::checked::checked_lanes; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; +//! Checked arithmetic for one primitive row. + +use std::ops::BitOrAssign; + use crate::dtype::NativePType; -use crate::dtype::PType; use crate::dtype::half::f16; -use crate::match_each_native_ptype; -use crate::scalar::NumericOperator; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; -use crate::validity::Validity; -struct CheckedAdd; +/// Checked addition, failing on integer overflow. +pub(super) struct CheckedAdd; -struct CheckedSub; +/// Checked subtraction, failing on integer overflow. +pub(super) struct CheckedSub; -struct CheckedMul; +/// Checked multiplication, failing on integer overflow. +pub(super) struct CheckedMul; -struct CheckedDiv; +/// Checked division, failing on integer division by zero and on `MIN / -1`. +pub(super) struct CheckedDiv; -trait CheckedPrimitiveOp: Sized { - /// Message for the error raised when this operation fails on a valid lane. - const ERROR: &'static str; +/// OR-reducible evidence that a row failed, with [`Default`] meaning success. +pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - /// Whether to check inside the value loop and exit early, rather than reducing evidence over - /// the whole batch and re-running only once some lane has flagged. - const CHECKED_VALUE_LOOP: bool = false; +impl Failure for T {} - /// How this operation reports a failing lane. See [`Failure`]. - type Failure: Failure; - - /// Compute a lane's value and its failure evidence together. - /// - /// Overflowing-style rather than `Option`, so the vectorizable kernels can write every - /// lane's value unconditionally and reduce the evidence separately. [`Self::checked`] is the - /// shape for the scalar and early-exit paths. - fn apply(lhs: T, rhs: T) -> (T, Self::Failure); +/// One arithmetic operator at one width, split into its value and failure evidence. +pub(super) trait CheckedPrimitiveOp: 'static + Sized { + /// The error reported for a batch in which some valid row failed. + const ERROR: &'static str; - /// [`Self::apply`] folded into the `Option` that the early-exit kernels take. - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - let (value, failed) = Self::apply(lhs, rhs); + /// How this operation reports a failing row. See [`Failure`]. + type Fail: Failure; - (failed == Self::Failure::default()).then_some(value) - } + /// The result of this operation, paired with evidence of whether the row failed. + fn apply(lhs: T, rhs: T) -> (T, Self::Fail); } impl CheckedPrimitiveOp for CheckedAdd { const ERROR: &'static str = "integer overflow in checked add"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.add_value(rhs), lhs.add_error(rhs)) } @@ -76,9 +51,9 @@ impl CheckedPrimitiveOp for CheckedAdd { impl CheckedPrimitiveOp for CheckedSub { const ERROR: &'static str = "integer overflow in checked sub"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.sub_value(rhs), lhs.sub_error(rhs)) } @@ -87,9 +62,9 @@ impl CheckedPrimitiveOp for CheckedSub { impl CheckedPrimitiveOp for CheckedMul { const ERROR: &'static str = "integer overflow in checked mul"; - type Failure = T::MulFailure; + type Fail = T::MulFailure; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, T::MulFailure) { (lhs.mul_value(rhs), lhs.mul_failure(rhs)) } @@ -97,16 +72,10 @@ impl CheckedPrimitiveOp for CheckedMul { impl CheckedPrimitiveOp for CheckedDiv { const ERROR: &'static str = "integer division by zero or overflow in checked div"; - // Integer division still lowers to scalar divides, so the split - // value/error-scan loop used to auto-vectorize add/sub/mul only adds a - // second full scan. Use the one-pass early-exit checked kernel for integer - // division, matching Arrow/Velox. Float division has no checked errors and - // stays on the split/vectorizable default path. - const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { let failed = lhs.div_error(rhs); let value = if failed { @@ -116,151 +85,13 @@ impl CheckedPrimitiveOp for CheckedDiv { }; (value, failed) } - - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - lhs.div_checked(rhs) - } -} - -/// Execute a numeric operation between two primitive-typed arrays. -pub(super) fn execute_numeric_primitive( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - - match_each_native_ptype!(ptype, |T| { - match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), - } - }) -} - -fn execute_checked_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType, - Op: CheckedPrimitiveOp, - Scalar: From, - Scalar: From>, -{ - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - let len = lhs.len(); - debug_assert_eq!(len, rhs.len()); - - let validity = lhs.validity().and(rhs.validity())?; - let valid_rows = validity.execute_mask(len, ctx)?; - - let values = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_op_lanes::<_, T, Op>( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - &valid_rows, - |(lhs, rhs)| (lhs, rhs), - ), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Capture the constant by value so it stays hoisted out of the lane loop. - let rhs = *rhs; - checked_op_lanes::<_, T, Op>(lhs.as_slice(), &valid_rows, move |lhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let lhs = *lhs; - checked_op_lanes::<_, T, Op>(rhs.as_slice(), &valid_rows, move |rhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let value = Op::checked(*lhs, *rhs) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - return Ok(constant_result_array(value, len, &result_dtype)); - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => Ok(Buffer::zeroed(len)), - } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - - primitive_result_array::(values, validity, &result_dtype) } -/// Run `Op` over the lanes of `source` in the loop shape it declares: the one-pass early-exit -/// kernel when `Op::CHECKED_VALUE_LOOP` is set, and the split value/failure kernel otherwise. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -fn checked_op_lanes( - source: S, - valid_rows: &Mask, - mut to_operands: impl FnMut(S::Item) -> (T, T), -) -> Result, usize> -where - S: IndexedSource + Copy, - T: NativePType, - Op: CheckedPrimitiveOp, -{ - if Op::CHECKED_VALUE_LOOP { - checked_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::checked(lhs, rhs) - }) - } else { - checked_apply_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::apply(lhs, rhs) - }) - } -} - -fn primitive_result_array( - values: Buffer, - validity: Validity, - dtype: &DType, -) -> VortexResult { - let array = PrimitiveArray::new(values, validity).into_array(); - if array.dtype() == dtype { - return Ok(array); - } - array.cast(dtype.clone()) -} - -fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef -where - T: NativePType, - Scalar: From + From>, -{ - if dtype.is_nullable() { - ConstantArray::new(Some(value), len).into_array() - } else { - ConstantArray::new(value, len).into_array() - } -} - -trait CheckedArithmetic: NativePType { - const DIV_CHECKS_IN_VALUE_LOOP: bool; - - /// How multiplication reports a failing lane, which is width-dependent in a way the other - /// operations are not. `impl_checked_unsigned!` and `impl_checked_signed!` say why each family - /// reports what it reports. +/// Per-width checked arithmetic. Every value method **must** be total over stored lane values. +pub(super) trait CheckedArithmetic: NativePType { + /// How multiplication reports a failing row. + /// + /// This may be a word rather than `bool` when narrowing evidence would block vectorization. type MulFailure: Failure; fn add_value(self, rhs: Self) -> Self; @@ -271,16 +102,9 @@ trait CheckedArithmetic: NativePType { fn mul_failure(self, rhs: Self) -> Self::MulFailure; fn div_value(self, rhs: Self) -> Self; fn div_error(self, rhs: Self) -> bool; - fn div_checked(self, rhs: Self) -> Option; } -/// The integer arithmetic every width shares, given the two things that actually differ between -/// them: how multiplication reports a failing lane, and how add/sub/div detect one. -/// -/// Only `mul_failure` genuinely varies per width, so it is the macro's parameter and everything -/// else is written once. The `$mul_failure_ty` and body are supplied by the caller because the -/// choice between a widened high half and a `bool` is exactly the vectorization decision described -/// on [`Failure`]. +/// Generate the shared integer operations from their failure predicates. macro_rules! impl_checked_integer { ( $ty:ty, @@ -291,67 +115,57 @@ macro_rules! impl_checked_integer { = |$mf_lhs:ident, $mf_rhs:ident| $mul_failure:expr, ) => { impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = true; - type MulFailure = $mul_failure_ty; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self.wrapping_add(rhs) } - #[inline(always)] + #[inline] fn add_error(self, rhs: Self) -> bool { let ($add_lhs, $add_rhs) = (self, rhs); $add_error } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self.wrapping_sub(rhs) } - #[inline(always)] + #[inline] fn sub_error(self, rhs: Self) -> bool { let ($sub_lhs, $sub_rhs) = (self, rhs); $sub_error } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self.wrapping_mul(rhs) } - #[inline(always)] + #[inline] $(#[$mul_failure_attr])* fn mul_failure(self, rhs: Self) -> $mul_failure_ty { let ($mf_lhs, $mf_rhs) = (self, rhs); $mul_failure } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, rhs: Self) -> bool { let ($div_lhs, $div_rhs) = (self, rhs); $div_error } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - self.checked_div(rhs) - } } }; } -/// The unsigned widths. `add`, `sub` and `div` are written once here, and `widening_mul` covers -/// every width including the 64-bit one, which widens into `u128`: the discarded high half of the -/// widened product is the failure evidence, and costs none of the comparison LLVM folds into -/// `umul.with.overflow`. +/// Unsigned multiplication reports its discarded high half as failure evidence. macro_rules! impl_checked_unsigned { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_integer!( @@ -364,12 +178,7 @@ macro_rules! impl_checked_unsigned { }; } -/// The signed widths, on the same principle. `widening_mul` is the shorthand for the narrow widths, -/// whose two-sided range check over a wider product is not the shape LLVM folds into an overflow -/// intrinsic, so they vectorize while reporting a plain `bool`. The 64-bit width cannot: deriving a -/// `bool` there costs the comparison that scalarizes the loop, so `high_half_mul` hands back the -/// discarded high half as a word. Both arms derive every shift and bound from `$ty`, so -/// instantiating one at a new width cannot silently keep another width's constants. +/// Signed widths use a range check or discarded high-half evidence. macro_rules! impl_checked_signed { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_signed!($ty, mul_failure: bool = |lhs, rhs| { @@ -377,9 +186,6 @@ macro_rules! impl_checked_signed { product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide }); }; - // Zero exactly when the product fits: a signed multiply overflows iff the high half of the true - // product differs from the sign extension of the half that was kept, so the two XOR to zero on - // the lanes that fit. `tests::test_multiply_overflow_boundaries` pins the boundaries. ($ty:ty, high_half_mul: $wide:ty => $failure:ty) => { impl_checked_signed!($ty, mul_failure: #[expect( clippy::cast_possible_truncation, @@ -395,7 +201,7 @@ macro_rules! impl_checked_signed { ( $ty:ty, mul_failure: $(#[$mul_failure_attr:meta])* $mul_failure_ty:ty - = |$l:ident, $r:ident| $mul_failure:expr + = |$lhs:ident, $rhs:ident| $mul_failure:expr ) => { impl_checked_integer!( $ty, @@ -408,7 +214,7 @@ macro_rules! impl_checked_signed { ((lhs ^ rhs) & (lhs ^ value)) < 0 }, div_error: |lhs, rhs| rhs == 0 || (lhs == <$ty>::MIN && rhs == -1), - mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$l, $r| $mul_failure, + mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$lhs, $rhs| $mul_failure, ); }; } @@ -417,54 +223,47 @@ macro_rules! impl_checked_float { ($($ty:ty),+ $(,)?) => { $( impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = false; - type MulFailure = bool; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self + rhs } - #[inline(always)] + #[inline] fn add_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self - rhs } - #[inline(always)] + #[inline] fn sub_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self * rhs } - #[inline(always)] + #[inline] fn mul_failure(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, _rhs: Self) -> bool { false } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - Some(self / rhs) - } } )+ }; @@ -484,30 +283,25 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; - /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change - /// pivots, and a spread of magnitudes that straddles the 64-bit split. const PROBES: &[i64] = &[ - 0, // - 1, // - -1, // - 2, // - -2, // - 3, // - i64::MIN, // - i64::MIN + 1, // - i64::MAX, // - i64::MAX - 1, // - 1 << 31, // - 1 << 32, // - 1 << 62, // - -(1 << 62), // - 0x7FFF_FFFF, // - -0x8000_0000, // + 0, + 1, + -1, + 2, + -2, + 3, + i64::MIN, + i64::MIN + 1, + i64::MAX, + i64::MAX - 1, + 1 << 31, + 1 << 32, + 1 << 62, + -(1 << 62), + 0x7FFF_FFFF, + -0x8000_0000, ]; - /// Every `mul_failure` impl is either a bit trick or a two-sided range check, so hold each - /// against the obvious reference: `reference` is the width's own `checked_mul`, whose `None` - /// _is_ the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { let failed = lhs.mul_failure(rhs) != ::default(); @@ -522,14 +316,11 @@ mod tests { assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); let (lhs, rhs) = (lhs as u64, rhs as u64); - assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); } } } - /// The 8-bit widths are cheap enough to check exhaustively, which pins the shift of the - /// unsigned formula and the range check of the signed one against every product that exists. #[test] fn mul_failure_agrees_with_checked_mul_exhaustively_at_8_bits() { for lhs in u8::MIN..=u8::MAX { diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs new file mode 100644 index 00000000000..f0efaca929d --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive arithmetic execution through [`RowFn`]. +//! +//! `Binary` keeps its registered contract; [`NumericBinary`] is only an execution helper. Decimal +//! arithmetic remains on its existing columnar path. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::primitive::CheckedAdd; +use super::primitive::CheckedArithmetic; +use super::primitive::CheckedDiv; +use super::primitive::CheckedMul; +use super::primitive::CheckedPrimitiveOp; +use super::primitive::CheckedSub; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::NumericOperator; +use crate::scalar_fn::BorrowedExecutionArgs; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::RowVisitor; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::execute_rows; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::row::InitializedElement; +use crate::scalar_fn::row::UninitElementSink; + +pub(super) fn execute_numeric_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let inputs = [lhs.clone(), rhs.clone()]; + let args = BorrowedExecutionArgs::new(&inputs, lhs.len()); + + execute_rows(&NumericBinary, &op, &args, ctx) +} + +/// Internal row execution for the primitive arithmetic operators. +#[derive(Clone)] +struct NumericBinary; + +impl RowFn for NumericBinary { + type Options = NumericOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + // Fallibility is queried without input dtypes, so this conservatively covers integer widths. + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + // `NumericBinary` is a private implementation detail of `Binary`: it is never registered or + // serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. + ScalarFnVTable::id(&Binary) + } + + fn dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = PType::try_from( + args.first() + .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?, + )?; + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_div::(visitor), + } + }) + } +} + +fn visit_checked(visitor: V) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + V: RowVisitor, +{ + visitor.visit_deferred::<(T, T), T, Op::Fail>( + |(lhs, rhs)| Op::apply(lhs, rhs), + |failure| { + if failure != ::default() { + return Err(numeric_error(Op::ERROR)); + } + + Ok(()) + }, + ) +} + +fn visit_div(visitor: V) -> VortexResult +where + T: CheckedArithmetic, + V: RowVisitor, +{ + if T::PTYPE.is_float() { + return visit_checked::(visitor); + } + + // Integer division is scalar and expensive, so deferring its cheap failure check preserves no + // vectorization. Check each divide immediately and stop at the first failure. + // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when + // they need to skip invalid rows. + visitor.visit_into::<(T, T), UninitElementSink, _>(|(lhs, rhs), output| { + let (value, failed) = CheckedDiv::apply(lhs, rhs); + if failed { + return Err(numeric_error(>::ERROR)); + } + + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) + }) +} + +/// Keep rich error construction out of row closures so the closures remain inlineable. +#[cold] +#[inline(never)] +fn numeric_error(message: &'static str) -> VortexError { + vortex_err!(InvalidArgument: "{message}") +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 7ee98ae717e..3813c8612b3 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -201,8 +201,7 @@ fn test_integer_array_array_errors_on_valid_lanes() { assert!(result.is_err()); } -/// Multiply two non-nullable lanes of `lhs` by two of `rhs`, expecting `Some(product)` where the -/// product fits and `None` where the checked kernel must report overflow. +/// Assert one checked multiplication through the complete array execution path. #[track_caller] fn assert_multiply(lhs: T, rhs: T, expected: Option) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -297,13 +296,11 @@ fn test_multiply_overflow_on_null_lane_ignored( Ok(()) } -/// The hot pass OR-reduces evidence across whole 64-lane chunks before anything looks at it, so an -/// overflow in a late chunk must still be caught, and must still be suppressed when its lane is -/// null. Every other test here fits in a single chunk and cannot show either. +/// An overflow late in the batch must still be reported, unless its row is null. #[rstest] #[case::reported(true)] #[case::suppressed_by_null(false)] -fn test_multiply_overflow_survives_chunk_reduction( +fn test_multiply_overflow_survives_batch_reduction( #[case] lane_is_valid: bool, ) -> VortexResult<()> { const LEN: u32 = 1000; From f95784799b8c414d4307406d545573baf2a91639 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:48:58 -0400 Subject: [PATCH 112/160] Tighten checked arithmetic contracts Signed-off-by: Connor Tsui --- .../scalar_fn/fns/binary/numeric/primitive.rs | 53 ++++++++++++------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 42fe3fd3e03..7cbe88c7c8e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -23,7 +23,11 @@ pub(super) struct CheckedDiv; /// OR-reducible evidence that a row failed, with [`Default`] meaning success. pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} -impl Failure for T {} +impl Failure for bool {} +impl Failure for u8 {} +impl Failure for u16 {} +impl Failure for u32 {} +impl Failure for u64 {} /// One arithmetic operator at one width, split into its value and failure evidence. pub(super) trait CheckedPrimitiveOp: 'static + Sized { @@ -87,7 +91,10 @@ impl CheckedPrimitiveOp for CheckedDiv { } } -/// Per-width checked arithmetic. Every value method **must** be total over stored lane values. +/// Per-width arithmetic used to compute values and failure evidence. +/// +/// The add, subtract, and multiply value methods **must** be total over every stored lane value. +/// [`Self::div_value`] may assume that [`Self::div_error`] returned `false` for the same operands. pub(super) trait CheckedArithmetic: NativePType { /// How multiplication reports a failing row. /// @@ -100,7 +107,11 @@ pub(super) trait CheckedArithmetic: NativePType { fn sub_error(self, rhs: Self) -> bool; fn mul_value(self, rhs: Self) -> Self; fn mul_failure(self, rhs: Self) -> Self::MulFailure; + + /// Divide operands that [`Self::div_error`] accepted. fn div_value(self, rhs: Self) -> Self; + + /// Return whether [`Self::div_value`] would trap for these operands. fn div_error(self, rhs: Self) -> bool; } @@ -195,6 +206,10 @@ macro_rules! impl_checked_signed { let kept = wide as $ty; let discarded = (wide >> <$ty>::BITS) as $ty; + // A product fits exactly when its discarded half is the sign extension of the kept + // half. XOR reduces that comparison to zero evidence for success and nonzero evidence + // for overflow without converting the wide product to a branch. + (discarded ^ (kept >> (<$ty>::BITS - 1))) as $failure }); }; @@ -283,23 +298,25 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; + /// Values around zero, signed extrema, and 32- and 64-bit boundaries where the discarded + /// multiplication half or its sign extension changes. const PROBES: &[i64] = &[ - 0, - 1, - -1, - 2, - -2, - 3, - i64::MIN, - i64::MIN + 1, - i64::MAX, - i64::MAX - 1, - 1 << 31, - 1 << 32, - 1 << 62, - -(1 << 62), - 0x7FFF_FFFF, - -0x8000_0000, + 0, // Additive identity. + 1, // Smallest positive value. + -1, // All sign bits set. + 2, // Small positive power of two. + -2, // Small negative power of two. + 3, // Small non-power of two. + i64::MIN, // Minimum signed value. + i64::MIN + 1, // Minimum signed value's neighbor. + i64::MAX, // Maximum signed value. + i64::MAX - 1, // Maximum signed value's neighbor. + 1 << 31, // First positive value outside i32. + 1 << 32, // First value with bit 32 set. + 1 << 62, // Largest positive power of two in i64. + -(1 << 62), // Negative counterpart of the largest power of two. + 0x7FFF_FFFF, // Maximum i32 represented as i64. + -0x8000_0000, // Minimum i32 represented as i64. ]; #[track_caller] From 8c203cf2b067aa2fda3e92775409db0169918402 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:07:50 -0400 Subject: [PATCH 113/160] Benchmark primitive RowFn execution shapes Signed-off-by: Connor Tsui --- vortex-array/benches/binary_ops.rs | 50 ++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 3bd466da0b1..89048fcefd4 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -38,10 +38,60 @@ static SESSION: LazyLock = LazyLock::new(array_session); const LEN: usize = 32_768; +const ROWFN_MATRIX_CASES: &[(usize, RowFnShape)] = &[ + (128, RowFnShape::VaryingVarying), + (128, RowFnShape::VaryingConstant), + (128, RowFnShape::ConstantVarying), + (128, RowFnShape::VaryingNullableConstant), + (LEN, RowFnShape::VaryingVarying), + (LEN, RowFnShape::VaryingConstant), + (LEN, RowFnShape::ConstantVarying), + (LEN, RowFnShape::VaryingNullableConstant), +]; + +#[derive(Clone, Copy, Debug)] +enum RowFnShape { + VaryingVarying, + VaryingConstant, + ConstantVarying, + VaryingNullableConstant, +} + /// Decimal Mul and Div cost far more per lane than Add, so they run over a shorter array to keep /// the instrumented CodSpeed runs quick. const DECIMAL_MUL_DIV_LEN: usize = 8_192; +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_add(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Add); +} + +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_subtract(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Sub); +} + +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_multiply(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Mul); +} + +fn bench_rowfn_shape(bencher: Bencher, len: usize, shape: RowFnShape, operator: Operator) { + let varying = + || PrimitiveArray::from_iter((0..len).map(|index| (index % 1_024) as i64 + 1)).into_array(); + let constant = || ConstantArray::new(17_i64, len).into_array(); + let nullable_constant = || ConstantArray::new(Some(17_i64), len).into_array(); + + let (lhs, rhs) = match shape { + RowFnShape::VaryingVarying => (varying(), varying()), + RowFnShape::VaryingConstant => (varying(), constant()), + RowFnShape::ConstantVarying => (constant(), varying()), + RowFnShape::VaryingNullableConstant => (varying(), nullable_constant()), + }; + + bench_primitive(bencher, lhs, rhs, operator); +} + #[divan::bench] fn add_i64_nonnull(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); From 09f9721508c6b399fe8b453277b61860254e1346 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 08:13:14 -0400 Subject: [PATCH 114/160] Use generic execution arguments for numeric rows Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/numeric/row.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index f0efaca929d..bc38dd0e5eb 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -23,11 +23,11 @@ use crate::dtype::NativePType; use crate::dtype::PType; use crate::match_each_native_ptype; use crate::scalar::NumericOperator; -use crate::scalar_fn::BorrowedExecutionArgs; use crate::scalar_fn::RowFn; use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::execute_rows; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::row::InitializedElement; @@ -39,8 +39,7 @@ pub(super) fn execute_numeric_primitive( op: NumericOperator, ctx: &mut ExecutionCtx, ) -> VortexResult { - let inputs = [lhs.clone(), rhs.clone()]; - let args = BorrowedExecutionArgs::new(&inputs, lhs.len()); + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); execute_rows(&NumericBinary, &op, &args, ctx) } From 6f891c2ad69df4cc10f10b2a19a57a156941b6ef Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:08:17 -0400 Subject: [PATCH 115/160] Use per-row terminology in numeric benchmarks Signed-off-by: Connor Tsui --- vortex-array/benches/binary_ops.rs | 34 +++++++++++++++--------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 89048fcefd4..ccd440ff939 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -39,22 +39,22 @@ static SESSION: LazyLock = LazyLock::new(array_session); const LEN: usize = 32_768; const ROWFN_MATRIX_CASES: &[(usize, RowFnShape)] = &[ - (128, RowFnShape::VaryingVarying), - (128, RowFnShape::VaryingConstant), - (128, RowFnShape::ConstantVarying), - (128, RowFnShape::VaryingNullableConstant), - (LEN, RowFnShape::VaryingVarying), - (LEN, RowFnShape::VaryingConstant), - (LEN, RowFnShape::ConstantVarying), - (LEN, RowFnShape::VaryingNullableConstant), + (128, RowFnShape::PerRowPerRow), + (128, RowFnShape::PerRowConstant), + (128, RowFnShape::ConstantPerRow), + (128, RowFnShape::PerRowNullableConstant), + (LEN, RowFnShape::PerRowPerRow), + (LEN, RowFnShape::PerRowConstant), + (LEN, RowFnShape::ConstantPerRow), + (LEN, RowFnShape::PerRowNullableConstant), ]; #[derive(Clone, Copy, Debug)] enum RowFnShape { - VaryingVarying, - VaryingConstant, - ConstantVarying, - VaryingNullableConstant, + PerRowPerRow, + PerRowConstant, + ConstantPerRow, + PerRowNullableConstant, } /// Decimal Mul and Div cost far more per lane than Add, so they run over a shorter array to keep @@ -77,16 +77,16 @@ fn rowfn_multiply(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { } fn bench_rowfn_shape(bencher: Bencher, len: usize, shape: RowFnShape, operator: Operator) { - let varying = + let per_row = || PrimitiveArray::from_iter((0..len).map(|index| (index % 1_024) as i64 + 1)).into_array(); let constant = || ConstantArray::new(17_i64, len).into_array(); let nullable_constant = || ConstantArray::new(Some(17_i64), len).into_array(); let (lhs, rhs) = match shape { - RowFnShape::VaryingVarying => (varying(), varying()), - RowFnShape::VaryingConstant => (varying(), constant()), - RowFnShape::ConstantVarying => (constant(), varying()), - RowFnShape::VaryingNullableConstant => (varying(), nullable_constant()), + RowFnShape::PerRowPerRow => (per_row(), per_row()), + RowFnShape::PerRowConstant => (per_row(), constant()), + RowFnShape::ConstantPerRow => (constant(), per_row()), + RowFnShape::PerRowNullableConstant => (per_row(), nullable_constant()), }; bench_primitive(bencher, lhs, rhs, operator); From b7eb750cc5d4758d93b12fc35bccff7426967c4f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:34:43 -0400 Subject: [PATCH 116/160] Opt numeric operators into unstable RowFn Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs | 6 +++--- vortex-array/src/scalar_fn/fns/binary/numeric/row.rs | 10 +++++----- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index c7ae86b93c9..a0c427b142e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,9 +4,9 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! -//! The primitive widths are computed by a [`RowFn`](crate::scalar_fn::RowFn), which owns null -//! handling, constants, and validity for them; see [`row`]. Decimal keeps its own columnar -//! implementation in [`decimal`]. +//! The primitive widths are computed by a +//! [`RowFn`](crate::scalar_fn::unstable::row::RowFn), which owns null handling, constants, and +//! validity for them; see [`row`]. Decimal keeps its own columnar implementation in [`decimal`]. //! //! [`Binary`]: super::Binary diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index bc38dd0e5eb..b6f565d8ff4 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -23,15 +23,15 @@ use crate::dtype::NativePType; use crate::dtype::PType; use crate::match_each_native_ptype; use crate::scalar::NumericOperator; -use crate::scalar_fn::RowFn; -use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; -use crate::scalar_fn::execute_rows; use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::row::InitializedElement; -use crate::scalar_fn::row::UninitElementSink; +use crate::scalar_fn::unstable::row::InitializedElement; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::UninitElementSink; +use crate::scalar_fn::unstable::row::execute_rows; pub(super) fn execute_numeric_primitive( lhs: &ArrayRef, From f6780e25bafde74568b79c9ca2f0beb850c2a717 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 21:51:32 -0400 Subject: [PATCH 117/160] Execute primitive comparisons with RowFn Use RowFn for primitive comparisons while retaining fused x86 bit-packing for the measured wide ordered cases where LLVM generates faster code. Signed-off-by: Connor Tsui --- encodings/runend/src/trace_tests.rs | 8 + .../src/scalar_fn/fns/binary/compare/mod.rs | 8 +- .../scalar_fn/fns/binary/compare/primitive.rs | 166 ++++++++---------- .../fns/binary/compare/primitive/columnar.rs | 120 +++++++++++++ .../primitive/operand.rs} | 19 +- vortex-array/src/scalar_fn/fns/binary/mod.rs | 1 - vortex-array/src/test_harness/trace/tests.rs | 8 + vortex-btrblocks/src/trace_tests.rs | 16 ++ 8 files changed, 248 insertions(+), 98 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs rename vortex-array/src/scalar_fn/fns/binary/{primitive_operand.rs => compare/primitive/operand.rs} (78%) diff --git a/encodings/runend/src/trace_tests.rs b/encodings/runend/src/trace_tests.rs index 96f2afff50a..be8fd0ba762 100644 --- a/encodings/runend/src/trace_tests.rs +++ b/encodings/runend/src/trace_tests.rs @@ -73,6 +73,14 @@ fn trace_compare_on_runend() -> VortexResult<()> { iter 0 current=vortex.runend(bool, len=9) builder_active=false execute_until target=AnyCanonical root=vortex.binary(bool, len=3) iter 0 current=vortex.binary(bool, len=3) builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 1 current=vortex.bool(bool, len=3) builder_active=false return output=vortex.bool(bool, len=3) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index d25a652ee57..a36d0a22bde 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -4,9 +4,9 @@ //! Native comparison kernels. //! //! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every -//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from -//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise -//! comparator for nested types. There is no Arrow fallback. +//! comparison directly over Vortex canonical arrays: bit buffers for booleans, row or fused lane +//! kernels for primitives, lane kernels for decimals, binary views for strings and bytes, and a +//! row-wise comparator for nested types. There is no Arrow fallback. //! //! Floating point values compare with Vortex's total ordering (`NaN` is the largest value, //! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics. @@ -211,7 +211,7 @@ fn compare_arrays( ) .into_array()), DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx), - DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), + DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, ctx), DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Map(..) => { diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 1247358dce1..f93475292fe 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,27 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays via bit-packing lane kernels. +//! Primitive comparison execution through [`RowFn`]. + +#[cfg(target_arch = "x86_64")] +mod columnar; +#[cfg(target_arch = "x86_64")] +mod operand; -use vortex_buffer::BitBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_bail; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::dtype::NativePType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::compare::collect_bits; -use crate::scalar_fn::fns::binary::compare::collect_zip_bits; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; +use crate::scalar_fn::BorrowedExecutionArgs; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::RowVisitor; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::execute_rows; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; /// Compare two primitive arrays of the same [`PType`]. @@ -32,99 +34,79 @@ pub(super) fn compare_primitive( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - match_each_native_ptype!(ptype, |T| { - compare_primitive_typed::(lhs, rhs, op, nullability, ctx) - }) + #[cfg(target_arch = "x86_64")] + if use_columnar_comparison(lhs, rhs, op)? { + return columnar::compare_primitive(lhs, rhs, op, ctx); + } + + let inputs = [lhs.clone(), rhs.clone()]; + let args = BorrowedExecutionArgs::new(&inputs, lhs.len()); + + execute_rows(&PrimitiveCompare, &op, &args, ctx) } -fn compare_primitive_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: CompareOperator, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let len = lhs.len(); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - if lhs.len() != rhs.len() { - vortex_bail!( - "compare operator requires equal lengths, got {} and {}", - lhs.len(), - rhs.len() - ); +/// Internal row execution for primitive comparison operators. +#[derive(Clone)] +struct PrimitiveCompare; + +impl RowFn for PrimitiveCompare { + type Options = CompareOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + ScalarFnVTable::id(&Binary) } - let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; - - let bits = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slices(lhs, rhs, op), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => compare_slice_constant(lhs, *rhs, op), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slice_constant(rhs, *lhs, op.swap()), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Unreachable through `execute_compare` (constant-constant is folded there), but - // cheap to answer anyway. - BitBuffer::full(apply_op(*lhs, *rhs, op), len) - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { - return Ok( - ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) - .into_array(), - ); - } - }; - - Ok(BoolArray::try_new(bits, validity)?.into_array()) -} + fn dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = + PType::try_from(args.first().ok_or_else(|| { + vortex_err!("a comparison operator takes two operands, got none") + })?)?; -#[inline(always)] -fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { - match op { - CompareOperator::Eq => lhs.is_eq(rhs), - CompareOperator::NotEq => !lhs.is_eq(rhs), - CompareOperator::Gt => lhs.is_gt(rhs), - CompareOperator::Gte => lhs.is_ge(rhs), - CompareOperator::Lt => lhs.is_lt(rhs), - CompareOperator::Lte => lhs.is_le(rhs), + match_each_native_ptype!(ptype, |T| { visit_compare::(*op, visitor) }) } } -fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { - // Dispatch the operator outside the lane loop so each instantiation vectorizes a single - // branch-free predicate. - match op { - CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)), - CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)), - CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), - CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), - CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), - CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), +#[cfg(target_arch = "x86_64")] +fn use_columnar_comparison( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, +) -> VortexResult { + if matches!(op, CompareOperator::Eq | CompareOperator::NotEq) { + return Ok(false); } + + let ptype = PType::try_from(lhs.dtype())?; + Ok(match ptype { + // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit + // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. + PType::I64 | PType::F64 => true, + // LLVM vectorizes varying u64 inputs, but not the mixed-constant RowFn loop. + PType::U64 => lhs.as_constant().is_some() || rhs.as_constant().is_some(), + _ => false, + }) } -fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { +fn visit_compare(op: CompareOperator, visitor: V) -> VortexResult +where + T: NativePType, + V: RowVisitor, +{ match op { - CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), - CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), - CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)), - CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)), - CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)), - CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), + CompareOperator::Eq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)), + CompareOperator::NotEq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| !lhs.is_eq(rhs)), + CompareOperator::Gt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_gt(rhs)), + CompareOperator::Gte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_ge(rhs)), + CompareOperator::Lt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_lt(rhs)), + CompareOperator::Lte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_le(rhs)), } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs new file mode 100644 index 00000000000..ccd7ffb8719 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -0,0 +1,120 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fused comparison and bit-packing for wide x86 lanes. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use super::operand::PrimitiveOperand; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// Compare primitive operands with one fused comparison and bit-packing loop. +pub(super) fn compare_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match PType::try_from(lhs.dtype())? { + PType::I64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::U64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::F64 => compare_primitive_typed::(lhs, rhs, op, ctx), + ptype => vortex_bail!("columnar comparison is not selected for {ptype}"), + } +} + +fn compare_primitive_typed( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable()); + let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; + let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + let bits = match (&lhs, &rhs) { + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slices(lhs, rhs, op), + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => compare_slice_constant(lhs, *rhs, op), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slice_constant(rhs, *lhs, op.swap()), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => BitBuffer::full(apply_op(*lhs, *rhs, op), len), + (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) + .into_array(), + ); + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +#[inline(always)] +fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { + match op { + CompareOperator::Eq => lhs.is_eq(rhs), + CompareOperator::NotEq => !lhs.is_eq(rhs), + CompareOperator::Gt => lhs.is_gt(rhs), + CompareOperator::Gte => lhs.is_ge(rhs), + CompareOperator::Lt => lhs.is_lt(rhs), + CompareOperator::Lte => lhs.is_le(rhs), + } +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + } +} + +fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(lhs, |lhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_bits(lhs, |lhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_bits(lhs, |lhs: T| lhs.is_gt(rhs)), + CompareOperator::Gte => collect_bits(lhs, |lhs: T| lhs.is_ge(rhs)), + CompareOperator::Lt => collect_bits(lhs, |lhs: T| lhs.is_lt(rhs)), + CompareOperator::Lte => collect_bits(lhs, |lhs: T| lhs.is_le(rhs)), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs similarity index 78% rename from vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs rename to vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs index 71d1122fc79..55d81153b1f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Decoding shared by primitive binary operators. +//! Operand decoding for the fused primitive comparison path. use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -15,19 +15,33 @@ use crate::validity::Validity; /// A materialized primitive column, a non-null constant, or an all-null constant. pub(super) enum PrimitiveOperand { + /// A varying primitive column and its validity. Array { + /// The materialized values. values: Buffer, + + /// The validity of the values. validity: Validity, }, + + /// A non-null value repeated for every row. Constant { + /// The repeated value. value: T, + + /// The number of repeated rows. len: usize, + + /// The validity implied by the constant's dtype. validity: Validity, }, + + /// An all-null constant with this row count. Null(usize), } impl PrimitiveOperand { + /// Decode an operand once for the fused comparison loop. pub(super) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { if let Some(constant) = array.as_opt::() { return Ok( @@ -49,9 +63,11 @@ impl PrimitiveOperand { let array = array.clone().execute::(ctx)?; let validity = array.validity()?; let values = array.into_buffer::(); + Ok(Self::Array { values, validity }) } + /// Return the logical row count. pub(super) fn len(&self) -> usize { match self { Self::Array { values, .. } => values.len(), @@ -59,6 +75,7 @@ impl PrimitiveOperand { } } + /// Return the operand validity. pub(super) fn validity(&self) -> Validity { match self { Self::Array { validity, .. } => validity.clone(), diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 80faff20e0a..a5b9fe70539 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -43,7 +43,6 @@ mod compare; pub use compare::*; mod numeric; pub(crate) use numeric::*; -mod primitive_operand; use crate::scalar::NumericOperator; use crate::scalar::Scalar; diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index 98043caec13..fd9685dcb35 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -685,6 +685,14 @@ fn trace_compare_on_dict() -> VortexResult<()> { iter 0 current=vortex.dict(bool, len=5) builder_active=false ExecuteSlot slot=1 parent=vortex.dict(bool, len=5) child=vortex.binary(bool, len=3) iter 1 current=vortex.binary(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=5) diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index d779757bc77..bf2b00f563e 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -209,6 +209,14 @@ fn trace_scan_compare_on_compressed_shipdate() -> VortexResult<()> { Done array=vortex.primitive(i32, len=4096) iter 1 current=vortex.primitive(i32, len=4096) builder_active=false return output=vortex.primitive(i32, len=4096) + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=4096) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=4096) iter 2 current=vortex.bool(bool, len=4096) builder_active=false return output=vortex.bool(bool, len=4096) @@ -267,6 +275,14 @@ fn trace_scan_compare_on_compressed_quantity() -> VortexResult<()> { Done array=vortex.primitive(i16, len=50) iter 1 current=vortex.primitive(i16, len=50) builder_active=false return output=vortex.primitive(i16, len=50) + optimize root=vortex.slice(i16, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i16, len=1) child=vortex.constant(i16, len=50) -> vortex.constant(i16, len=1) + done output=vortex.constant(i16, len=1) + execute_until target=AnyCanonical root=vortex.constant(i16, len=1) + iter 0 current=vortex.constant(i16, len=1) builder_active=false + Done array=vortex.primitive(i16, len=1) + iter 1 current=vortex.primitive(i16, len=1) builder_active=false + return output=vortex.primitive(i16, len=1) Done array=vortex.bool(bool, len=50) iter 6 current=vortex.bool(bool, len=50) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=4096) From 0728082a4c0099b23abef3e439cf028f750e25df Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:44:18 -0400 Subject: [PATCH 118/160] Exercise both primitive comparison paths Signed-off-by: Connor Tsui --- .../scalar_fn/fns/binary/compare/primitive.rs | 62 ++++++-- .../fns/binary/compare/primitive/columnar.rs | 3 +- .../src/scalar_fn/fns/binary/compare/tests.rs | 141 ++++++++++++++++++ 3 files changed, 192 insertions(+), 14 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index f93475292fe..18f0c10faea 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -3,9 +3,7 @@ //! Primitive comparison execution through [`RowFn`]. -#[cfg(target_arch = "x86_64")] mod columnar; -#[cfg(target_arch = "x86_64")] mod operand; use vortex_error::VortexResult; @@ -36,8 +34,49 @@ pub(super) fn compare_primitive( op: CompareOperator, ctx: &mut ExecutionCtx, ) -> VortexResult { - #[cfg(target_arch = "x86_64")] - if use_columnar_comparison(lhs, rhs, op)? { + compare_primitive_with_path(lhs, rhs, op, PrimitiveComparisonPath::Auto, ctx) +} + +/// Selects automatic production dispatch or a forced implementation in tests. +#[derive(Clone, Copy)] +pub(super) enum PrimitiveComparisonPath { + /// Use the architecture and operand-specific production policy. + Auto, + + /// Force row execution. + #[cfg(test)] + Row, + + /// Force fused columnar execution. + #[cfg(test)] + Columnar, +} + +/// Compare primitives through the selected implementation. +pub(super) fn compare_primitive_with_path( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + path: PrimitiveComparisonPath, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let use_columnar = match path { + PrimitiveComparisonPath::Auto => { + #[cfg(target_arch = "x86_64")] + { + use_columnar_comparison(lhs, rhs, op)? + } + #[cfg(not(target_arch = "x86_64"))] + { + false + } + } + #[cfg(test)] + PrimitiveComparisonPath::Row => false, + #[cfg(test)] + PrimitiveComparisonPath::Columnar => true, + }; + if use_columnar { return columnar::compare_primitive(lhs, rhs, op, ctx); } @@ -75,23 +114,22 @@ impl RowFn for PrimitiveCompare { } } -#[cfg(target_arch = "x86_64")] fn use_columnar_comparison( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, ) -> VortexResult { - if matches!(op, CompareOperator::Eq | CompareOperator::NotEq) { - return Ok(false); - } - let ptype = PType::try_from(lhs.dtype())?; - Ok(match ptype { + Ok(match (ptype, op) { + // Equality bit-packs efficiently for every type supported by the columnar path. + (PType::I64 | PType::U64 | PType::F64, CompareOperator::Eq | CompareOperator::NotEq) => { + true + } // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. - PType::I64 | PType::F64 => true, + (PType::I64 | PType::F64, _) => true, // LLVM vectorizes varying u64 inputs, but not the mixed-constant RowFn loop. - PType::U64 => lhs.as_constant().is_some() || rhs.as_constant().is_some(), + (PType::U64, _) => lhs.as_constant().is_some() || rhs.as_constant().is_some(), _ => false, }) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs index ccd7ffb8719..2a8ee71cfb0 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Fused comparison and bit-packing for wide x86 lanes. +//! Fused comparison and bit-packing for wide primitive lanes. use vortex_buffer::BitBuffer; use vortex_error::VortexResult; @@ -85,7 +85,6 @@ fn compare_primitive_typed( Ok(BoolArray::try_new(bits, validity)?.into_array()) } -#[inline(always)] fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { match op { CompareOperator::Eq => lhs.is_eq(rhs), diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index 9831a963354..d0f7a9b5e57 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -12,7 +12,9 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::VTable; use crate::array_session; +use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; @@ -21,6 +23,7 @@ use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::ScalarFn; use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; @@ -39,6 +42,8 @@ use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; use crate::scalar::DecimalValue; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::primitive::PrimitiveComparisonPath; +use crate::scalar_fn::fns::binary::compare::primitive::compare_primitive_with_path; use crate::scalar_fn::fns::binary::scalar_cmp; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -429,6 +434,142 @@ fn float_total_order() { ); } +#[rstest] +#[case::row_eq(PrimitiveComparisonPath::Row, CompareOperator::Eq)] +#[case::row_not_eq(PrimitiveComparisonPath::Row, CompareOperator::NotEq)] +#[case::row_lt(PrimitiveComparisonPath::Row, CompareOperator::Lt)] +#[case::columnar_eq(PrimitiveComparisonPath::Columnar, CompareOperator::Eq)] +#[case::columnar_not_eq(PrimitiveComparisonPath::Columnar, CompareOperator::NotEq)] +#[case::columnar_lt(PrimitiveComparisonPath::Columnar, CompareOperator::Lt)] +fn test_primitive_comparison_paths_preserve_semantics_and_encoding( + #[case] path: PrimitiveComparisonPath, + #[case] op: CompareOperator, +) -> VortexResult<()> { + let lhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::NAN, // Null on the left. + -0.0, // Signed zero ordering. + 1.0, // A finite value below NaN. + f64::NAN, // Null on the right. + ], + Validity::from_iter([ + true, // + false, // + true, // + true, // + true, // + ]), + ) + .into_array(); + let rhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::INFINITY, // Null on the left. + 0.0, // Signed zero ordering. + f64::NAN, // A finite value below NaN. + 1.0, // Null on the right. + ], + Validity::from_iter([ + true, // + true, // + true, // + true, // + false, // + ]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let actual = compare_primitive_with_path(&lhs, &rhs, op, path, &mut ctx)?; + let expected = match op { + CompareOperator::Eq => [ + Some(true), // Equal NaNs. + None, // Null on the left. + Some(false), // Distinct signed zeroes. + Some(false), // A finite value and NaN. + None, // Null on the right. + ], + CompareOperator::NotEq | CompareOperator::Lt => [ + Some(false), // Equal NaNs. + None, // Null on the left. + Some(true), // Distinct signed zeroes. + Some(true), // A finite value and NaN. + None, // Null on the right. + ], + _ => unreachable!(), + }; + let expected = BoolArray::from_iter(expected); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(actual.dtype(), &DType::Bool(Nullability::Nullable)); + + // This encoding difference is intentional: the fused path materializes bits and validity + // together, while the RowFn path keeps masking lazy. + match path { + PrimitiveComparisonPath::Columnar => assert_eq!(actual.encoding_id(), Bool.id()), + PrimitiveComparisonPath::Row => assert!(actual.as_opt::().is_some()), + PrimitiveComparisonPath::Auto => unreachable!(), + } + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[rstest] +#[case::i64_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::i64_not_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::u64_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::u64_not_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::f64_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::f64_not_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +fn test_primitive_equality_auto_uses_columnar_for_supported_ptype( + #[case] lhs: ArrayRef, + #[case] rhs: ArrayRef, + #[case] op: CompareOperator, + #[case] expected: [bool; 3], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let actual = + compare_primitive_with_path(&lhs, &rhs, op, PrimitiveComparisonPath::Auto, &mut ctx)?; + + assert_eq!(actual.encoding_id(), Bool.id()); + assert_arrays_eq!(actual, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +} + #[rstest] #[case(Operator::Eq, [true, false, true, true])] #[case(Operator::Lt, [false, true, false, false])] From 415deb1c935f2b904748a752148b97e9d2e9a2a7 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 18:59:22 -0400 Subject: [PATCH 119/160] Benchmark primitive comparison shapes Cover lane widths, equality, nullability, and both constant operand positions for primitive comparison dispatch. Signed-off-by: Connor Tsui --- vortex-array/benches/compare.rs | 115 ++++++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 4a399760dc2..9e6dd3e4e5b 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -4,6 +4,7 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use divan::counter::ItemsCount; use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; @@ -38,6 +39,7 @@ const ARRAY_SIZE: usize = 65_536; fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); bencher + .counter(ItemsCount::new(ARRAY_SIZE)) .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input @@ -49,6 +51,31 @@ fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { }); } +fn u8_array(offset: u8) -> ArrayRef { + (0u8..=u8::MAX) + .cycle() + .take(ARRAY_SIZE) + .map(|value| value.wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn i32_array(offset: i32) -> ArrayRef { + (0i32..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn u64_array(offset: u64) -> ArrayRef { + (0u64..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + fn bool_array(rng: &mut StdRng) -> ArrayRef { BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.5))).into_array() } @@ -87,6 +114,13 @@ fn float_array(rng: &mut StdRng) -> ArrayRef { .into_array() } +fn f32_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f32..1.0)) + .collect::>() + .into_array() +} + fn string_array(rng: &mut StdRng) -> ArrayRef { VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { let len = rng.random_range(1usize..24); @@ -153,6 +187,14 @@ fn compare_int_constant(bencher: Bencher) { bench_compare(bencher, arr, constant, Operator::Gte); } +#[divan::bench] +fn compare_int_constant_lhs(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let constant = ConstantArray::new(50_000_000i64, ARRAY_SIZE).into_array(); + let arr = int_array(&mut rng); + bench_compare(bencher, constant, arr, Operator::Gte); +} + #[divan::bench] fn compare_int_eq(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -161,6 +203,55 @@ fn compare_int_eq(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Eq); } +#[divan::bench] +fn compare_i32(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = i32_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_i32_constant(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = ConstantArray::new(1_000_000i32, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = u8_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8_constant(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = ConstantArray::new(127u8, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_constant(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = ConstantArray::new(1_000_000u64, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_eq(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Eq); +} + #[divan::bench] fn compare_float(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -169,6 +260,30 @@ fn compare_float(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Gte); } +#[divan::bench] +fn compare_float_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = float_array(&mut rng); + let arr2 = float_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_f32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_f32_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + #[divan::bench] fn compare_decimal(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); From 507d0bd4cccce01821c855756121bb158c93fc36 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:46:45 -0400 Subject: [PATCH 120/160] Document primitive comparison dispatch contracts Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs | 3 +++ .../src/scalar_fn/fns/binary/compare/primitive/columnar.rs | 3 +++ 2 files changed, 6 insertions(+) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 18f0c10faea..a3f50f907fd 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -96,6 +96,9 @@ impl RowFn for PrimitiveCompare { const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; fn id(&self) -> ScalarFnId { + // `PrimitiveCompare` is a private implementation detail of `Binary`: it is never registered + // or serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. ScalarFnVTable::id(&Binary) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs index 2a8ee71cfb0..6728437e6a8 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -2,6 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Fused comparison and bit-packing for wide primitive lanes. +//! +//! Production uses this implementation only for measured x86 paths. Keeping it portable lets the +//! semantic tests exercise the RowFn and fused paths on every target. use vortex_buffer::BitBuffer; use vortex_error::VortexResult; From 29a01b5d8d04250c16c6aad87f9687c3ab035f07 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 08:13:57 -0400 Subject: [PATCH 121/160] Use generic execution arguments for comparisons Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index a3f50f907fd..f40ac2f6a90 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -15,11 +15,11 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::scalar_fn::BorrowedExecutionArgs; use crate::scalar_fn::RowFn; use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::execute_rows; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; @@ -80,8 +80,7 @@ pub(super) fn compare_primitive_with_path( return columnar::compare_primitive(lhs, rhs, op, ctx); } - let inputs = [lhs.clone(), rhs.clone()]; - let args = BorrowedExecutionArgs::new(&inputs, lhs.len()); + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); execute_rows(&PrimitiveCompare, &op, &args, ctx) } From f4feded373610bcff17f302cd4a50608d43c0b26 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:08:16 -0400 Subject: [PATCH 122/160] Avoid cloning primitive comparison constants Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs | 5 +++-- .../src/scalar_fn/fns/binary/compare/primitive/operand.rs | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index f40ac2f6a90..ac631075394 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::arrays::Constant; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::PType; @@ -130,8 +131,8 @@ fn use_columnar_comparison( // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. (PType::I64 | PType::F64, _) => true, - // LLVM vectorizes varying u64 inputs, but not the mixed-constant RowFn loop. - (PType::U64, _) => lhs.as_constant().is_some() || rhs.as_constant().is_some(), + // LLVM vectorizes per-row u64 inputs, but not the mixed-constant RowFn loop. + (PType::U64, _) => lhs.is::() || rhs.is::(), _ => false, }) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs index 55d81153b1f..1563b8e68e6 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs @@ -15,7 +15,7 @@ use crate::validity::Validity; /// A materialized primitive column, a non-null constant, or an all-null constant. pub(super) enum PrimitiveOperand { - /// A varying primitive column and its validity. + /// A per-row primitive column and its validity. Array { /// The materialized values. values: Buffer, From 30e02d3bddb854284c195d642db9067dde1ce47f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:36:09 -0400 Subject: [PATCH 123/160] Document primitive RowFn dispatch policy Signed-off-by: Connor Tsui --- .../src/scalar_fn/fns/binary/compare/primitive.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index ac631075394..979d8d5c038 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -16,14 +16,14 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::scalar_fn::RowFn; -use crate::scalar_fn::RowVisitor; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::VecExecutionArgs; -use crate::scalar_fn::execute_rows; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; /// Compare two primitive arrays of the same [`PType`]. /// @@ -131,7 +131,9 @@ fn use_columnar_comparison( // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. (PType::I64 | PType::F64, _) => true, - // LLVM vectorizes per-row u64 inputs, but not the mixed-constant RowFn loop. + // LLVM 22 vectorizes the mixed-constant RowFn loop at 16 CGUs without LTO. However, the + // fused comparison and bit-packing path is still about 38% faster in + // `compare_u64_constant`. Recheck that benchmark before changing this dispatch. (PType::U64, _) => lhs.is::() || rhs.is::(), _ => false, }) From 67e132877e00908049659437ade7b5c139a51975 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:20:20 -0400 Subject: [PATCH 124/160] Declare primitive comparison fallibility Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 979d8d5c038..b68acd4d0d5 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -94,6 +94,7 @@ impl RowFn for PrimitiveCompare { type Options = CompareOperator; const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { // `PrimitiveCompare` is a private implementation detail of `Binary`: it is never registered From 8c215b8dc1ff0c3f95aaedb6cb8e07ca315f786d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:36:22 -0400 Subject: [PATCH 125/160] Compile columnar comparison policy only on x86 Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index b68acd4d0d5..93afcf538ed 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; +#[cfg(target_arch = "x86_64")] use crate::arrays::Constant; use crate::dtype::DType; use crate::dtype::NativePType; @@ -118,6 +119,7 @@ impl RowFn for PrimitiveCompare { } } +#[cfg(target_arch = "x86_64")] fn use_columnar_comparison( lhs: &ArrayRef, rhs: &ArrayRef, From a2b4dc03bb6116d4043d5826d4ca5e110c04f9dd Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 12:14:21 -0400 Subject: [PATCH 126/160] Add encoding-aware RowFn reductions Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 5 + .../scalar_fn/unstable/row/batch/execution.rs | 82 +++++++- .../src/scalar_fn/unstable/row/batch/tests.rs | 197 +++++++++++++++++- .../src/scalar_fn/unstable/row/row_fn.rs | 39 +++- .../src/scalar_fn/unstable/row/vtable.rs | 1 + 5 files changed, 306 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index 781d15711be..922e3ee5095 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -55,6 +55,11 @@ impl<'a> BorrowedExecutionArgs<'a> { } } + /// Return the concrete arrays used by encoding-aware execution. + pub(crate) fn arrays(&self) -> &'a [ArrayRef] { + self.arrays + } + /// Return the original input dtypes used to select the row implementation. pub(crate) fn dtypes(&self) -> &'a [DType] { self.dtypes diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs index 93aad88810e..b14b1c71452 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -3,11 +3,12 @@ //! Applies columnar semantics around one typed row kernel invocation. //! -//! [`Batch`] owns strict null propagation, constant broadcasting, execution strategy selection, and -//! output validation. The row kernel therefore handles only decoded values and its selected output -//! capability. +//! [`Batch`] owns strict null propagation, encoded reductions, constant broadcasting, execution +//! strategy selection, and output validation. The row kernel therefore handles only decoded values +//! and its selected output capability. use smallvec::SmallVec; +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -121,12 +122,17 @@ impl Batch { }) } - /// Apply constant folding and null handling around `kernel`. + /// Apply encoded reductions, constant folding, and null handling around `kernel`. /// - /// For a mixed validity mask, `try_unfiltered` may avoid filtering; `Ok(None)` selects - /// filter-and-scatter. Every kernel result is checked against the planned shape and dtype. + /// `reduce` receives the original inputs before constant broadcasting. For a mixed validity + /// mask, `try_unfiltered` may avoid filtering; `Ok(None)` selects filter-and-scatter. Every + /// kernel result is checked against the planned shape and dtype. pub fn execute( &self, + reduce: impl FnOnce( + BorrowedExecutionArgs<'_>, + &mut ExecutionCtx, + ) -> VortexResult>, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_unfiltered: impl FnOnce( BorrowedExecutionArgs<'_>, @@ -147,6 +153,20 @@ impl Batch { return Ok(self.all_null()); } + // An empty mask is both all-true and all-false, so deferred encoded evidence cannot be + // attributed to an observable row. Let the ordinary policy construct the typed empty + // output instead. + if self.row_count > 0 + && let Some(execution) = reduce(self.execution_args(&self.inputs, self.row_count), ctx)? + { + match execution { + RowExecution::Output(values) => return self.finalize_reduced(values, ctx), + RowExecution::DeferredError(error) => { + return self.resolve_reduced_error(error, kernel, try_unfiltered, ctx); + } + } + } + // All inputs constant, and their conjoined validity proves every row non-null. This sees // through extension and masked wrappers just like argument decoding does. if self.row_count > 0 @@ -331,6 +351,56 @@ impl Batch { ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() } + /// Reconcile an encoding-aware result and apply the batch's strict input validity. + fn finalize_reduced(&self, values: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + validate_output(self.id, &self.result_dtype, self.row_count, &values)?; + + let input_valid = self.validity.execute_mask(self.row_count, ctx)?; + let output_valid = values.validity()?.execute_mask(self.row_count, ctx)?; + vortex_ensure!( + input_valid.bitand_not(&output_valid).all_false(), + "the {} encoded reduction produced nulls for valid rows", + self.id, + ); + + let values = match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => values, + Validity::Array(valid) => values.mask(valid)?, + // Handled before the encoding-aware hook runs. + Validity::AllInvalid => return Ok(self.all_null()), + }; + + cast_output_nullability(&self.result_dtype, values) + } + + /// Resolve deferred evidence from the encoded path by executing only observable rows. + fn resolve_reduced_error( + &self, + error: VortexError, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + /// Pair an input view with this batch's planning metadata. fn execution_args<'b>( &'b self, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 6c92fd52c40..e8e34082990 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -31,7 +31,6 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::InputElement; @@ -59,7 +58,13 @@ struct AddShort; struct ShortDecodeI64; #[derive(Clone)] -struct Identity; +struct OriginalInputReducer; + +#[derive(Clone)] +struct InvalidEncodedReduction; + +#[derive(Clone)] +struct DeferredOriginalReducer; #[derive(Clone)] struct SinkOptions; @@ -308,16 +313,103 @@ impl RowFn for RetryConstantAdd { }, ) } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 1 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(0u8, args[0].len()).into_array(), + ))); + } + + Ok(None) + } } -impl RowFn for Identity { +impl RowFn for OriginalInputReducer { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("test.identity"); + static ID: CachedId = CachedId::new("test.original_input_reducer"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 3 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(42_i64, 3).into_array(), + ))); + } + + Ok(None) + } +} + +impl RowFn for InvalidEncodedReduction { + type Options = usize; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_encoded_reduction"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + null_index: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::Output( + PrimitiveArray::new( + vec![10_i64, 20], + Validity::from_iter((0..2).map(|index| index != *null_index)), + ) + .into_array(), + ))) + } +} + +impl RowFn for DeferredOriginalReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.deferred_original_reducer"); *ID } @@ -329,6 +421,17 @@ impl RowFn for Identity { ) -> VortexResult { visitor.visit::<(i64,), i64>(|(value,)| value) } + + fn reduce_encoded( + &self, + _options: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::DeferredError(vortex_err!( + InvalidArgument: "encoded payload failed" + )))) + } } impl RowFn for SinkOptions { @@ -461,7 +564,7 @@ fn test_short_decode_beside_constant_is_rejected() -> VortexResult<()> { } #[test] -fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> { +fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { let lhs = PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); let rhs = ConstantArray::new(1u8, 2).into_array(); @@ -489,13 +592,89 @@ fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { Ok(()) } +#[test] +fn test_reduce_encoded_defers_errors_behind_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_empty_batch_skips_deferred_encoded_error() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(Vec::::new()).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 0); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid(Validity::AllValid)] +#[case::mixed(Validity::from_iter([true, false]))] +fn test_reduce_encoded_rejects_nulls_on_valid_rows(#[case] validity: Validity) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![10_i64, 20], validity).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&InvalidEncodedReduction, &0, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an encoded reduction introduced a null on a valid row"), + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_encoded_reduction"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("encoded reduction produced nulls for valid rows"), + "the boundary error must identify invalid reduced output, got {error}", + ); + Ok(()) +} + +#[test] +fn test_reduce_encoded_preserves_input_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&InvalidEncodedReduction, &1, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + let expected = ConstantArray::new(42_i64, 3).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + #[test] fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { let input = ConstantArray::new(7_i64, 2).into_array(); let args = VecExecutionArgs::new(vec![input.clone()], 2); let mut ctx = array_session().create_execution_ctx(); - let actual = execute_rows(&Identity, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; assert_arrays_eq!(&actual, &input, &mut ctx); Ok(()) @@ -519,7 +698,8 @@ fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResul let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; @@ -547,7 +727,8 @@ fn test_valid_only_filters_and_scatters() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 8a982c4fb37..0b813e2e039 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -4,8 +4,8 @@ //! The [`RowFn`] contract for scalar functions whose natural kernel computes one row at a time. //! //! Implementations declare their arity and fallibility, then use [`RowFn::dispatch`] to select the -//! typed row signature for each supported dtype combination. Optional methods provide -//! serialization without putting persistence plumbing in the row kernel. +//! typed row signature for each supported dtype combination. Optional hooks provide serialization +//! and encoding-aware execution without putting columnar plumbing in the row kernel. use std::fmt::Debug; use std::fmt::Display; @@ -16,8 +16,11 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use super::visitor::RowVisitor; +use crate::ArrayRef; +use crate::ExecutionCtx; use crate::dtype::DType; use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::unstable::row::RowExecution; /// A scalar function computed one row at a time. /// @@ -35,12 +38,14 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// The arguments in display order. Its length is the function's exact arity. const ARG_NAMES: &'static [&'static str]; - /// Whether any dispatch can raise a semantic error. + /// Whether any dispatch or encoded reduction can raise a semantic error. /// /// See [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible) for a /// more detailed explanation of semantic errors. /// - /// The framework checks dispatched element and result types. A conservative `true` is allowed. + /// The framework checks dispatched element and result types, but cannot inspect + /// [`reduce_encoded`](Self::reduce_encoded). Set this to `true` when that hook can return a + /// semantic error or [`RowExecution::DeferredError`]. A conservative `true` is allowed. const FALLIBLE: bool; /// Returns the ID of the scalar function. @@ -71,4 +76,30 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { args: &[DType], visitor: V, ) -> VortexResult; + + /// Try an encoding-aware implementation before decoding the inputs into row elements. + /// + /// `None` continues to the row loop. [`Output`](RowExecution::Output) may remain encoded or + /// lazy. [`DeferredError`](RowExecution::DeferredError) retries only valid rows. Batch execution + /// calls this hook at most once with the original nonempty inputs. Nullary functions, empty + /// batches, slices, and compacted retries skip it. + /// + /// Like a dense row closure, this hook must be total over every stored payload, including + /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or + /// retried through the row layer. + /// + /// # Requirements + /// + /// - `output.len()` **must** equal `args[0].len()`. + /// - The output dtype **must** match the planned dtype when ignoring nullability. + /// - The output **must not** introduce a null where every input is valid. + fn reduce_encoded( + &self, + options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + _ = (options, args, ctx); + Ok(None) + } } diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 9675219d547..98e43c3dd9c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -132,6 +132,7 @@ pub fn execute_rows( let batch = prepare_batch(function, options, args)?; batch.execute( + |args, ctx| function.reduce_encoded(options, args.arrays(), ctx), |args, ctx| execute_row_kernel(function, options, args, ctx), |args, valid, ctx| try_execute_rows_unfiltered(function, options, args, valid, ctx), ctx, From 7eb107845cca3c9b35bbd17d554cdf6be9988325 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 20:58:24 -0400 Subject: [PATCH 127/160] Execute tensor L2 norm with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/benches/l2_norm.rs | 36 +- .../src/scalar_fns/cosine_similarity.rs | 7 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 401 ++++-------------- vortex-tensor/src/scalar_fns/mod.rs | 4 + vortex-tensor/src/scalar_fns/row.rs | 165 +++++++ vortex-tensor/src/scalar_fns/tests/l2_norm.rs | 318 ++++++++++++++ vortex-tensor/src/scalar_fns/tests/mod.rs | 7 + vortex-tensor/src/scalar_fns/tests/row.rs | 98 +++++ vortex-tensor/src/utils.rs | 59 +++ 9 files changed, 773 insertions(+), 322 deletions(-) create mode 100644 vortex-tensor/src/scalar_fns/row.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/l2_norm.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/mod.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/row.rs diff --git a/vortex-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index 6f597084113..bf8832f2520 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -15,11 +15,19 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::l2_norm::L2Norm; @@ -54,13 +62,25 @@ fn vectors(width: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { let session = vortex_array::array_session(); bencher .counter(ItemsCount::new(input.len())) .with_inputs(|| { ( - L2Norm::try_new_array(input.clone()).unwrap().into_array(), + L2Norm + .try_new_array(input.len(), EmptyOptions, [input.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -80,3 +100,17 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_l2_norm(bencher, input); } + +#[divan::bench(args = WIDTHS)] +fn constant(bencher: Bencher, width: usize) { + bench_l2_norm(bencher, constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let input = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_l2_norm(bencher, input); +} diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index ca8fcf0efd4..76d266a2721 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -10,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; @@ -144,8 +145,8 @@ impl ScalarFnVTable for CosineSimilarity { let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm::try_new_array(lhs_ref.clone())?; - let norm_rhs_arr = L2Norm::try_new_array(rhs_ref.clone())?; + let norm_lhs_arr = L2Norm.try_new_array(len, EmptyOptions, [lhs_ref.clone()])?; + let norm_rhs_arr = L2Norm.try_new_array(len, EmptyOptions, [rhs_ref.clone()])?; let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; // Execute to get the inner product and norms of the arrays. We only fully decompress @@ -288,7 +289,7 @@ impl CosineSimilarity { let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - let norm_arr = L2Norm::try_new_array(plain_ref.clone())?; + let norm_arr = L2Norm.try_new_array(len, EmptyOptions, [plain_ref.clone()])?; let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index b7e9060ed3f..f8349d3431b 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -3,50 +3,44 @@ //! L2 norm expression for tensor-like types. -use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowExecution; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::UninitElementSink; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::Normalized; -use crate::matcher::AnyTensor; -use crate::utils::extract_flat_elements; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::extract_normalized_children; +use crate::utils::l2_norm_row; use crate::utils::validate_tensor_float_input; /// L2 norm (Euclidean norm) of a tensor or vector column. @@ -62,142 +56,116 @@ use crate::utils::validate_tensor_float_input; /// of the storage contract, not a separate lossy-compute mode. /// /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct L2Norm; impl L2Norm { /// Creates a new [`TypedScalarFnInstance`] wrapping the L2 norm operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(L2Norm, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the L2 norm over `child`. /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtype is + /// unsupported. pub fn try_new_array(child: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(L2Norm::new().erased(), vec![child]) + ScalarFnArray::try_new(Self::new().erased(), vec![child]) } } -impl ScalarFnVTable for L2Norm { +impl RowFn for L2Norm { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["input"]; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.l2_norm"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(1) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("input"), - _ => unreachable!("L2Norm must have exactly one child"), - } - } - - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let input_dtype = &arg_dtypes[0]; - let tensor_match = validate_tensor_float_input(input_dtype)?; - let ptype = tensor_match.element_ptype(); - - let nullability = Nullability::from(input_dtype.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn execute( + fn dispatch>( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let input_ref = args.get(0)?; - let row_count = args.row_count(); - - let ext = input_ref.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let tensor_flat_size = tensor_match.list_size() as usize; - let element_ptype = tensor_match.element_ptype(); - - let norm_dtype = DType::Primitive(element_ptype, ext.nullability()); - - // L2Norm over a `Normalized`-encoded column is defined to read back the authoritative stored - // norms. Callers of lossy encodings opt into that storage semantics instead of forcing a - // decode-and-recompute path here. - if input_ref.is::() { - let (_, norms) = extract_normalized_children(&input_ref); - vortex_ensure_eq!(norms.dtype(), &norm_dtype); - return Ok(norms); - } - - // Optimize for the constant array case. - if let Some(array) = input_ref.as_opt::() { - let scalar = array.scalar().as_extension().to_storage_scalar(); - - let Some(elements) = scalar.as_list().elements() else { - return Ok(ConstantArray::new(Scalar::null(norm_dtype), row_count).into_array()); - }; - - let norm_scalar = match_each_float_ptype!(element_ptype, |T| { - let values: Vec = elements - .iter() - .map(|s| { - s.as_primitive() - .as_::() - .vortex_expect("element was somehow not the correct float") - }) - .collect(); - let norm = l2_norm_row::(&values); - - Scalar::try_new(norm_dtype, Some(norm.into())) - })?; - - let norms = ConstantArray::new(norm_scalar, row_count).into_array(); - return Ok(norms); - } - - let input: ExtensionArray = input_ref.execute(ctx)?; - let validity = input.as_ref().validity()?; - - let storage = input.storage_array(); - let flat = extract_flat_elements(storage, tensor_flat_size, ctx)?; - - match_each_float_ptype!(flat.ptype(), |T| { - let buffer: Buffer = (0..row_count) - .map(|i| l2_norm_row(flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow,), UninitElementSink, _>(|(row,), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l2_norm_row(row)) } + }) }) } - fn validity( + /// `L2Norm` over a [`Normalized`]-encoded column is defined to read back the authoritative + /// stored norms. Callers of lossy encodings opt into that storage semantics instead of forcing + /// a decode-and-recompute path here. + fn reduce_encoded( &self, _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if the input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let input = &args[0]; + if input.is::() { + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let (_, norms) = extract_normalized_children(input); + vortex_ensure!( + norms.dtype().is_primitive(), + "normalized norms must be primitive, got {}", + norms.dtype(), + ); + vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); + return Ok(Some(RowExecution::Output(norms))); + } - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + let Some(constant) = input.as_opt::() else { + return Ok(None); + }; + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let norm_dtype = + DType::Primitive(element_ptype, input.dtype().as_extension().nullability()); + let storage = constant.scalar().as_extension().to_storage_scalar(); + + let Some(elements) = storage.as_list().elements() else { + let output = ConstantArray::new(Scalar::null(norm_dtype), input.len()); + return Ok(Some(RowExecution::Output(output.into_array()))); + }; + + let norm = match_each_float_ptype!(element_ptype, |T| { + let values: Vec = elements + .iter() + .map(|element| { + element + .as_primitive() + .as_::() + .vortex_expect("tensor element must match its declared ptype") + }) + .collect(); + Scalar::try_new(norm_dtype, Some(l2_norm_row::(&values).into())) + })?; + let output = ConstantArray::new(norm, input.len()); + Ok(Some(RowExecution::Output(output.into_array()))) } } +vortex_array::impl_row_fn_vtable!(L2Norm); + /// Metadata for a serialized [`L2Norm`] array: the single `input` child's [`DType`], which carries /// the extension type (`FixedShapeTensor` vs `Vector`), dimension, and nullability that are not /// recoverable from the parent's primitive-float output. @@ -240,206 +208,3 @@ impl ScalarFnArrayVTable for L2Norm { }) } } - -/// Computes the L2 norm (Euclidean norm) of a float slice. -/// -/// Returns `sqrt(sum(v_i^2))`. A zero-length or all-zero input produces `0.0`. -fn l2_norm_row(v: &[T]) -> T { - let mut sum_sq = T::zero(); - for &x in v { - sum_sq = sum_sq + x * x; - } - sum_sq.sqrt() -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::EmptyMetadata; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Constant; - use vortex_array::arrays::ConstantArray; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::dtype::extension::ExtDType; - use vortex_array::scalar::Scalar; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::scalar_fns::l2_norm::L2Norm; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::literal_vector_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. - fn eval_l2_norm(input: ArrayRef) -> VortexResult> { - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[rstest] - #[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] - #[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] - #[case::single_element(&[1], &[7.0], &[7.0])] - #[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] - fn known_norms( - #[case] shape: &[usize], - #[case] elements: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let arr = tensor_array(shape, elements)?; - assert_close(&eval_l2_norm(arr)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let arr = tensor_array( - &[3], - &[ - 3.0, 4.0, 0.0, // norm = 5.0 - 0.0, 0.0, 0.0, // norm = 0.0 - 1.0, 1.0, 1.0, // norm = sqrt(3) - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); - Ok(()) - } - - #[test] - fn vector_multiple_rows() -> VortexResult<()> { - let arr = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // norm = 1.0 - 3.0, 4.0, 0.0, // norm = 5.0 - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 is masked as null. - let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; - let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: norm = 5.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is a non-null tensor should short-circuit to a - /// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so - /// execution stops at the [`Constant`] encoding instead of canonicalizing into a - /// [`PrimitiveArray`]. - #[test] - fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { - let input = literal_vector_array(&[3.0f64, 4.0], 4); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let output = result.execute_until::(&mut ctx)?; - - let constant = output - .as_opt::() - .expect("L2Norm over a constant input must produce a constant output"); - assert_eq!(constant.len(), 4); - let norm = constant - .scalar() - .as_primitive() - .as_::() - .expect("norm scalar must be a non-null primitive"); - assert_close(&[norm], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of - /// the correct primitive dtype and length. - #[test] - fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { - let storage_dtype = DType::FixedSizeList( - DType::Primitive(PType::F64, Nullability::NonNullable).into(), - 2, - Nullability::Nullable, - ); - let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); - let null_scalar = Scalar::null(DType::Extension(ext_dtype)); - let input = ConstantArray::new(null_scalar, 3).into_array(); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let output = result.execute_until::(&mut ctx)?; - - let constant = output - .as_opt::() - .expect("null constant input must produce a constant output"); - assert_eq!(constant.len(), 3); - assert!(constant.scalar().is_null()); - assert_eq!( - constant.dtype(), - &DType::Primitive(PType::F64, Nullability::Nullable) - ); - Ok(()) - } - - #[rstest] - #[case::fixed_shape_tensor(l2_norm_tensor_child())] - #[case::vector(l2_norm_vector_child())] - fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { - let original = L2Norm::try_new_array(child.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(L2Norm); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("L2Norm serialize must produce metadata"); - - let children = vec![child]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn l2_norm_tensor_child() -> ArrayRef { - tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") - } - - fn l2_norm_vector_child() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } -} diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 68f10ca6b01..da9b8950e7a 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -6,3 +6,7 @@ pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; +pub(crate) mod row; + +#[cfg(test)] +mod tests; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs new file mode 100644 index 00000000000..6f7e4d64d9a --- /dev/null +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the tensor scalar functions add to the row-function machinery: an element type that reads a +//! tensor row and the width rule they share. + +use std::marker::PhantomData; + +use num_traits::Float; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::Masked; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::masked::MaskedArraySlotsExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::scalar_fn::InputElement; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::utils::extract_flat_elements; +use crate::utils::validate_tensor_float_input; +use crate::utils::validate_tensor_float_inputs; + +/// The width rule the tensor scalar functions share: every argument is the same float tensor dtype, +/// and the width is its element ptype. +pub(crate) fn tensor_element_ptype(args: &[DType]) -> VortexResult { + Ok(validate_tensor_float_inputs(args)?.element_ptype()) +} + +/// Marker for tensor-valued input elements: accepts any tensor-like extension column whose +/// elements are `T`, and presents each row as its flat elements, `&[T]`. +pub struct TensorRow(PhantomData); + +/// The decoded form of a [`TensorRow`] column: one flat typed buffer plus the stride to read it at. +/// +/// Typed at decode time rather than per row. `FlatElements::row` re-derives its typed slice on every +/// call, which costs a ptype check and a buffer downcast per row; a row loop reads every row, so it +/// pays that once here instead. +pub struct TensorRows { + /// Every row's elements, back to back. + elements: Buffer, + + /// Number of logical tensor rows, stored so zero-width tensors retain their length. + rows: usize, + + /// Elements per row, the length of each row slice. + list_size: usize, + + /// `list_size` for a full column and `0` for constant-backed storage, so `index * stride` pins a + /// constant to its single materialized row without a branch in the loop. + stride: usize, +} + +// SAFETY: `TensorRows` records the row count validated during decode, and both checked and +// unchecked access use the same stride and row width. +unsafe impl InputElement for TensorRow { + type Column = TensorRows; + type Varying<'a> = &'a TensorRows; + type Elem<'a> = &'a [T]; + + // Tensor storage is a fully materialized non-nullable primitive buffer, so the elements behind + // a null row are arbitrary values rather than an unresolvable reference. + const DENSE_SAFE: bool = true; + // Tensor storage is a primitive buffer; reading it cannot fail on account of its values. + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + let tensor_match = validate_tensor_float_input(dtype)?; + let expected = T::PTYPE; + vortex_ensure_eq!( + tensor_match.element_ptype(), + expected, + "expected a tensor of {expected} elements, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Dense batch execution owns the mask and restores it on the result. Decode the values + // directly so a nullable tensor does not rebuild its extension storage under that mask. + let array = match array.as_opt::() { + Some(masked) => masked.child().clone(), + None => array, + }; + + let rows = array.len(); + let list_size = validate_tensor_float_input(array.dtype())?.list_size() as usize; + let ext: ExtensionArray = array.execute(ctx)?; + let flat = extract_flat_elements(ext.storage_array(), list_size, ctx)?; + let list_size = flat.list_size(); + let stride = flat.row_stride(); + let elements = flat.into_buffer::(); + + let expected_elements = if stride == 0 { + list_size + } else { + vortex_ensure_eq!( + stride, + list_size, + "varying tensor row stride must equal its width, got {stride}", + ); + let Some(expected_elements) = rows.checked_mul(stride) else { + vortex_bail!( + "tensor row storage length must fit usize, got {rows} rows of width {stride}", + ); + }; + + expected_elements + }; + vortex_ensure_eq!( + elements.len(), + expected_elements, + "tensor row storage must contain {expected_elements} elements, got {}", + elements.len(), + ); + + Ok(TensorRows { + elements, + rows, + list_size, + stride, + }) + } + + fn get(column: &Self::Column, index: usize) -> &[T] { + let start = index * column.stride; + &column.elements.as_slice()[start..start + column.list_size] + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.rows + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + Self::get(column, index) + } + + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + let start = index * column.stride; + + // SAFETY: decode established one complete stored row for stride 0, or `rows` contiguous + // `list_size`-element rows otherwise. The caller guarantees `index < rows`. + unsafe { + std::slice::from_raw_parts( + column.elements.as_slice().as_ptr().add(start), + column.list_size, + ) + } + } +} diff --git a/vortex-tensor/src/scalar_fns/tests/l2_norm.rs b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs new file mode 100644 index 00000000000..930ef423d47 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs @@ -0,0 +1,318 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::l2_norm::L2Norm; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. +fn eval_l2_norm(input: ArrayRef) -> VortexResult> { + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = L2Norm::new(); + let array = L2Norm::try_new_array(tensor_array(&[1], &[3.0])?)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + assert_close( + &eval_l2_norm(zero_width_vector_array::(3)?)?, + &[0.0, 0.0, 0.0], + ); + assert!(eval_l2_norm(vector_array(2, &[] as &[f64])?)?.is_empty()); + + let constant = Vector::constant_array::(&[], 3)?; + assert_close(&eval_l2_norm(constant)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +#[rstest] +#[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] +#[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] +#[case::single_element(&[1], &[7.0], &[7.0])] +#[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] +fn known_norms( + #[case] shape: &[usize], + #[case] elements: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let arr = tensor_array(shape, elements)?; + assert_close(&eval_l2_norm(arr)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let arr = tensor_array( + &[3], + &[ + 3.0, 4.0, 0.0, // norm = 5.0 + 0.0, 0.0, 0.0, // norm = 0.0 + 1.0, 1.0, 1.0, // norm = sqrt(3) + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); + Ok(()) +} + +#[test] +fn vector_multiple_rows() -> VortexResult<()> { + let arr = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // norm = 1.0 + 3.0, 4.0, 0.0, // norm = 5.0 + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 is masked as null. + let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: norm = 5.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is a non-null tensor should short-circuit to a +/// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so +/// execution stops at the [`Constant`] encoding instead of canonicalizing into a +/// [`PrimitiveArray`]. +#[test] +fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { + let input = literal_vector_array(&[3.0f64, 4.0], 4); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over a constant input must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// An extension array over constant storage is folded just like a top-level constant instead of +/// recomputing the same norm once per row. +#[test] +fn extension_backed_constant_yields_constant_output() -> VortexResult<()> { + let input = Vector::constant_array(&[3.0f64, 4.0], 4)?; + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over constant-backed extension storage must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of +/// the correct primitive dtype and length. +#[test] +fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { + let storage_dtype = DType::FixedSizeList( + DType::Primitive(PType::F64, Nullability::NonNullable).into(), + 2, + Nullability::Nullable, + ); + let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); + let null_scalar = Scalar::null(DType::Extension(ext_dtype)); + let input = ConstantArray::new(null_scalar, 3).into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("null constant input must produce a constant output"); + assert_eq!(constant.len(), 3); + assert!(constant.scalar().is_null()); + assert_eq!( + constant.dtype(), + &DType::Primitive(PType::F64, Nullability::Nullable) + ); + Ok(()) +} + +/// An `f32` column must dispatch at `f32` and produce an `f32` result, which is the property that +/// makes width polymorphism load-bearing rather than decorative. +#[rstest] +#[case::f32(&[3.0f32, 4.0], PType::F32)] +#[case::f64(&[3.0f64, 4.0], PType::F64)] +fn dispatches_at_input_width( + #[case] elements: &[T], + #[case] expected: PType, +) -> VortexResult<()> { + let arr = tensor_array(&[2], elements)?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L2Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .execute(&mut ctx)?; + assert_eq!(prim.ptype(), expected); + Ok(()) +} + +/// `L2Norm(Normalized(normalized, norms))` reads back the authoritative stored norms rather than +/// recomputing over decoded coordinates. The normalized child here is deliberately *not* +/// unit-norm, mimicking lossy storage, so readthrough and recompute disagree: row 0 decodes to +/// `[6, 8]` (norm `10`) and row 1 to `[6, 0]` (norm `6`), while the stored norms are `5` and `2`. +#[test] +fn normalized_readthrough_returns_stored_norms() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 2.0]).into_array(); + // SAFETY: A focused test of the lossy storage contract: the stored norms are authoritative + // even though this normalized child violates the unit-norm invariant. + let denorm = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); + + assert_close(&eval_l2_norm(denorm)?, &[5.0, 2.0]); + Ok(()) +} + +/// The readthrough must survive a partially-null column. +/// +/// This pins the dense policy the row contract derives. Filtering could hand `reduce_encoded` a +/// filtered input, which is no longer an `ExactScalarFn`, silently falling back to +/// decode-and-recompute. For a lossy child that changes the answer: row 0 below would come back as +/// `10` (recomputed from `[6, 8]`) instead of the authoritative stored `5`. +#[test] +fn normalized_readthrough_survives_null_rows() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + // SAFETY: Intentionally lossy, as in `normalized_readthrough_returns_stored_norms`, so that + // a recompute fallback is observable. + let denorm = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// The readthrough must still propagate nulls carried by the `norms` child. +#[test] +fn normalized_readthrough_propagates_null_norms() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let denorm = Normalized::try_new(normalized, norms, &mut ctx)?.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +#[rstest] +#[case::fixed_shape_tensor(l2_norm_tensor_child())] +#[case::vector(l2_norm_vector_child())] +fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { + let original = L2Norm.try_new_array(child.len(), EmptyOptions, [child.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(L2Norm); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("L2Norm serialize must produce metadata"); + + let children = vec![child]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn l2_norm_tensor_child() -> ArrayRef { + tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") +} + +fn l2_norm_vector_child() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/mod.rs b/vortex-tensor/src/scalar_fns/tests/mod.rs new file mode 100644 index 00000000000..5447772adda --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/mod.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for the tensor scalar functions. + +mod l2_norm; +mod row; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs new file mode 100644 index 00000000000..77248329ae5 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::Float; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::match_each_float_ptype; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::tensor_array; + +/// The marginal cost of a new tensor scalar function is this entire definition. Everything else +/// (null propagation, constants, validity, f16/f32/f64 dispatch, dtype checks, and constructors) is +/// derived. +#[derive(Clone, Debug, Default)] +struct L1Norm; + +impl RowFn for L1Norm { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.l1_norm"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow,), UninitElementSink, _>(|(row,), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l1_norm_row(row)) } + }) + }) + } +} + +vortex_array::impl_row_fn_vtable!(L1Norm); + +fn l1_norm_row(row: &[T]) -> T { + row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) +} + +#[test] +fn derived_fn_executes_with_nulls() -> VortexResult<()> { + let arr = tensor_array(&[2], &[3.0, -4.0, 1.0, 1.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L1Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[7.0]); + Ok(()) +} + +/// A kernel written once serves every float width. +#[test] +fn derived_fn_dispatches_at_input_width() -> VortexResult<()> { + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + + let f32_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f32, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f32_result.ptype(), PType::F32); + assert_eq!(f32_result.as_slice::(), &[7.0f32]); + + let f64_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f64, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f64_result.ptype(), PType::F64); + Ok(()) +} diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 488694bd47f..df40bdea44a 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Shared helpers for tensor scalar functions. + use half::f16; +use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -20,6 +23,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::dtype::proto::dtype as pb; use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -58,6 +62,16 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { SAFETY_FACTOR as f64 * machine_epsilon * dimensions_root } +/// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. +pub(crate) fn l2_norm_row(row: &[T]) -> T { + let mut sum_squared = T::zero(); + for &element in row { + sum_squared = sum_squared + element * element; + } + + sum_squared.sqrt() +} + /// Extracts the `(normalized, norms)` children of a [`Normalized`]-encoded array. /// /// # Panics @@ -110,6 +124,22 @@ pub fn validate_binary_tensor_float_inputs<'a>( validate_tensor_float_input(lhs) } +/// Validates that every argument has the same float tensor dtype, ignoring nullability. +pub fn validate_tensor_float_inputs(args: &[DType]) -> VortexResult> { + let (first, rest) = args + .split_first() + .ok_or_else(|| vortex_err!("tensor expression expects at least one input"))?; + + for arg in rest { + vortex_ensure!( + first.eq_ignore_nullability(arg), + "tensor expression expects inputs to have the same dtype, got {first} and {arg}" + ); + } + + validate_tensor_float_input(first) +} + /// The flat primitive elements of a tensor storage array, with typed row access. /// /// This struct hides the stride detail that arises from the [`ConstantArray`] optimization: a @@ -138,6 +168,23 @@ impl FlatElements { let slice = self.elems.as_slice::(); &slice[row_idx * self.list_size..][..self.list_size] } + + /// Returns the number of elements in each row. + #[must_use] + pub fn list_size(&self) -> usize { + self.list_size + } + + /// Returns the physical distance between rows, or zero when every row uses one stored value. + #[must_use] + pub fn row_stride(&self) -> usize { + if self.is_constant { 0 } else { self.list_size } + } + + /// Returns the elements as a typed buffer, performing the ptype check once for the batch. + pub fn into_buffer(self) -> Buffer { + self.elems.into_buffer::() + } } /// Extracts the flat primitive elements from a tensor storage array (FixedSizeList). @@ -343,6 +390,18 @@ pub mod test_helpers { Vector::try_new_vector_array(flat_fsl(elements, dim)) } + /// Builds `rows` zero-width vectors over an empty typed element buffer. + pub fn zero_width_vector_array(rows: usize) -> VortexResult { + let storage = FixedSizeListArray::new( + Buffer::::empty().into_array(), + 0, + Validity::NonNullable, + rows, + ) + .into_array(); + Vector::try_new_vector_array(storage) + } + /// Builds a [`FixedShapeTensor`] extension array whose storage is a [`ConstantArray`], /// representing a single query tensor broadcast to `len` rows. pub fn constant_tensor_array>( From 8d707ad8157bd69e6b9997a8ade27f269d910b34 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:08:55 -0400 Subject: [PATCH 128/160] Update tensor L2 for the RowFn API Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/l2_norm.rs | 2 -- vortex-tensor/src/scalar_fns/row.rs | 23 ++++++++++------------- vortex-tensor/src/scalar_fns/tests/row.rs | 2 -- 3 files changed, 10 insertions(+), 17 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index f8349d3431b..2a4fb0368cf 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -164,8 +164,6 @@ impl RowFn for L2Norm { } } -vortex_array::impl_row_fn_vtable!(L2Norm); - /// Metadata for a serialized [`L2Norm`] array: the single `input` child's [`DType`], which carries /// the extension type (`FixedShapeTensor` vs `Vector`), dimension, and nullability that are not /// recoverable from the parent's primitive-float output. diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 6f7e4d64d9a..8340e97160a 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -60,7 +60,7 @@ pub struct TensorRows { // unchecked access use the same stride and row width. unsafe impl InputElement for TensorRow { type Column = TensorRows; - type Varying<'a> = &'a TensorRows; + type View<'a> = &'a TensorRows; type Elem<'a> = &'a [T]; // Tensor storage is a fully materialized non-nullable primitive buffer, so the elements behind @@ -102,7 +102,7 @@ unsafe impl InputElement for TensorRow { vortex_ensure_eq!( stride, list_size, - "varying tensor row stride must equal its width, got {stride}", + "per-row tensor stride must equal its width, got {stride}", ); let Some(expected_elements) = rows.checked_mul(stride) else { vortex_bail!( @@ -132,34 +132,31 @@ unsafe impl InputElement for TensorRow { &column.elements.as_slice()[start..start + column.list_size] } - fn varying(column: &Self::Column) -> Self::Varying<'_> { + fn view(column: &Self::Column) -> Self::View<'_> { column } - fn varying_len(column: &Self::Varying<'_>) -> usize { - column.rows + fn view_len(view: &Self::View<'_>) -> usize { + view.rows } - fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] where Self: 'a, { - Self::get(column, index) + Self::get(view, index) } - unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> &'a [T] + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] where Self: 'a, { - let start = index * column.stride; + let start = index * view.stride; // SAFETY: decode established one complete stored row for stride 0, or `rows` contiguous // `list_size`-element rows otherwise. The caller guarantees `index < rows`. unsafe { - std::slice::from_raw_parts( - column.elements.as_slice().as_ptr().add(start), - column.list_size, - ) + std::slice::from_raw_parts(view.elements.as_slice().as_ptr().add(start), view.list_size) } } } diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index 77248329ae5..b88a22838aa 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -57,8 +57,6 @@ impl RowFn for L1Norm { } } -vortex_array::impl_row_fn_vtable!(L1Norm); - fn l1_norm_row(row: &[T]) -> T { row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) } From 0beffe06a1aa6bcae5db29dbaea56a02caa82a4e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:37:56 -0400 Subject: [PATCH 129/160] Opt tensor L2 into unstable RowFn Signed-off-by: Connor Tsui --- vortex-tensor/Cargo.toml | 2 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 10 +++++----- vortex-tensor/src/scalar_fns/row.rs | 2 +- vortex-tensor/src/scalar_fns/tests/row.rs | 8 ++++---- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/vortex-tensor/Cargo.toml b/vortex-tensor/Cargo.toml index abdca676775..9706102f6d6 100644 --- a/vortex-tensor/Cargo.toml +++ b/vortex-tensor/Cargo.toml @@ -17,7 +17,7 @@ version = { workspace = true } workspace = true [dependencies] -vortex-array = { workspace = true } +vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 2a4fb0368cf..d0e99576f83 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -20,13 +20,13 @@ use vortex_array::dtype::proto::dtype as pb; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowExecution; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; use vortex_error::VortexExpect; use vortex_error::VortexResult; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 8340e97160a..9a9600bd6f9 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -16,7 +16,7 @@ use vortex_array::arrays::masked::MaskedArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; -use vortex_array::scalar_fn::InputElement; +use vortex_array::scalar_fn::unstable::row::InputElement; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index b88a22838aa..aff3c2d0bb1 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -12,11 +12,11 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_session::registry::CachedId; From 35e9945512a66f192e0986d5cba6eb4da71235df Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:21:25 -0400 Subject: [PATCH 130/160] Declare tensor norm fallibility Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/l2_norm.rs | 1 + vortex-tensor/src/scalar_fns/tests/row.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index d0e99576f83..78cabc961b0 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -80,6 +80,7 @@ impl RowFn for L2Norm { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.l2_norm"); diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs index aff3c2d0bb1..8a7fa23edd4 100644 --- a/vortex-tensor/src/scalar_fns/tests/row.rs +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -36,6 +36,7 @@ impl RowFn for L1Norm { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.test.l1_norm"); From 6c84e6ed1f449f8b2745a52019a74d58ac7e13a6 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:25:53 -0400 Subject: [PATCH 131/160] Opt tensor rows into null-tolerant decoding Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/row.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index 9a9600bd6f9..a99ad76ea6e 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -127,6 +127,10 @@ unsafe impl InputElement for TensorRow { }) } + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + fn get(column: &Self::Column, index: usize) -> &[T] { let start = index * column.stride; &column.elements.as_slice()[start..start + column.list_size] From 4ddec3c9f7ffc6b57c6c37d555fc1bff647ad0b9 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:00:46 -0400 Subject: [PATCH 132/160] Execute tensor product functions with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/benches/cosine_similarity.rs | 35 +- vortex-tensor/benches/inner_product.rs | 43 +- vortex-tensor/src/encodings/normalized/mod.rs | 1 - .../src/scalar_fns/cosine_similarity.rs | 921 +++++------------- vortex-tensor/src/scalar_fns/inner_product.rs | 497 ++-------- vortex-tensor/src/scalar_fns/row.rs | 17 + .../src/scalar_fns/tests/cosine_similarity.rs | 607 ++++++++++++ .../src/scalar_fns/tests/inner_product.rs | 282 ++++++ vortex-tensor/src/scalar_fns/tests/mod.rs | 2 + vortex-tensor/src/utils.rs | 18 +- vortex-tensor/src/vector_search.rs | 4 +- 11 files changed, 1297 insertions(+), 1130 deletions(-) create mode 100644 vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/inner_product.rs diff --git a/vortex-tensor/benches/cosine_similarity.rs b/vortex-tensor/benches/cosine_similarity.rs index 6cc5eb867ef..49551fbf701 100644 --- a/vortex-tensor/benches/cosine_similarity.rs +++ b/vortex-tensor/benches/cosine_similarity.rs @@ -21,11 +21,14 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity; @@ -41,15 +44,11 @@ fn main() { } /// Total `f64` elements per operand, held constant across widths: the row count is -/// `ELEMENTS / width`. This budget is a quarter of the one the other tensor benches use, because -/// the constant arms recompute the broadcast vector's norm per row and cost roughly ten times the -/// column arms per element. It is what keeps every arm inside the 1 ms per-iteration limit from -/// `docs/developer-guide/benchmarking.md`, measured against CodSpeed's CPU simulation. +/// `ELEMENTS / width`. The smaller budget keeps the wider cosine kernels inside the 1 ms +/// per-iteration limit from `docs/developer-guide/benchmarking.md` under CodSpeed simulation. const ELEMENTS: usize = 2_048; -/// Widths chosen to separate the two costs, as in `l2_norm.rs`: the redundant norm pass is -/// `O(rows * width)`, one third of the closure's arithmetic, so wide tensors show the hoist -/// while a narrow one is dominated by per-row framework costs. +/// Widths that expose both fixed row-framework costs and the `O(width)` kernel work. const WIDTHS: &[usize] = &[2, 32, 256]; /// `ELEMENTS / width` vectors of `width` `f64` elements, non-nullable. `seed` offsets the values so @@ -85,9 +84,9 @@ fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { bencher .with_inputs(|| { ( - CosineSimilarity::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + CosineSimilarity + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -106,6 +105,22 @@ fn column_x_constant(bencher: Bencher, width: usize) { bench_cosine(bencher, vectors(width, 0), constant_vector(width)); } +/// The lhs is a broadcast query vector, whose norm is the same in every row. +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_cosine(bencher, constant_vector(width), vectors(width, 31)); +} + +/// A nullable broadcast rhs exercises constant preparation and output validity together. +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_cosine(bencher, vectors(width, 0), rhs); +} + /// One query vector represented as an extension array over constant storage. fn extension_constant_vector(width: usize) -> ArrayRef { let ext_dtype = vectors(width, 0).dtype().as_extension().clone(); diff --git a/vortex-tensor/benches/inner_product.rs b/vortex-tensor/benches/inner_product.rs index 796e9b648d6..5c4adf1c7ec 100644 --- a/vortex-tensor/benches/inner_product.rs +++ b/vortex-tensor/benches/inner_product.rs @@ -15,11 +15,19 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::inner_product::InnerProduct; @@ -56,15 +64,25 @@ fn vectors(width: usize, seed: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_inner_product(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { let session = vortex_array::array_session(); bencher .counter(ItemsCount::new(lhs.len())) .with_inputs(|| { ( - InnerProduct::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + InnerProduct + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -84,3 +102,22 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_inner_product(bencher, lhs, vectors(width, 31)); } + +#[divan::bench(args = WIDTHS)] +fn column_x_constant(bencher: Bencher, width: usize) { + bench_inner_product(bencher, vectors(width, 0), constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_inner_product(bencher, constant_vector(width), vectors(width, 31)); +} + +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_inner_product(bencher, vectors(width, 0), rhs); +} diff --git a/vortex-tensor/src/encodings/normalized/mod.rs b/vortex-tensor/src/encodings/normalized/mod.rs index 545236bba7d..2c9a72d2988 100644 --- a/vortex-tensor/src/encodings/normalized/mod.rs +++ b/vortex-tensor/src/encodings/normalized/mod.rs @@ -31,7 +31,6 @@ pub use array::NormalizedSlots; mod compress; pub use compress::NormalizedScheme; pub use compress::normalize; -pub(crate) use compress::try_build_constant_normalized; mod execute; diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 76d266a2721..6de8240d41d 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Cosine similarity expression for tensor-like types. +//! Cosine similarity between two tensor columns. +use num_traits::Float; use num_traits::Zero; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -14,36 +15,39 @@ use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; +use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowExecution; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::UninitElementSink; use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::encodings::normalized::try_build_constant_normalized; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_norm::L2Norm; +use crate::scalar_fns::row::TensorRow; +#[cfg(test)] +use crate::scalar_fns::row::probe; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; +use crate::utils::l2_norm_row; /// Cosine similarity between two columns. /// /// Computes `dot(a, b) / (||a|| * ||b||)` over the flat backing buffer of each tensor or vector. /// The shape and permutation do not affect the result because cosine similarity only depends on the -/// element values, not their logical arrangement. +/// element values, not their logical arrangement. A zero norm on either side yields `0.0`. /// /// Both inputs must be tensor-like extension arrays ([`FixedShapeTensor`] or [`Vector`]) with the /// same dtype and a float element type. The output is a float column of the same float type. @@ -56,13 +60,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct CosineSimilarity; impl CosineSimilarity { /// Creates a new [`TypedScalarFnInstance`] wrapping the cosine similarity operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(CosineSimilarity, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the cosine similarity between `lhs` and @@ -70,132 +74,94 @@ impl CosineSimilarity { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(CosineSimilarity::new().erased(), vec![lhs, rhs]) + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) } } -impl ScalarFnVTable for CosineSimilarity { +impl RowFn for CosineSimilarity { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.cosine_similarity"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("CosineSimilarity must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_prepared_into::<(TensorRow, TensorRow), UninitElementSink, _, _>( + |(lhs, rhs)| { + #[cfg(test)] + probe::record(lhs.is_some(), rhs.is_some()); + ConstantNorms { + lhs: lhs.map(l2_norm_row), + rhs: rhs.map(l2_norm_row), + } + }, + |norms, (lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { + InitializedElement::write( + output, + cosine_similarity_row_prepared(norms, lhs, rhs), + ) + } + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands make the _stored_ norms and normalized children + /// authoritative: `cos(D(x, s), D(y, t)) = dot(x, y)` and `cos(D(x, s), y) = dot(x, y) / + /// ||y||`, in both cases forced to `0.0` on rows where any authoritative norm is `0.0` (even + /// for lossy children whose decoded coordinates are nonzero). + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, + args: &[ArrayRef], ctx: &mut ExecutionCtx, - ) -> VortexResult { - let mut lhs_ref = args.get(0)?; - let mut rhs_ref = args.get(1)?; - let len = args.row_count(); - - // If either side is a constant tensor-like extension array, eagerly normalize the single - // stored row and re-encode it as an `Normalized` whose children are both `ConstantArray`s. - // The `Normalized` fast path below then picks it up. - if let Some(normalized_array) = try_build_constant_normalized(&lhs_ref, len, ctx)? { - lhs_ref = normalized_array.into_array(); - } - if let Some(normalized_array) = try_build_constant_normalized(&rhs_ref, len, ctx)? { - rhs_ref = normalized_array.into_array(); - } + ) -> VortexResult> { + let lhs = args[0].clone(); + let rhs = args[1].clone(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { - NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); - } + match NormalizedOrientation::classify(&lhs, &rhs) { + NormalizedOrientation::Both { lhs, rhs } => cosine_both_normalized(lhs, rhs, ctx) + .map(|output| Some(RowExecution::Output(output))), NormalizedOrientation::One { normalized_array, plain, - } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); - } - NormalizedOrientation::Neither => {} + } => cosine_one_normalized(normalized_array, plain, ctx) + .map(|output| Some(RowExecution::Output(output))), + NormalizedOrientation::Neither => Ok(None), } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm.try_new_array(len, EmptyOptions, [lhs_ref.clone()])?; - let norm_rhs_arr = L2Norm.try_new_array(len, EmptyOptions, [rhs_ref.clone()])?; - let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; - - // Execute to get the inner product and norms of the arrays. We only fully decompress - // because we need to perform special logic (guard against 0) during division. - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - let norm_l: PrimitiveArray = norm_lhs_arr.into_array().execute(ctx)?; - let norm_r: PrimitiveArray = norm_rhs_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norm_l.as_slice::(); - let norms_r = norm_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - let denom = norms_l[i] * norms_r[i]; - - if denom == T::zero() { - T::zero() - } else { - dots[i] / denom - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false } } +vortex_array::impl_row_fn_vtable!(CosineSimilarity); + impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( &self, @@ -222,578 +188,177 @@ impl ScalarFnArrayVTable for CosineSimilarity { } } -impl CosineSimilarity { - /// Both sides are [`Normalized`]-encoded: treat the normalized children as authoritative, so - /// `cosine_similarity = dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - // `Normalized` makes the normalized children authoritative, so their dot product is the - // cosine similarity even for lossy storage wrappers, except that a zero stored norm still - // represents a zero vector. - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norms_l.as_slice::(); - let norms_r = norms_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if norms_l[i] == T::zero() || norms_r[i] == T::zero() { - T::zero() - } else { - dots[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: treat the normalized child as authoritative, so - /// `cosine_similarity = dot(n, b) / ||b||`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, normalized_norms) = extract_normalized_children(normalized_ref); - - let dot_arr = InnerProduct::try_new_array(normalized, plain_ref.clone())?; - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - - let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - - let norm_arr = L2Norm.try_new_array(len, EmptyOptions, [plain_ref.clone()])?; - let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let normalized_norms = normalized_norms.as_slice::(); - let plain_norms = plain_norm.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if normalized_norms[i] == T::zero() || plain_norms[i] == T::zero() { - T::zero() - } else { - dots[i] / plain_norms[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } +/// Per-batch state for the cosine row kernel: the L2 norm of each operand that is constant for +/// the batch. +/// +/// A broadcast query vector holds the same elements in every row, so its norm is the same in +/// every row too. Computing it in the prepare step hoists an `O(width)` pass and a `sqrt` per row +/// out of the row loop. `None` marks an operand that varies by row, whose norm the row closure +/// computes exactly as it did before the hoist. +struct ConstantNorms { + /// The norm of the lhs when it is batch-constant. + lhs: Option, + + /// The norm of the rhs when it is batch-constant. + rhs: Option, } -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::cosine_similarity::CosineSimilarity; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::constant_tensor_array; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. - fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[test] - fn unit_vectors_1d() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 0.0, 1.0, 0.0, // Tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 1.0, 0.0, 0.0, // Tensor 2 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - /// Single-row cosine similarity for various vector pairs. - #[rstest] - // Antiparallel -> -1.0. - #[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] - // dot=24, both magnitudes=5 -> 24/25 = 0.96. - #[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] - // Zero vector -> guarded to 0.0. - #[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); - Ok(()) - } - - /// Self-similarity across various tensor shapes should always produce 1.0. - #[rstest] - // 2x3 matrix, flattened to 6 elements. - #[case::matrix_2d( - &[2, 3], - &[ - 1.0, 0.0, 0.0, // row 0 - 0.0, 0.0, 0.0, // row 1 - ], - )] - // 2x2x2 tensor, 8 elements. - #[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] - fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { - let lhs = tensor_array(shape, elements)?; - let rhs = tensor_array(shape, elements)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn scalar_0d() -> VortexResult<()> { - // 0-dimensional tensor: each "tensor" is a single scalar value. - let lhs = tensor_array(&[], &[5.0, 3.0])?; - let rhs = tensor_array(&[], &[5.0, -3.0])?; - - // Same sign -> 1.0, opposite sign -> -1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); - Ok(()) - } - - #[test] - fn many_rows() -> VortexResult<()> { - // 5 tensors of shape [4] compared against themselves -> all 1.0. - let lhs = tensor_array( - &[4], - &[ - 1.0, 2.0, 3.0, 4.0, // tensor 0 - 0.0, 1.0, 0.0, 0.0, // tensor 1 - 5.0, 0.0, 5.0, 0.0, // tensor 2 - 1.0, 1.0, 1.0, 1.0, // tensor 3 - 0.0, 0.0, 0.0, 7.0, // tensor 4 - ], - )?; - let rhs = lhs.clone(); - - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0, 1.0, 1.0, 1.0, 1.0], - ); - Ok(()) - } - - #[test] - fn constant_query_tensor() -> VortexResult<()> { - // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. - let data = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 0.0, 1.0, 0.0, // tensor 1 - 0.0, 0.0, 1.0, // tensor 2 - 1.0, 0.0, 0.0, // tensor 3 - ], - )?; - let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn vector_unit_vectors() -> VortexResult<()> { - let lhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 1.0, 0.0, 0.0, // vector 1 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn vector_constant_query() -> VortexResult<()> { - let data = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - 0.0, 0.0, 1.0, // vector 2 - 1.0, 0.0, 0.0, // vector 3 - ], - )?; - let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. - let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; - let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; - let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: self-similarity = 1.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_self_similarity() -> VortexResult<()> { - // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. - // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Self-similarity should always be 1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); - Ok(()) - } - - #[test] - fn both_normalized_orthogonal() -> VortexResult<()> { - // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. - // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn both_normalized_zero_norm() -> VortexResult<()> { - // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS is plain [3.0, 4.0]. - // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[3.0, 4.0])?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 0.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); - Ok(()) - } - - #[test] - fn both_normalized_null_norms() -> VortexResult<()> { - // Row 0: valid, row 1: null (via nullable norms on rhs). - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_r = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); - let rhs = Normalized::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine - // similarity for that row must be `0.0` even though the dot product of the normalized - // children is nonzero. - let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; - let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally violates the unit-norm invariant by pairing a nonzero normalized row - // with a stored norm of `0.0`, mimicking lossy storage. - // SAFETY: The children are structurally valid. - let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l) }.into_array(); - - let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; - let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Same as above for the rhs operand. - // SAFETY: The children are structurally valid. - let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r) }.into_array(); - - // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both - // `0.0`, so cosine similarity must be `0.0`. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. The plain side is a normal nonzero - // tensor with positive norm. cosine similarity must still be `0.0` because the - // authoritative stored norm on the normalized_array side is `0.0`. - let normalized = tensor_array(&[2], &[0.6, 0.8])?; - let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally pairs a nonzero normalized row with a stored norm of `0.0`, mimicking - // lossy storage where the stored norm is authoritative. - // SAFETY: The children are structurally valid. - let normalized_array = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); - - let plain = tensor_array(&[2], &[1.0, 0.0])?; - - // Normalized encoding on the lhs: `One { normalized_array: lhs, plain: rhs }`. - assert_close( - &eval_cosine_similarity(normalized_array.clone(), plain.clone())?, - &[0.0], - ); - - // Normalized encoding on the rhs: `One { normalized_array: rhs, plain: lhs }`. The same - // zero-norm guard must fire regardless of operand order. - assert_close(&eval_cosine_similarity(plain, normalized_array)?, &[0.0]); - Ok(()) - } - - #[test] - fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { - // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. - // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. - let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 - 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 - 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 - 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { - // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn both_constant_tensors() -> VortexResult<()> { - // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). - let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; - let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; - let expected = 1.0 / 2.0_f64.sqrt(); - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[expected, expected, expected], - ); - Ok(()) - } - - #[test] - fn constant_zero_norm_query() -> VortexResult<()> { - // A zero-norm constant query must produce `0.0` for every row via the zero-norm guard in - // `execute_one_normalized` and `execute_both_normalized`. - let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 2.0, 3.0, // - 4.0, 5.0, 6.0, // - 7.0, 8.0, 9.0, // - ], - )?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); - Ok(()) - } - - #[test] - fn constant_self_similarity_nonunit() -> VortexResult<()> { - // A non-unit constant query compared to itself must produce `1.0`. This exercises the - // helper's division: after normalization, both sides must be exactly unit so the - // Normalized fast path's inner product yields 1. - let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); - Ok(()) - } - - #[test] - fn vector_constant_matches_plain() -> VortexResult<()> { - // Exercise the `Vector` extension variant through the new pre-pass. - let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[rstest] - #[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] - #[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = CosineSimilarity::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("CosineSimilarity serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) +/// Computes the cosine similarity of one row, taking any hoisted norm from `norms` and computing +/// the rest exactly as [`cosine_similarity_row`] does. +/// +/// Each arm accumulates the same values in the same order as [`cosine_similarity_row`], and the +/// denominator keeps its lhs-times-rhs order, so the result is bit-identical whether a norm was +/// hoisted or not. The match costs one predictable branch per row: the arm is the same for the +/// whole batch. +fn cosine_similarity_row_prepared( + norms: &ConstantNorms, + lhs: &[T], + rhs: &[T], +) -> T { + match (norms.lhs, norms.rhs) { + (None, None) => cosine_similarity_row(lhs, rhs), + (Some(lhs_norm), None) => { + let mut dot = T::zero(); + let mut rhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + rhs_norm_squared = rhs_norm_squared + rhs_element * rhs_element; + } + cosine_from_parts(dot, lhs_norm * rhs_norm_squared.sqrt()) + } + (None, Some(rhs_norm)) => { + let mut dot = T::zero(); + let mut lhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + lhs_norm_squared = lhs_norm_squared + lhs_element * lhs_element; + } + cosine_from_parts(dot, lhs_norm_squared.sqrt() * rhs_norm) + } + (Some(lhs_norm), Some(rhs_norm)) => { + let mut dot = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + } + cosine_from_parts(dot, lhs_norm * rhs_norm) + } } +} - fn cosine_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") - } +/// Computes the cosine similarity of two equal-length float slices. +/// +/// Returns `dot(a, b) / (||a|| * ||b||)`, or `0.0` when either norm is zero. +fn cosine_similarity_row(lhs: &[T], rhs: &[T]) -> T { + let mut dot = T::zero(); + let mut lhs_norm_squared = T::zero(); + let mut rhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + lhs_norm_squared = lhs_norm_squared + lhs_element * lhs_element; + rhs_norm_squared = rhs_norm_squared + rhs_element * rhs_element; + } + + cosine_from_parts(dot, lhs_norm_squared.sqrt() * rhs_norm_squared.sqrt()) +} - fn cosine_vector_rhs() -> ArrayRef { - vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +/// The shared tail of every cosine arm: `dot / denominator`, guarded to `0.0` when it is +/// zero. +fn cosine_from_parts(dot: T, denominator: T) -> T { + if denominator == T::zero() { + T::zero() + } else { + dot / denominator } +} - fn cosine_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") - } +/// Both sides are [`Normalized`]-encoded: the normalized children are authoritative, so their dot +/// product is the cosine similarity, except that a row with a zero _stored_ norm is a zero vector. +/// +/// Unlike [`InnerProduct::reduce_encoded`], which composes lazy `Mul` arrays over the norm columns, +/// this executes and materializes. The zero-norm guard is a conditional per row rather than an +/// arithmetic factor, so there is no lazy array that expresses it; the norm columns are one value +/// per row rather than one per coordinate, so materializing them is cheap next to the decode this +/// avoids. +/// +/// [`InnerProduct::reduce_encoded`]: InnerProduct::reduce_encoded +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_both_normalized( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized_l, normalized_r])? + .execute(ctx)?; + let norms_l: PrimitiveArray = norms_l.execute(ctx)?; + let norms_r: PrimitiveArray = norms_r.execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let norms_l = norms_l.as_slice::(); + let norms_r = norms_r.as_slice::(); + // Zipped rather than indexed by `0..len`: one bounds check per iterator instead of three + // per row. A length disagreement between the children shortens the result, which the + // lifting reports against the batch row count rather than panicking mid-loop. + let buffer: Buffer = dots + .iter() + .zip(norms_l) + .zip(norms_r) + .map(|((&dot, &norm_l), &norm_r)| { + if norm_l.is_zero() || norm_r.is_zero() { + T::zero() + } else { + dot + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) +} - fn cosine_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") - } +/// One side is [`Normalized`]-encoded: `cos = dot(normalized, plain) / ||plain||`, forced to `0.0` +/// on rows where the stored norm or the plain norm is `0.0`. +/// +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_one_normalized( + normalized_array: &ArrayRef, + plain: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = normalized_array.len(); + let (normalized, normalized_norms) = extract_normalized_children(normalized_array); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized, plain.clone()])? + .execute(ctx)?; + let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; + let plain_norm: PrimitiveArray = L2Norm + .try_new_array(len, EmptyOptions, [plain.clone()])? + .execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let normalized_norms = normalized_norms.as_slice::(); + let plain_norms = plain_norm.as_slice::(); + // Zipped for the same reason as [`cosine_both_normalized`]. + let buffer: Buffer = dots + .iter() + .zip(normalized_norms) + .zip(plain_norms) + .map(|((&dot, &stored_norm), &plain_norm)| { + if stored_norm.is_zero() || plain_norm.is_zero() { + T::zero() + } else { + dot / plain_norm + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) } diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 53ae82eb4a2..a25e36f9aab 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -6,40 +6,34 @@ use num_traits::Float; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowExecution; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::matcher::AnyTensor; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; -use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; /// Inner product (dot product) between two columns. /// @@ -52,13 +46,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct InnerProduct; impl InnerProduct { /// Creates a new [`TypedScalarFnInstance`] wrapping the inner product operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(InnerProduct, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the inner product between `lhs` and @@ -66,119 +60,93 @@ impl InnerProduct { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(InnerProduct::new().erased(), vec![lhs, rhs]) + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) } } -impl ScalarFnVTable for InnerProduct { +impl RowFn for InnerProduct { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.inner_product"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("InnerProduct must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - // TODO(connor): relax the float-only gate once integer tensors are supported. - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow, TensorRow), UninitElementSink, _>( + |(lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, inner_product_row(lhs, rhs)) } + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands factor through their stored norms: with `D(x, s)` denoting + /// `x * s` rowwise, `dot(D(x, s), D(y, t)) = s * t * dot(x, y)` and + /// `dot(D(x, s), y) = s * dot(x, y)`. The rewrite is expressed with lazy [`Operator::Mul`] + /// arrays over the (much smaller) norm columns, so no denormalized coordinates are decoded. + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let lhs_ref = args.get(0)?; - let rhs_ref = args.get(1)?; - let len = args.row_count(); + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let len = args[0].len(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { + Ok(match NormalizedOrientation::classify(&args[0], &args[1]) { NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized_l, normalized_r])?; + Some( + dot.binary(norms_l, Operator::Mul)? + .binary(norms_r, Operator::Mul)?, + ) } NormalizedOrientation::One { normalized_array, plain, } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); + let (normalized, norms) = extract_normalized_children(normalized_array); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized, plain.clone()])?; + Some(dot.binary(norms, Operator::Mul)?) } - NormalizedOrientation::Neither => {} + NormalizedOrientation::Neither => None, } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Canonicalize so we can perform the math directly. - let lhs: ExtensionArray = lhs_ref.execute(ctx)?; - let rhs: ExtensionArray = rhs_ref.execute(ctx)?; - - // We validated that both inputs have the same type. - let ext = lhs.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let dimensions = tensor_match.list_size() as usize; - - // Extract the storage array from each extension input. We pass the storage (FSL) rather - // than the extension array to avoid canonicalizing the extension wrapper. - let lhs_storage = lhs.storage_array(); - let rhs_storage = rhs.storage_array(); - - let lhs_flat = extract_flat_elements(lhs_storage, dimensions, ctx)?; - let rhs_flat = extract_flat_elements(rhs_storage, dimensions, ctx)?; - - match_each_float_ptype!(lhs_flat.ptype(), |T| { - let buffer: Buffer = (0..len) - .map(|i| inner_product_row(lhs_flat.row::(i), rhs_flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + .map(RowExecution::Output)) } } +vortex_array::impl_row_fn_vtable!(InnerProduct); + impl ScalarFnArrayVTable for InnerProduct { fn serialize( &self, @@ -205,329 +173,12 @@ impl ScalarFnArrayVTable for InnerProduct { } } -impl InnerProduct { - /// Both sides are [`Normalized`]-encoded: `inner_product = s_l * s_r * dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let nl = norms_l.as_slice::(); - let nr = norms_r.as_slice::(); - let buffer: Buffer = (0..len).map(|i| nl[i] * nr[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: `inner_product = s * dot(n, other)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, norms) = extract_normalized_children(normalized_ref); - let normalized_norms: PrimitiveArray = norms.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized, plain_ref.clone())? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let ns = normalized_norms.as_slice::(); - let buffer: Buffer = (0..len).map(|i| ns[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } -} - /// Computes the inner product (dot product) of two equal-length float slices. /// /// Returns `sum(a_i * b_i)`. -fn inner_product_row(a: &[T], b: &[T]) -> T { - a.iter() - .zip(b.iter()) - .map(|(&x, &y)| x * y) - .fold(T::zero(), |acc, v| acc + v) -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::inner_product::InnerProduct; - use crate::tests::SESSION; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates inner product between two tensor arrays and returns the result as `Vec`. - fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - /// Single-row inner product for various vector pairs. - #[rstest] - // Orthogonal: [1, 0] . [0, 1] = 0. - #[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] - // Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. - #[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] - // Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. - #[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] - // Scaled: [2, 0] . [3, 0] = 6. - #[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_inner_product(lhs, rhs)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 3.0, 4.0, 0.0, // tensor 1 - 1.0, 1.0, 1.0, // tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 0.0, 1.0, 0.0, // tensor 0: dot = 0 - 3.0, 4.0, 0.0, // tensor 1: dot = 25 - 2.0, 2.0, 2.0, // tensor 2: dot = 6 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); - Ok(()) - } - - #[test] - fn vector_inner_product() -> VortexResult<()> { - let lhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0 - 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0: dot = 25 - 0.0, 1.0, // vector 1: dot = 0 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. - let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; - let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; - let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert!(prim.is_valid(2, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[23.0]); - assert_close(&[prim.as_slice::()[2]], &[127.0]); - Ok(()) - } - - #[test] - fn rejects_non_extension_dtype() { - let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); - let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - } - - #[test] - fn rejects_mismatched_dtypes() -> VortexResult<()> { - let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; - let rhs = vector_array(2, &[3.0_f64, 4.0])?; - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn both_normalized() -> VortexResult<()> { - // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). - // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). - // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; - - // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. - assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); - Ok(()) - } - - #[test] - fn both_normalized_multiple_rows() -> VortexResult<()> { - // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. - // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS: plain [1.0, 2.0]. - // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[1.0, 2.0])?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS: plain [1.0, 2.0]. - // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 2.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn both_normalized_null_norms() -> VortexResult<()> { - // Row 0: valid, row 1: null (via nullable norms on lhs). - let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_l = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - let lhs = Normalized::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[25.0]); - Ok(()) - } - - #[rstest] - #[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] - #[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = InnerProduct::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(InnerProduct); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("InnerProduct serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn inner_product_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } - - fn inner_product_vector_rhs() -> ArrayRef { - vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") - } - - fn inner_product_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") - } - - fn inner_product_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") - } +fn inner_product_row(lhs: &[T], rhs: &[T]) -> T { + lhs.iter() + .zip(rhs) + .map(|(&lhs_element, &rhs_element)| lhs_element * rhs_element) + .fold(T::zero(), |sum, product| sum + product) } diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index a99ad76ea6e..0b3cf3634ee 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -164,3 +164,20 @@ unsafe impl InputElement for TensorRow { } } } + +/// Records which operands a test's prepare step received as batch constants. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + thread_local! { + /// Bit 0 records the lhs and bit 1 records the rhs. Thread-local storage prevents + /// concurrent tests from racing; row execution remains on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Records whether each operand was constant for the current batch. + pub(crate) fn record(lhs_constant: bool, rhs_constant: bool) { + SEEN_CONSTANTS.set(u8::from(lhs_constant) | (u8::from(rhs_constant) << 1)); + } +} diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs new file mode 100644 index 00000000000..0e2687690a9 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -0,0 +1,607 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::assert_arrays_eq; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::cosine_similarity::CosineSimilarity; +use crate::scalar_fns::row::probe; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::constant_tensor_array; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. +fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = CosineSimilarity::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = CosineSimilarity::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_cosine_similarity(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +/// Like [`eval_cosine_similarity`], but returns the executed array for exact array comparisons. +fn eval_cosine_similarity_array( + lhs: ArrayRef, + rhs: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + Ok(result + .into_array() + .execute::(ctx)? + .into_array()) +} + +#[test] +fn unit_vectors_1d() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 0.0, 1.0, 0.0, // Tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 1.0, 0.0, 0.0, // Tensor 2 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +/// Single-row cosine similarity for various vector pairs. +#[rstest] +// Antiparallel -> -1.0. +#[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] +// dot=24, both magnitudes=5 -> 24/25 = 0.96. +#[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] +// Zero vector -> guarded to 0.0. +#[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); + Ok(()) +} + +/// Self-similarity across various tensor shapes should always produce 1.0. +#[rstest] +// 2x3 matrix, flattened to 6 elements. +#[case::matrix_2d( + &[2, 3], + &[ + 1.0, 0.0, 0.0, // row 0 + 0.0, 0.0, 0.0, // row 1 + ], +)] +// 2x2x2 tensor, 8 elements. +#[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] +fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { + let lhs = tensor_array(shape, elements)?; + let rhs = tensor_array(shape, elements)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn scalar_0d() -> VortexResult<()> { + // 0-dimensional tensor: each "tensor" is a single scalar value. + let lhs = tensor_array(&[], &[5.0, 3.0])?; + let rhs = tensor_array(&[], &[5.0, -3.0])?; + + // Same sign -> 1.0, opposite sign -> -1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); + Ok(()) +} + +#[test] +fn many_rows() -> VortexResult<()> { + // 5 tensors of shape [4] compared against themselves -> all 1.0. + let lhs = tensor_array( + &[4], + &[ + 1.0, 2.0, 3.0, 4.0, // tensor 0 + 0.0, 1.0, 0.0, 0.0, // tensor 1 + 5.0, 0.0, 5.0, 0.0, // tensor 2 + 1.0, 1.0, 1.0, 1.0, // tensor 3 + 0.0, 0.0, 0.0, 7.0, // tensor 4 + ], + )?; + let rhs = lhs.clone(); + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0, 1.0, 1.0, 1.0, 1.0], + ); + Ok(()) +} + +#[test] +fn constant_query_tensor() -> VortexResult<()> { + // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. + let data = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 0.0, 1.0, 0.0, // tensor 1 + 0.0, 0.0, 1.0, // tensor 2 + 1.0, 0.0, 0.0, // tensor 3 + ], + )?; + let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn vector_unit_vectors() -> VortexResult<()> { + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 1.0, 0.0, 0.0, // vector 1 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn vector_constant_query() -> VortexResult<()> { + let data = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + 0.0, 0.0, 1.0, // vector 2 + 1.0, 0.0, 0.0, // vector 3 + ], + )?; + let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. + let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; + let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; + let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: self-similarity = 1.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_self_similarity() -> VortexResult<()> { + // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. + // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Self-similarity should always be 1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); + Ok(()) +} + +#[test] +fn both_normalized_orthogonal() -> VortexResult<()> { + // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. + // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn both_normalized_zero_norm() -> VortexResult<()> { + // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS is plain [3.0, 4.0]. + // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[3.0, 4.0])?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 0.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); + Ok(()) +} + +#[test] +fn both_normalized_null_norms() -> VortexResult<()> { + // Row 0: valid, row 1: null (via nullable norms on rhs). + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_r = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let rhs = Normalized::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine + // similarity for that row must be `0.0` even though the dot product of the normalized + // children is nonzero. + let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; + let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally violates the unit-norm invariant by + // pairing a nonzero normalized row with a stored norm of `0.0`, mimicking lossy storage. + let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l) }.into_array(); + + let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; + let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: Same as above for the rhs operand. + let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r) }.into_array(); + + // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both + // `0.0`, so cosine similarity must be `0.0`. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. The plain side is a normal nonzero + // tensor with positive norm. cosine similarity must still be `0.0` because the + // authoritative stored norm on the denorm side is `0.0`. + let normalized = tensor_array(&[2], &[0.6, 0.8])?; + let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally pairs a nonzero normalized row with a + // stored norm of `0.0`, mimicking lossy storage where the stored norm is authoritative. + let denorm = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); + + let plain = tensor_array(&[2], &[1.0, 0.0])?; + + // Denorm on the lhs: `One { denorm: lhs, plain: rhs }`. + assert_close( + &eval_cosine_similarity(denorm.clone(), plain.clone())?, + &[0.0], + ); + + // Denorm on the rhs: `One { denorm: rhs, plain: lhs }`. The same zero-norm guard must + // fire regardless of operand order. + assert_close(&eval_cosine_similarity(plain, denorm)?, &[0.0]); + Ok(()) +} + +#[test] +fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { + // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. + // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. + let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 + 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 + 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 + 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 + ], + )?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { + // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn both_constant_tensors() -> VortexResult<()> { + // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). + let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; + let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[expected, expected, expected], + ); + Ok(()) +} + +#[test] +fn constant_zero_norm_query() -> VortexResult<()> { + // A zero-norm constant query must produce `0.0` through the prepared row kernel's + // zero-denominator guard. + let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 2.0, 3.0, // + 4.0, 5.0, 6.0, // + 7.0, 8.0, 9.0, // + ], + )?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +#[test] +fn constant_self_similarity_nonunit() -> VortexResult<()> { + // The prepared path hoists both norms and computes the same dot product for every row. + let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); + Ok(()) +} + +/// An extension array over constant storage (what [`Vector::constant_array`] builds) is a batch +/// constant like any other. The row layer sees through the wrapper, so `prepare` hoists its norm. +#[test] +fn vector_constant_matches_plain() -> VortexResult<()> { + let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + assert_eq!( + probe::SEEN_CONSTANTS.get(), + 0b01, + "the extension-over-constant lhs must reach prepare as a batch constant", + ); + Ok(()) +} + +/// Both literal and extension-wrapped constant storage reach the prepared row path. The probe +/// ensures that the literal query remains a batch constant instead of becoming a varying column. +/// +/// [`ConstantArray`]: vortex_array::arrays::ConstantArray +#[test] +fn literal_constant_rhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(lhs.clone(), literal_vector_array(&query, 4), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(lhs, vector_array(3, &query.repeat(4))?, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// The mirror of [`literal_constant_rhs_matches_expanded_column`], exercising the hoisted-lhs arm +/// of the prepared kernel. +#[test] +fn literal_constant_lhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(literal_vector_array(&query, 4), rhs.clone(), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(vector_array(3, &query.repeat(4))?, rhs, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// A zero-norm literal constant query must be guarded to `0.0` on every row by the prepared row +/// kernel, exactly as the unprepared kernel guards it. +#[test] +fn literal_constant_zero_norm_query_yields_zero() -> VortexResult<()> { + let lhs = vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = literal_vector_array(&[0.0f64, 0.0, 0.0], 2); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0]); + Ok(()) +} + +/// Two literal constants are folded to a single-row execution by the row lifting, and that row +/// still runs the prepared kernel with both norms hoisted. +#[test] +fn both_literal_constants() -> VortexResult<()> { + let lhs = literal_vector_array(&[1.0f64, 0.0, 0.0], 3); + let rhs = literal_vector_array(&[1.0f64, 1.0, 0.0], 3); + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[expected; 3]); + Ok(()) +} + +#[rstest] +#[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] +#[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("CosineSimilarity serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn cosine_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_vector_rhs() -> ArrayRef { + vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn cosine_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/inner_product.rs b/vortex-tensor/src/scalar_fns/tests/inner_product.rs new file mode 100644 index 00000000000..9a5a60f5dda --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/inner_product.rs @@ -0,0 +1,282 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::inner_product::InnerProduct; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates inner product between two tensor arrays and returns the result as `Vec`. +fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = InnerProduct::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = InnerProduct::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_inner_product(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +/// Single-row inner product for various vector pairs. +#[rstest] +// Orthogonal: [1, 0] . [0, 1] = 0. +#[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] +// Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. +#[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] +// Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. +#[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] +// Scaled: [2, 0] . [3, 0] = 6. +#[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_inner_product(lhs, rhs)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 3.0, 4.0, 0.0, // tensor 1 + 1.0, 1.0, 1.0, // tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 0.0, 1.0, 0.0, // tensor 0: dot = 0 + 3.0, 4.0, 0.0, // tensor 1: dot = 25 + 2.0, 2.0, 2.0, // tensor 2: dot = 6 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); + Ok(()) +} + +#[test] +fn vector_inner_product() -> VortexResult<()> { + let lhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0 + 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0: dot = 25 + 0.0, 1.0, // vector 1: dot = 0 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. + let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; + let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert!(prim.is_valid(2, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[23.0]); + assert_close(&[prim.as_slice::()[2]], &[127.0]); + Ok(()) +} + +#[test] +fn rejects_non_extension_dtype() { + let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); + let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); +} + +#[test] +fn rejects_mismatched_dtypes() -> VortexResult<()> { + let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; + let rhs = vector_array(2, &[3.0_f64, 4.0])?; + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn both_normalized() -> VortexResult<()> { + // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). + // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). + // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; + + // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. + assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); + Ok(()) +} + +#[test] +fn both_normalized_multiple_rows() -> VortexResult<()> { + // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. + // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS: plain [1.0, 2.0]. + // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[1.0, 2.0])?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS: plain [1.0, 2.0]. + // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 2.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn both_normalized_null_norms() -> VortexResult<()> { + // Row 0: valid, row 1: null (via nullable norms on lhs). + let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_l = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let lhs = Normalized::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[25.0]); + Ok(()) +} + +#[rstest] +#[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] +#[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(InnerProduct); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("InnerProduct serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn inner_product_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} + +fn inner_product_vector_rhs() -> ArrayRef { + vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") +} + +fn inner_product_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn inner_product_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/mod.rs b/vortex-tensor/src/scalar_fns/tests/mod.rs index 5447772adda..bb3726e9329 100644 --- a/vortex-tensor/src/scalar_fns/tests/mod.rs +++ b/vortex-tensor/src/scalar_fns/tests/mod.rs @@ -3,5 +3,7 @@ //! Tests for the tensor scalar functions. +mod cosine_similarity; +mod inner_product; mod l2_norm; mod row; diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index df40bdea44a..62f0713319b 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -63,6 +63,9 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { } /// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. +/// +/// L2 norm and cosine similarity share this implementation so prepared constant norms use the +/// same accumulation order as varying rows. pub(crate) fn l2_norm_row(row: &[T]) -> T { let mut sum_squared = T::zero(); for &element in row { @@ -111,19 +114,6 @@ pub fn validate_tensor_float_input(input_dtype: &DType) -> VortexResult( - lhs: &'a DType, - rhs: &DType, -) -> VortexResult> { - vortex_ensure!( - lhs.eq_ignore_nullability(rhs), - "binary tensor expression expects inputs to have the same dtype, got {lhs} and {rhs}" - ); - validate_tensor_float_input(lhs) -} - /// Validates that every argument has the same float tensor dtype, ignoring nullability. pub fn validate_tensor_float_inputs(args: &[DType]) -> VortexResult> { let (first, rest) = args @@ -322,7 +312,7 @@ impl BinaryTensorOpMetadata { let lhs_dtype = DType::from_proto(lhs_pb, session)?; let rhs_dtype = DType::from_proto(rhs_pb, session)?; - validate_binary_tensor_float_inputs(&lhs_dtype, &rhs_dtype)?; + validate_tensor_float_inputs(&[lhs_dtype.clone(), rhs_dtype.clone()])?; let lhs = children.get(0, &lhs_dtype, len)?; let rhs = children.get(1, &rhs_dtype, len)?; diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index ad3b96d1bff..492bc837b89 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -35,11 +35,13 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; @@ -79,7 +81,7 @@ pub fn build_similarity_search_tree>( let num_rows = data.len(); let query_vec = Vector::constant_array(query, num_rows)?; - let cosine = CosineSimilarity::try_new_array(data, query_vec)?.into_array(); + let cosine = CosineSimilarity.try_new_array(num_rows, EmptyOptions, [data, query_vec])?; let threshold_scalar = Scalar::primitive(threshold, Nullability::NonNullable); let threshold_array = ConstantArray::new(threshold_scalar, num_rows).into_array(); From 87bd382973efe4264d749e718cdbbda139f603a0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:00 -0400 Subject: [PATCH 133/160] Update tensor products for the RowFn API Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 2 -- vortex-tensor/src/scalar_fns/inner_product.rs | 2 -- vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs | 2 +- vortex-tensor/src/utils.rs | 2 +- 4 files changed, 2 insertions(+), 6 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 6de8240d41d..a2788cdaae1 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -160,8 +160,6 @@ impl RowFn for CosineSimilarity { } } -vortex_array::impl_row_fn_vtable!(CosineSimilarity); - impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index a25e36f9aab..976fbe3cb18 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -145,8 +145,6 @@ impl RowFn for InnerProduct { } } -vortex_array::impl_row_fn_vtable!(InnerProduct); - impl ScalarFnArrayVTable for InnerProduct { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs index 0e2687690a9..ecdbe23f835 100644 --- a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -490,7 +490,7 @@ fn vector_constant_matches_plain() -> VortexResult<()> { } /// Both literal and extension-wrapped constant storage reach the prepared row path. The probe -/// ensures that the literal query remains a batch constant instead of becoming a varying column. +/// ensures that the literal query remains a batch constant instead of becoming a per-row column. /// /// [`ConstantArray`]: vortex_array::arrays::ConstantArray #[test] diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 62f0713319b..c05dcf5c1b6 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -65,7 +65,7 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { /// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. /// /// L2 norm and cosine similarity share this implementation so prepared constant norms use the -/// same accumulation order as varying rows. +/// same accumulation order as rows from per-row inputs. pub(crate) fn l2_norm_row(row: &[T]) -> T { let mut sum_squared = T::zero(); for &element in row { From 659aa973a47f2f501a464c59063da76b34f14772 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:38:43 -0400 Subject: [PATCH 134/160] Opt tensor products into unstable RowFn Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 10 +++++----- vortex-tensor/src/scalar_fns/inner_product.rs | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index a2788cdaae1..e3defc3432f 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -18,13 +18,13 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowExecution; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_buffer::Buffer; diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 976fbe3cb18..d6a8be4d333 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -16,14 +16,14 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowExecution; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; use vortex_error::VortexResult; use vortex_session::VortexSession; From c3ed052b5045ee1afb8d4e6679d05fe8cb978238 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:22:42 -0400 Subject: [PATCH 135/160] Declare tensor product fallibility Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 1 + vortex-tensor/src/scalar_fns/inner_product.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index e3defc3432f..1582ef9c7d0 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -85,6 +85,7 @@ impl RowFn for CosineSimilarity { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.cosine_similarity"); diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index d6a8be4d333..a31cce5ddda 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -71,6 +71,7 @@ impl RowFn for InnerProduct { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.inner_product"); From 324fb6ba14a4c10cace2948def6df0d7917b65c2 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:00:39 -0400 Subject: [PATCH 136/160] Execute spatial distance with RowFn Signed-off-by: Connor Tsui --- vortex-spatial/src/extension/mod.rs | 41 +++++++++ vortex-spatial/src/extension/point.rs | 18 ++++ vortex-spatial/src/extension/polygon.rs | 18 ++++ vortex-spatial/src/scalar_fn/distance.rs | 107 +++++++---------------- vortex-spatial/src/scalar_fn/mod.rs | 1 + vortex-spatial/src/scalar_fn/row.rs | 92 +++++++++++++++++++ 6 files changed, 201 insertions(+), 76 deletions(-) create mode 100644 vortex-spatial/src/scalar_fn/row.rs diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index f24f31f02aa..7e47ec0f71d 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -178,6 +178,47 @@ pub(crate) fn geometries( } } +/// The geometry a null row decodes to under [`geometries_null_tolerant`]. Arbitrary: the caller +/// guarantees null rows are never read. +pub(crate) fn placeholder_geometry() -> Geometry { + Geometry::Point(geo_types::Point::new(0.0, 0.0)) +} + +/// Decode a native geometry column that may contain null rows, writing [`placeholder_geometry`] +/// into their slots. The caller guarantees null rows are never read. +/// +/// `Ok(None)` means this geometry type has no null-tolerant decode yet (`Point` and `Polygon` are +/// covered), and the caller falls back to the filter strategy, which never decodes a null row. A +/// column with definitely no nulls delegates to the ordinary [`geometries`] for any type. +pub(crate) fn geometries_null_tolerant( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>>> { + if array.validity()?.definitely_no_nulls() { + return geometries(array, ctx).map(Some); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + + if ext.is::() { + point_geometries_null_tolerant(&storage, ctx).map(Some) + } else if ext.is::() { + polygon_geometries_null_tolerant(&storage, ctx).map(Some) + } else { + Ok(None) + } +} + /// Decode a constant operand scalar to one geometry, a constant of any /// supported geometry type is decoded exactly like a column. pub(crate) fn single_geometry( diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index e6a00fe8fea..b774f624c4c 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -50,6 +50,7 @@ use super::coordinate::coordinate_from_struct; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A single location: `geoarrow.point`, stored as `Struct` of non-nullable `f64`. @@ -150,6 +151,23 @@ pub(crate) fn point_geometries( .collect() } +/// Like [`point_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn point_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + point_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + impl ArrowExportVTable for Point { fn arrow_ext_id(&self) -> Id { *ARROW_POINT diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index a4c88b07b22..bce33efe6e5 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -53,6 +53,7 @@ use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A polygon: `geoarrow.polygon`, stored as `List>>` (rings of vertices). @@ -153,6 +154,23 @@ pub(crate) fn polygon_geometries( .collect() } +/// Like [`polygon_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn polygon_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + polygon_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + /// Build a geoarrow `PolygonArray` from a `Polygon`'s `List>` storage. fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let polygon_type = polygon_type( diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index dfd3d09ed23..751e402f118 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -6,43 +6,20 @@ use geo::Distance; use geo::Euclidean; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Distance`. -fn validate_distance_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: distance requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: distance operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -60,69 +37,46 @@ impl SpatialDistance { } } -impl ScalarFnVTable for SpatialDistance { +impl RowFn for SpatialDistance { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.distance"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { - Ok(EmptyOptions) - } - - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("distance has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_distance_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Primitive(PType::F64, nullability)) - } - - fn execute( + fn deserialize( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Distance is a value, not a verdict: no bounding-rect test can decide it. - execute_binary_geo_types(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx) + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn validity( + fn dispatch>( &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } - - fn is_strict(&self, _: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _: &Self::Options) -> bool { - false + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(GeometryRow, GeometryRow), UninitElementSink, _>( + |(a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, Euclidean.distance(a, b)) } + }, + ) } } +vortex_array::impl_row_fn_vtable!(SpatialDistance); + #[cfg(test)] mod tests { use vortex_array::ArrayRef; @@ -196,8 +150,9 @@ mod tests { Ok(()) } - /// Distance passes no bounding-rect rejection: a point far outside a constant polygon's - /// bounding rect still gets its true distance, alongside an inside point at distance zero. + /// Distance is a value rather than a verdict, so no bounding-rect rejection may fire for it: a + /// point far outside a constant polygon's rect still gets its true distance. Carried over from + /// #9076, which added the rejection to the predicates but deliberately not to this function. #[test] fn distance_to_constant_polygon_is_exact() -> VortexResult<()> { let session = vortex_array::array_session(); diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index 99fe5d28528..6291075246a 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -13,3 +13,4 @@ mod execute; pub mod intersects; pub mod length; pub mod make_line; +pub(crate) mod row; diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs new file mode 100644 index 00000000000..575d0793ce6 --- /dev/null +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the geo scalar functions add to the row-function machinery: an element type that decodes a +//! native geometry column into `geo_types` geometries. + +use geo_types::Geometry; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::InputElement; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::extension::geometries; +use crate::extension::geometries_null_tolerant; +use crate::extension::is_native_geometry; + +/// Marker for native geometry input elements: accepts any native geometry column and presents each +/// row as a decoded `geo_types` geometry. +/// +/// The two operands of a binary geo function need not share a geometry type, since distance, +/// containment and intersection across types are all meaningful, so this validates only that the +/// column is *some* native geometry. +pub(crate) struct GeometryRow; + +// SAFETY: [`varying`](InputElement::varying) returns the decoded geometry slice and +// [`varying_len`](InputElement::varying_len) reports that slice's exact length. +unsafe impl InputElement for GeometryRow { + type Column = Vec>; + type Varying<'a> = &'a [Geometry]; + type Elem<'a> = &'a Geometry; + + // A geometry row is decoded from its coordinate storage, which behind a null row holds arbitrary + // coordinates that need not describe a well-formed geometry. + const DENSE_SAFE: bool = false; + // Decoding builds a geometry from stored coordinates, and a malformed one in a *valid* row is a + // domain error rather than an infrastructural failure. + const DECODE_FALLIBLE: bool = true; + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + is_native_geometry(dtype), + "spatial: operand {dtype} is not a native geometry type" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + geometries(&array, ctx) + } + + fn get(column: &Self::Column, index: usize) -> &Geometry { + &column[index] + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column.as_slice() + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a Geometry + where + Self: 'a, + { + &column[index] + } + + unsafe fn get_varying_unchecked<'a>( + column: &Self::Varying<'a>, + index: usize, + ) -> &'a Geometry + where + Self: 'a, + { + // SAFETY: The caller established that `index` is below the slice length returned by + // `varying_len` for this exact view. + unsafe { column.get_unchecked(index) } + } + + /// Null rows decode to a placeholder geometry that the branch-and-skip row loop never reads. + /// `Point` and `Polygon` columns are covered; other geometry types return `Ok(None)` and the + /// batch falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + geometries_null_tolerant(&array, ctx) + } +} From 365392e5a0afe0679dce5600d813d6d1a32b5866 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:43 -0400 Subject: [PATCH 137/160] Update spatial distance for RowFn views Signed-off-by: Connor Tsui --- vortex-spatial/src/scalar_fn/distance.rs | 2 -- vortex-spatial/src/scalar_fn/row.rs | 27 +++++++++++------------- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index 751e402f118..6f6b7e7bd93 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -75,8 +75,6 @@ impl RowFn for SpatialDistance { } } -vortex_array::impl_row_fn_vtable!(SpatialDistance); - #[cfg(test)] mod tests { use vortex_array::ArrayRef; diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index 575d0793ce6..1f5e2f2d0ec 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -21,14 +21,14 @@ use crate::extension::is_native_geometry; /// /// The two operands of a binary geo function need not share a geometry type, since distance, /// containment and intersection across types are all meaningful, so this validates only that the -/// column is *some* native geometry. +/// column is _some_ native geometry. pub(crate) struct GeometryRow; -// SAFETY: [`varying`](InputElement::varying) returns the decoded geometry slice and -// [`varying_len`](InputElement::varying_len) reports that slice's exact length. +// SAFETY: [`view`](InputElement::view) returns the decoded geometry slice and +// [`view_len`](InputElement::view_len) reports that slice's exact length. unsafe impl InputElement for GeometryRow { type Column = Vec>; - type Varying<'a> = &'a [Geometry]; + type View<'a> = &'a [Geometry]; type Elem<'a> = &'a Geometry; // A geometry row is decoded from its coordinate storage, which behind a null row holds arbitrary @@ -53,31 +53,28 @@ unsafe impl InputElement for GeometryRow { &column[index] } - fn varying(column: &Self::Column) -> Self::Varying<'_> { + fn view(column: &Self::Column) -> Self::View<'_> { column.as_slice() } - fn varying_len(column: &Self::Varying<'_>) -> usize { - column.len() + fn view_len(view: &Self::View<'_>) -> usize { + view.len() } - fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> &'a Geometry + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry where Self: 'a, { - &column[index] + &view[index] } - unsafe fn get_varying_unchecked<'a>( - column: &Self::Varying<'a>, - index: usize, - ) -> &'a Geometry + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry where Self: 'a, { // SAFETY: The caller established that `index` is below the slice length returned by - // `varying_len` for this exact view. - unsafe { column.get_unchecked(index) } + // `view_len` for this exact view. + unsafe { view.get_unchecked(index) } } /// Null rows decode to a placeholder geometry that the branch-and-skip row loop never reads. From 2c24de405c117fbdc741860edad51f1567834aae Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:39:38 -0400 Subject: [PATCH 138/160] Opt spatial distance into unstable RowFn Signed-off-by: Connor Tsui --- vortex-spatial/Cargo.toml | 2 +- vortex-spatial/src/scalar_fn/distance.rs | 8 ++++---- vortex-spatial/src/scalar_fn/row.rs | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 3be2b2d9d66..cd306089325 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -22,7 +22,7 @@ geo-types = { workspace = true } geoarrow = { workspace = true } geoarrow-cast = { workspace = true } prost = { workspace = true } -vortex-array = { workspace = true } +vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-edition = { workspace = true } diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index 6f6b7e7bd93..e41999338a6 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -9,12 +9,12 @@ use vortex_array::ArrayRef; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index 1f5e2f2d0ec..127f34779b5 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -8,7 +8,7 @@ use geo_types::Geometry; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::dtype::DType; -use vortex_array::scalar_fn::InputElement; +use vortex_array::scalar_fn::unstable::row::InputElement; use vortex_error::VortexResult; use vortex_error::vortex_ensure; From d86f86f4258579c8b22c3152aaa3a5e62a4147f0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:27:19 -0400 Subject: [PATCH 139/160] Preflight null-tolerant geometry decoding Signed-off-by: Connor Tsui --- vortex-spatial/src/extension/mod.rs | 16 ++++++++++++++++ vortex-spatial/src/scalar_fn/row.rs | 5 +++++ 2 files changed, 21 insertions(+) diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index 7e47ec0f71d..8a670bb0e12 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -184,6 +184,22 @@ pub(crate) fn placeholder_geometry() -> Geometry { Geometry::Point(geo_types::Point::new(0.0, 0.0)) } +/// Whether [`geometries_null_tolerant`] supports this array without filtering null rows first. +pub(crate) fn can_decode_geometries_null_tolerant(array: &ArrayRef) -> VortexResult { + if array.validity()?.definitely_no_nulls() { + return Ok(true); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + + Ok(ext.is::() || ext.is::()) +} + /// Decode a native geometry column that may contain null rows, writing [`placeholder_geometry`] /// into their slots. The caller guarantees null rows are never read. /// diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index 127f34779b5..b7353e10c1c 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -12,6 +12,7 @@ use vortex_array::scalar_fn::unstable::row::InputElement; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use crate::extension::can_decode_geometries_null_tolerant; use crate::extension::geometries; use crate::extension::geometries_null_tolerant; use crate::extension::is_native_geometry; @@ -49,6 +50,10 @@ unsafe impl InputElement for GeometryRow { geometries(&array, ctx) } + fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { + can_decode_geometries_null_tolerant(array) + } + fn get(column: &Self::Column, index: usize) -> &Geometry { &column[index] } From 7c571caeac9749aaa05a4627859e95aa3bfca10e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:05:11 -0400 Subject: [PATCH 140/160] Execute spatial predicates with RowFn Signed-off-by: Connor Tsui --- Cargo.toml | 8 +- vortex-spatial/src/scalar_fn/contains.rs | 648 +++++++++++++++--- vortex-spatial/src/scalar_fn/execute.rs | 7 +- .../src/scalar_fn/execute/binary.rs | 253 +------ .../src/scalar_fn/execute/geo_types.rs | 24 - vortex-spatial/src/scalar_fn/intersects.rs | 246 +++++-- vortex-spatial/src/scalar_fn/row.rs | 94 ++- 7 files changed, 837 insertions(+), 443 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 65452f18300..fc38587579e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -162,7 +162,13 @@ flatbuffers = "25.2.10" fsst-rs = "0.6.0" futures = { version = "0.3.31", default-features = false } fuzzy-matcher = "0.3" -geo = "0.31.0" +# `vortex-spatial`'s `contains_route` transcribes geo's `impl_contains_from_relate!` dispatch +# table, so any bump that moves a row silently changes containment verdicts. The tests stay green +# wherever relate and the direct algorithm agree. Pinned exactly so that taking any new geo, +# patch releases included, is a deliberate edit of this line that re-verifies the table; a caret +# requirement would let `cargo update` (or automated lockfile maintenance) take 0.31.x with no diff +# to review. See `vortex-spatial/src/scalar_fn/contains.rs`. +geo = "=0.31.0" geo-traits = "0.3.0" geo-types = "0.7.19" geoarrow = "0.8.0" diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 8850e59f751..7e841081f9f 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -3,44 +3,31 @@ //! `ST_Contains`: OGC containment test between two native geometries. +use std::cell::OnceCell; + +use geo::BoundingRect; use geo::Contains; +use geo::PreparedGeometry; +use geo::Relate; +use geo_types::Geometry; +use geo_types::Rect; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Contains`. -fn validate_contains_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: contains requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: contains operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; +#[cfg(test)] +use crate::scalar_fn::row::probe; /// OGC `ST_Contains` between two native geometry operands, each a column or a constant /// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone @@ -59,83 +46,316 @@ impl SpatialContains { } } -impl ScalarFnVTable for SpatialContains { +impl RowFn for SpatialContains { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.contains"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { Ok(EmptyOptions) } - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) + /// Containment is not symmetric, so `a` is always the container and `b` the contained. + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), UninitElementSink, _, _>( + |(a, b)| { + #[cfg(test)] + probe::record(a.is_some(), b.is_some()); + ConstOperands { + a: a.map(PreparedOperand::new), + b: b.map(PreparedOperand::new), + } + }, + |operands, (a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, contains_row_prepared(operands, a, b)) } + }, + ) } +} + +vortex_array::impl_row_fn_vtable!(SpatialContains); + +/// Per-batch state for the contains row kernel: the prepared form of whichever operand is +/// constant for the batch. `None` marks an operand that varies by row. +struct ConstOperands { + /// Operand `a` (the container) when it is batch-constant. + a: Option, - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("contains has exactly two children"), + /// Operand `b` (the contained) when it is batch-constant. + b: Option, +} + +/// One batch-constant operand: its bounding rectangle and the [`PreparedGeometry`] built on the +/// first row whose pairing routes through relate. +/// +/// The build is lazy because preparation (self-noding the topology graph plus an R*-tree over the +/// edges) costs `O(edges log edges)` and pays off only on relate-routed pairings; a batch of +/// point rows against a constant polygon never touches it, and preparing a large constant eagerly +/// would charge such a batch for nothing. +struct PreparedOperand { + /// The constant's bounding rectangle, folded once for conservative row rejection. + bbox: Option>, + + /// The constant's prepared form, initialized only when a relate route needs it. + prepared: OnceCell, f64>>, +} + +impl PreparedOperand { + fn new(geometry: &Geometry) -> Self { + Self { + bbox: finite_bounding_rect(geometry), + prepared: OnceCell::new(), } } - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_contains_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Bool(nullability)) + /// Return the prepared geometry, cloning the decoded constant only on first use. + /// + /// `geometry` **must** be the constant represented by this state. The row kernel maintains + /// that relationship by passing the operand from the same decoded constant column that + /// produced this [`PreparedOperand`]. + fn get(&self, geometry: &Geometry) -> &PreparedGeometry<'static, Geometry, f64> { + self.prepared + .get_or_init(|| PreparedGeometry::from(geometry.clone())) } +} - fn execute( - &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Containment is not symmetric: `a` is always the container and `b` the contained. A - // container's rect must cover the contained's rect (`Rect::contains` is the closed - // test), so a contained rect poking outside proves the row false. - execute_binary_geo_types( - &a, - &b, - |a, b| a.contains(b), - Some(|ra, rb| (!ra.contains(rb)).then_some(false)), - ctx, - ) - } +/// Returns a bounding rectangle only when ordered comparisons can conservatively reject a row. +/// +/// Geo permits non-finite coordinates. A rectangle containing NaN cannot prove non-containment, +/// because its ordered comparisons can return false even when the exact algorithm accepts the +/// geometry. +fn finite_bounding_rect(geometry: &Geometry) -> Option> { + let bbox = geometry.bounding_rect()?; + let min = bbox.min(); + let max = bbox.max(); + + [min.x, min.y, max.x, max.y] + .into_iter() + .all(f64::is_finite) + .then_some(bbox) +} - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) +/// How geo's `a.contains(b)` computes its verdict for a pairing. +enum ContainsRoute { + /// `a.relate(b).is_contains()`. + ForwardRelate, + + /// `b.relate(a).is_within()`, how geo phrases relate for `MultiPolygon` containers. + ReversedRelate, + + /// A direct algorithm (coordinate position, point arithmetic); nothing to prepare. + Direct, +} + +/// The route geo 0.31's `Contains` dispatch takes for `a.contains(b)`. +/// +/// The prepared substitution in [`contains_row_prepared`] **must** run relate exactly where geo +/// runs relate, with the same argument order, because geo's direct algorithms are not everywhere +/// bit-identical to a relate matrix query (they resolve degenerate and boundary cases with +/// different arithmetic). The relate rows below transcribe geo's `impl_contains_from_relate!` +/// lists per container type; everything else, notably every `Point`/`MultiPoint` contained side +/// and every `Point` container, is direct. +/// +/// **This table is coupled to the geo version.** It transcribes a dispatch that geo is free to +/// reshuffle in any release, and a wrong row is a silently wrong verdict rather than a build error. +/// The workspace therefore pins `geo = "=0.31.0"`: taking any new geo, patch releases included, is +/// a deliberate edit of that line, and the edit must re-verify this table against +/// `impl_contains_from_relate!`. +/// +/// `constant_operands_agree_with_columns` is the mechanical check, and it is **not** complete: it +/// compares the prepared route against plain `a.contains(b)` only for the container types it has +/// cases for. `routes_agree_with_geo_for_every_container` covers the rest, one representative +/// pairing per container variant, and is the one to extend when geo grows a geometry type. Both +/// stay green wherever relate and the direct algorithm agree, so neither replaces the pin. +fn contains_route(a: &Geometry, b: &Geometry) -> ContainsRoute { + use Geometry as G; + + match (a, b) { + // Line contains [Polygon, MultiLineString, MultiPolygon, GeometryCollection, Rect, + // Triangle]. + ( + G::Line(_), + G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // LineString contains [Polygon, MultiPoint, MultiLineString, MultiPolygon, + // GeometryCollection, Rect, Triangle]. + | ( + G::LineString(_), + G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // MultiLineString contains everything except Point. + | ( + G::MultiLineString(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // MultiPoint contains [Line, LineString, Polygon, MultiLineString, MultiPolygon, + // GeometryCollection, Rect, Triangle]. + | ( + G::MultiPoint(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // Polygon contains everything except Point and MultiPoint. + | ( + G::Polygon(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // Rect contains [Line, LineString, MultiPoint, MultiLineString, MultiPolygon, + // GeometryCollection, Triangle]; Rect contains Rect and Polygon are direct. + | ( + G::Rect(_), + G::Line(_) + | G::LineString(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Triangle(_), + ) + // Triangle and GeometryCollection contain everything except Point. + | ( + G::Triangle(_) | G::GeometryCollection(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) => ContainsRoute::ForwardRelate, + + // MultiPolygon contains everything except Point and MultiPoint, phrased reversed. + ( + G::MultiPolygon(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) => ContainsRoute::ReversedRelate, + + _ => ContainsRoute::Direct, } +} - fn is_strict(&self, _: &Self::Options) -> bool { - true +/// Computes one row of contains, substituting a prepared graph for a constant operand on the +/// pairings geo itself answers through relate. +/// +/// [`PreparedGeometry`] carries the operand's self-noded topology graph and edge R*-tree, so a +/// relate against it skips rebuilding both and reads its bounding rect from cache; geo asserts +/// the cached graph equal to a freshly built one (its `swap_arg_index` test), which is what makes +/// the substitution result-preserving. Before dispatch, a disjoint constant-side bounding rect +/// conservatively rejects the row, matching the columnar implementation's #9076 optimization. +/// All other rows delegate to the same direct or relate route as `a.contains(b)`. +fn contains_row_prepared(operands: &ConstOperands, a: &Geometry, b: &Geometry) -> bool { + let rejected = match (&operands.a, &operands.b) { + (None, None) => false, + (Some(const_a), Some(const_b)) => const_a + .bbox + .zip(const_b.bbox) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + (Some(const_a), None) => const_a + .bbox + .zip(finite_bounding_rect(b)) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + (None, Some(const_b)) => finite_bounding_rect(a) + .zip(const_b.bbox) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + }; + + if rejected { + return false; } - fn is_fallible(&self, _: &Self::Options) -> bool { - false + match contains_route(a, b) { + ContainsRoute::Direct => a.contains(b), + ContainsRoute::ForwardRelate => match (&operands.a, &operands.b) { + (Some(const_a), Some(const_b)) => const_a.get(a).relate(const_b.get(b)).is_contains(), + (Some(const_a), None) => const_a.get(a).relate(b).is_contains(), + (None, Some(const_b)) => a.relate(const_b.get(b)).is_contains(), + (None, None) => a.contains(b), + }, + ContainsRoute::ReversedRelate => match (&operands.a, &operands.b) { + (Some(const_a), Some(const_b)) => const_b.get(b).relate(const_a.get(a)).is_within(), + (Some(const_a), None) => b.relate(const_a.get(a)).is_within(), + (None, Some(const_b)) => const_b.get(b).relate(a).is_within(), + (None, None) => a.contains(b), + }, } } #[cfg(test)] mod tests { + use geo::Contains; + use geo_types::Coord; use geo_types::Geometry; + use geo_types::GeometryCollection; + use geo_types::Line; use geo_types::LineString; + use geo_types::MultiLineString; + use geo_types::MultiPoint; + use geo_types::MultiPolygon; use geo_types::Point; use geo_types::Polygon; + use geo_types::Rect; + use geo_types::Triangle; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -144,6 +364,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::MaskedArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -158,10 +379,15 @@ mod tests { use vortex_error::vortex_err; use wkb::writer::WriteOptions; + use super::ConstOperands; + use super::PreparedOperand; use super::SpatialContains; + use super::contains_row_prepared; + use crate::scalar_fn::row::probe::assert_prepared_agrees_with_columns; use crate::test_harness::linestring_column; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; + use crate::test_harness::polygon_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { @@ -218,6 +444,21 @@ mod tests { assert_contains(container, other, [expected; 3]) } + /// A non-finite bounding rectangle cannot reject a containment that the exact geometry + /// algorithm accepts. + #[test] + fn nan_bounding_rect_does_not_reject_containment() { + let container = multipoint(vec![(f64::NAN, f64::NAN), (1.0, 1.0)]); + let contained = point(1.0, 1.0); + let operands = ConstOperands { + a: Some(PreparedOperand::new(&container)), + b: Some(PreparedOperand::new(&contained)), + }; + + assert!(container.contains(&contained)); + assert!(contains_row_prepared(&operands, &container, &contained)); + } + /// Partially overlapping polygons contain each other in neither direction. #[test] fn overlapping_polygons_contain_neither_way() -> VortexResult<()> { @@ -246,6 +487,20 @@ mod tests { assert_contains(container, points, [true, false, false]) } + /// Constant container vs a linestring column: a row whose bounding rect pokes outside the + /// container's is not contained, while one wholly inside is. Carried over from the columnar + /// bounding-rect rejection in #9076, since it constrains the verdict rather than the mechanism. + #[test] + fn constant_container_vs_row_rect_poking_outside() -> VortexResult<()> { + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let lines = linestring_column(vec![ + vec![(1.0, 1.0), (3.0, 3.0)], + vec![(1.0, 1.0), (9.0, 1.0)], + vec![(5.0, 5.0), (9.0, 9.0)], + ])?; + assert_contains(container, lines, [true, false, false]) + } + /// Polygon column vs constant point: only the polygon around the point contains it. #[test] fn polygon_column_vs_constant_point() -> VortexResult<()> { @@ -266,20 +521,6 @@ mod tests { assert_contains(away, point, [false; 2]) } - /// Constant container vs a linestring column: a row whose bounding rect pokes outside the - /// container's rect is proven false by the rect pre-check alone; a fully inside row still - /// needs (and passes) the exact test. - #[test] - fn constant_container_vs_row_rect_poking_outside() -> VortexResult<()> { - let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; - let lines = linestring_column(vec![ - vec![(1.0, 1.0), (3.0, 3.0)], - vec![(1.0, 1.0), (9.0, 1.0)], - vec![(5.0, 5.0), (9.0, 9.0)], - ])?; - assert_contains(container, lines, [true, false, false]) - } - /// Column vs column pairs rows: each polygon row is tested against the point row at the /// same position. #[test] @@ -410,6 +651,83 @@ mod tests { Ok(()) } + /// A nullable polygon column: unit squares at `centers`, the rows where `nulls` is true + /// masked out, spelled as `Masked` over non-nullable storage. + fn nullable_squares(centers: &[(f64, f64)], nulls: &[bool]) -> VortexResult { + let squares = centers + .iter() + .map(|&(x, y)| { + vec![vec![ + (x - 1.0, y - 1.0), + (x + 1.0, y - 1.0), + (x + 1.0, y + 1.0), + (x - 1.0, y + 1.0), + (x - 1.0, y - 1.0), + ]] + }) + .collect(); + let polygons = polygon_column(squares)?; + + Ok( + MaskedArray::try_new(polygons, Validity::from_iter(nulls.iter().map(|n| !n)))? + .into_array(), + ) + } + + /// Nullable geometry operands conjoin their validity before computing containment. + #[test] + fn contains_nullable_geometries_conjoins_validity() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let centers = [(0.0, 0.0), (5.0, 5.0), (0.5, -0.2), (9.0, 9.0), (0.0, 1.0)]; + let nulls = [false, true, false, false, true]; + let polygons = nullable_squares(¢ers, &nulls)?; + let points = nullable_point_column(vec![ + Some((0.0, 0.0)), + Some((5.0, 5.0)), + None, + Some((0.0, 0.0)), + Some((0.0, 1.0)), + ])?; + + let actual = SpatialContains::try_new_array(polygons, points)? + .into_array() + .execute::(&mut ctx)? + .into_array(); + let expected = BoolArray::from_iter([Some(true), None, None, Some(false), None]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// Geometry types without a null-tolerant decode fall back to filtering valid rows. + #[test] + fn contains_unsupported_geometry_falls_back_to_filter() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let validity = Validity::from_iter([true, false, true, true]); + let lines = linestring_column(vec![ + vec![(0.0, 0.0), (4.0, 4.0)], + vec![(0.0, 0.0), (1.0, 1.0)], + vec![(2.0, 2.0), (3.0, 3.0)], + vec![(0.0, 4.0), (4.0, 0.0)], + ])?; + let nullable_lines = MaskedArray::try_new(lines.clone(), validity.clone())?.into_array(); + let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 4)?; + + let expected = SpatialContains::try_new_array(lines, point.clone())?.into_array(); + let expected = MaskedArray::try_new(expected, validity)?.into_array(); + let actual = SpatialContains::try_new_array(nullable_lines, point)? + .into_array() + .execute::(&mut ctx)? + .into_array(); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + /// A non-geometry operand dtype is rejected up front, before execution. #[test] fn non_geometry_operand_is_rejected() -> VortexResult<()> { @@ -419,4 +737,166 @@ mod tests { assert!(result.is_err()); Ok(()) } + + // The prepared-vs-expanded agreement grid: every constant arrangement of a pairing must + // return exactly what the fully expanded columns return. + + /// A point geometry. + fn point(x: f64, y: f64) -> Geometry { + Geometry::Point(Point::new(x, y)) + } + + /// A linestring geometry through `coords`. + fn line(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::LineString(LineString::from(coords)) + } + + /// A multipoint geometry over `coords`. + fn multipoint(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::MultiPoint(MultiPoint::from(coords)) + } + + /// A two-point line segment geometry, the `Line` container variant. + fn line_geometry(start: (f64, f64), end: (f64, f64)) -> Geometry { + Geometry::Line(Line::new( + Coord { + x: start.0, + y: start.1, + }, + Coord { x: end.0, y: end.1 }, + )) + } + + /// A multilinestring geometry over one linestring per entry of `parts`. + fn multilinestring(parts: Vec>) -> Geometry { + Geometry::MultiLineString(MultiLineString::new( + parts.into_iter().map(LineString::from).collect(), + )) + } + + /// A geometry collection wrapping `parts`. + fn collection(parts: Vec) -> Geometry { + Geometry::GeometryCollection(GeometryCollection::from(parts)) + } + + /// An axis-aligned rectangle geometry, the `Rect` container variant. + fn rect_geometry(x0: f64, y0: f64, x1: f64, y1: f64) -> Geometry { + Geometry::Rect(Rect::new(Coord { x: x0, y: y0 }, Coord { x: x1, y: y1 })) + } + + /// A triangle geometry large enough to contain the small test polygons. + fn triangle_geometry() -> Geometry { + Geometry::Triangle(Triangle::new( + Coord { x: 0.0, y: 0.0 }, + Coord { x: 8.0, y: 0.0 }, + Coord { x: 0.0, y: 8.0 }, + )) + } + + /// A two-part multipolygon: `4x4` squares at the origin and at `(10, 10)`. + fn two_part_multipolygon() -> Geometry { + Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 4.0, 4.0), + rect_polygon(10.0, 10.0, 14.0, 14.0), + ])) + } + + /// Every container variant `contains_route` distinguishes, checked against plain + /// `a.contains(b)` in all four constant arrangements. + /// + /// Every case is a containment geo answers `true`, which the test asserts: a pairing that is + /// false regardless of route (a lower-dimensional container, say) also agrees regardless of + /// route, and pins nothing. A true case fails when the prepared substitution diverges from + /// geo: a table row whose relate phrasing disagrees with geo's dispatch on this input, or a + /// bounding-rect prescreen that wrongly rejects a contained row. It is **not** a version + /// tripwire: a geo release that reshuffles its dispatch stays green wherever relate and the + /// direct algorithm agree, which is why the workspace pins `geo` exactly. + /// + /// This is the table's own regression, and the one to extend when geo grows a geometry type: + /// `constant_operands_agree_with_columns` below goes through real arrays and so is the better + /// end-to-end check, but it only covers the container types it has cases for, and WKB decoding + /// limits which types those can be. The MultiPoint and Line containers route relate only for + /// contained types a MultiPoint or Line can rarely contain, so their true cases lean on + /// `GeometryCollection` membership and collinear `MultiLineString` parts respectively. + #[rstest] + #[case::point(point(1.0, 1.0), point(1.0, 1.0))] + #[case::line(line_geometry((0.0, 0.0), (4.0, 4.0)), point(2.0, 2.0))] + #[case::line_x_multilinestring(line_geometry((0.0, 0.0), (4.0, 4.0)), multilinestring(vec![vec![(1.0, 1.0), (2.0, 2.0)]]))] + #[case::linestring(line(vec![(0.0, 0.0), (4.0, 4.0)]), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::multipoint(multipoint(vec![(0.0, 0.0), (2.0, 2.0), (4.0, 4.0)]), collection(vec![point(2.0, 2.0)]))] + #[case::multilinestring(multilinestring(vec![vec![(0.0, 0.0), (4.0, 4.0)]]), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipolygon(two_part_multipolygon(), rect_polygon(1.0, 1.0, 3.0, 3.0).into())] + #[case::geometrycollection(collection(vec![rect_polygon(0.0, 0.0, 8.0, 8.0).into()]), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::rect(rect_geometry(0.0, 0.0, 8.0, 8.0), line(vec![(2.0, 2.0), (4.0, 4.0)]))] + #[case::triangle(triangle_geometry(), rect_polygon(1.0, 1.0, 2.0, 2.0).into())] + fn routes_agree_with_geo_for_every_container(#[case] a: Geometry, #[case] b: Geometry) { + let expected = a.contains(&b); + assert!( + expected, + "route cases must be containments geo answers true, or every route agrees vacuously", + ); + + let arrangements = [ + (None, None), + (Some(PreparedOperand::new(&a)), None), + (None, Some(PreparedOperand::new(&b))), + ( + Some(PreparedOperand::new(&a)), + Some(PreparedOperand::new(&b)), + ), + ]; + + for (index, (const_a, const_b)) in arrangements.into_iter().enumerate() { + let operands = ConstOperands { + a: const_a, + b: const_b, + }; + assert_eq!( + contains_row_prepared(&operands, &a, &b), + expected, + "arrangement {index} disagrees with geo's own contains", + ); + } + } + + /// Constant arrangements agree with expanded columns across the routes the prepared kernel + /// distinguishes: forward relate (polygon, linestring and multipoint containers), reversed + /// relate (multipolygon containers), and the direct pairings (a point on either side, + /// multipoint over multipoint, polygon over multipoint), including boundary contact, + /// crossing, disjoint and empty cases. + #[rstest] + #[case::polygon_nested_polygon(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::polygon_touching_from_inside(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(0.0, 2.0, 2.0, 4.0).into())] + #[case::polygon_overlapping_polygon(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(2.0, 2.0, 6.0, 6.0).into())] + #[case::polygon_disjoint_polygon(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::polygon_x_point_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(2.0, 2.0))] + #[case::polygon_x_point_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(0.0, 2.0))] + #[case::polygon_x_point_outside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(20.0, 20.0))] + #[case::polygon_x_nan_point(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(f64::NAN, 2.0))] + #[case::polygon_x_linestring_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon_x_linestring_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(0.0, 1.0), (0.0, 3.0)]))] + #[case::polygon_x_linestring_crossing(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(-2.0, 2.0), (2.0, 2.0)]))] + #[case::polygon_x_empty_linestring(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![]))] + #[case::polygon_x_multipoint_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon_x_multipoint_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), multipoint(vec![(0.0, 1.0), (0.0, 3.0)]))] + #[case::linestring_x_multipoint_on_line(line(vec![(0.0, 0.0), (4.0, 4.0)]), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipoint_x_multipoint_subset(multipoint(vec![(0.0, 0.0), (2.0, 2.0), (4.0, 4.0)]), multipoint(vec![(2.0, 2.0)]))] + #[case::multipoint_x_linestring_between_points(multipoint(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipolygon_x_polygon_in_one_part(two_part_multipolygon(), rect_polygon(1.0, 1.0, 3.0, 3.0).into())] + #[case::multipolygon_x_polygon_straddling(two_part_multipolygon(), rect_polygon(3.0, 3.0, 11.0, 11.0).into())] + #[case::multipolygon_x_polygon_disjoint(two_part_multipolygon(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::multipolygon_x_point_inside(two_part_multipolygon(), point(11.0, 11.0))] + #[case::point_x_point_equal(point(1.0, 1.0), point(1.0, 1.0))] + #[case::point_x_polygon(point(2.0, 2.0), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + fn constant_operands_agree_with_columns( + #[case] a: Geometry, + #[case] b: Geometry, + ) -> VortexResult<()> { + assert_prepared_agrees_with_columns( + SpatialContains::try_new_array, + geometry_constant(&a, 3)?, + geometry_constant(&b, 3)?, + ) + } } diff --git a/vortex-spatial/src/scalar_fn/execute.rs b/vortex-spatial/src/scalar_fn/execute.rs index 3a7494bcb39..2acdce76e36 100644 --- a/vortex-spatial/src/scalar_fn/execute.rs +++ b/vortex-spatial/src/scalar_fn/execute.rs @@ -7,17 +7,14 @@ //! propagation without prescribing how a kernel represents geometries or builds its output. //! Native columnar kernels such as `ST_MakeLine` use these dispatchers directly. //! -//! [`execute_unary_geo_types`] and [`execute_binary_geo_types`] are convenience adapters for -//! row-oriented algorithms from the `geo` ecosystem. They decode valid inputs into -//! `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`], such as an `f64` or -//! boolean array. +//! [`execute_unary_geo_types`] adapts row-oriented algorithms from the `geo` ecosystem. It decodes +//! valid inputs into `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`]. mod binary; mod geo_types; mod unary; pub(crate) use binary::dispatch_binary; -pub(crate) use binary::execute_binary_geo_types; pub(crate) use unary::dispatch_unary; pub(crate) use unary::execute_unary_geo_types; use vortex_array::ArrayRef; diff --git a/vortex-spatial/src/scalar_fn/execute/binary.rs b/vortex-spatial/src/scalar_fn/execute/binary.rs index f2c03bd1beb..5cf639461b0 100644 --- a/vortex-spatial/src/scalar_fn/execute/binary.rs +++ b/vortex-spatial/src/scalar_fn/execute/binary.rs @@ -1,28 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Binary pairwise dispatch, plus an adapter for row-oriented `geo_types` kernels. +//! Binary constant-and-column operand dispatch. -use geo::BoundingRect; -use geo_types::Geometry; -use geo_types::Rect; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_mask::Mask; use super::Execution; use super::Operand; -use super::geo_types::GeoTypesOutput; -use super::geo_types::eval_column; -use super::geo_types::eval_column_pair; -use crate::extension::single_geometry; /// Dispatch a binary strict geometry kernel over constants and columns. /// @@ -80,6 +72,7 @@ where if len != 0 && valid.all_false() { return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); } + kernel( Execution { operands: [left, right], @@ -90,245 +83,3 @@ where ctx, ) } - -/// A bounding-rectangle pre-check for [`execute_binary_geo_types`]'s one-constant paths. -/// -/// Called per row with rectangles in operand order, it returns `Some(result)` when they prove the -/// result and `None` when the exact kernel must run. -pub(crate) type BboxPrecheck = fn(&Rect, &Rect) -> Option; - -/// Run a binary row-oriented kernel whose inputs are decoded to `geo_types::Geometry`. -/// -/// The `geo_types` name describes the values passed to `compute`, not the output. `T` is converted -/// into a Vortex array before this function returns. Nulls propagate from either operand. With -/// exactly one constant operand, `bbox_precheck` may prove a result from the fixed constant -/// bounding rectangle and the current row's rectangle before the exact kernel runs. -pub(crate) fn execute_binary_geo_types( - left: &ArrayRef, - right: &ArrayRef, - compute: F, - bbox_precheck: Option>, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T + Copy, -{ - let nullability = Nullability::from(left.dtype().is_nullable() || right.dtype().is_nullable()); - dispatch_binary( - left, - right, - T::dtype(nullability), - |execution, ctx| match execution.operands { - [Operand::Constant(left), Operand::Constant(right)] => { - let left = single_geometry(&left, ctx)?; - let right = single_geometry(&right, ctx)?; - Ok(ConstantArray::new( - compute(&left, &right).into_scalar(execution.nullability), - execution.len, - ) - .into_array()) - } - [Operand::Constant(left), Operand::Column(right)] => { - let left = single_geometry(&left, ctx)?; - let prescreen = bbox_precheck.zip(left.bounding_rect()); - eval_column( - &right, - &execution.valid, - |right| { - prescreen - .and_then(|(precheck, fixed)| precheck(&fixed, &right.bounding_rect()?)) - .unwrap_or_else(|| compute(&left, right)) - }, - execution.nullability, - ctx, - ) - } - [Operand::Column(left), Operand::Constant(right)] => { - let right = single_geometry(&right, ctx)?; - let prescreen = bbox_precheck.zip(right.bounding_rect()); - eval_column( - &left, - &execution.valid, - |left| { - prescreen - .and_then(|(precheck, fixed)| precheck(&left.bounding_rect()?, &fixed)) - .unwrap_or_else(|| compute(left, &right)) - }, - execution.nullability, - ctx, - ) - } - [Operand::Column(left), Operand::Column(right)] => eval_column_pair( - &left, - &right, - &execution.valid, - compute, - execution.nullability, - ctx, - ), - }, - ctx, - ) -} - -#[cfg(test)] -mod tests { - use std::cell::Cell; - - use geo::Contains; - use geo::Intersects; - use geo_types::Geometry; - use vortex_array::ArrayRef; - use vortex_array::ExecutionCtx; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::ConstantArray; - use vortex_array::assert_arrays_eq; - use vortex_array::validity::Validity; - use vortex_buffer::BitBuffer; - use vortex_error::VortexResult; - - use super::BboxPrecheck; - use super::execute_binary_geo_types; - use crate::test_harness::linestring_column; - use crate::test_harness::nullable_point_column; - use crate::test_harness::point_column; - use crate::test_harness::polygon_column; - - const DISJOINT_PRECHECK: BboxPrecheck = - |left, right| (!left.intersects(right)).then_some(false); - - fn triangle_constant(len: usize, ctx: &mut ExecutionCtx) -> VortexResult { - let ring = vec![(0.0, 0.0), (10.0, 0.0), (0.0, 10.0), (0.0, 0.0)]; - let scalar = polygon_column(vec![vec![ring]])?.execute_scalar(0, ctx)?; - Ok(ConstantArray::new(scalar, len).into_array()) - } - - fn counting_intersects( - counter: &Cell, - ) -> impl Fn(&Geometry, &Geometry) -> bool + Copy { - move |left, right| { - counter.set(counter.get() + 1); - left.intersects(right) - } - } - - #[test] - fn bbox_precheck_skips_exact_test() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(3, &mut ctx)?; - let probes = point_column(vec![50.0, 8.0, 2.0], vec![50.0, 8.0, 2.0])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - #[test] - fn bbox_precheck_leaves_nulls_alone() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(3, &mut ctx)?; - let probes = nullable_point_column(vec![Some((50.0, 50.0)), None, Some((2.0, 2.0))])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - let expected = BoolArray::new( - BitBuffer::from_iter([false, false, true]), - Validity::from_iter([true, false, true]), - ) - .into_array(); - - assert_arrays_eq!(result, expected, &mut ctx); - assert_eq!(exact_runs.get(), 1); - Ok(()) - } - - #[test] - fn bbox_precheck_sees_rects_in_operand_order() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let triangle = triangle_constant(2, &mut ctx)?; - let exact_runs = Cell::new(0); - let counted = |left: &Geometry, right: &Geometry| { - exact_runs.set(exact_runs.get() + 1); - left.contains(right) - }; - - let result = execute_binary_geo_types( - &probes, - &triangle, - counted, - Some(|left, right| (!left.contains(right)).then_some(false)), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 0); - Ok(()) - } - - #[test] - fn empty_constant_falls_through_to_exact() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let scalar = linestring_column(vec![vec![]])?.execute_scalar(0, &mut ctx)?; - let empty = ConstantArray::new(scalar, 2).into_array(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &empty, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - #[test] - fn bbox_precheck_matches_exact_results() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(6, &mut ctx)?; - let probes = nullable_point_column(vec![ - Some((50.0, 50.0)), - Some((8.0, 8.0)), - Some((2.0, 2.0)), - None, - Some((0.0, 0.0)), - Some((10.0, 0.0)), - ])?; - let exact = |left: &Geometry, right: &Geometry| left.intersects(right); - - let with_precheck = - execute_binary_geo_types(&triangle, &probes, exact, Some(DISJOINT_PRECHECK), &mut ctx)?; - let exact_only = execute_binary_geo_types(&triangle, &probes, exact, None, &mut ctx)?; - - assert_arrays_eq!(with_precheck, exact_only, &mut ctx); - Ok(()) - } -} diff --git a/vortex-spatial/src/scalar_fn/execute/geo_types.rs b/vortex-spatial/src/scalar_fn/execute/geo_types.rs index 038aca46502..7007f02cfc6 100644 --- a/vortex-spatial/src/scalar_fn/execute/geo_types.rs +++ b/vortex-spatial/src/scalar_fn/execute/geo_types.rs @@ -118,27 +118,3 @@ where let values = decoded.iter().map(compute).collect(); Ok(T::build_array(len, valid, values, nullability)) } - -/// Evaluate a decoded kernel over rows where both geometry columns are valid. -pub(super) fn eval_column_pair( - left: &ArrayRef, - right: &ArrayRef, - valid: &Mask, - compute: F, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T, -{ - let len = left.len(); - let left = geometries(&left.filter(valid.clone())?, ctx)?; - let right = geometries(&right.filter(valid.clone())?, ctx)?; - let values = left - .iter() - .zip(&right) - .map(|(left, right)| compute(left, right)) - .collect(); - Ok(T::build_array(len, valid, values, nullability)) -} diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index 77d33886ff3..160b384cd84 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -3,44 +3,27 @@ //! `ST_Intersects`: OGC intersection test between two native geometries. +use geo::BoundingRect; use geo::Intersects; +use geo_types::Geometry; +use geo_types::Rect; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Intersects`. -fn validate_intersects_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: intersects requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: intersects operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; +#[cfg(test)] +use crate::scalar_fn::row::probe; /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry /// operands, each a column or a constant literal. @@ -58,74 +41,102 @@ impl SpatialIntersects { } } -impl ScalarFnVTable for SpatialIntersects { +impl RowFn for SpatialIntersects { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.intersects"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { Ok(EmptyOptions) } - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("intersects has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_intersects_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Bool(nullability)) - } - - fn execute( + fn dispatch>( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Disjoint bounding rects prove the geometries disjoint; rect contact (closed test) - // falls through to the exact test. - execute_binary_geo_types( - &a, - &b, - |x, y| x.intersects(y), - Some(|ra, rb| (!ra.intersects(rb)).then_some(false)), - ctx, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), UninitElementSink, _, _>( + |(a, b)| { + #[cfg(test)] + probe::record(a.is_some(), b.is_some()); + ConstBboxes::new(a, b) + }, + |bboxes, (a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, intersects_row_prepared(bboxes, a, b)) } + }, ) } +} - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } +vortex_array::impl_row_fn_vtable!(SpatialIntersects); + +/// Per-batch state for the intersects row kernel: the bounding rect of each operand that is +/// constant for the batch. +/// +/// geo opens many intersects pairings with `has_disjoint_bboxes`, an early-out that folds +/// [`bounding_rect`] over both operands. For a batch-constant operand that fold recomputes the +/// same rect every row, so it is hoisted here and [`intersects_row_prepared`] replays the +/// comparison with the hoisted value. `None` marks an operand that varies by row or has no +/// bounding rect (an empty geometry); both mean no early-out, exactly as `has_disjoint_bboxes` +/// treats a missing rect. +/// +/// [`bounding_rect`]: BoundingRect::bounding_rect +struct ConstBboxes { + /// The bounding rect of operand `a` when it is batch-constant. + a: Option>, + + /// The bounding rect of operand `b` when it is batch-constant. + b: Option>, +} - fn is_strict(&self, _: &Self::Options) -> bool { - true +impl ConstBboxes { + fn new(a: Option<&Geometry>, b: Option<&Geometry>) -> Self { + Self { + a: a.and_then(BoundingRect::bounding_rect), + b: b.and_then(BoundingRect::bounding_rect), + } } +} - fn is_fallible(&self, _: &Self::Options) -> bool { - false - } +/// Computes one row of intersects, spending any bounding rect hoisted into `bboxes`. +/// +/// Disjoint bounding rectangles conservatively prove that the geometries do not intersect. The +/// fall-through delegates to the unchanged `a.intersects(b)`, which refolds both rects internally, +/// so a batch where every row overlaps pays one extra `bounding_rect` fold over the row operand; +/// the win concentrates where most rows are disjoint, the usual spatial-filter shape. +fn intersects_row_prepared(bboxes: &ConstBboxes, a: &Geometry, b: &Geometry) -> bool { + let disjoint = match (bboxes.a, bboxes.b) { + (None, None) => false, + (Some(bbox_a), Some(bbox_b)) => !bbox_a.intersects(&bbox_b), + (Some(bbox_a), None) => b + .bounding_rect() + .is_some_and(|bbox_b| !bbox_a.intersects(&bbox_b)), + (None, Some(bbox_b)) => a + .bounding_rect() + .is_some_and(|bbox_a| !bbox_a.intersects(&bbox_b)), + }; + + if disjoint { + return false; + } + + a.intersects(b) } #[cfg(test)] @@ -133,7 +144,9 @@ mod tests { use geo_types::Coord; use geo_types::Geometry; use geo_types::LineString; + use geo_types::MultiPoint; use geo_types::MultiPolygon; + use geo_types::Point; use geo_types::Polygon; use rstest::rstest; use vortex_array::ArrayRef; @@ -158,8 +171,10 @@ mod tests { use wkb::writer::WriteOptions; use super::SpatialIntersects; + use crate::scalar_fn::row::probe::assert_prepared_agrees_with_columns; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; + use crate::test_harness::rect_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { @@ -441,4 +456,85 @@ mod tests { assert!(result.is_err()); Ok(()) } + + // The prepared-vs-expanded agreement grid: every constant arrangement of a pairing must + // return exactly what the fully expanded columns return. + + /// A point geometry. + fn point(x: f64, y: f64) -> Geometry { + Geometry::Point(Point::new(x, y)) + } + + /// A linestring geometry through `coords`. + fn line(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::LineString(LineString::from(coords)) + } + + /// A multipoint geometry over `coords`. + fn multipoint(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::MultiPoint(MultiPoint::from(coords)) + } + + /// Constant arrangements agree with expanded columns across the pairing classes the prepared + /// kernel treats differently: bbox-prechecked pairs (polygon x polygon, linestring x + /// anything, multipolygon blankets), direct pairs (points), the excluded `MultiPoint` route, + /// and an empty geometry whose bounding rect does not exist. + #[rstest] + #[case::polygons_overlapping(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(2.0, 2.0, 6.0, 6.0).into())] + #[case::polygons_touching_edge(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(4.0, 0.0, 8.0, 4.0).into())] + #[case::polygons_touching_corner(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(4.0, 4.0, 8.0, 8.0).into())] + #[case::polygons_disjoint(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::polygons_nested(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::polygon_x_point_inside(donut(), point(2.0, 2.0))] + #[case::polygon_x_point_on_boundary(donut(), point(0.0, 5.0))] + #[case::polygon_x_point_in_hole(donut(), point(5.0, 5.0))] + #[case::point_outside_x_polygon(point(20.0, 20.0), donut())] + #[case::nan_point_x_polygon(point(f64::NAN, 2.0), donut())] + #[case::polygon_x_nan_point(donut(), point(f64::NAN, 2.0))] + #[case::points_equal(point(1.0, 1.0), point(1.0, 1.0))] + #[case::points_distinct(point(1.0, 1.0), point(2.0, 1.0))] + #[case::linestring_crossing_polygon(line(vec![(-2.0, -2.0), (2.0, 2.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::linestring_disjoint_polygon(line(vec![(-2.0, -2.0), (-6.0, -6.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::linestring_in_polygon_hole(line(vec![(4.5, 4.5), (5.5, 5.5)]), donut())] + #[case::linestrings_crossing(line(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(0.0, 4.0), (4.0, 0.0)]))] + #[case::linestrings_disjoint(line(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(10.0, 10.0), (14.0, 14.0)]))] + #[case::empty_linestring_x_polygon(line(vec![]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipoint_straddling_polygon(multipoint(vec![(2.0, 2.0), (20.0, 20.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipoint_outside_polygon(multipoint(vec![(20.0, 20.0), (30.0, 30.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipolygon_disjoint_polygon( + Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 2.0, 2.0), + rect_polygon(10.0, 10.0, 12.0, 12.0), + ])), + rect_polygon(20.0, 20.0, 24.0, 24.0).into() + )] + fn constant_operands_agree_with_columns( + #[case] a: Geometry, + #[case] b: Geometry, + ) -> VortexResult<()> { + assert_prepared_agrees_with_columns( + SpatialIntersects::try_new_array, + geometry_constant(&a, 3)?, + geometry_constant(&b, 3)?, + ) + } + + /// `Rect` has no WKB form, so its constant comes from a one-row rect column; its conservative + /// bbox early-out and exact fall-through must agree with the expanded form like the rest. + #[test] + fn rect_operand_agrees_with_columns() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let rect_scalar = rect_column(vec![(0.0, 0.0, 4.0, 4.0)])?.execute_scalar(0, &mut ctx)?; + let rect_constant = ConstantArray::new(rect_scalar, 3).into_array(); + let polygon_constant = + geometry_constant(&Geometry::Polygon(rect_polygon(2.0, 2.0, 6.0, 6.0)), 3)?; + + assert_prepared_agrees_with_columns( + SpatialIntersects::try_new_array, + rect_constant, + polygon_constant, + ) + } } diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index b7353e10c1c..750497e5f32 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -32,10 +32,10 @@ unsafe impl InputElement for GeometryRow { type View<'a> = &'a [Geometry]; type Elem<'a> = &'a Geometry; - // A geometry row is decoded from its coordinate storage, which behind a null row holds arbitrary - // coordinates that need not describe a well-formed geometry. + // A geometry row is decoded from its coordinate storage, which behind a null row holds + // arbitrary coordinates that need not describe a well-formed geometry. const DENSE_SAFE: bool = false; - // Decoding builds a geometry from stored coordinates, and a malformed one in a *valid* row is a + // Decoding builds a geometry from stored coordinates, and a malformed one in a _valid_ row is a // domain error rather than an infrastructural failure. const DECODE_FALLIBLE: bool = true; fn validate(dtype: &DType) -> VortexResult<()> { @@ -92,3 +92,91 @@ unsafe impl InputElement for GeometryRow { geometries_null_tolerant(&array, ctx) } } + +/// Test-only support for the prepared geo row kernels: a probe recording which operands a +/// `prepare` step saw as batch-constant, and the shared prepared-vs-expanded agreement check +/// built on it. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::ScalarFnArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + + thread_local! { + /// Which operands the last `prepare` saw as constant, as a bitmask (bit 0 for `a`, bit 1 + /// for `b`). Thread-local rather than a process global so concurrent tests in one process + /// (plain `cargo test`) cannot race it; execution runs on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Record which operands `prepare` saw as constant. + pub(crate) fn record(a_constant: bool, b_constant: bool) { + SEEN_CONSTANTS.set(u8::from(a_constant) | (u8::from(b_constant) << 1)); + } + + /// Execute `build(a, b)` and assert that `prepare` saw exactly `expect_seen` as its constant + /// operands, so the test knows which decode path the inputs took. + fn run_probed( + build: &impl Fn(ArrayRef, ArrayRef) -> VortexResult, + a: ArrayRef, + b: ArrayRef, + expect_seen: u8, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + SEEN_CONSTANTS.set(u8::MAX); + let result = build(a, b)? + .into_array() + .execute::(ctx)? + .into_array(); + + assert_eq!( + SEEN_CONSTANTS.get(), + expect_seen, + "prepare saw the wrong constant operands", + ); + Ok(result) + } + + /// Assert that every constant-operand arrangement of `build(a, b)` returns exactly what the + /// fully expanded columns return, and that each arrangement's constness really reached + /// `prepare` (so the constants exercised the stride-0 path rather than a decoded column). + /// + /// Arrangements: `a` constant, `b` constant, and both constant with `a` masked. A plain + /// constant pair folds to a single-row execution before the row loop, so masking one side is + /// what drives the both-hoisted arm across rows; that run is compared against the same mask + /// over the expanded column. + pub(crate) fn assert_prepared_agrees_with_columns( + build: impl Fn(ArrayRef, ArrayRef) -> VortexResult, + const_a: ArrayRef, + const_b: ArrayRef, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let col_a = const_a.clone().execute::(&mut ctx)?.into_array(); + let col_b = const_b.clone().execute::(&mut ctx)?.into_array(); + + let baseline = run_probed(&build, col_a.clone(), col_b.clone(), 0b00, &mut ctx)?; + let a_hoisted = run_probed(&build, const_a.clone(), col_b.clone(), 0b01, &mut ctx)?; + let b_hoisted = run_probed(&build, col_a.clone(), const_b.clone(), 0b10, &mut ctx)?; + assert_arrays_eq!(a_hoisted, baseline, &mut ctx); + assert_arrays_eq!(b_hoisted, baseline, &mut ctx); + + let validity = Validity::from_iter((0..col_a.len()).map(|row| row != 1)); + let masked_const_a = MaskedArray::try_new(const_a, validity.clone())?.into_array(); + let masked_col_a = MaskedArray::try_new(col_a, validity)?.into_array(); + let both_hoisted = run_probed(&build, masked_const_a, const_b, 0b11, &mut ctx)?; + let masked_baseline = run_probed(&build, masked_col_a, col_b, 0b00, &mut ctx)?; + assert_arrays_eq!(both_hoisted, masked_baseline, &mut ctx); + + Ok(()) + } +} From 4774ba5d399bd0d254d09c04e6d3adf23725a4a7 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:45 -0400 Subject: [PATCH 141/160] Use the blanket RowFn vtable for spatial predicates Signed-off-by: Connor Tsui --- vortex-spatial/src/scalar_fn/contains.rs | 2 -- vortex-spatial/src/scalar_fn/intersects.rs | 2 -- 2 files changed, 4 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 7e841081f9f..c0a3a2aa071 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -93,8 +93,6 @@ impl RowFn for SpatialContains { } } -vortex_array::impl_row_fn_vtable!(SpatialContains); - /// Per-batch state for the contains row kernel: the prepared form of whichever operand is /// constant for the batch. `None` marks an operand that varies by row. struct ConstOperands { diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index 160b384cd84..e9cd6f34866 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -84,8 +84,6 @@ impl RowFn for SpatialIntersects { } } -vortex_array::impl_row_fn_vtable!(SpatialIntersects); - /// Per-batch state for the intersects row kernel: the bounding rect of each operand that is /// constant for the batch. /// From c0f5e8e4b37af05ef291116ee65793fce49ca4f0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:40:48 -0400 Subject: [PATCH 142/160] Benchmark prepared contains reuse Signed-off-by: Connor Tsui --- vortex-spatial/benches/binary_predicates.rs | 26 +++++++++++++++++---- vortex-spatial/src/scalar_fn/contains.rs | 8 +++---- vortex-spatial/src/scalar_fn/intersects.rs | 8 +++---- 3 files changed, 29 insertions(+), 13 deletions(-) diff --git a/vortex-spatial/benches/binary_predicates.rs b/vortex-spatial/benches/binary_predicates.rs index b84ab67b17c..281578e4e16 100644 --- a/vortex-spatial/benches/binary_predicates.rs +++ b/vortex-spatial/benches/binary_predicates.rs @@ -12,11 +12,6 @@ //! column-x-column arms are the control: no operand is constant, so a prepared path has nothing to //! hoist and must not regress them. //! -//! `contains` has no all-overlapping arm. One `contains(query polygon, contained square)` row -//! builds a topology graph over the constant's 128 edges, which CodSpeed's CPU simulation charges -//! around 120 µs, so no row count both fits the per-iteration budget and exercises the row loop. -//! [`intersects::polygons_overlapping_x_constant`] covers the never-rejects case instead. -//! //! Run with `cargo bench -p vortex-spatial --bench binary_predicates`. #![expect(clippy::unwrap_used)] @@ -64,6 +59,10 @@ const ROWS: usize = 1 << 7; /// pairwise predicate. It needs a smaller fixture than [`ROWS`] to stay inside the same budget. const OVERLAPPING_POLYGON_ROWS: usize = 1 << 5; +/// Containment builds a topology graph for each polygon pair. Four rows fit the benchmark budget +/// while exercising construction followed by reuse of the prepared constant geometry. +const CONTAINED_POLYGON_ROWS: usize = 4; + /// Deterministic pseudo-random value in `[0, 1)`. fn unit(i: usize) -> f64 { ((i.wrapping_mul(2654435761) >> 8) % 10_000) as f64 / 10_000.0 @@ -225,6 +224,23 @@ mod contains { }); } + /// Constant container against contained polygons: every bbox check passes, the first row + /// prepares the constant geometry, and the remaining rows reuse it for the full predicate. + #[divan::bench] + fn constant_x_polygons_overlapping(bencher: Bencher) { + let mut ctx = SESSION.create_execution_ctx(); + let query = query_constant(&mut ctx, CONTAINED_POLYGON_ROWS); + let polygons = squares_mostly_overlapping(CONTAINED_POLYGON_ROWS); + bencher + .counter(ItemsCount::new(CONTAINED_POLYGON_ROWS)) + .bench_local(|| { + execute( + SpatialContains::try_new_array(query.clone(), polygons.clone()), + &mut ctx, + ) + }); + } + /// Constant container against a point column with one null row in eight. #[divan::bench] fn constant_x_nullable_points(bencher: Bencher) { diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index c0a3a2aa071..9b22a471204 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -15,12 +15,12 @@ use vortex_array::ArrayRef; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index e9cd6f34866..9a3a198e838 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -11,12 +11,12 @@ use vortex_array::ArrayRef; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; From 0ceaf7defaa6d077376bd0e47df82a04be7e9aec Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:02:12 -0400 Subject: [PATCH 143/160] Automate RowFn benchmark comparisons Signed-off-by: Connor Tsui --- scripts/benchmark-rowfn.sh | 362 ++++++++++++++++++++++++++ scripts/rowfn_benchmark.py | 356 +++++++++++++++++++++++++ scripts/tests/test_rowfn_benchmark.py | 134 ++++++++++ 3 files changed, 852 insertions(+) create mode 100755 scripts/benchmark-rowfn.sh create mode 100755 scripts/rowfn_benchmark.py create mode 100644 scripts/tests/test_rowfn_benchmark.py diff --git a/scripts/benchmark-rowfn.sh b/scripts/benchmark-rowfn.sh new file mode 100755 index 00000000000..4c6ae4152cc --- /dev/null +++ b/scripts/benchmark-rowfn.sh @@ -0,0 +1,362 @@ +#!/usr/bin/env bash + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +set -Eeu -o pipefail + +script_directory=$(dirname "$(realpath "${BASH_SOURCE[0]}")") + +usage() { + cat >&2 <<'EOF' +Usage: benchmark-rowfn.sh [OPTIONS] + +Options: + --suite NAME Select a preset or benchmark label. Repeatable; defaults to full. + --filter PATTERN Pass a Divan benchmark filter. Repeatable. + --config NAME primary (1 CGU/fat LTO) or repository (16 CGUs/no LTO). + --target-root PATH New directory for separate baseline and candidate Cargo targets. + --codegen-units N Override the selected configuration. + --lto VALUE Override LTO with false, thin, or fat. + --rustflags FLAGS Override RUSTFLAGS; defaults to -C target-cpu=native. + --build-jobs N Jobs per concurrent revision build; defaults to 8 and cannot exceed 8. + --bench-cpu N Logical CPU used for every timed process; defaults to 4. + --warm-runs N Warm runs per revision; defaults to 2. + --measured-pairs N Alternating measured pairs; defaults to 7. + --sample-count N Divan sample count; defaults to 100. + --min-time SECONDS Divan minimum time; defaults to 0.25. + --max-time SECONDS Divan maximum time; defaults to 0.5. + --lock-file PATH Global timed-run lock; defaults to /tmp/vortex-rowfn-benchmark.lock. + --list-suites Print presets and benchmark labels, then exit. +EOF +} + +suite_catalog=( + "array-binary_ops|vortex-array|binary_ops|array,numeric,design-a-matrix,full" + "array-compare|vortex-array|compare|array,compare,full" + "array-row_fn_executor|vortex-array|row_fn_executor|array,framework,full" + "array-strict_validity|vortex-array|strict_validity|array,framework,full" + "array-like|vortex-array|like|array,full" + "array-take_filter|vortex-array|take_filter|array,full" + "array-varbinview_compact|vortex-array|varbinview_compact|array,full" + "tensor-l2_norm|vortex-tensor|l2_norm|tensor,full" + "tensor-inner_product|vortex-tensor|inner_product|tensor,full" + "tensor-cosine_similarity|vortex-tensor|cosine_similarity|tensor,full" + "tensor-normalized|vortex-tensor|normalized|tensor,full" + "spatial-binary_predicates|vortex-spatial|binary_predicates|spatial,full" + "spatial-distance|vortex-spatial|distance|spatial,full" + "spatial-envelope|vortex-spatial|envelope|spatial,full" + "spatial-predicate_bbox|vortex-spatial|predicate_bbox|spatial,full" +) + +requested_suites=() +filters=() +configuration=primary +target_root= +codegen_units_override= +lto_override= +rustflags_override= +build_jobs=8 +bench_cpu=4 +warm_runs=2 +measured_pairs=7 +sample_count=100 +min_time=0.25 +max_time=0.5 +lock_file=/tmp/vortex-rowfn-benchmark.lock + +while [[ $# -gt 0 ]]; do + case $1 in + --suite) requested_suites+=("$2"); shift 2 ;; + --filter) filters+=("$2"); shift 2 ;; + --config) configuration=$2; shift 2 ;; + --target-root) target_root=$2; shift 2 ;; + --codegen-units) codegen_units_override=$2; shift 2 ;; + --lto) lto_override=$2; shift 2 ;; + --rustflags) rustflags_override=$2; shift 2 ;; + --build-jobs) build_jobs=$2; shift 2 ;; + --bench-cpu) bench_cpu=$2; shift 2 ;; + --warm-runs) warm_runs=$2; shift 2 ;; + --measured-pairs) measured_pairs=$2; shift 2 ;; + --sample-count) sample_count=$2; shift 2 ;; + --min-time) min_time=$2; shift 2 ;; + --max-time) max_time=$2; shift 2 ;; + --lock-file) lock_file=$2; shift 2 ;; + --list-suites) + echo "Presets: full array framework numeric design-a-matrix compare tensor spatial" + printf '%s\n' "${suite_catalog[@]}" | cut -d '|' -f 1 + exit 0 + ;; + -h|--help) usage; exit 0 ;; + --*) echo "Unknown option: $1" >&2; usage; exit 1 ;; + *) break ;; + esac +done + +if [[ $# -ne 3 ]]; then + usage + exit 1 +fi +if [[ $(uname -m) != x86_64 ]]; then + echo "RowFn native performance decisions require an x86_64 host." >&2 + exit 1 +fi +if ((build_jobs < 1 || build_jobs > 8)); then + echo "--build-jobs must be between 1 and 8 so two builds cannot exceed 16 jobs." >&2 + exit 1 +fi +command -v flock >/dev/null || { echo "benchmark-rowfn.sh requires flock." >&2; exit 1; } + +baseline=$(realpath "$1") +candidate=$(realpath "$2") +output=$(realpath -m "$3") +if [[ -e $output ]]; then + echo "Output path already exists: $output" >&2 + exit 1 +fi + +case $configuration in + primary) codegen_units=1; lto=fat ;; + repository) codegen_units=16; lto=false ;; + *) echo "Unknown configuration: $configuration" >&2; exit 1 ;; +esac +codegen_units=${codegen_units_override:-$codegen_units} +lto=${lto_override:-$lto} +rustflags=${rustflags_override:--C target-cpu=native} + +if ((${#requested_suites[@]} == 0)); then + requested_suites=(full) +fi +selected_suites=() +declare -A selected_labels=() +for request in "${requested_suites[@]}"; do + matched=false + for entry in "${suite_catalog[@]}"; do + IFS='|' read -r label _ _ groups <<<"$entry" + if [[ $request == "$label" || ,$groups, == *,$request,* ]]; then + matched=true + if [[ -z ${selected_labels[$label]:-} ]]; then + selected_suites+=("$entry") + selected_labels[$label]=1 + fi + fi + done + if [[ $matched == false ]]; then + echo "Unknown suite or benchmark label: $request" >&2 + exit 1 + fi +done + +common_suites=() +skipped_suites=() +for entry in "${selected_suites[@]}"; do + IFS='|' read -r label package bench _ <<<"$entry" + baseline_source="$baseline/$package/benches/$bench.rs" + candidate_source="$candidate/$package/benches/$bench.rs" + + if [[ -f $baseline_source && -f $candidate_source ]]; then + common_suites+=("$entry") + elif [[ -f $baseline_source ]]; then + skipped_suites+=("$label (baseline only)") + elif [[ -f $candidate_source ]]; then + skipped_suites+=("$label (candidate only)") + else + skipped_suites+=("$label (missing from both revisions)") + fi +done +if ((${#common_suites[@]} == 0)); then + echo "No requested benchmark targets exist in both revisions; no comparison is possible." >&2 + printf 'Skipped: %s\n' "${skipped_suites[@]}" >&2 + exit 1 +fi +selected_suites=("${common_suites[@]}") +if ((${#skipped_suites[@]} != 0)); then + printf 'Skipping one-sided benchmark target: %s\n' "${skipped_suites[@]}" >&2 +fi + +common_git_dir=$(git -C "$candidate" rev-parse --path-format=absolute --git-common-dir) +repository_root=$(dirname "$common_git_dir") +if [[ -z $target_root ]]; then + target_root="$repository_root/target/rowfn-benchmark/$(basename "$output")" +fi +target_root=$(realpath -m "$target_root") +if [[ -e $target_root ]]; then + echo "Target root already exists: $target_root" >&2 + exit 1 +fi + +mkdir -p "$output/build" "$output/warm" "$output/measured" "$target_root" +baseline_target="$target_root/baseline" +candidate_target="$target_root/candidate" +parser="$script_directory/rowfn_benchmark.py" + +{ + echo "RowFn benchmark machine record" + echo "Date: $(date --iso-8601=seconds)" + echo "Host: $(hostname)" + echo "Kernel: $(uname -srvmo)" + echo "Benchmark CPU: $bench_cpu" + echo "Configuration: $configuration" + echo "Cargo profile: bench, $codegen_units codegen units, LTO $lto" + echo "RUSTFLAGS: $rustflags" + echo "Warm runs: $warm_runs" + echo "Measured pairs: $measured_pairs" + echo "Divan: TSC timer, $sample_count samples, min $min_time s, max $max_time s" + if ((${#skipped_suites[@]} == 0)); then + echo "Skipped one-sided benchmark targets: none" + else + printf 'Skipped one-sided benchmark target: %s\n' "${skipped_suites[@]}" + fi + echo + rustc -vV + cargo -V + echo + lscpu + echo + rg -m1 '^microcode' /proc/cpuinfo || true + for path in \ + /sys/devices/system/cpu/cpu"$bench_cpu"/cpufreq/scaling_governor \ + /sys/devices/system/cpu/cpu"$bench_cpu"/cpufreq/energy_performance_preference \ + /sys/devices/system/cpu/cpufreq/boost; do + [[ -r $path ]] && echo "$path: $(<"$path")" + done +} >"$output/machine.txt" + +build_revision() { + local worktree=$1 + local target=$2 + local log=$3 + + ( + cd "$worktree" + export CARGO_TARGET_DIR=$target + export CARGO_PROFILE_BENCH_CODEGEN_UNITS=$codegen_units + export CARGO_PROFILE_BENCH_LTO=$lto + export RUSTFLAGS=$rustflags + for entry in "${selected_suites[@]}"; do + IFS='|' read -r _ package bench _ <<<"$entry" + cargo bench --no-run -j "$build_jobs" -p "$package" --bench "$bench" + done + ) >"$log" 2>&1 +} + +echo "Building baseline and candidate with $build_jobs jobs each." +build_revision "$baseline" "$baseline_target" "$output/build/baseline.txt" & +baseline_pid=$! +build_revision "$candidate" "$candidate_target" "$output/build/candidate.txt" & +candidate_pid=$! +baseline_status=0 +candidate_status=0 +wait "$baseline_pid" || baseline_status=$? +wait "$candidate_pid" || candidate_status=$? +if [[ $baseline_status -ne 0 || $candidate_status -ne 0 ]]; then + echo "Benchmark build failed; see $output/build/." >&2 + exit 1 +fi + +find_benchmark() { + local target=$1 + local name=$2 + local binary + + binary=$(find "$target/release/deps" -maxdepth 1 -type f -executable -name "$name-*" \ + -printf '%T@ %p\n' | sort -nr | head -n 1 | cut -d ' ' -f 2-) + [[ -n $binary ]] || { echo "Cannot find benchmark $name under $target." >&2; exit 1; } + echo "$binary" +} + +declare -A baseline_binaries=() +declare -A candidate_binaries=() +manifest_args=( + manifest + --output "$output/manifest.json" + --machine-record "$output/machine.txt" + --baseline-worktree "$baseline" + --candidate-worktree "$candidate" + --baseline-target "$baseline_target" + --candidate-target "$candidate_target" + --setting "configuration=$configuration" + --setting "codegen_units=$codegen_units" + --setting "lto=$lto" + --setting "rustflags=$rustflags" + --setting "bench_cpu=$bench_cpu" + --setting "warm_runs=$warm_runs" + --setting "measured_pairs=$measured_pairs" + --setting "sample_count=$sample_count" + --setting "min_time=$min_time" + --setting "max_time=$max_time" +) +for filter in "${filters[@]}"; do + manifest_args+=(--filter "$filter") +done +for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ bench _ <<<"$entry" + baseline_binaries[$label]=$(find_benchmark "$baseline_target" "$bench") + candidate_binaries[$label]=$(find_benchmark "$candidate_target" "$bench") + manifest_args+=( + --suite "$label" + --baseline-binary "$label=${baseline_binaries[$label]}" + --candidate-binary "$label=${candidate_binaries[$label]}" + ) +done +python3 "$parser" "${manifest_args[@]}" + +run_suite() { + local revision=$1 + local label=$2 + local destination=$3 + local binary + local command + + if [[ $revision == baseline ]]; then + binary=${baseline_binaries[$label]} + else + binary=${candidate_binaries[$label]} + fi + command=( + taskset -c "$bench_cpu" "$binary" + --bench --timer tsc --sample-count "$sample_count" + --min-time "$min_time" --max-time "$max_time" --color never + "${filters[@]}" + ) + echo "Running $label ($revision) -> $destination" + "${command[@]}" >"$destination" 2>&1 +} + +echo "Waiting for the global timed benchmark lock: $lock_file" +exec {benchmark_lock}>"$lock_file" +flock "$benchmark_lock" +if pgrep -x cargo >/dev/null || pgrep -x rustc >/dev/null; then + echo "Cargo or rustc is active after acquiring the benchmark lock; refusing to measure." >&2 + exit 1 +fi + +for ((round = 1; round <= warm_runs; round++)); do + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + if ((round % 2 == 1)); then + run_suite baseline "$label" "$output/warm/$label-baseline-$round.txt" + run_suite candidate "$label" "$output/warm/$label-candidate-$round.txt" + else + run_suite candidate "$label" "$output/warm/$label-candidate-$round.txt" + run_suite baseline "$label" "$output/warm/$label-baseline-$round.txt" + fi + done +done + +for ((pair = 1; pair <= measured_pairs; pair++)); do + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + if ((pair % 2 == 1)); then + run_suite baseline "$label" "$output/measured/$label-baseline-$pair.txt" + run_suite candidate "$label" "$output/measured/$label-candidate-$pair.txt" + else + run_suite candidate "$label" "$output/measured/$label-candidate-$pair.txt" + run_suite baseline "$label" "$output/measured/$label-baseline-$pair.txt" + fi + done +done + +python3 "$parser" summarize "$output" +echo "Raw results: $output" +echo "Summary: $output/summary.md" diff --git a/scripts/rowfn_benchmark.py b/scripts/rowfn_benchmark.py new file mode 100755 index 00000000000..b2f778ea3f7 --- /dev/null +++ b/scripts/rowfn_benchmark.py @@ -0,0 +1,356 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Capture and summarize evidence from ``benchmark-rowfn.sh`` runs.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import re +import statistics +import subprocess +from collections.abc import Iterable +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from pathlib import Path + +RESULT_FILE = re.compile(r"^(?P.+)-(?Pbaseline|candidate)-(?P\d+)\.txt$") +TREE_ROW = re.compile(r"^(?P(?:│ | )*)(?:├─ |╰─ )(?P.*)$") +TIMING = re.compile(r"(?P\d+(?:\.\d+)?)\s*(?Pps|ns|µs|us|ms|s)\s*$") +UNIT_TO_NS = { + "ps": 0.001, + "ns": 1.0, + "µs": 1_000.0, + "us": 1_000.0, + "ms": 1_000_000.0, + "s": 1_000_000_000.0, +} + + +@dataclass(frozen=True) +class BenchmarkSummary: + suite: str + benchmark: str + pairs: int + baseline_median_ns: float + candidate_median_ns: float + median_ratio: float + minimum_ratio: float + maximum_ratio: float + ratio_mad: float + + +def run_git(worktree: Path, *args: str, binary: bool = False) -> str | bytes: + """Run one read-only Git command in ``worktree``.""" + + result = subprocess.run( + ["git", "-C", str(worktree), *args], + check=True, + capture_output=True, + text=not binary, + ) + return result.stdout if binary else result.stdout.strip() + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + + return digest.hexdigest() + + +def revision_record(worktree: Path, target: Path, binaries: Iterable[str]) -> dict[str, object]: + """Describe the exact revision, dirty patch, targets, and benchmark executables.""" + + status = str(run_git(worktree, "status", "--short")).splitlines() + diff = run_git(worktree, "diff", "--binary", "HEAD", binary=True) + assert isinstance(diff, bytes) + + untracked = run_git(worktree, "ls-files", "--others", "--exclude-standard", "-z", binary=True) + assert isinstance(untracked, bytes) + dirty_digest = hashlib.sha256(diff) + dirty_digest.update(untracked) + for relative_path in filter(None, untracked.decode().split("\0")): + path = worktree / relative_path + if path.is_file(): + dirty_digest.update(relative_path.encode()) + dirty_digest.update(bytes.fromhex(sha256_file(path))) + + executable_records: dict[str, object] = {} + for entry in binaries: + label, separator, raw_path = entry.partition("=") + if not separator: + raise ValueError(f"expected LABEL=PATH for benchmark binary, got {entry!r}") + path = Path(raw_path).resolve() + executable_records[label] = { + "path": str(path), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + + return { + "worktree": str(worktree.resolve()), + "head": run_git(worktree, "rev-parse", "HEAD"), + "changed_paths": status, + "tracked_diff_sha256": hashlib.sha256(diff).hexdigest(), + "dirty_state_sha256": dirty_digest.hexdigest(), + "target": str(target.resolve()), + "binaries": executable_records, + } + + +def write_manifest(args: argparse.Namespace) -> None: + settings = dict(setting.split("=", 1) for setting in args.setting) + manifest = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + "settings": settings, + "suites": args.suite, + "filters": args.filter, + "machine_record": str(Path(args.machine_record).resolve()), + "baseline": revision_record( + Path(args.baseline_worktree), + Path(args.baseline_target), + args.baseline_binary, + ), + "candidate": revision_record( + Path(args.candidate_worktree), + Path(args.candidate_target), + args.candidate_binary, + ), + } + output = Path(args.output) + output.write_text(f"{json.dumps(manifest, indent=2, sort_keys=True)}\n", encoding="utf-8") + + +def timing_ns(field: str) -> float: + match = TIMING.search(field.strip()) + if match is None: + raise ValueError(f"cannot parse Divan timing from {field!r}") + + return float(match.group("value")) * UNIT_TO_NS[match.group("unit")] + + +def parse_divan(path: Path) -> dict[str, float]: + """Return benchmark paths and median nanoseconds from one Divan table.""" + + parents: dict[int, str] = {} + timings: dict[str, float] = {} + for line in path.read_text(encoding="utf-8").splitlines(): + fields = re.split(r"\s+│\s+", line) + tree_match = TREE_ROW.match(fields[0]) + if tree_match is None: + continue + + depth = len(tree_match.group("prefix")) // 3 + body = tree_match.group("body").rstrip() + timing_match = TIMING.search(body) + name = body[: timing_match.start()].rstrip() if timing_match else body.strip() + parents = {level: parent for level, parent in parents.items() if level < depth} + + if timing_match is None: + parents[depth] = name + continue + if len(fields) < 3: + raise ValueError(f"timed Divan row has no median column in {path}: {line}") + + components = [parents[level] for level in sorted(parents) if level < depth] + benchmark = "/".join([*components, name]) + if benchmark in timings: + raise ValueError(f"duplicate benchmark {benchmark!r} in {path}") + timings[benchmark] = timing_ns(fields[2]) + + if not timings: + raise ValueError(f"no Divan benchmark timings found in {path}") + + return timings + + +def read_measurements(directory: Path) -> dict[tuple[str, str, int, str], float]: + measurements: dict[tuple[str, str, int, str], float] = {} + for path in sorted(directory.glob("*.txt")): + match = RESULT_FILE.match(path.name) + if match is None: + continue + suite = match.group("suite") + revision = match.group("revision") + pair = int(match.group("pair")) + for benchmark, median_ns in parse_divan(path).items(): + measurements[suite, revision, pair, benchmark] = median_ns + + if not measurements: + raise ValueError(f"no measured result files found in {directory}") + + return measurements + + +def summarize(measurements: dict[tuple[str, str, int, str], float]) -> list[BenchmarkSummary]: + inventories: dict[tuple[str, str], set[str]] = {} + for suite, revision, _, benchmark in measurements: + inventories.setdefault((suite, revision), set()).add(benchmark) + + suites = {suite for suite, _ in inventories} + comparable = { + (suite, benchmark) + for suite in suites + for benchmark in inventories.get((suite, "baseline"), set()) & inventories.get((suite, "candidate"), set()) + } + groups = { + (suite, pair, benchmark) for suite, _, pair, benchmark in measurements if (suite, benchmark) in comparable + } + incomplete = [ + group + for group in groups + if (group[0], "baseline", group[1], group[2]) not in measurements + or (group[0], "candidate", group[1], group[2]) not in measurements + ] + if incomplete: + raise ValueError(f"unpaired benchmark measurements: {sorted(incomplete)!r}") + if not groups: + raise ValueError("unpaired benchmark measurements: no comparable benchmarks") + + by_benchmark: dict[tuple[str, str], list[tuple[float, float]]] = {} + for suite, pair, benchmark in sorted(groups): + baseline = measurements[suite, "baseline", pair, benchmark] + candidate = measurements[suite, "candidate", pair, benchmark] + by_benchmark.setdefault((suite, benchmark), []).append((baseline, candidate)) + + summaries = [] + for (suite, benchmark), pairs in sorted(by_benchmark.items()): + baseline_values = [baseline for baseline, _ in pairs] + candidate_values = [candidate for _, candidate in pairs] + ratios = [candidate / baseline for baseline, candidate in pairs] + median_ratio = statistics.median(ratios) + summaries.append( + BenchmarkSummary( + suite=suite, + benchmark=benchmark, + pairs=len(pairs), + baseline_median_ns=statistics.median(baseline_values), + candidate_median_ns=statistics.median(candidate_values), + median_ratio=median_ratio, + minimum_ratio=min(ratios), + maximum_ratio=max(ratios), + ratio_mad=statistics.median(abs(ratio - median_ratio) for ratio in ratios), + ) + ) + + return summaries + + +def inventory_differences( + measurements: dict[tuple[str, str, int, str], float], +) -> list[tuple[str, str, str]]: + """Return benchmarks that exist in only one revision.""" + + inventories: dict[tuple[str, str], set[str]] = {} + for suite, revision, _, benchmark in measurements: + inventories.setdefault((suite, revision), set()).add(benchmark) + + differences = [] + for suite in sorted({suite for suite, _ in inventories}): + baseline = inventories.get((suite, "baseline"), set()) + candidate = inventories.get((suite, "candidate"), set()) + differences.extend((suite, "baseline only", benchmark) for benchmark in baseline - candidate) + differences.extend((suite, "candidate only", benchmark) for benchmark in candidate - baseline) + + return sorted(differences) + + +def format_ns(value: float) -> str: + for divisor, unit in ((1_000_000_000, "s"), (1_000_000, "ms"), (1_000, "µs")): + if value >= divisor: + return f"{value / divisor:.3f} {unit}" + + return f"{value:.3f} ns" + + +def write_summary( + output_directory: Path, + summaries: list[BenchmarkSummary], + differences: Iterable[tuple[str, str, str]] = (), +) -> None: + csv_path = output_directory / "ratios.csv" + with csv_path.open("w", encoding="utf-8", newline="") as file: + writer = csv.DictWriter(file, fieldnames=list(asdict(summaries[0]))) + writer.writeheader() + writer.writerows(asdict(summary) for summary in summaries) + + markdown = [ + "# RowFn benchmark comparison", + "", + "Ratios are paired candidate/baseline medians. Lower is faster.", + "", + "| Suite | Benchmark | Pairs | Baseline | Candidate | Ratio | Change | MAD |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for summary in sorted(summaries, key=lambda result: result.median_ratio, reverse=True): + change = (summary.median_ratio - 1.0) * 100.0 + markdown.append( + f"| {summary.suite} | `{summary.benchmark}` | {summary.pairs} " + f"| {format_ns(summary.baseline_median_ns)} " + f"| {format_ns(summary.candidate_median_ns)} " + f"| {summary.median_ratio:.6f} | {change:+.2f}% | {summary.ratio_mad:.6f} |" + ) + differences = list(differences) + if differences: + markdown.extend( + [ + "", + "## Unpaired benchmark inventory", + "", + "These benchmarks were recorded for only one revision and are excluded from ratios.", + "", + ] + ) + markdown.extend(f"- `{suite}/{benchmark}`: {revision}." for suite, revision, benchmark in differences) + markdown.append("") + (output_directory / "summary.md").write_text("\n".join(markdown), encoding="utf-8") + + +def summarize_directory(args: argparse.Namespace) -> None: + output_directory = Path(args.output_directory) + measurements = read_measurements(output_directory / "measured") + summaries = summarize(measurements) + write_summary(output_directory, summaries, inventory_differences(measurements)) + + +def argument_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(required=True) + + manifest = subparsers.add_parser("manifest", help="capture revisions and executable hashes") + manifest.add_argument("--output", required=True) + manifest.add_argument("--machine-record", required=True) + manifest.add_argument("--baseline-worktree", required=True) + manifest.add_argument("--candidate-worktree", required=True) + manifest.add_argument("--baseline-target", required=True) + manifest.add_argument("--candidate-target", required=True) + manifest.add_argument("--setting", action="append", default=[]) + manifest.add_argument("--suite", action="append", default=[]) + manifest.add_argument("--filter", action="append", default=[]) + manifest.add_argument("--baseline-binary", action="append", default=[]) + manifest.add_argument("--candidate-binary", action="append", default=[]) + manifest.set_defaults(function=write_manifest) + + summary = subparsers.add_parser("summarize", help="write ratios.csv and summary.md") + summary.add_argument("output_directory") + summary.set_defaults(function=summarize_directory) + + return parser + + +def main() -> None: + args = argument_parser().parse_args() + args.function(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/tests/test_rowfn_benchmark.py b/scripts/tests/test_rowfn_benchmark.py new file mode 100644 index 00000000000..308580f8d68 --- /dev/null +++ b/scripts/tests/test_rowfn_benchmark.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +SCRIPT = REPO_ROOT / "scripts" / "rowfn_benchmark.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("rowfn_benchmark", SCRIPT) + assert spec is not None + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def write_divan(path: Path, rows: list[str]) -> None: + path.write_text( + "\n".join( + [ + "Timer precision: 20 ns", + "bench fastest │ slowest │ median │ mean │ samples │ iters", + *rows, + "", + ] + ), + encoding="utf-8", + ) + + +class RowFnBenchmarkTest(unittest.TestCase): + def setUp(self) -> None: + self.module = load_module() + self.temporary_directory = tempfile.TemporaryDirectory() + self.directory = Path(self.temporary_directory.name) + + def tearDown(self) -> None: + self.temporary_directory.cleanup() + + def test_parse_divan_preserves_nested_benchmark_names_and_converts_units(self) -> None: + output = self.directory / "result.txt" + write_divan( + output, + [ + "├─ non_nullable │ │ │ │ │", + "│ ├─ 2 17.18 µs │ 18 µs │ 17.33 µs │ 17.4 µs │ 100 │ 100", + "│ ╰─ 32 6.709 µs │ 8 µs │ 6.829 µs │ 7 µs │ 100 │ 100", + "╰─ nullable │ │ │ │ │", + " ╰─ 2 799.7 ns │ 1 µs │ 979.7 ns │ 986 ns │ 100 │ 100", + ], + ) + + self.assertEqual( + self.module.parse_divan(output), + { + "non_nullable/2": 17_330.0, + "non_nullable/32": 6_829.0, + "nullable/2": 979.7, + }, + ) + + def test_summarize_writes_paired_ratios_and_slowest_first_markdown(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + results = { + "numeric-baseline-1.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-1.txt": ["├─ add 12 ns │ 12 ns │ 12 ns │ 12 ns │ 100 │ 100"], + "numeric-baseline-2.txt": ["├─ add 20 ns │ 20 ns │ 20 ns │ 20 ns │ 100 │ 100"], + "numeric-candidate-2.txt": ["├─ add 18 ns │ 18 ns │ 18 ns │ 18 ns │ 100 │ 100"], + "numeric-baseline-3.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-3.txt": ["├─ add 11 ns │ 11 ns │ 11 ns │ 11 ns │ 100 │ 100"], + "numeric-baseline-4.txt": ["├─ mul 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-4.txt": ["├─ mul 9 ns │ 9 ns │ 9 ns │ 9 ns │ 100 │ 100"], + } + for name, rows in results.items(): + write_divan(measured / name, rows) + + summaries = self.module.summarize(self.module.read_measurements(measured)) + self.module.write_summary(self.directory, summaries) + + add = next(summary for summary in summaries if summary.benchmark == "add") + self.assertEqual(add.pairs, 3) + self.assertAlmostEqual(add.median_ratio, 1.1) + self.assertAlmostEqual(add.ratio_mad, 0.1) + + csv_output = (self.directory / "ratios.csv").read_text(encoding="utf-8") + markdown = (self.directory / "summary.md").read_text(encoding="utf-8") + self.assertIn("suite,benchmark,pairs", csv_output) + self.assertLess(markdown.index("`add`"), markdown.index("`mul`")) + + def test_summarize_rejects_unpaired_measurements(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + write_divan( + measured / "numeric-baseline-1.txt", + ["╰─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + ) + + with self.assertRaisesRegex(ValueError, "unpaired benchmark measurements"): + self.module.summarize(self.module.read_measurements(measured)) + + def test_summarize_excludes_and_reports_revision_only_benchmarks(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + results = { + "numeric-baseline-1.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-1.txt": [ + "├─ add 11 ns │ 11 ns │ 11 ns │ 11 ns │ 100 │ 100", + "╰─ candidate 5 ns │ 5 ns │ 5 ns │ 5 ns │ 100 │ 100", + ], + } + for name, rows in results.items(): + write_divan(measured / name, rows) + + measurements = self.module.read_measurements(measured) + summaries = self.module.summarize(measurements) + differences = self.module.inventory_differences(measurements) + self.module.write_summary(self.directory, summaries, differences) + + self.assertEqual([summary.benchmark for summary in summaries], ["add"]) + self.assertEqual(differences, [("numeric", "candidate only", "candidate")]) + markdown = (self.directory / "summary.md").read_text(encoding="utf-8") + self.assertIn("`numeric/candidate`: candidate only.", markdown) + + +if __name__ == "__main__": + unittest.main() From e24d2a8a2948b697bc9b0dbd3d9093e07bd0ead9 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 08:08:32 -0400 Subject: [PATCH 144/160] Add focused RowFn framework benchmarks Signed-off-by: Connor Tsui --- vortex-array/Cargo.toml | 8 + vortex-array/benches/row_fn_executor.rs | 287 ++++++++++++++++++++++++ vortex-array/benches/strict_validity.rs | 217 ++++++++++++++++++ 3 files changed, 512 insertions(+) create mode 100644 vortex-array/benches/row_fn_executor.rs create mode 100644 vortex-array/benches/strict_validity.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 2af2eacf238..68bd189ef11 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -134,6 +134,10 @@ harness = false name = "binary_ops" harness = false +[[bench]] +name = "row_fn_executor" +harness = false + [[bench]] name = "kleene_bool" harness = false @@ -213,6 +217,10 @@ harness = false name = "validity_is_valid" harness = false +[[bench]] +name = "strict_validity" +harness = false + [[bench]] name = "dict_unreferenced_mask" harness = false diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs new file mode 100644 index 00000000000..bff6a29d2da --- /dev/null +++ b/vortex-array/benches/row_fn_executor.rs @@ -0,0 +1,287 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares owned-output, sink-writing, and hand-written primitive row loops. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::OutputSink; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const ROWS: usize = 65_536; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +#[derive(Clone)] +struct RowWrappingAdd; + +impl RowFn for RowWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_wrapping_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64), i64>(|(lhs, rhs)| lhs.wrapping_add(rhs)) + } +} + +#[derive(Clone)] +struct RowCheckedAdd; + +impl RowFn for RowCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_checked_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(i64, i64), i64, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(checked_add_error()); + } + Ok(()) + }, + ) + } +} + +/// Keep error construction out of the benchmarked success path. +#[cold] +#[inline(never)] +fn checked_add_error() -> VortexError { + vortex_err!("integer overflow in row checked add") +} + +/// A benchmark sink that writes one `i64` per row. +struct I64Sink( + /// The output values written by the row loop. + BufferMut, +); + +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +#[derive(Clone)] +struct RowSinkWrappingAdd; + +impl RowFn for RowSinkWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_wrapping_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64, i64), I64Sink, _>(|(lhs, rhs), out| { + *out = lhs.wrapping_add(rhs); + }) + } +} + +vortex_array::impl_row_fn_vtable!(RowWrappingAdd); +vortex_array::impl_row_fn_vtable!(RowCheckedAdd); +vortex_array::impl_row_fn_vtable!(RowSinkWrappingAdd); + +fn inputs() -> (ArrayRef, ArrayRef) { + let lhs = (0..ROWS) + .map(|index| index as i64) + .collect::>() + .into_array(); + let rhs = (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>() + .into_array(); + (lhs, rhs) +} + +fn constant_inputs() -> (ArrayRef, ArrayRef) { + let (lhs, _) = inputs(); + let rhs = ConstantArray::new(Scalar::from(7i64), ROWS).into_array(); + (lhs, rhs) +} + +fn nullable_inputs() -> (ArrayRef, ArrayRef) { + let lhs = PrimitiveArray::new( + (0..ROWS).map(|index| index as i64).collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(5))), + ) + .into_array(); + let rhs = PrimitiveArray::new( + (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(7))), + ) + .into_array(); + (lhs, rhs) +} + +fn bench_row_fn(bencher: Bencher, function: F, make_inputs: fn() -> (ArrayRef, ArrayRef)) +where + F: RowFn + ScalarFnVTable, +{ + bencher + .with_inputs(make_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + function + .clone() + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, inputs); +} + +#[divan::bench] +fn handrolled_sink_wrapping_add(bencher: Bencher) { + bencher + .with_inputs(inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = lhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let rhs = rhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let mut output = BufferMut::zeroed(ROWS); + for ((out, lhs), rhs) in output + .as_mut_slice() + .iter_mut() + .zip(lhs.as_slice()) + .zip(rhs.as_slice()) + { + *out = lhs.wrapping_add(*rhs); + } + PrimitiveArray::new(output.freeze(), Validity::NonNullable).into_array() + }); +} + +#[divan::bench] +fn row_checked_add(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, inputs); +} + +#[divan::bench] +fn row_wrapping_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, nullable_inputs); +} + +#[divan::bench] +fn row_wrapping_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, nullable_inputs); +} diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs new file mode 100644 index 00000000000..e0a8b6a565e --- /dev/null +++ b/vortex-array/benches/strict_validity.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks how the row lifting applies validity to a dense kernel. +//! +//! Both arms run the same kernel over the same nullable input and differ only in how the output's +//! nulls are restored: +//! +//! - `lazy` is what the lifting behind [`RowFn`] does: conjoin the input validities (itself lazy) +//! and hand the resulting boolean array straight to `mask`. +//! - `eager` is the strategy it replaced: materialize the conjunction into a `Mask`, copy it into a +//! `BitBuffer`, wrap that in a `BoolArray`, and mask with it. +//! +//! `chain` composes three of the same function, which is where a per-call materialization +//! compounds. + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowExecution; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const SIZES: &[usize] = &[65_536, 1 << 20]; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +/// The shared kernel: double every lane, ignoring validity. +fn doubled(input: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let input = input.clone().execute::(ctx)?; + let values: Buffer = input + .as_slice::() + .iter() + .map(|value| value.wrapping_mul(2)) + .collect(); + Ok(PrimitiveArray::new(values, Validity::NonNullable).into_array()) +} + +/// Dense row function: the lifting decides how validity is applied. +/// +/// Its [`reduce_encoded`](RowFn::reduce_encoded) answers with [`doubled`] before the row loop is +/// reached, so this arm runs exactly the kernel [`EagerDouble`] runs and the two differ only in +/// validity. The row closure below is what makes it a [`RowFn`] at all, and never executes. +#[derive(Clone)] +struct LazyDouble; + +impl RowFn for LazyDouble { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.lazy_double"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i32,), i32>(|(value,)| value.wrapping_mul(2)) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + doubled(&args[0], ctx).map(|output| Some(RowExecution::Output(output))) + } +} + +vortex_array::impl_row_fn_vtable!(LazyDouble); + +/// The same function, applying validity the way the adapter used to: materialize a mask first. +#[derive(Clone)] +struct EagerDouble; + +impl ScalarFnVTable for EagerDouble { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.eager_double"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { + ChildName::from("input") + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + Ok(DType::Primitive(PType::I32, args[0].nullability())) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let valid = input.validity()?.execute_mask(args.row_count(), ctx)?; + let values = doubled(&input, ctx)?; + + if valid.all_true() { + return values.cast(DType::Primitive(PType::I32, input.dtype().nullability())); + } + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + values.mask(mask) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// A nullable i32 column with array-backed validity (~10% nulls). +fn nullable_input(len: usize) -> ArrayRef { + PrimitiveArray::new( + (0..len as i32).collect::>(), + Validity::from_iter((0..len).map(|index| !index.is_multiple_of(10))), + ) + .into_array() +} + +fn bench_depth(bencher: Bencher, function: F, len: usize, depth: usize) +where + F: ScalarFnVTable + Clone, +{ + bencher + .with_inputs(|| nullable_input(len)) + .bench_local_values(|input| { + let mut ctx = SESSION.create_execution_ctx(); + let mut array = input; + for _ in 0..depth { + array = function + .clone() + .try_new_array(len, EmptyOptions, [array]) + .unwrap() + .into_array(); + } + array.execute::(&mut ctx).unwrap() + }); +} + +#[divan::bench(args = SIZES)] +fn lazy(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn eager(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn lazy_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 3); +} + +#[divan::bench(args = SIZES)] +fn eager_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 3); +} From 5c211e199e2ac5231ae95189958f1fa8b80b32c2 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 10:56:03 -0400 Subject: [PATCH 145/160] Reuse RowFn benchmark builds safely Signed-off-by: Connor Tsui --- scripts/benchmark-rowfn.sh | 180 ++++++++++++++++++++++---- scripts/rowfn_benchmark.py | 130 ++++++++++++++++++- scripts/tests/test_rowfn_benchmark.py | 63 +++++++++ 3 files changed, 344 insertions(+), 29 deletions(-) diff --git a/scripts/benchmark-rowfn.sh b/scripts/benchmark-rowfn.sh index 4c6ae4152cc..790b10462f1 100755 --- a/scripts/benchmark-rowfn.sh +++ b/scripts/benchmark-rowfn.sh @@ -14,8 +14,13 @@ Usage: benchmark-rowfn.sh [OPTIONS] &2 + exit 1 + fi + run_measure=false + shift + ;; + --measure-only) + if [[ $run_measure == false ]]; then + echo "--build-only and --measure-only are mutually exclusive." >&2 + exit 1 + fi + run_build=false + shift + ;; --config) configuration=$2; shift 2 ;; --target-root) target_root=$2; shift 2 ;; + --baseline-target) baseline_target_override=$2; shift 2 ;; + --candidate-target) candidate_target_override=$2; shift 2 ;; --codegen-units) codegen_units_override=$2; shift 2 ;; --lto) lto_override=$2; shift 2 ;; --rustflags) rustflags_override=$2; shift 2 ;; @@ -105,7 +132,9 @@ if ((build_jobs < 1 || build_jobs > 8)); then echo "--build-jobs must be between 1 and 8 so two builds cannot exceed 16 jobs." >&2 exit 1 fi -command -v flock >/dev/null || { echo "benchmark-rowfn.sh requires flock." >&2; exit 1; } +if [[ $run_measure == true ]]; then + command -v flock >/dev/null || { echo "benchmark-rowfn.sh requires flock." >&2; exit 1; } +fi baseline=$(realpath "$1") candidate=$(realpath "$2") @@ -176,18 +205,28 @@ fi common_git_dir=$(git -C "$candidate" rev-parse --path-format=absolute --git-common-dir) repository_root=$(dirname "$common_git_dir") +if [[ -n $target_root && (-n $baseline_target_override || -n $candidate_target_override) ]]; then + echo "--target-root cannot be combined with revision-specific target paths." >&2 + exit 1 +fi if [[ -z $target_root ]]; then target_root="$repository_root/target/rowfn-benchmark/$(basename "$output")" fi target_root=$(realpath -m "$target_root") -if [[ -e $target_root ]]; then - echo "Target root already exists: $target_root" >&2 +baseline_target=$(realpath -m "${baseline_target_override:-$target_root/baseline}") +candidate_target=$(realpath -m "${candidate_target_override:-$target_root/candidate}") +if [[ $baseline_target == "$candidate_target" ]]; then + echo "Baseline and candidate must use different Cargo target directories." >&2 exit 1 fi -mkdir -p "$output/build" "$output/warm" "$output/measured" "$target_root" -baseline_target="$target_root/baseline" -candidate_target="$target_root/candidate" +mkdir -p "$output" +if [[ $run_build == true ]]; then + mkdir -p "$output/build" "$baseline_target" "$candidate_target" +fi +if [[ $run_measure == true ]]; then + mkdir -p "$output/warm" "$output/measured" +fi parser="$script_directory/rowfn_benchmark.py" { @@ -208,8 +247,11 @@ parser="$script_directory/rowfn_benchmark.py" printf 'Skipped one-sided benchmark target: %s\n' "${skipped_suites[@]}" fi echo - rustc -vV - cargo -V + echo "Baseline toolchain:" + (cd "$baseline" && rustc -vV && cargo -V) + echo + echo "Candidate toolchain:" + (cd "$candidate" && rustc -vV && cargo -V) echo lscpu echo @@ -233,25 +275,37 @@ build_revision() { export CARGO_PROFILE_BENCH_CODEGEN_UNITS=$codegen_units export CARGO_PROFILE_BENCH_LTO=$lto export RUSTFLAGS=$rustflags - for entry in "${selected_suites[@]}"; do - IFS='|' read -r _ package bench _ <<<"$entry" - cargo bench --no-run -j "$build_jobs" -p "$package" --bench "$bench" + for package in vortex-array vortex-tensor vortex-spatial; do + local command=(cargo bench --no-run -j "$build_jobs" -p "$package") + local has_bench=false + for entry in "${selected_suites[@]}"; do + IFS='|' read -r _ suite_package bench _ <<<"$entry" + if [[ $suite_package == "$package" ]]; then + command+=(--bench "$bench") + has_bench=true + fi + done + if [[ $has_bench == true ]]; then + "${command[@]}" + fi done ) >"$log" 2>&1 } -echo "Building baseline and candidate with $build_jobs jobs each." -build_revision "$baseline" "$baseline_target" "$output/build/baseline.txt" & -baseline_pid=$! -build_revision "$candidate" "$candidate_target" "$output/build/candidate.txt" & -candidate_pid=$! -baseline_status=0 -candidate_status=0 -wait "$baseline_pid" || baseline_status=$? -wait "$candidate_pid" || candidate_status=$? -if [[ $baseline_status -ne 0 || $candidate_status -ne 0 ]]; then - echo "Benchmark build failed; see $output/build/." >&2 - exit 1 +if [[ $run_build == true ]]; then + echo "Building baseline and candidate with $build_jobs jobs each." + build_revision "$baseline" "$baseline_target" "$output/build/baseline.txt" & + baseline_pid=$! + build_revision "$candidate" "$candidate_target" "$output/build/candidate.txt" & + candidate_pid=$! + baseline_status=0 + candidate_status=0 + wait "$baseline_pid" || baseline_status=$? + wait "$candidate_pid" || candidate_status=$? + if [[ $baseline_status -ne 0 || $candidate_status -ne 0 ]]; then + echo "Benchmark build failed; see $output/build/." >&2 + exit 1 + fi fi find_benchmark() { @@ -267,6 +321,73 @@ find_benchmark() { declare -A baseline_binaries=() declare -A candidate_binaries=() +build_settings=( + --setting "configuration=$configuration" + --setting "codegen_units=$codegen_units" + --setting "lto=$lto" + --setting "rustflags=$rustflags" +) + +record_build() { + local revision=$1 + local worktree=$2 + local target=$3 + local metadata="$target/rowfn-benchmark-build.json" + local arguments=( + record-build + --output "$metadata" + --worktree "$worktree" + --target "$target" + "${build_settings[@]}" + ) + + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ bench _ <<<"$entry" + local binary + binary=$(find_benchmark "$target" "$bench") + arguments+=(--binary "$label=$binary") + done + python3 "$parser" "${arguments[@]}" + echo "Recorded $revision build metadata: $metadata" +} + +if [[ $run_build == true ]]; then + record_build baseline "$baseline" "$baseline_target" + record_build candidate "$candidate" "$candidate_target" +fi + +load_binaries() { + local worktree=$1 + local target=$2 + local metadata="$target/rowfn-benchmark-build.json" + local arguments=( + validate-build + --metadata "$metadata" + --worktree "$worktree" + --target "$target" + "${build_settings[@]}" + ) + + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + arguments+=(--suite "$label") + done + python3 "$parser" "${arguments[@]}" +} + +baseline_binary_output=$(load_binaries "$baseline" "$baseline_target") +candidate_binary_output=$(load_binaries "$candidate" "$candidate_target") +mapfile -t baseline_binary_records <<<"$baseline_binary_output" +mapfile -t candidate_binary_records <<<"$candidate_binary_output" +for record in "${baseline_binary_records[@]}"; do + label=${record%%=*} + baseline_binaries[$label]=${record#*=} +done +for record in "${candidate_binary_records[@]}"; do + label=${record%%=*} + candidate_binaries[$label]=${record#*=} +done + manifest_args=( manifest --output "$output/manifest.json" @@ -290,9 +411,7 @@ for filter in "${filters[@]}"; do manifest_args+=(--filter "$filter") done for entry in "${selected_suites[@]}"; do - IFS='|' read -r label _ bench _ <<<"$entry" - baseline_binaries[$label]=$(find_benchmark "$baseline_target" "$bench") - candidate_binaries[$label]=$(find_benchmark "$candidate_target" "$bench") + IFS='|' read -r label _ _ _ <<<"$entry" manifest_args+=( --suite "$label" --baseline-binary "$label=${baseline_binaries[$label]}" @@ -301,6 +420,13 @@ for entry in "${selected_suites[@]}"; do done python3 "$parser" "${manifest_args[@]}" +if [[ $run_measure == false ]]; then + echo "Build evidence: $output" + echo "Baseline target: $baseline_target" + echo "Candidate target: $candidate_target" + exit 0 +fi + run_suite() { local revision=$1 local label=$2 diff --git a/scripts/rowfn_benchmark.py b/scripts/rowfn_benchmark.py index b2f778ea3f7..35c08918d75 100755 --- a/scripts/rowfn_benchmark.py +++ b/scripts/rowfn_benchmark.py @@ -66,6 +66,22 @@ def sha256_file(path: Path) -> str: return digest.hexdigest() +def toolchain_record(worktree: Path) -> dict[str, str]: + """Capture the tools selected from a revision's working directory.""" + + def version(*command: str) -> str: + result = subprocess.run( + command, + cwd=worktree, + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + return {"rustc": version("rustc", "-vV"), "cargo": version("cargo", "-V")} + + def revision_record(worktree: Path, target: Path, binaries: Iterable[str]) -> dict[str, object]: """Describe the exact revision, dirty patch, targets, and benchmark executables.""" @@ -106,6 +122,96 @@ def revision_record(worktree: Path, target: Path, binaries: Iterable[str]) -> di } +def build_identity(worktree: Path, target: Path, settings: dict[str, str]) -> dict[str, object]: + revision = revision_record(worktree, target, []) + revision.pop("binaries") + return { + "settings": settings, + "toolchain": toolchain_record(worktree), + "revision": revision, + } + + +def binary_records(entries: Iterable[str]) -> dict[str, object]: + records: dict[str, object] = {} + for entry in entries: + label, separator, raw_path = entry.partition("=") + if not separator: + raise ValueError(f"expected LABEL=PATH for benchmark binary, got {entry!r}") + path = Path(raw_path).resolve() + records[label] = { + "path": str(path), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + return records + + +def write_build_record(args: argparse.Namespace) -> None: + output = Path(args.output) + worktree = Path(args.worktree) + target = Path(args.target) + settings = dict(setting.split("=", 1) for setting in args.setting) + identity = build_identity(worktree, target, settings) + binaries = binary_records(args.binary) + + if output.exists(): + previous = json.loads(output.read_text(encoding="utf-8")) + previous_identity = {key: previous.get(key) for key in identity} + if previous_identity == identity: + binaries = {**previous.get("binaries", {}), **binaries} + + record = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + **identity, + "binaries": binaries, + } + output.write_text(f"{json.dumps(record, indent=2, sort_keys=True)}\n", encoding="utf-8") + + +def validated_build_binaries(args: argparse.Namespace) -> dict[str, str]: + metadata = Path(args.metadata) + if not metadata.is_file(): + raise ValueError(f"build metadata does not exist: {metadata}") + + record = json.loads(metadata.read_text(encoding="utf-8")) + if record.get("schema_version") != 1: + raise ValueError(f"unsupported build metadata schema in {metadata}") + + settings = dict(setting.split("=", 1) for setting in args.setting) + identity = build_identity(Path(args.worktree), Path(args.target), settings) + mismatches = [key for key in identity if record.get(key) != identity[key]] + if mismatches: + fields = ", ".join(mismatches) + raise ValueError(f"stale benchmark build metadata ({fields} changed): {metadata}") + + binaries = record.get("binaries", {}) + resolved: dict[str, str] = {} + for suite in args.suite: + stored = binaries.get(suite) + if stored is None: + raise ValueError(f"benchmark suite {suite!r} was not recorded in {metadata}") + path = Path(stored["path"]) + if not path.is_file(): + raise ValueError(f"recorded benchmark binary does not exist: {path}") + current = { + "path": str(path.resolve()), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + if current != stored: + raise ValueError(f"recorded benchmark binary changed: {path}") + resolved[suite] = str(path.resolve()) + + return resolved + + +def validate_build_record(args: argparse.Namespace) -> None: + for suite, path in validated_build_binaries(args).items(): + print(f"{suite}={path}") + + def write_manifest(args: argparse.Namespace) -> None: settings = dict(setting.split("=", 1) for setting in args.setting) manifest = { @@ -340,6 +446,22 @@ def argument_parser() -> argparse.ArgumentParser: manifest.add_argument("--candidate-binary", action="append", default=[]) manifest.set_defaults(function=write_manifest) + record_build = subparsers.add_parser("record-build", help="record reusable benchmark binaries") + record_build.add_argument("--output", required=True) + record_build.add_argument("--worktree", required=True) + record_build.add_argument("--target", required=True) + record_build.add_argument("--setting", action="append", default=[]) + record_build.add_argument("--binary", action="append", default=[]) + record_build.set_defaults(function=write_build_record) + + validate_build = subparsers.add_parser("validate-build", help="validate a reusable build") + validate_build.add_argument("--metadata", required=True) + validate_build.add_argument("--worktree", required=True) + validate_build.add_argument("--target", required=True) + validate_build.add_argument("--setting", action="append", default=[]) + validate_build.add_argument("--suite", action="append", default=[]) + validate_build.set_defaults(function=validate_build_record) + summary = subparsers.add_parser("summarize", help="write ratios.csv and summary.md") summary.add_argument("output_directory") summary.set_defaults(function=summarize_directory) @@ -348,8 +470,12 @@ def argument_parser() -> argparse.ArgumentParser: def main() -> None: - args = argument_parser().parse_args() - args.function(args) + parser = argument_parser() + args = parser.parse_args() + try: + args.function(args) + except (OSError, subprocess.CalledProcessError, ValueError) as error: + parser.error(str(error)) if __name__ == "__main__": diff --git a/scripts/tests/test_rowfn_benchmark.py b/scripts/tests/test_rowfn_benchmark.py index 308580f8d68..ace0797c6a3 100644 --- a/scripts/tests/test_rowfn_benchmark.py +++ b/scripts/tests/test_rowfn_benchmark.py @@ -6,6 +6,8 @@ import tempfile import unittest from pathlib import Path +from types import SimpleNamespace +from unittest import mock REPO_ROOT = Path(__file__).resolve().parents[2] SCRIPT = REPO_ROOT / "scripts" / "rowfn_benchmark.py" @@ -129,6 +131,67 @@ def test_summarize_excludes_and_reports_revision_only_benchmarks(self) -> None: markdown = (self.directory / "summary.md").read_text(encoding="utf-8") self.assertIn("`numeric/candidate`: candidate only.", markdown) + def test_build_record_validates_identity_and_executable(self) -> None: + target = self.directory / "target" + target.mkdir() + binary = target / "binary_ops-123" + binary.write_bytes(b"first binary") + metadata = target / "rowfn-benchmark-build.json" + identity = { + "settings": {"codegen_units": "1", "lto": "fat"}, + "toolchain": {"rustc": "rustc 1.97.1", "cargo": "cargo 1.97.1"}, + "revision": {"head": "abc123", "dirty_state_sha256": "clean"}, + } + arguments = SimpleNamespace( + output=str(metadata), + worktree=str(self.directory), + target=str(target), + setting=["codegen_units=1", "lto=fat"], + binary=[f"numeric={binary}"], + ) + + with mock.patch.object(self.module, "build_identity", return_value=identity): + self.module.write_build_record(arguments) + + validation = SimpleNamespace( + metadata=str(metadata), + worktree=str(self.directory), + target=str(target), + setting=["codegen_units=1", "lto=fat"], + suite=["numeric"], + ) + with mock.patch.object(self.module, "build_identity", return_value=identity): + self.assertEqual( + self.module.validated_build_binaries(validation), + {"numeric": str(binary.resolve())}, + ) + + changed_identities = { + "settings": {**identity, "settings": {"codegen_units": "16", "lto": "false"}}, + "toolchain": { + **identity, + "toolchain": {"rustc": "rustc 1.98.0", "cargo": "cargo 1.98.0"}, + }, + "revision": { + **identity, + "revision": {"head": "def456", "dirty_state_sha256": "changed"}, + }, + } + for field, changed_identity in changed_identities.items(): + with ( + self.subTest(field=field), + mock.patch.object(self.module, "build_identity", return_value=changed_identity), + self.assertRaisesRegex(ValueError, f"{field} changed"), + ): + self.module.validated_build_binaries(validation) + + binary.write_bytes(b"second binary") + with ( + mock.patch.object(self.module, "build_identity", return_value=identity), + self.assertRaisesRegex(ValueError, "binary changed"), + ): + self.module.validated_build_binaries(validation) + if __name__ == "__main__": unittest.main() From 9850c8933199a0e322d7030edf2991ffd226098b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:47 -0400 Subject: [PATCH 146/160] Use the blanket RowFn vtable in benchmarks Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 7 +------ vortex-array/benches/strict_validity.rs | 2 -- 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index bff6a29d2da..df9bbb5aa8c 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -24,7 +24,6 @@ use vortex_array::scalar_fn::OutputSink; use vortex_array::scalar_fn::RowFn; use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -168,10 +167,6 @@ impl RowFn for RowSinkWrappingAdd { } } -vortex_array::impl_row_fn_vtable!(RowWrappingAdd); -vortex_array::impl_row_fn_vtable!(RowCheckedAdd); -vortex_array::impl_row_fn_vtable!(RowSinkWrappingAdd); - fn inputs() -> (ArrayRef, ArrayRef) { let lhs = (0..ROWS) .map(|index| index as i64) @@ -208,7 +203,7 @@ fn nullable_inputs() -> (ArrayRef, ArrayRef) { fn bench_row_fn(bencher: Bencher, function: F, make_inputs: fn() -> (ArrayRef, ArrayRef)) where - F: RowFn + ScalarFnVTable, + F: RowFn, { bencher .with_inputs(make_inputs) diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs index e0a8b6a565e..8064c8582e8 100644 --- a/vortex-array/benches/strict_validity.rs +++ b/vortex-array/benches/strict_validity.rs @@ -106,8 +106,6 @@ impl RowFn for LazyDouble { } } -vortex_array::impl_row_fn_vtable!(LazyDouble); - /// The same function, applying validity the way the adapter used to: materialize a mask first. #[derive(Clone)] struct EagerDouble; From 5ad8fb88e5e2b601da3a3ca7f9d72c90c93b5d87 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:42:17 -0400 Subject: [PATCH 147/160] Cover repository RowFn benchmark paths Signed-off-by: Connor Tsui --- scripts/benchmark-rowfn.sh | 4 +- vortex-array/benches/row_fn_executor.rs | 58 ++++++++++++++++++++++--- vortex-array/benches/strict_validity.rs | 6 +-- 3 files changed, 58 insertions(+), 10 deletions(-) diff --git a/scripts/benchmark-rowfn.sh b/scripts/benchmark-rowfn.sh index 790b10462f1..7536e37ec87 100755 --- a/scripts/benchmark-rowfn.sh +++ b/scripts/benchmark-rowfn.sh @@ -16,7 +16,7 @@ Options: --filter PATTERN Pass a Divan benchmark filter. Repeatable. --build-only Build and record benchmark executables without measuring. --measure-only Measure previously recorded benchmark executables without building. - --config NAME primary (1 CGU/fat LTO) or repository (16 CGUs/no LTO). + --config NAME repository (16 CGUs/no LTO, default) or primary (1 CGU/fat LTO). --target-root PATH Parent for reusable baseline and candidate Cargo targets. --baseline-target PATH Reusable Cargo target for the baseline revision. --candidate-target PATH @@ -58,7 +58,7 @@ requested_suites=() filters=() run_build=true run_measure=true -configuration=primary +configuration=repository target_root= baseline_target_override= candidate_target_override= diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index df9bbb5aa8c..18c1ba87e63 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -20,10 +20,12 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::OutputSink; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::OutputSink; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -133,8 +135,9 @@ unsafe impl OutputSink for I64Sink { rows.len() == row_count } - fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { - &mut rows[index] + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } } unsafe fn finish(self) -> VortexResult { @@ -167,6 +170,36 @@ impl RowFn for RowSinkWrappingAdd { } } +#[derive(Clone)] +struct RowSinkCheckedAdd; + +impl RowFn for RowSinkCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_checked_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64, i64), UninitElementSink, _>( + |(lhs, rhs), output| -> VortexResult { + let value = lhs.checked_add(rhs).ok_or_else(checked_add_error)?; + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) + }, + ) + } +} + fn inputs() -> (ArrayRef, ArrayRef) { let lhs = (0..ROWS) .map(|index| index as i64) @@ -229,6 +262,16 @@ fn row_sink_wrapping_add(bencher: Bencher) { bench_row_fn(bencher, RowSinkWrappingAdd, inputs); } +#[divan::bench] +fn row_sink_wrapping_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, constant_inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, nullable_inputs); +} + #[divan::bench] fn handrolled_sink_wrapping_add(bencher: Bencher) { bencher @@ -276,6 +319,11 @@ fn row_checked_add_nullable(bencher: Bencher) { bench_row_fn(bencher, RowCheckedAdd, nullable_inputs); } +#[divan::bench] +fn row_sink_checked_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowSinkCheckedAdd, nullable_inputs); +} + #[divan::bench] fn row_wrapping_add_nullable(bencher: Bencher) { bench_row_fn(bencher, RowWrappingAdd, nullable_inputs); diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs index 8064c8582e8..f4993cd6d42 100644 --- a/vortex-array/benches/strict_validity.rs +++ b/vortex-array/benches/strict_validity.rs @@ -38,11 +38,11 @@ use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ExecutionArgs; -use vortex_array::scalar_fn::RowExecution; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; From b416a1fcb7ac6526c05a74f3c23162a16d45f5f5 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:25:04 -0400 Subject: [PATCH 148/160] Adapt RowFn benchmarks to explicit contracts Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 6 ++++-- vortex-array/benches/strict_validity.rs | 1 + 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index 18c1ba87e63..62b1b08b795 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -51,6 +51,7 @@ impl RowFn for RowWrappingAdd { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("bench.row_wrapping_add"); @@ -119,11 +120,11 @@ unsafe impl OutputSink for I64Sink { type Row<'a> = &'a mut i64; type WriteToken = (); - fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(rows: usize) -> VortexResult { Ok(Self(BufferMut::zeroed(rows))) } @@ -152,6 +153,7 @@ impl RowFn for RowSinkWrappingAdd { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("bench.row_sink_wrapping_add"); diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs index f4993cd6d42..9a713220fab 100644 --- a/vortex-array/benches/strict_validity.rs +++ b/vortex-array/benches/strict_validity.rs @@ -81,6 +81,7 @@ impl RowFn for LazyDouble { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("bench.lazy_double"); From 9c0dac2bd89e22f4360c6f5564bbbd71bc4db6e8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:28:23 -0400 Subject: [PATCH 149/160] Use sink row counts in benchmarks Signed-off-by: Connor Tsui --- vortex-array/benches/row_fn_executor.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs index 62b1b08b795..84a2e029412 100644 --- a/vortex-array/benches/row_fn_executor.rs +++ b/vortex-array/benches/row_fn_executor.rs @@ -132,8 +132,8 @@ unsafe impl OutputSink for I64Sink { self.0.as_mut_slice() } - fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { - rows.len() == row_count + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { From 479a06f1c8b84f2eb1603830f94fe1993fed41c4 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 12:48:59 -0400 Subject: [PATCH 150/160] Implement the RowFn execution backend Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 91 ++ .../scalar_fn/unstable/row/batch/execution.rs | 475 +++++++++++ .../src/scalar_fn/unstable/row/batch/mod.rs | 25 + .../src/scalar_fn/unstable/row/batch/tests.rs | 793 ++++++++++++++++++ .../src/scalar_fn/unstable/row/execute/mod.rs | 19 + .../scalar_fn/unstable/row/execute/outcome.rs | 43 + .../scalar_fn/unstable/row/execute/owned.rs | 128 +++ .../scalar_fn/unstable/row/execute/sink.rs | 380 +++++++++ .../src/scalar_fn/unstable/row/mod.rs | 11 + .../src/scalar_fn/unstable/row/row_fn.rs | 7 +- .../unstable/row/types/element/input.rs | 2 +- .../unstable/row/types/element/mod.rs | 1 + .../row/types/element/tuple/element_tuple.rs | 41 +- .../unstable/row/types/element/tuple/mod.rs | 1 + .../unstable/row/types/element/tuple/tests.rs | 2 +- .../src/scalar_fn/unstable/row/types/mod.rs | 1 + .../src/scalar_fn/unstable/row/types/sink.rs | 12 +- .../scalar_fn/unstable/row/visitor/execute.rs | 306 +++++++ .../src/scalar_fn/unstable/row/visitor/mod.rs | 7 + .../scalar_fn/unstable/row/visitor/plan.rs | 7 +- .../unstable/row/visitor/row_visitor.rs | 2 +- .../src/scalar_fn/unstable/row/vtable.rs | 198 ++++- 22 files changed, 2511 insertions(+), 41 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/args.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execution.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/tests.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/owned.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/sink.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs new file mode 100644 index 00000000000..781d15711be --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution arguments paired with the metadata selected during planning. +//! +//! [`BorrowedExecutionArgs`] can point at original, sliced, or filtered arrays while retaining the +//! dtypes, output dtype, and null policy of the original batch plan. + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::visitor::RowPolicy; + +/// A borrowed [`ExecutionArgs`] view with the planning metadata selected for its row kernel. +/// +/// `arrays` may be filtered or sliced, while `dtypes` and `output_dtype` always describe the +/// original planned batch. Keeping them together prevents an execution path from pairing an input +/// view with unrelated planning metadata. +#[derive(Clone, Copy)] +pub(crate) struct BorrowedExecutionArgs<'a> { + /// The input arrays for this kernel invocation. + arrays: &'a [ArrayRef], + + /// The number of rows in this kernel invocation. + row_count: usize, + + /// The original input dtypes used to select the row implementation. + dtypes: &'a [DType], + + /// The non-nullable dtype built by the selected output capability. + output_dtype: &'a DType, + + /// The nullable execution policy selected during planning. + policy: RowPolicy, +} + +impl<'a> BorrowedExecutionArgs<'a> { + /// Pair one input view with the planning metadata selected for its batch. + pub(crate) fn new( + arrays: &'a [ArrayRef], + row_count: usize, + dtypes: &'a [DType], + output_dtype: &'a DType, + policy: RowPolicy, + ) -> Self { + Self { + arrays, + row_count, + dtypes, + output_dtype, + policy, + } + } + + /// Return the original input dtypes used to select the row implementation. + pub(crate) fn dtypes(&self) -> &'a [DType] { + self.dtypes + } + + /// Return the non-nullable dtype built by the selected output capability. + pub(crate) fn output_dtype(&self) -> &'a DType { + self.output_dtype + } + + /// Return the nullable execution policy selected during planning. + pub(crate) fn policy(&self) -> RowPolicy { + self.policy + } +} + +impl ExecutionArgs for BorrowedExecutionArgs<'_> { + fn get(&self, index: usize) -> VortexResult { + self.arrays.get(index).cloned().ok_or_else(|| { + vortex_err!( + "row-function input index must be less than {}, got {index}", + self.arrays.len(), + ) + }) + } + + fn num_inputs(&self) -> usize { + self.arrays.len() + } + + fn row_count(&self) -> usize { + self.row_count + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs new file mode 100644 index 00000000000..8a0e10d79df --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -0,0 +1,475 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Applies columnar semantics around typed row-kernel invocations. +//! +//! [`Batch`] owns strict null propagation, constant broadcasting, execution strategy selection, and +//! output validation. The row kernel therefore handles only decoded values and its selected output +//! capability. + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::BatchPlan; +use super::RowPolicy; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::types::batch_constant; +use crate::validity::Validity; + +/// The result of resolving batch validity. +enum ResolvedValidity { + /// The output for an all-valid or all-null batch. + Output(ArrayRef), + + /// A mask with both valid and invalid rows. + PartiallyValid(Mask), +} + +/// One batch of inputs and the metadata needed before its row kernel runs. +pub(crate) struct Batch { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The number of rows in the original execution scope. + row_count: usize, + + /// The input columns, collected once: constant folding inspects them and the filter strategy + /// filters them. + inputs: SmallVec<[ArrayRef; 4]>, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: SmallVec<[DType; 4]>, + + /// The conjoined input validity, so a row of the output is valid iff it is valid in every + /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. + validity: Validity, + + /// The dtype the function declares for these inputs, which the kernel's output is reconciled + /// against. Already widened to nullable if any input is nullable. + result_dtype: DType, + + /// The non-nullable dtype the dispatched output capability builds, computed while planning. + output_dtype: DType, + + /// How the concrete dispatch executes nullable rows. + policy: RowPolicy, +} + +impl Batch { + /// Collect the inputs and derive their dtypes, validity, and execution policy. + /// + /// **Not** for a nullary function: with no inputs there is no validity to propagate and no + /// per-row work to fold, and the all-constant check below would vacuously pass. + pub(crate) fn new( + id: ScalarFnId, + args: &dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let row_count = args.row_count(); + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>()?; + + for (index, input) in inputs.iter().enumerate() { + vortex_ensure_eq!( + input.len(), + row_count, + "the {id} input {index} must have {row_count} rows, got {}", + input.len(), + ); + } + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let result_dtype = plan.result_dtype(&arg_dtypes); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + row_count, + inputs, + arg_dtypes, + validity, + result_dtype, + output_dtype: plan.output_dtype, + policy: plan.policy, + }) + } + + /// Apply constant folding and null handling around `kernel`. + /// + /// When the mask contains valid and invalid rows, `try_unfiltered` may avoid filtering. + /// `Ok(None)` filters the valid rows and scatters the output back. Every result is checked + /// against the planned shape and dtype. + pub(crate) fn execute( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.validity, Validity::AllInvalid) + || self.inputs.iter().any(|input| { + input + .as_opt::() + .is_some_and(|constant| constant.scalar().is_null()) + }) + { + return Ok(self.all_null()); + } + + // All inputs constant, and their conjoined validity proves every row non-null. This sees + // through extension and masked wrappers just like argument decoding does. + if self.row_count > 0 + && self.validity.definitely_no_nulls() + && self + .inputs + .iter() + .all(|input| batch_constant(input).is_some()) + { + return self.broadcast_one_row(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, false, ctx), + RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), + RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx), + } + } + + /// Evaluate one row of constant inputs and broadcast the validated result. + fn broadcast_one_row( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one_row: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.slice(0..1)) + .collect::>()?; + + let result = VortexResult::from(kernel(self.execution_args(&one_row, 1), ctx)?)?; + let result = self.validate_kernel_output(result, 1, ctx)?; + let result = self.finalize_output(result, 1)?; + let scalar = result.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.row_count).into_array()) + } + + /// Run every stored payload, then attach the input validity without materializing its mask. + fn execute_dense( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + retry_deferred_error: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let values = match kernel(self.execution_args(&self.inputs, self.row_count), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Unlike `resolve_validity`, all-true preserves the deferred error and all-false + // suppresses evidence that came entirely from null rows. An empty loop cannot + // produce deferred evidence, so the ambiguous empty mask cannot reach this arm. + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + // Deferred retry receives only the dense kernel. Filter first so the retry cannot + // evaluate null rows again. + return self.filter_and_scatter(kernel, &valid, ctx); + } + RowExecution::DeferredError(error) => return Err(error), + }; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.row_count) + } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // Handled by the guard above, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } + + /// Materialize validity and handle all-valid or all-null batches. + fn resolve_validity( + &self, + kernel: &impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Check all-true before all-false: an empty mask is both, and must not be treated as + // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). + if valid.all_true() { + let values = VortexResult::from(kernel( + self.execution_args(&self.inputs, self.row_count), + ctx, + )?)?; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = self.finalize_output(values, self.row_count)?; + + return Ok(ResolvedValidity::Output(values)); + } + + if valid.all_false() { + return Ok(ResolvedValidity::Output(self.all_null())); + } + + Ok(ResolvedValidity::PartiallyValid(valid)) + } + + /// Resolve validity, try unfiltered execution, then fall back to filtering. + fn execute_valid_only( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedValidity::Output(output) => return Ok(output), + ResolvedValidity::PartiallyValid(valid) => valid, + }; + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Try execution against the original inputs, then mask a returned full-length result. + fn try_execute_unfiltered( + &self, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(execution) = try_unfiltered( + self.execution_args(&self.inputs, self.row_count), + valid, + ctx, + )? + else { + return Ok(None); + }; + let values = VortexResult::from(execution)?; + let values = self.validate_kernel_output(values, valid.len(), ctx)?; + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + self.finalize_output(values.mask(mask)?, valid.len()) + .map(Some) + } + + /// Filter to valid rows, run the kernel, then scatter into a null-padded output. + fn filter_and_scatter( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(valid.clone())) + .collect::>()?; + + let values = VortexResult::from(kernel( + self.execution_args(&filtered, valid.true_count()), + ctx, + )?)?; + let values = self.validate_kernel_output(values, valid.true_count(), ctx)?; + + self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) + } + + fn all_null(&self) -> ArrayRef { + ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() + } + + /// Pair an input view with this batch's planning metadata. + fn execution_args<'b>( + &'b self, + arrays: &'b [ArrayRef], + row_count: usize, + ) -> BorrowedExecutionArgs<'b> { + BorrowedExecutionArgs::new( + arrays, + row_count, + &self.arg_dtypes, + &self.output_dtype, + self.policy, + ) + } + + fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { + reconcile_output(self.id, &self.result_dtype, expected_len, values) + } + + /// Validate the output from a row kernel before batch validity is attached. + fn validate_kernel_output( + &self, + values: ArrayRef, + expected_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx) + } + + /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set + /// bits, producing an array of length `valid.len()` that is null at every unset position. + fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult { + vortex_ensure_eq!( + values.len(), + valid.true_count(), + "the {} kernel output must contain {} filtered rows, got {}", + self.id, + valid.true_count(), + values.len(), + ); + + let AllOr::Some(slices) = valid.slices() else { + // The caller handles the all-true and all-false masks. + vortex_bail!( + "scatter_valid requires valid and invalid rows, got an all-valid or all-invalid mask" + ); + }; + + // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index + // 0, and any in-bounds index would do since they are masked out below (values is non-empty + // here). + let mut indices = vec![0u64; valid.len()]; + let mut rank = 0u64; + for &(start, end) in slices { + for index in &mut indices[start..end] { + *index = rank; + rank += 1; + } + } + let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array(); + + let scattered = values.take(indices)?; + + // A nullable gathered array cannot be wrapped because a `Masked` child must be all valid. + // The general masking pass unions its nulls with the batch validity instead. + if scattered.dtype().is_nullable() { + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + return scattered.mask(mask); + } + + // The gathered values are all valid, so attaching validity is sufficient. + Ok(MaskedArray::try_new( + scattered, + Validity::from_mask(valid.clone(), Nullability::Nullable), + )? + .into_array()) + } +} + +/// Validate the output produced directly by a row kernel. +/// +/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring +/// nullability, and every produced row **must** be valid. Batch execution owns strict null +/// propagation and attaches input-derived validity only after this boundary. +pub(crate) fn finalize_kernel_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + vortex_ensure!( + values.all_valid(ctx)?, + "the {id} row kernel must produce only valid rows, got at least one null row", + ); + + cast_output_nullability(result_dtype, values) +} + +/// Reconcile an output with the function's declared shape and nullability. +fn reconcile_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + + cast_output_nullability(result_dtype, values) +} + +/// Validate an output's shape and logical dtype without executing a nullability cast. +fn validate_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: &ArrayRef, +) -> VortexResult<()> { + vortex_ensure_eq!( + values.len(), + expected_len, + "the {id} kernel output must contain {expected_len} rows, got {}", + values.len(), + ); + vortex_ensure!( + values.dtype().eq_ignore_nullability(result_dtype), + "the {id} kernel output dtype must match {result_dtype} ignoring nullability, got {}", + values.dtype(), + ); + + Ok(()) +} + +/// Cast only the output nullability after its shape, dtype, and validity are accepted. +fn cast_output_nullability(result_dtype: &DType, values: ArrayRef) -> VortexResult { + if values.dtype() == result_dtype { + Ok(values) + } else { + values.cast(result_dtype.clone()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs new file mode 100644 index 00000000000..b501965fa61 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batch execution around a strict row kernel. +//! +//! A row kernel handles typed values for one row. This module adds the columnar concerns around it: +//! planning the output and null strategy, preserving batch constants, propagating strict validity, +//! selecting an execution strategy, and validating the finished output. +//! +//! [`BatchPlan`] carries the nullable execution strategy selected by a concrete dispatch. [`Batch`] +//! applies that strategy, and [`BorrowedExecutionArgs`] pairs each kernel invocation with its +//! planning metadata. + +mod args; +pub(super) use args::BorrowedExecutionArgs; + +mod execution; +pub(super) use execution::Batch; +pub(super) use execution::finalize_kernel_output; + +pub(super) use super::visitor::BatchPlan; +pub(super) use super::visitor::RowPolicy; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs new file mode 100644 index 00000000000..7ad4b69c4e7 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -0,0 +1,793 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::super::execute::RowExecution; +use super::Batch; +use super::BatchPlan; +use super::RowPolicy; +use super::finalize_kernel_output; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; +use crate::scalar_fn::unstable::row::row_fn_return_dtype; +use crate::validity::Validity; + +#[derive(Clone)] +struct RetryConstantAdd; + +#[derive(Clone)] +struct NullarySeven; + +#[derive(Clone)] +struct AddThree; + +#[derive(Clone)] +struct AddShort(ShortVisit); + +/// An element whose decode drops the last row, standing in for an invalid element implementation. +struct ShortDecodeI64; + +#[derive(Clone)] +struct Identity; + +#[derive(Clone)] +struct ValidOnlyIdentity; + +#[derive(Clone)] +struct SinkOptions; + +struct OptionsCheckingSink; + +#[derive(Clone)] +struct InvalidKernelOutput; + +/// Deliberately violates [`OutputElement::build`] to test validation at the public boundary. +struct NullProducingI64(i64); + +#[derive(Clone)] +struct PreparedAdd { + visit: PreparedVisit, + prepares: Arc, +} + +#[derive(Clone, Copy)] +enum PreparedVisit { + Owned, + Sink, + Deferred, +} + +#[derive(Clone, Copy)] +enum ShortVisit { + Owned, + Sink, +} + +// SAFETY: the view and unchecked access delegate to the `i64` implementation. The implementation +// deliberately returns a short column so the executor's pre-loop length guard can be tested. +unsafe impl InputElement for ShortDecodeI64 { + type Column = Buffer; + type View<'a> = &'a [i64]; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = false; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let column = ::decode(array, ctx)?; + + Ok(column.slice(0..column.len().saturating_sub(1))) + } + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + ::view(column) + } + + fn view_len(view: &Self::View<'_>) -> usize { + ::view_len(view) + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_from_view(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { ::get_from_view_unchecked(view, index) } + } +} + +// SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or +// `finish` through the executor. The row-initialization requirements are therefore vacuous. +unsafe impl OutputSink for OptionsCheckingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn output_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { + if !enabled { + vortex_bail!(InvalidArgument: "the test sink is disabled"); + } + + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize) -> VortexResult { + vortex_bail!("the planning-only test sink must not be allocated") + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 + } + + unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} + + unsafe fn finish(self) -> VortexResult { + vortex_bail!("the planning-only test sink must not finish") + } +} + +impl OutputElement for NullProducingI64 { + fn element_dtype() -> DType { + DType::from(i64::PTYPE) + } + + fn build(values: Vec) -> ArrayRef { + let values: Vec<_> = values.into_iter().map(|value| value.0).collect(); + let validity = Validity::from_iter((0..values.len()).map(|index| index != 0)); + + PrimitiveArray::new(values, validity).into_array() + } +} + +struct I64Sink(BufferMut); + +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +impl RowFn for NullarySeven { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.nullary_seven"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), I64Sink, _>(|(), output| { + *output = 7; + }) + } +} + +impl RowFn for AddThree { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["first", "second", "third"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.add_three"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64, i64), i64>(|(first, second, third)| first + second + third) + } +} + +impl RowFn for AddShort { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.add_short"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + match self.0 { + ShortVisit::Owned => { + visitor.visit::<(ShortDecodeI64, i64), i64>(|(lhs, rhs)| lhs + rhs) + } + ShortVisit::Sink => { + visitor.visit_into::<(ShortDecodeI64, i64), I64Sink, _>(|(lhs, rhs), output| { + *output = lhs + rhs; + }) + } + } + } +} + +impl RowFn for RetryConstantAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.retry_constant_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(u8, u8), u8, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "checked add overflowed")); + } + + Ok(()) + }, + ) + } +} + +impl RowFn for Identity { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.identity"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } +} + +impl RowFn for ValidOnlyIdentity { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.valid_only_identity"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| { + *output = value; + Ok(()) + }) + } +} + +impl RowFn for SinkOptions { + type Options = bool; + + const ARG_NAMES: &'static [&'static str] = &[]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.sink_options"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), OptionsCheckingSink, _>(|(), ()| ()) + } +} + +impl RowFn for InvalidKernelOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_kernel_output"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), NullProducingI64>(|(value,)| NullProducingI64(value)) + } +} + +impl RowFn for PreparedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.prepared_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + let prepares = Arc::clone(&self.prepares); + let prepare = move |(_lhs, rhs): (Option, Option)| { + prepares.fetch_add(1, Ordering::Relaxed); + rhs + }; + + match self.visit { + PreparedVisit::Owned => visitor + .visit_prepared::<(i64, i64), i64, _>(prepare, |constant_rhs, (lhs, rhs)| { + lhs.wrapping_add(constant_rhs.unwrap_or(rhs)) + }), + PreparedVisit::Sink => visitor.visit_prepared_into::<(i64, i64), I64Sink, _, ()>( + prepare, + |constant_rhs, (lhs, rhs), output| { + *output = lhs.wrapping_add(constant_rhs.unwrap_or(rhs)); + }, + ), + PreparedVisit::Deferred => visitor.visit_prepared_deferred::<(i64, i64), i64, _, bool>( + prepare, + |constant_rhs, (lhs, rhs)| lhs.overflowing_add(constant_rhs.unwrap_or(rhs)), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "prepared add overflowed")); + } + + Ok(()) + }, + ), + } + } +} + +#[test] +fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.row_batch"); + + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let result = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::Dense, + }) + }); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_short_decode_beside_constant_is_rejected() -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter(0..64_i64).into_array(); + let rhs = ConstantArray::new(10_i64, 64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 64); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort(ShortVisit::Sink), &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short decoded column passed the pre-loop length check"), + }; + + assert!( + error + .to_string() + .contains("does not address exactly 64 rows"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[rstest] +#[case::owned(ShortVisit::Owned)] +#[case::sink(ShortVisit::Sink)] +fn test_short_constant_decode_is_rejected(#[case] visit: ShortVisit) -> VortexResult<()> { + let lhs = ConstantArray::new(10_i64, 64).into_array(); + let rhs = PrimitiveArray::from_iter(0..64_i64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 64); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort(visit), &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short constant decode passed the pre-loop length check"), + }; + + assert!( + error + .to_string() + .contains("batch-constant input must contain exactly 1 row, got 0"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[test] +fn test_short_constant_null_tolerant_decode_is_rejected() -> VortexResult<()> { + let lhs = MaskedArray::try_new( + ConstantArray::new(10_i64, 4).into_array(), + Validity::from_iter([true, false, true, false]), + )? + .into_array(); + let rhs = PrimitiveArray::from_iter(0..4_i64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 4); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort(ShortVisit::Sink), &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short null-tolerant constant decode passed validation"), + }; + + assert!( + error + .to_string() + .contains("decoded batch-constant input must contain exactly 1 row, got 0"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[test] +fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("valid-row overflow must remain observable"), + }; + + assert!( + error.to_string().contains("checked add overflowed"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[test] +fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![1, u8::MAX], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1_u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::new(vec![2_u8, 0], Validity::from_iter([true, false])); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} + +#[test] +fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&Identity, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid([true, true])] +#[case::all_invalid([false, false])] +fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.resolve_validity"); + + let validity = Validity::Array(BoolArray::from_iter(validity).into_array()); + let input = PrimitiveArray::new(vec![4_i64, 5], validity).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_filters_and_scatters() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.filter_and_scatter"); + + let input = PrimitiveArray::new( + vec![10_i64, 20, 30, 40], + Validity::from_iter([true, false, true, false]), + ) + .into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 4); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.finalize_kernel_output"); + + let values = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let result_dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); + let mut ctx = array_session().create_execution_ctx(); + + let actual = finalize_kernel_output(*ID, &result_dtype, 2, values.clone(), &mut ctx)?; + let expected = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + assert_eq!(actual.dtype(), &result_dtype); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + assert!(finalize_kernel_output(*ID, &result_dtype, 3, values, &mut ctx).is_err()); + + let bools = BoolArray::from_iter([true, false]).into_array(); + assert!(finalize_kernel_output(*ID, &result_dtype, 2, bools, &mut ctx).is_err()); + Ok(()) +} + +#[test] +fn test_output_dtype_receives_function_options() -> VortexResult<()> { + assert_eq!( + row_fn_return_dtype(&SinkOptions, &true, &[])?, + DType::from(i64::PTYPE) + ); + assert!(row_fn_return_dtype(&SinkOptions, &false, &[]).is_err()); + Ok(()) +} + +#[rstest] +#[case::nonnullable(Validity::NonNullable)] +#[case::all_valid(Validity::AllValid)] +fn test_kernel_output_rejects_nulls_at_function_boundary( + #[case] validity: Validity, +) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + let execution = execute_rows(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); + let error = match execution { + Err(error) => error, + Ok(output) => match output.execute::(&mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an invalid row kernel output passed boundary validation"), + }, + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_kernel_output"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("row kernel must produce only valid rows"), + "the boundary error must identify invalid row output, got {error}", + ); + Ok(()) +} + +#[rstest] +#[case::owned_constant(PreparedVisit::Owned, true)] +#[case::owned_per_row(PreparedVisit::Owned, false)] +#[case::sink_constant(PreparedVisit::Sink, true)] +#[case::sink_per_row(PreparedVisit::Sink, false)] +#[case::deferred_constant(PreparedVisit::Deferred, true)] +#[case::deferred_per_row(PreparedVisit::Deferred, false)] +fn test_prepared_visits( + #[case] visit: PreparedVisit, + #[case] constant_rhs: bool, +) -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let rhs = if constant_rhs { + ConstantArray::new(3_i64, 2).into_array() + } else { + PrimitiveArray::from_iter([3_i64, 4]).into_array() + }; + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let prepares = Arc::new(AtomicUsize::new(0)); + let function = PreparedAdd { + visit, + prepares: Arc::clone(&prepares), + }; + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; + let expected = if constant_rhs { + PrimitiveArray::from_iter([4_i64, 5]).into_array() + } else { + PrimitiveArray::from_iter([4_i64, 6]).into_array() + }; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(prepares.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[test] +fn test_nullary_row_function_broadcasts() -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([7_i64, 7, 7]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_empty_batch_preserves_nonnullable_dtype() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(std::iter::empty::()).into_array(); + let args = VecExecutionArgs::new(vec![input], 0); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&ValidOnlyIdentity, &EmptyOptions, &args, &mut ctx)?; + + assert_eq!(actual.len(), 0); + assert_eq!(actual.dtype(), &DType::from(i64::PTYPE)); + Ok(()) +} + +#[test] +fn test_owned_execution_traverses_three_per_row_inputs() -> VortexResult<()> { + let args = VecExecutionArgs::new( + vec![ + PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(), + PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(), + PrimitiveArray::from_iter([100_i64, 200, 300]).into_array(), + ], + 3, + ); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&AddThree, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([111_i64, 222, 333]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs new file mode 100644 index 00000000000..73722549f41 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -0,0 +1,19 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Row-loop execution for owned outputs and output sinks. +//! +//! [`owned`] stores one independent value per row and reduces compact failure evidence. [`sink`] +//! drives output builders whose row handles may share batch state. Both return [`RowExecution`], +//! which distinguishes a completed array from a deferred error that batch validity may suppress. + +mod owned; +pub(super) use owned::execute_owned; +pub(super) use owned::execute_owned_infallible; + +mod outcome; +pub use outcome::RowExecution; + +mod sink; +pub(super) use sink::execute_sink; +pub(super) use sink::execute_sink_valid_rows; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs new file mode 100644 index 00000000000..fc013e7a317 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The result of a completed row loop before batch-level null handling. +//! +//! [`RowExecution`] preserves deferred failure evidence until batch execution can determine whether +//! the failing payload belonged to a valid row. + +use vortex_error::VortexError; +use vortex_error::VortexResult; + +use crate::ArrayRef; + +/// The outcome of a row loop before batch execution decides whether an error is observable. +/// +/// Together with the surrounding [`VortexResult`], this represents three outcomes: +/// +/// - `Err(error)` is a non-retryable execution or immediate row error. +/// - [`Output`](Self::Output) is a successful row loop. +/// - [`DeferredError`](Self::DeferredError) is failure evidence from a completed row loop. +/// +/// A dense loop can evaluate null payloads, so its deferred error is not always observable. Batch +/// execution can retry only valid rows to discard errors caused by null payloads. A plain +/// `VortexResult` cannot distinguish these errors from failures that a retry cannot fix. +/// +/// Once execution is known to contain only valid rows, converting this outcome into a +/// `VortexResult` turns [`DeferredError`](Self::DeferredError) into an ordinary error. +pub enum RowExecution { + /// The successfully built, full-length output column. + Output(ArrayRef), + + /// An error constructed from failure evidence reduced across a completed row loop. + DeferredError(VortexError), +} + +impl From for VortexResult { + fn from(execution: RowExecution) -> Self { + match execution { + RowExecution::Output(output) => Ok(output), + RowExecution::DeferredError(error) => Err(error), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs new file mode 100644 index 00000000000..e3e64194cae --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -0,0 +1,128 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Executes row kernels that return one independent owned value per row. +//! +//! [`execute_owned`] decodes inputs once, prepares constant state, writes into spare vector +//! capacity, and reduces compact failure evidence without putting error construction in the hot +//! loop. [`execute_owned_infallible`] removes that failure path for infallible kernels. + +use std::ops::BitOrAssign; + +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; + +/// Zero-sized evidence used to erase failure reduction from infallible owned visits. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + +/// Decode every input column for one kernel invocation, then store one infallible output per row. +pub(crate) fn execute_owned_infallible( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, +{ + execute_owned::( + args, + ctx, + prepare, + move |prepared, args| (apply(prepared, args), NoFailure), + |_| Ok(()), + ) +} + +/// Decode every input column for one kernel invocation, then store outputs and reduce failures. +pub(crate) fn execute_owned( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, +{ + const { assert_owned_output_needs_no_drop::() }; + + // Keep the vector length at zero until every row succeeds. An unwind then abandons partially + // initialized spare capacity without treating it as initialized output. The no-drop assertion + // above proves that no initialized value requires its destructor to run. + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let failure; + + { + let output = &mut values.spare_capacity_mut()[..row_count]; + + // When every input stores one value per row, the indexed source removes argument-shape + // dispatch from the hot loop and lets the lane kernel optimize the traversal as one + // operation. Keep view construction and its length proof in this branch. Hoisting them + // through the shared validation helper changed add, subtract, and multiply with + // batch-constant and per-row arguments from 9.219, 9.229, and 18.94 us to 30.46, 31.11, + // and 37.73 us on a Ryzen 9 7950X with rustc 1.91.0 and LLVM 21.1.2. + // Restoring this placement recovered the fast code under the 16-CGU, no-LTO bench profile. + if let Some(views) = Args::per_row_views(&columns) { + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows + // immediately above. + failure = unsafe { Args::indexed_source(views, row_count) } + .map_checked_into(output, |elements| apply(&prepared, elements)); + } else { + // A batch-constant input was collapsed to one row during decoding. This path reads that + // row repeatedly while indexing only the per-row inputs. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // Keep the output-slot iterator as the loop bound. `row_count` is address-taken by the + // validation error formatting above. With rustc 1.97.1 and LLVM 22.1.6 under 16 CGUs + // without LTO, indexing `output` by a `0..row_count` range retains an early-exit bounds + // check and prevents vectorization with batch-constant and per-row arguments. Recheck + // the optimized IR and those benchmarks before restoring that range loop. + let mut accumulated = Fail::default(); + for (index, slot) in output.iter_mut().enumerate() { + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + slot.write(value); + accumulated |= row_failure; + } + failure = accumulated; + } + } + + // SAFETY: normal completion of either loop initializes every slot in `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + // Failure evidence is reduced inside the loop so its richer error construction stays cold. + // Preserve that provenance so batch execution may retry over only valid rows. + match finish_failure(failure) { + Ok(()) => Ok(RowExecution::Output(Out::build(values))), + Err(error) => Ok(RowExecution::DeferredError(error)), + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs new file mode 100644 index 00000000000..ec121b21b44 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -0,0 +1,380 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Executes row kernels that write through an [`OutputSink`]. +//! +//! Dense execution visits every row. Skip-invalid execution can instead initialize omitted output +//! positions and visit only rows that are valid in every input, falling back when either the input +//! representation or sink lacks that capability. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::SinkResult; + +fn ensure_decoded_lengths( + columns: &Args::Columns, + views: Option<&Args::Views<'_>>, + row_count: usize, +) -> VortexResult<()> { + let lengths_match = match views { + Some(views) => Args::view_lens_match(views, row_count), + None => Args::decoded_lens_match(columns, row_count), + }; + vortex_ensure!( + lengths_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + Ok(()) +} + +/// Decode every input column and allocate one sink for one kernel invocation. +/// +/// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state +/// does not need to be captured by the closure. +pub(crate) fn execute_sink( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + let row_count = args.row_count(); + let mut sink = >::with_capacity(row_count)?; + let columns = Args::decode(args, ctx)?; + let constants = Args::constants(&columns); + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + let prepared = prepare(constants); + + { + // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This + // scope releases the borrow before `finish_sink` consumes the sink. + let mut rows = >::rows(&mut sink); + let sink_row_count = >::row_count(&rows); + vortex_ensure_eq!( + sink_row_count, + row_count, + "the output sink must address exactly {row_count} rows, got {sink_row_count}", + ); + + // The all-per-row representation removes argument-shape dispatch from the hot loop. The + // constant-and-per-row path instead reads collapsed batch constants at row zero. + if let Some(views) = views { + for index in 0..row_count { + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before + // the loop. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + // SAFETY: the sink row-count check above proved every loop index is in bounds. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + apply(&prepared, elements, output).into_result()?; + } + } else { + for index in 0..row_count { + // SAFETY: the sink row-count check above proved every loop index is in bounds. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + apply(&prepared, Args::get(&columns, index), output).into_result()?; + } + } + } + + finish_sink::(sink) +} + +/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +pub(crate) fn execute_sink_valid_rows( + args: &dyn ExecutionArgs, + valid: &Mask, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult> +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + // Decline before input decoding or sink allocation when this sink cannot initialize rows that + // the mask skips. The capability and the operation are the same function pointer. + let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + else { + return Ok(None); + }; + + // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An + // element representation may decline when it cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + let constants = Args::constants(&columns); + let row_count = args.row_count(); + let mut sink = >::with_capacity(row_count)?; + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid) = valid.bit_buffer() else { + vortex_bail!( + "execute_sink_valid_rows requires valid and invalid rows, got an all-valid or all-invalid mask" + ); + }; + vortex_ensure_eq!( + valid.len(), + row_count, + "the validity mask must address exactly {row_count} rows, got {}", + valid.len(), + ); + + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + let prepared = prepare(constants); + + { + let mut rows = >::rows(&mut sink); + + // Initialize every slot before skipping rows. Recheck addressability afterward because the + // initializer mutably borrows the row representation. + initialize_skipped_rows(&mut rows); + let initialized_row_count = >::row_count(&rows); + vortex_ensure_eq!( + initialized_row_count, + row_count, + "the initialized output sink must address exactly {row_count} rows, got {initialized_row_count}", + ); + + // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first + // immediate error, turn later callbacks into no-ops, and return before finishing the sink. + let mut error = None; + valid.for_each_set_index(|index| { + if error.is_some() { + return; + } + + // SAFETY: the post-initialization row-count check proved that the sink addresses every + // mask index, which is below the mask's validated `row_count`. + let output = unsafe { >::row_unchecked(&mut rows, index) }; + let result = match &views { + Some(views) => { + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows, and + // mask indices are below `row_count`. + let elements = unsafe { Args::get_from_views_unchecked(views, index) }; + apply(&prepared, elements, output) + } + None => apply(&prepared, Args::get(&columns, index), output), + }; + if let Err(row_error) = result.into_result() { + error = Some(row_error); + } + }); + + if let Some(error) = error { + return Err(error); + } + } + + finish_sink::(sink).map(Some) +} + +fn finish_sink(sink: Sink) -> VortexResult +where + Sink: OutputSink, +{ + // SAFETY: callers reach this helper only after every completed callback returned the sink's + // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. + // The sink contract defines how that evidence establishes initialization of its row storage. + unsafe { >::finish(sink) }.map(RowExecution::Output) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_error::vortex_err; + use vortex_mask::Mask; + + use super::RowExecution; + use super::execute_sink_valid_rows; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::NativePType; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::InitializedElement; + use crate::scalar_fn::unstable::row::OutputSink; + use crate::scalar_fn::unstable::row::UninitElementSink; + use crate::validity::Validity; + + struct NonSkippingSink; + + struct ShrinkingSink(Vec); + + // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or + // `finish` through the executor. The row-initialization requirements are therefore vacuous. + unsafe impl OutputSink for NonSkippingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize) -> VortexResult { + Err(vortex_err!( + "a non-skipping sink must decline before allocation" + )) + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 + } + + unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> { + } + + unsafe fn finish(self) -> VortexResult { + Err(vortex_err!("a non-skipping sink must not finish")) + } + } + + // SAFETY: the initializer deliberately shrinks the row collection to exercise the executor's + // post-initialization length check. If execution incorrectly continues, safe indexing in + // `row_unchecked` panics instead of accessing invalid memory. + unsafe impl OutputSink for ShrinkingSink { + type Rows<'a> = &'a mut Vec; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + rows.pop(); + }) + } + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(vec![0; rows])) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.0 + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::from_iter(self.0).into_array()) + } + } + + #[test] + fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([true, false]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( + &args, + &valid, + &mut ctx, + |_| (), + |_, _, _| (), + )?; + + assert!(execution.is_none()); + Ok(()) + } + + #[test] + fn test_skip_invalid_sink_initializes_and_writes_addressed_rows() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let valid = Mask::from_iter([true, false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::< + (i64,), + (), + UninitElementSink, + InitializedElement, + EmptyOptions, + >( + &args, + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + // SAFETY: `output` is the row supplied to this callback. + unsafe { InitializedElement::write(output, value * 2) } + }, + )?; + let Some(RowExecution::Output(actual)) = execution else { + vortex_bail!("the skip-invalid sink must produce an output"); + }; + let expected = PrimitiveArray::from_iter([20_i64, 0, 60]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) + } + + #[test] + fn test_skip_invalid_sink_rechecks_rows_after_initialization() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>( + &args, + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + *output = value; + }, + ); + + let error = match result { + Err(error) => error, + Ok(_) => vortex_bail!("the sink must reject rows changed by its initializer"), + }; + assert!( + error + .to_string() + .contains("initialized output sink must address exactly 2 rows, got 1"), + "unexpected error: {error}", + ); + Ok(()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index bcb3a008488..91f225ca071 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -11,10 +11,21 @@ //! [`RowFn::dispatch`] implementation uses a [`RowVisitor`] to select an [`ElementTuple`] and //! either an [`OutputElement`] or [`OutputSink`] for each supported dtype combination. //! +//! Unlike a general strict function, a [`RowFn`] cannot produce null from valid inputs. +//! +//! A _partially valid_ batch contains both valid and invalid rows. _Skip-invalid_ runs the kernel +//! only for valid rows without changing row positions. _Filter-and-scatter_ compacts valid rows, +//! runs the kernel, and restores their positions. +//! //! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits //! reduce compact failure evidence in that loop and retry only valid rows when null payloads may //! have caused the failure. +mod execute; +pub use execute::RowExecution; + +mod batch; + mod row_fn; pub use row_fn::RowFn; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 8a982c4fb37..831cc25f251 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -19,7 +19,12 @@ use super::visitor::RowVisitor; use crate::dtype::DType; use crate::scalar_fn::ScalarFnId; -/// A scalar function computed one row at a time. +/// A strict scalar function whose row kernel cannot produce null from valid inputs. +/// +/// This is stronger than +/// [`ScalarFnVTable::is_strict`](crate::scalar_fn::ScalarFnVTable::is_strict), which requires null +/// propagation but permits valid inputs to produce null. The framework derives output validity +/// only from input validity. /// /// Declare argument names and use [`dispatch`](Self::dispatch) to select element and output types. /// Every implementation receives the standard [`ScalarFnVTable`]. A public type that needs custom diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 2764840da7a..c9140c00b8c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -68,7 +68,7 @@ pub unsafe trait InputElement: 'static { /// cannot decode this particular array. /// /// Override this for a non-dense-safe representation that can still place safe placeholders in - /// null slots. The skip-invalid executor never reads those slots. + /// null slots. Valid-row execution never reads those slots. fn decode_null_tolerant( array: ArrayRef, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 51d66594332..7a8a8e92e41 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -20,3 +20,4 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; +pub use tuple::batch_constant; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index b0a3e709696..e6f63628e88 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -27,20 +27,31 @@ pub struct ArgColumn( ); enum ArgColumnKind { + /// One decoded value per batch row; executors validate the exact length before traversal. PerRow(T::Column), + + /// Exactly one decoded row, established by [`ArgColumn::try_from_constant`]. Constant(T::Column), } impl ArgColumn { + fn try_from_constant(column: T::Column) -> VortexResult { + let decoded_len = T::view_len(&T::view(&column)); + vortex_ensure_eq!( + decoded_len, + 1, + "a decoded batch-constant input must contain exactly 1 row, got {decoded_len}", + ); + + Ok(Self(ArgColumnKind::Constant(column))) + } + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // An empty input has no row 0 to slice, and its row loop runs zero times either way. if let Some(constant) = batch_constant(&array) && !array.is_empty() { - return Ok(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?))); + return Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?); } Ok(Self(ArgColumnKind::PerRow(T::decode(array, ctx)?))) @@ -52,10 +63,7 @@ impl ArgColumn { if let Some(constant) = batch_constant(&array) && !array.is_empty() { - return Ok(Some(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?)))); + return Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?).map(Some); } Ok(T::decode_null_tolerant(array, ctx)? @@ -88,14 +96,14 @@ impl ArgColumn { } fn addresses_rows(&self, row_count: usize) -> bool { - // A constant is always read at index zero, so it addresses any batch length. + // A constant is validated when constructed and is always read at index zero. match &self.0 { ArgColumnKind::PerRow(column) => T::view_len(&T::view(column)) == row_count, ArgColumnKind::Constant(_) => true, } } - fn constant(&self) -> Option> { + fn constant_value(&self) -> Option> { match &self.0 { ArgColumnKind::PerRow(_) => None, ArgColumnKind::Constant(column) => Some(T::get(column, 0)), @@ -170,8 +178,8 @@ pub trait ElementTuple: 'static + private::Sealed { /// Decode every input column once while tolerating null rows. /// - /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid - /// strategy calls this once per batch. + /// Return `Ok(None)` when an argument has no null-tolerant representation. Valid-row execution + /// calls this once per batch. fn decode_null_tolerant( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, @@ -195,9 +203,10 @@ pub trait ElementTuple: 'static + private::Sealed { /// Whether every per-row argument contains exactly `row_count` rows. /// - /// This is the mixed-shape equivalent of [`view_lens_match`](Self::view_lens_match) when - /// [`per_row_views`](Self::per_row_views) declines. It runs once before the hot loop for the - /// same LLVM optimization. A batch constant is exempt because decoding collapsed it to one row. + /// This is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include + /// batch constants. It runs once before the hot loop for the same LLVM optimization. A batch + /// constant is exempt because its [`ArgColumn`] constructor already validated the one-row + /// representation produced by decoding. fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; /// Read one row from borrowed views. @@ -372,7 +381,7 @@ macro_rules! element_tuple { } fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { - ($(columns.$idx.constant(),)+) + ($(columns.$idx.constant_value(),)+) } } }; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index a2c143704a0..69b5cf686f6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -8,6 +8,7 @@ mod element_tuple; pub use element_tuple::ElementTuple; +pub use element_tuple::batch_constant; mod indexed; pub use indexed::IndexedElementTuple; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index 044ad15fd4b..560e1e48f4f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_bail; use vortex_mask::Mask; use super::ElementTuple; -use super::element_tuple::batch_constant; +use super::batch_constant; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs index ce119f32915..e47f195410c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -12,6 +12,7 @@ pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; +pub(super) use element::batch_constant; mod result; pub use result::SinkResult; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index 6cc3ce06d30..22c90b8da02 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -21,8 +21,7 @@ use crate::scalar_fn::unstable::row::OutputElement; /// batch state. The executor passes each row slot into an [`Fn`] closure. /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once. -/// Skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an -/// initializer. +/// Execution can omit invalid rows when [`skipped_rows_initializer`] returns an initializer. /// /// # Errors /// @@ -73,10 +72,11 @@ pub unsafe trait OutputSink: 'static + Sized { /// **must not** be able to construct one without establishing the invariant. type WriteToken: 'static; - /// The operation that initializes every output position before skip-invalid execution. + /// The operation that initializes every output position before + /// [skip-invalid execution](crate::scalar_fn::unstable::row). /// - /// `Some` enables skip-invalid execution. The initializer **must** make every row safe to - /// finish. Callbacks overwrite valid rows, and batch execution masks skipped rows. + /// `Some` enables this strategy. The initializer **must** make every row safe to finish. + /// Callbacks overwrite valid rows, and batch execution masks skipped rows. /// /// `None` makes the executor fall back to filtering the inputs. fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { @@ -153,7 +153,7 @@ impl InitializedElement { /// The row closure must return the [`InitializedElement`] from [`InitializedElement::write`] on /// success. The token is zero-sized, so the proof adds no runtime row state. /// -/// Skip-invalid execution initializes placeholders before omitting rows. Errors and unwinds are +/// When execution omits invalid rows, it initializes placeholders first. Errors and unwinds are /// safe because `values` keeps length zero until `finish`. The `T: Copy` bound means that /// initialized spare-capacity elements require no destruction. pub struct UninitElementSink { diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs new file mode 100644 index 00000000000..b003f70a395 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -0,0 +1,306 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visitors that execute dense and skip-invalid row loops. +//! +//! Each visit revalidates its concrete signature and checks that its output dtype and null policy +//! match the plan before entering a row loop. [`ExecuteValidRows`] can decline, so the batch layer +//! filters the inputs and retries with [`ExecuteRows`]. + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; + +use super::RowPolicy; +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; +use super::row_visitor::private; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::execute::execute_owned; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_sink; +use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; + +/// The runtime visit that decodes every column once and runs the selected row loop. +pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { + /// The inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { + pub(crate) fn new( + args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, + output_dtype: &'args DType, + policy: RowPolicy, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + dtypes, + options, + output_dtype, + policy, + ctx, + } + } +} + +impl private::Sealed for ExecuteRows<'_, '_, F> {} + +impl RowVisitor for ExecuteRows<'_, '_, F> { + type VisitResult = RowExecution; + + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; + + execute_owned_infallible::(self.args, self.ctx, prepare, apply) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; + + execute_sink::( + self.args, self.ctx, prepare, apply, + ) + } + + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; + + execute_owned::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} + +/// The runtime visit that executes valid rows over the original input columns. +/// +/// Only output sinks have a contract for skipped output positions. Owned visits therefore decline +/// so batch execution can filter the valid inputs and scatter the output back. +pub(crate) struct ExecuteValidRows<'args, 'ctx, F: RowFn> { + /// The original inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The conjoined validity, containing both valid and invalid rows. + valid: &'args Mask, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'args, 'ctx, F: RowFn> ExecuteValidRows<'args, 'ctx, F> { + pub(crate) fn new( + args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, + output_dtype: &'args DType, + policy: RowPolicy, + valid: &'args Mask, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + dtypes, + options, + output_dtype, + policy, + valid, + ctx, + } + } +} + +impl private::Sealed for ExecuteValidRows<'_, '_, F> {} + +impl RowVisitor for ExecuteValidRows<'_, '_, F> { + type VisitResult = Option; + + fn visit_prepared( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; + + // Owned execution has no sink that can initialize skipped output positions. Decline so + // batch execution filters the inputs and retries with the dense visitor. + Ok(None) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; + + execute_sink_valid_rows::( + self.args, self.valid, self.ctx, prepare, apply, + ) + } + + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + _finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; + + // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. + Ok(None) + } +} + +fn ensure_plan( + planned_output: &DType, + planned_policy: RowPolicy, + actual_output: DType, + actual_policy: RowPolicy, +) -> VortexResult<()> { + vortex_ensure_eq!( + actual_policy, + planned_policy, + "row dispatch must select the planned nullable execution policy: planned {planned_policy:?}, got {actual_policy:?}", + ); + vortex_ensure_eq!( + actual_output, + *planned_output, + "row dispatch must select the planned output dtype: planned {planned_output}, got {actual_output}", + ); + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index 57da5f4691b..c7f9baf6a62 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -6,9 +6,16 @@ //! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch mod check; +pub(super) use check::assert_owned_output_needs_no_drop; + +mod execute; +pub(super) use execute::ExecuteRows; +pub(super) use execute::ExecuteValidRows; mod plan; +pub(super) use plan::BatchPlan; pub(super) use plan::BatchPlanner; +pub(super) use plan::RowPolicy; mod row_visitor; pub use row_visitor::RowVisitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index c522376d491..5360ced573b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -114,15 +114,12 @@ pub(crate) struct BatchPlan { pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. - // TODO(connor)[RowFn]: Remove this allowance when the execution backend from #9130 consumes - // this policy. - #[allow(dead_code)] pub(crate) policy: RowPolicy, } impl BatchPlan { /// Return the output dtype widened with strict input nullability. - pub(crate) fn result_dtype(self, args: &[DType]) -> DType { + pub(crate) fn result_dtype(&self, args: &[DType]) -> DType { let nullability = self.output_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); @@ -139,7 +136,7 @@ pub(crate) enum RowPolicy { /// Evaluate all rows, retrying only valid rows if a deferred error is raised. DenseWithRetry, - /// Execute only valid rows, trying skip-invalid execution before filtering. + /// Execute only valid rows over the original inputs before filtering. ValidOnly, } diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 0e7abaa4322..1a9e8628bc7 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -74,7 +74,7 @@ pub trait RowVisitor: private::Sealed + Sized { /// # Examples /// /// Test whether each string occurs in its allowed-values list. The prepare closure builds one - /// lookup table for a batch-constant list. The row closure scans a varying list directly. + /// lookup table for a batch-constant list. The row closure scans a per-row list directly. /// /// ```ignore /// visitor.visit_prepared::< diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index b6c14585c48..257ac3ca0b5 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -4,12 +4,13 @@ //! Adapts [`RowFn`] implementations to the scalar-function interface. //! //! The blanket [`ScalarFnVTable`] implementation supplies common arity, validity, fallibility, and -//! execution behavior. [`row_fn_return_dtype`] and [`execute_rows`] expose the same planning and -//! execution paths to public vtables that delegate to a private row kernel. +//! execution behavior. The visitor layer validates and executes the concrete signature selected by +//! dispatch. [`row_fn_return_dtype`] and [`execute_rows`] expose the same paths to public vtables +//! that delegate to a private row kernel. use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; use vortex_session::VortexSession; use super::row_fn::RowFn; @@ -24,6 +25,12 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::unstable::row::batch::Batch; +use crate::scalar_fn::unstable::row::batch::BorrowedExecutionArgs; +use crate::scalar_fn::unstable::row::batch::finalize_kernel_output; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::visitor::ExecuteRows; +use crate::scalar_fn::unstable::row::visitor::ExecuteValidRows; impl ScalarFnVTable for F { type Options = F::Options; @@ -69,6 +76,8 @@ impl ScalarFnVTable for F { union_child_validities(expression) } + // `RowFn` is stricter than `ScalarFnVTable::is_strict`: its kernel cannot produce null from + // valid inputs, so batch execution derives output validity only from input validity. fn is_strict(&self, _options: &Self::Options) -> bool { true } @@ -98,20 +107,39 @@ pub fn row_fn_return_dtype( /// delegate row execution to a private `RowFn` kernel through this function. pub fn execute_rows( function: &F, - _options: &F::Options, + options: &F::Options, args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { ensure_arity(function, args.num_inputs())?; - // TODO(connor)[RowFn]: Replace this temporary error with the execution backend in #9129. - vortex_bail!( - "Row function {} does not yet have an execution backend", - RowFn::id(function) + // Nullary functions have no input validity to propagate, so they skip batch execution. + if args.num_inputs() == 0 { + let plan = function.dispatch(options, &[], BatchPlanner::::new(&[], options))?; + let result_dtype = plan.result_dtype(&[]); + let nullary_args = + BorrowedExecutionArgs::new(&[], args.row_count(), &[], &plan.output_dtype, plan.policy); + + let execution = execute_row_kernel(function, options, nullary_args, ctx)?; + let values = VortexResult::from(execution)?; + + return finalize_kernel_output( + RowFn::id(function), + &result_dtype, + args.row_count(), + values, + ctx, + ); + } + + let batch = prepare_batch(function, options, args)?; + batch.execute( + |args, ctx| execute_row_kernel(function, options, args, ctx), + |args, valid, ctx| try_execute_rows_unfiltered(function, options, args, valid, ctx), + ctx, ) } -/// Validate the number of arguments before calling user-defined dispatch code. fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { let expected = F::ARG_NAMES.len(); vortex_ensure_eq!( @@ -124,26 +152,101 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } +fn execute_row_kernel( + function: &F, + options: &F::Options, + args: BorrowedExecutionArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + function.dispatch( + options, + args.dtypes(), + ExecuteRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + ctx, + ), + ) +} + +fn try_execute_rows_unfiltered( + function: &F, + options: &F::Options, + args: BorrowedExecutionArgs<'_>, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + function.dispatch( + options, + args.dtypes(), + ExecuteValidRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + valid, + ctx, + ), + ) +} + +fn prepare_batch( + function: &F, + options: &F::Options, + args: &dyn ExecutionArgs, +) -> VortexResult { + Batch::new(RowFn::id(function), args, |arg_dtypes| { + function.dispatch( + options, + arg_dtypes, + BatchPlanner::::new(arg_dtypes, options), + ) + }) +} + #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use super::execute_rows; use super::row_fn_return_dtype; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::RowVisitor; + use crate::validity::Validity; #[derive(Clone)] struct IndexingRowFn; + #[derive(Clone)] + struct ChangingDispatchRowFn { + dispatches: Arc, + change: DispatchChange, + } + + #[derive(Clone, Copy)] + enum DispatchChange { + Policy, + Element, + } + impl RowFn for IndexingRowFn { type Options = EmptyOptions; @@ -168,6 +271,35 @@ mod tests { } } + impl RowFn for ChangingDispatchRowFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.changing_dispatch_row_fn"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { + visitor.visit::<(i64,), i64>(|(value,)| value) + } else { + match self.change { + DispatchChange::Policy => visitor + .visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())), + DispatchChange::Element => visitor.visit::<(u64,), u64>(|(value,)| value), + } + } + } + } + #[test] fn test_return_dtype_rejects_wrong_arity_before_dispatch() { let error = row_fn_return_dtype(&IndexingRowFn, &EmptyOptions, &[]) @@ -186,6 +318,52 @@ mod tests { assert_arity_error(error); } + #[test] + fn test_execute_rejects_dispatch_that_changes_after_planning() { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Policy, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&function, &EmptyOptions, &args, &mut ctx) + .expect_err("dispatch must not change after planning"); + let message = error.to_string(); + + assert!( + message.contains("row dispatch must select the planned nullable execution policy"), + "unexpected error: {error}", + ); + assert!( + message.contains("planned Dense, got DenseWithRetry"), + "unexpected error: {error}", + ); + } + + #[test] + fn test_execute_revalidates_element_types_after_planning() -> VortexResult<()> { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Element, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&function, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_error::vortex_bail!("dispatch must preserve its planned element types"), + }; + + assert!( + error.to_string().contains("expected a u64 column"), + "unexpected error: {error}", + ); + Ok(()) + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( From b76cf19cc6017c3e48e800e422ea4c98c2489ee8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 17:35:57 -0400 Subject: [PATCH 151/160] Refine RowFn row execution Signed-off-by: "Connor Tsui" --- .../scalar_fn/unstable/row/execute/owned.rs | 101 ++++---- .../scalar_fn/unstable/row/execute/sink.rs | 219 +++++++++++------- .../row/types/element/tuple/element_tuple.rs | 31 +-- .../row/types/element/tuple/indexed.rs | 2 +- vortex-buffer/src/bit/buf.rs | 65 +++++- 5 files changed, 269 insertions(+), 149 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index e3e64194cae..d3809f939be 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -20,7 +20,7 @@ use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; -/// Zero-sized evidence used to erase failure reduction from infallible owned visits. +/// Zero-sized failure accumulator for infallible owned visits. #[derive(Clone, Copy, Default)] struct NoFailure; @@ -28,7 +28,7 @@ impl BitOrAssign for NoFailure { fn bitor_assign(&mut self, _rhs: Self) {} } -/// Decode every input column for one kernel invocation, then store one infallible output per row. +/// Decode every input column, then store one output per row from an infallible kernel. pub(crate) fn execute_owned_infallible( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, @@ -48,7 +48,7 @@ where ) } -/// Decode every input column for one kernel invocation, then store outputs and reduce failures. +/// Decode every input column, then store outputs and combine per-row failure evidence. pub(crate) fn execute_owned( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, @@ -61,66 +61,59 @@ where Out: OutputElement, Fail: Copy + Default + BitOrAssign, { + // The output vector stays at length zero until every slot is initialized so that an unwind + // abandons partially initialized spare capacity. This no-drop assertion proves that no + // initialized value requires a destructor to run. const { assert_owned_output_needs_no_drop::() }; - // Keep the vector length at zero until every row succeeds. An unwind then abandons partially - // initialized spare capacity without treating it as initialized output. The no-drop assertion - // above proves that no initialized value requires its destructor to run. - let row_count = args.row_count(); - let mut values = Vec::::with_capacity(row_count); let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); - let failure; - - { - let output = &mut values.spare_capacity_mut()[..row_count]; - - // When every input stores one value per row, the indexed source removes argument-shape - // dispatch from the hot loop and lets the lane kernel optimize the traversal as one - // operation. Keep view construction and its length proof in this branch. Hoisting them - // through the shared validation helper changed add, subtract, and multiply with - // batch-constant and per-row arguments from 9.219, 9.229, and 18.94 us to 30.46, 31.11, - // and 37.73 us on a Ryzen 9 7950X with rustc 1.91.0 and LLVM 21.1.2. - // Restoring this placement recovered the fast code under the 16-CGU, no-LTO bench profile. - if let Some(views) = Args::per_row_views(&columns) { - vortex_ensure!( - Args::view_lens_match(&views, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - - // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows - // immediately above. - failure = unsafe { Args::indexed_source(views, row_count) } - .map_checked_into(output, |elements| apply(&prepared, elements)); - } else { - // A batch-constant input was collapsed to one row during decoding. This path reads that - // row repeatedly while indexing only the per-row inputs. - vortex_ensure!( - Args::decoded_lens_match(&columns, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - - // Keep the output-slot iterator as the loop bound. `row_count` is address-taken by the - // validation error formatting above. With rustc 1.97.1 and LLVM 22.1.6 under 16 CGUs - // without LTO, indexing `output` by a `0..row_count` range retains an early-exit bounds - // check and prevents vectorization with batch-constant and per-row arguments. Recheck - // the optimized IR and those benchmarks before restoring that range loop. - let mut accumulated = Fail::default(); - for (index, slot) in output.iter_mut().enumerate() { - let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); - slot.write(value); - accumulated |= row_failure; - } - failure = accumulated; + + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let output = &mut values.spare_capacity_mut()[..row_count]; + + let failure = if let Some(views) = Args::views_no_constants(&columns) { + // Keep this validation beside the views so LLVM sees their common length here. + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows + // immediately above. + let source = unsafe { Args::indexed_source(views, row_count) }; + + source.map_checked_into(output, |elements| apply(&prepared, elements)) + } else { + // Keep this proof branch-local. Shared validation prevents LLVM from specializing this + // loop for each batch-constant arrangement, leaving it scalar under multiple CGUs without + // LTO. The exact pass interaction is unknown. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut accumulated = Fail::default(); + + // Iterate over `output` directly. A `0..row_count` range reuses the address-taken value + // from the validation error formatter and retains an output bounds check. + for (index, slot) in output.iter_mut().enumerate() { + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the loop. + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + + slot.write(value); + accumulated |= row_failure; } - } - // SAFETY: normal completion of either loop initializes every slot in `0..row_count` exactly + accumulated + }; + + // SAFETY: normal completion of either execution path initializes `0..row_count` exactly // once, and `values` was allocated with at least `row_count` capacity. unsafe { values.set_len(row_count) }; - // Failure evidence is reduced inside the loop so its richer error construction stays cold. - // Preserve that provenance so batch execution may retry over only valid rows. + // Defer failures so batch execution can retry with only valid rows. match finish_failure(failure) { Ok(()) => Ok(RowExecution::Output(Out::build(values))), Err(error) => Ok(RowExecution::DeferredError(error)), diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index ec121b21b44..4fd2b42476b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -3,10 +3,11 @@ //! Executes row kernels that write through an [`OutputSink`]. //! -//! Dense execution visits every row. Skip-invalid execution can instead initialize omitted output -//! positions and visit only rows that are valid in every input, falling back when either the input -//! representation or sink lacks that capability. +//! Dense execution visits every row. Skip-invalid execution initializes skipped output rows and +//! visits only rows that are valid in every input. Skip-invalid execution declines when either the +//! input representation or sink cannot support that path. +use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -21,7 +22,12 @@ use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::SinkResult; -fn ensure_decoded_lengths( +/// Verify that every decoded input addresses exactly `row_count` rows. +/// +/// Unlike the owned executor, the paths with and without batch constants can share this check +/// without losing sink-loop vectorization under multiple CGUs without LTO. The exact pass +/// interaction is unknown. +fn verify_lengths( columns: &Args::Columns, views: Option<&Args::Views<'_>>, row_count: usize, @@ -30,6 +36,7 @@ fn ensure_decoded_lengths( Some(views) => Args::view_lens_match(views, row_count), None => Args::decoded_lens_match(columns, row_count), }; + vortex_ensure!( lengths_match, "a decoded row input does not address exactly {row_count} rows", @@ -38,10 +45,11 @@ fn ensure_decoded_lengths( Ok(()) } -/// Decode every input column and allocate one sink for one kernel invocation. +/// Decode inputs once, then write one sink row for each input row. /// -/// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state -/// does not need to be captured by the closure. +/// The executor owns the sink and passes each output row to `apply`. This keeps `apply` as [`Fn`]. +/// Capturing the sink would require [`FnMut`] and put its buffer metadata behind loop-carried +/// mutable closure state, which can prevent LLVM from treating that metadata as loop-invariant. pub(crate) fn execute_sink( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, @@ -53,18 +61,22 @@ where Sink: OutputSink, ApplyResult: SinkResult>::WriteToken>, { - let row_count = args.row_count(); - let mut sink = >::with_capacity(row_count)?; let columns = Args::decode(args, ctx)?; + let views = Args::views_no_constants(&columns); + + let row_count = args.row_count(); + verify_lengths::(&columns, views.as_ref(), row_count)?; + let constants = Args::constants(&columns); - let views = Args::per_row_views(&columns); - ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; let prepared = prepare(constants); + let mut sink = >::with_capacity(row_count)?; + + // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. { - // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This - // scope releases the borrow before `finish_sink` consumes the sink. let mut rows = >::rows(&mut sink); + + // This equality proves to LLVM that `0..row_count` is in bounds for `rows`. let sink_row_count = >::row_count(&rows); vortex_ensure_eq!( sink_row_count, @@ -72,16 +84,14 @@ where "the output sink must address exactly {row_count} rows, got {sink_row_count}", ); - // The all-per-row representation removes argument-shape dispatch from the hot loop. The - // constant-and-per-row path instead reads collapsed batch constants at row zero. if let Some(views) = views { for index in 0..row_count { - // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before - // the loop. + // SAFETY: `verify_lengths` proved every view has `row_count` rows before the loop. let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; // SAFETY: the sink row-count check above proved every loop index is in bounds. let output = unsafe { >::row_unchecked(&mut rows, index) }; + apply(&prepared, elements, output).into_result()?; } } else { @@ -89,15 +99,23 @@ where // SAFETY: the sink row-count check above proved every loop index is in bounds. let output = unsafe { >::row_unchecked(&mut rows, index) }; + + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the + // loop. apply(&prepared, Args::get(&columns, index), output).into_result()?; } } } - finish_sink::(sink) + // SAFETY: every row callback completed successfully, so each returned the required write token. + unsafe { >::finish(sink) }.map(RowExecution::Output) } -/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +/// Write only the rows set in `valid`, or decline when the inputs or sink cannot support +/// skip-invalid execution. +/// +/// `Ok(None)` signals batch execution to filter every input to the valid rows, run the dense +/// kernel, and scatter the results back into a null-padded array. pub(crate) fn execute_sink_valid_rows( args: &dyn ExecutionArgs, valid: &Mask, @@ -110,45 +128,32 @@ where Sink: OutputSink, ApplyResult: SinkResult>::WriteToken>, { - // Decline before input decoding or sink allocation when this sink cannot initialize rows that - // the mask skips. The capability and the operation are the same function pointer. - let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + let Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + mut sink, + }) = setup_sink_valid_rows::(args, valid, ctx)? else { return Ok(None); }; - // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An - // element representation may decline when it cannot provide those values safely. - let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { - return Ok(None); - }; - let constants = Args::constants(&columns); - let row_count = args.row_count(); - let mut sink = >::with_capacity(row_count)?; - - // Batch execution resolves all-valid and all-null inputs before selecting this path. - let AllOr::Some(valid) = valid.bit_buffer() else { - vortex_bail!( - "execute_sink_valid_rows requires valid and invalid rows, got an all-valid or all-invalid mask" - ); - }; - vortex_ensure_eq!( - valid.len(), - row_count, - "the validity mask must address exactly {row_count} rows, got {}", - valid.len(), - ); + let views = Args::views_no_constants(&columns); + verify_lengths::(&columns, views.as_ref(), row_count)?; - let views = Args::per_row_views(&columns); - ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + let constants = Args::constants(&columns); let prepared = prepare(constants); + // Keep `rows` scoped so its borrow ends before `finish`. With multiple CGUs and no LTO, using + // `drop(rows)` duplicates `Args::get` in every sparse callback. { + // Initialize every slot before visiting only valid rows. let mut rows = >::rows(&mut sink); - - // Initialize every slot before skipping rows. Recheck addressability afterward because the - // initializer mutably borrows the row representation. initialize_skipped_rows(&mut rows); + + // The initializer can change addressability. Recheck it so LLVM can prove every mask + // index is in bounds. let initialized_row_count = >::row_count(&rows); vortex_ensure_eq!( initialized_row_count, @@ -156,47 +161,100 @@ where "the initialized output sink must address exactly {row_count} rows, got {initialized_row_count}", ); - // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first - // immediate error, turn later callbacks into no-ops, and return before finishing the sink. - let mut error = None; - valid.for_each_set_index(|index| { - if error.is_some() { - return; - } + if let Some(views) = views { + valid_rows.try_for_each_set_index(|index| { + // SAFETY: the post-initialization row-count check proved that the sink addresses + // every mask index, which is below the mask's validated `row_count`. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; - // SAFETY: the post-initialization row-count check proved that the sink addresses every - // mask index, which is below the mask's validated `row_count`. - let output = unsafe { >::row_unchecked(&mut rows, index) }; - let result = match &views { - Some(views) => { - // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows, and - // mask indices are below `row_count`. - let elements = unsafe { Args::get_from_views_unchecked(views, index) }; - apply(&prepared, elements, output) - } - None => apply(&prepared, Args::get(&columns, index), output), - }; - if let Err(row_error) = result.into_result() { - error = Some(row_error); - } - }); + // SAFETY: `verify_lengths` proved every view has `row_count` rows, and mask indices + // are below `row_count`. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; - if let Some(error) = error { - return Err(error); + apply(&prepared, elements, output).into_result() + })?; + } else { + valid_rows.try_for_each_set_index(|index| { + // SAFETY: the post-initialization row-count check proved that the sink addresses + // every mask index, which is below the mask's validated `row_count`. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + apply(&prepared, Args::get(&columns, index), output).into_result() + })?; } } - finish_sink::(sink).map(Some) + // SAFETY: the initializer completed before traversal, and every visited callback completed + // successfully and returned the required write token. + unsafe { >::finish(sink) } + .map(RowExecution::Output) + .map(Some) } -fn finish_sink(sink: Sink) -> VortexResult +/// State resolved before preparing the skip-invalid row loop. +struct ValidRowsSetup<'valid, Args, Sink, Options> where + Args: ElementTuple, Sink: OutputSink, { - // SAFETY: callers reach this helper only after every completed callback returned the sink's - // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. - // The sink contract defines how that evidence establishes initialization of its row storage. - unsafe { >::finish(sink) }.map(RowExecution::Output) + initialize_skipped_rows: for<'rows> fn(&mut >::Rows<'rows>), + columns: Args::Columns, + valid_rows: &'valid BitBuffer, + row_count: usize, + sink: Sink, +} + +/// Resolve the capabilities, inputs, sink, and validity mask for skip-invalid execution. +fn setup_sink_valid_rows<'valid, Args, Sink, Options>( + args: &dyn ExecutionArgs, + valid: &'valid Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult>> +where + Args: ElementTuple, + Sink: OutputSink, +{ + // The initializer both declares support for skipping rows and initializes those rows. + let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + else { + return Ok(None); + }; + + // Null-tolerant decoding exposes values behind nulls without filtering. Decline when any input + // cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + + let row_count = args.row_count(); + + // Keep allocation before the validity and length checks. With multiple CGUs and no LTO, + // moving it later inlines `Args::get` into every sparse callback, duplicating its bounds + // checks. + let sink = >::with_capacity(row_count)?; + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid_rows) = valid.bit_buffer() else { + vortex_bail!( + "execute_sink_valid_rows requires valid and invalid rows, got an all-valid or all-invalid mask" + ); + }; + vortex_ensure_eq!( + valid_rows.len(), + row_count, + "the validity mask must address exactly {row_count} rows, got {}", + valid_rows.len(), + ); + + Ok(Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + sink, + })) } #[cfg(test)] @@ -313,6 +371,7 @@ mod tests { )?; assert!(execution.is_none()); + Ok(()) } @@ -345,6 +404,7 @@ mod tests { let expected = PrimitiveArray::from_iter([20_i64, 0, 60]); assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) } @@ -375,6 +435,7 @@ mod tests { .contains("initialized output sink must address exactly 2 rows, got 1"), "unexpected error: {error}", ); + Ok(()) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index e6f63628e88..7c1baed168a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -138,7 +138,7 @@ pub trait ElementTuple: 'static + private::Sealed { /// The decoded column representations. type Columns; - /// Borrowed views of decoded columns when every argument stores one value per row. + /// Borrowed views of decoded columns with no batch constants. type Views<'a>; /// The borrowed row of element values. @@ -146,9 +146,9 @@ pub trait ElementTuple: 'static + private::Sealed { /// The batch-constant element values. /// - /// `Some` carries the value of a batch-constant argument. `None` marks a per-row argument. A - /// [`RowVisitor`] passes these values to its prepare closure so constant work can leave the row - /// loop. + /// `Some` carries the value of a batch-constant argument. `None` marks a non-constant argument. + /// A [`RowVisitor`] passes these values to its prepare closure so constant work can leave the + /// row loop. /// /// [`RowVisitor`]: crate::scalar_fn::unstable::row::RowVisitor type ConstElems<'a>; @@ -186,22 +186,25 @@ pub trait ElementTuple: 'static + private::Sealed { ) -> VortexResult>; /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + /// + /// Each argument selects either its batch-constant value or row `index`. Keep that selection + /// visible in the loop so LLVM can unswitch it before vectorizing. fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; - /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// Borrow the decoded columns when none is batch-constant. /// - /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple - /// gives the optimizer ordinary contiguous column access without a per-row constant check. - fn per_row_views(columns: &Self::Columns) -> Option>; + /// Returns `None` if any column is batch-constant. Otherwise, omitting [`ArgColumn`] from the + /// returned tuple removes constant checks from the row loop. + fn views_no_constants(columns: &Self::Columns) -> Option>; /// Whether every view contains exactly `row_count` rows. /// - /// The executor calls this once before the all-per-row hot loop. A successful check gives LLVM - /// a dominating equality between the loop bound and every source length, which lets it optimize - /// the tuple access as one fixed-length traversal. + /// The executor calls this once before the loop used when no input is batch-constant. A + /// successful check gives LLVM a dominating equality between the loop bound and every source + /// length, which lets it optimize the tuple access as one fixed-length traversal. fn view_lens_match(views: &Self::Views<'_>, row_count: usize) -> bool; - /// Whether every per-row argument contains exactly `row_count` rows. + /// Whether every non-constant argument contains exactly `row_count` rows. /// /// This is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include /// batch constants. It runs once before the hot loop for the same LLVM optimization. A batch @@ -266,7 +269,7 @@ impl ElementTuple for () { fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} - fn per_row_views(_columns: &Self::Columns) -> Option> { + fn views_no_constants(_columns: &Self::Columns) -> Option> { Some(()) } @@ -350,7 +353,7 @@ macro_rules! element_tuple { ($(columns.$idx.get(index),)+) } - fn per_row_views(columns: &Self::Columns) -> Option> { + fn views_no_constants(columns: &Self::Columns) -> Option> { Some(($($t::view(columns.$idx.per_row_column()?),)+)) } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs index d612d874935..832dbf9ee19 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -17,7 +17,7 @@ use crate::scalar_fn::unstable::row::InputElement; /// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's /// unchecked view access after batch execution validates every decoded column length once. pub trait IndexedElementTuple: ElementTuple { - /// The source shared execution uses for a dense all-per-row loop. + /// The source used when no input is batch-constant. /// /// Its length must be the common view length. For every valid index it must preserve row order, /// return the same value as [`ElementTuple::get_from_views`], and uphold the unchecked read diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index aac8ead42fe..1da0b037f2a 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -479,6 +479,33 @@ impl BitBuffer { } } + /// Fallible variant of [`for_each_set_index`](Self::for_each_set_index). + /// + /// Stops and returns the first error from `f`. + #[inline] + pub fn try_for_each_set_index(&self, mut f: F) -> Result<(), E> + where + F: FnMut(usize) -> Result<(), E>, + { + let mut base = 0usize; + for word in self.chunks().iter_padded() { + if word == u64::MAX { + for k in 0..64 { + f(base + k)?; + } + } else { + let mut w = word; + while w != 0 { + f(base + w.trailing_zeros() as usize)?; + w &= w - 1; + } + } + base += 64; + } + + Ok(()) + } + /// Created a new BitBuffer with offset reset to 0 pub fn sliced(&self) -> Self { if self.offset.is_multiple_of(8) { @@ -970,12 +997,21 @@ mod tests { #[case(65)] #[case(200)] #[case(1000)] - fn test_for_each_set_index_matches_set_indices(#[case] len: usize) { + fn test_set_index_visitors_match_set_indices(#[case] len: usize) { let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0); let expected: Vec = buf.set_indices().collect(); + let mut got = Vec::new(); buf.for_each_set_index(|i| got.push(i)); assert_eq!(got, expected); + + let mut fallible_got = Vec::new(); + let result = buf.try_for_each_set_index(|i| { + fallible_got.push(i); + Ok::<(), ()>(()) + }); + assert_eq!(result, Ok(())); + assert_eq!(fallible_got, expected); } #[rstest] @@ -998,6 +1034,33 @@ mod tests { assert_eq!(got, (0..130).collect::>()); } + #[test] + fn test_try_for_each_set_index_stops_on_error() { + for (buffer, stop) in [ + (BitBuffer::new_set(130), 65), + (BitBuffer::collect_bool(130, |i| i % 3 == 0), 66), + ] { + let mut visited = Vec::new(); + let result = buffer.try_for_each_set_index(|index| { + visited.push(index); + if index == stop { + return Err(index); + } + + Ok(()) + }); + + assert_eq!(result, Err(stop)); + assert_eq!( + visited, + buffer + .set_indices() + .take_while(|&i| i <= stop) + .collect::>() + ); + } + } + #[test] fn test_map_cmp_conditional() { // map_cmp with conditional logic based on index and bit value From 8d7edf34dd445e152fc91be8c9b73d027d0131ad Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 17:36:14 -0400 Subject: [PATCH 152/160] Split RowFn batch execution by strategy Signed-off-by: "Connor Tsui" --- .../unstable/row/batch/execute/constant.rs | 35 ++ .../unstable/row/batch/execute/dense.rs | 54 ++ .../row/batch/execute/filter_scatter.rs | 95 ++++ .../unstable/row/batch/execute/mod.rs | 76 +++ .../unstable/row/batch/execute/output.rs | 104 ++++ .../unstable/row/batch/execute/valid_only.rs | 104 ++++ .../scalar_fn/unstable/row/batch/execution.rs | 475 ------------------ .../src/scalar_fn/unstable/row/batch/mod.rs | 44 +- .../scalar_fn/unstable/row/batch/planning.rs | 77 +++ 9 files changed, 586 insertions(+), 478 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs delete mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execution.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/planning.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs new file mode 100644 index 00000000000..8cb89a14bf3 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::scalar_fn::unstable::row::execute::RowExecution; + +impl Batch { + /// Evaluate one row of constant inputs and broadcast the validated result. + pub(super) fn broadcast_one_row( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one_row: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.slice(0..1)) + .collect::>()?; + + let result = VortexResult::from(kernel(self.execution_args(&one_row, 1), ctx)?)?; + let result = self.validate_kernel_output(result, 1, ctx)?; + let result = self.finalize_output(result, 1)?; + let scalar = result.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.row_count).into_array()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs new file mode 100644 index 00000000000..d3e0ef103c2 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +impl Batch { + /// Run every stored payload, then attach the input validity without materializing its mask. + pub(super) fn execute_dense( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + retry_deferred_error: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let values = match kernel(self.execution_args(&self.inputs, self.row_count), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Unlike `resolve_validity`, all-true preserves the deferred error and all-false + // suppresses evidence that came entirely from null rows. An empty loop cannot + // produce deferred evidence, so the ambiguous empty mask cannot reach this arm. + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + // Deferred retry receives only the dense kernel. Filter first so the retry cannot + // evaluate null rows again. + return self.filter_and_scatter(kernel, &valid, ctx); + } + RowExecution::DeferredError(error) => return Err(error), + }; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.row_count) + } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // Handled by the guard in `Batch::execute`, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs new file mode 100644 index 00000000000..871c7fc1219 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +impl Batch { + /// Filter to valid rows, run the kernel, then scatter into a null-padded output. + pub(super) fn filter_and_scatter( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(valid.clone())) + .collect::>()?; + + let values = VortexResult::from(kernel( + self.execution_args(&filtered, valid.true_count()), + ctx, + )?)?; + let values = self.validate_kernel_output(values, valid.true_count(), ctx)?; + + self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) + } + + /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set + /// bits, producing an array of length `valid.len()` that is null at every unset position. + fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult { + vortex_ensure_eq!( + values.len(), + valid.true_count(), + "the {} kernel output must contain {} filtered rows, got {}", + self.id, + valid.true_count(), + values.len(), + ); + + let AllOr::Some(slices) = valid.slices() else { + // The caller handles the all-true and all-false masks. + vortex_bail!( + "scatter_valid requires valid and invalid rows, got an all-valid or all-invalid mask" + ); + }; + + // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index + // 0, and any in-bounds index would do since they are masked out below (values is non-empty + // here). + let mut indices = vec![0u64; valid.len()]; + let mut rank = 0u64; + for &(start, end) in slices { + for index in &mut indices[start..end] { + *index = rank; + rank += 1; + } + } + let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array(); + + let scattered = values.take(indices)?; + + // A nullable gathered array cannot be wrapped because a `Masked` child must be all valid. + // The general masking pass unions its nulls with the batch validity instead. + if scattered.dtype().is_nullable() { + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + return scattered.mask(mask); + } + + // The gathered values are all valid, so attaching validity is sufficient. + Ok(MaskedArray::try_new( + scattered, + Validity::from_mask(valid.clone(), Nullability::Nullable), + )? + .into_array()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs new file mode 100644 index 00000000000..f29fbcaad27 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Selects a batch execution strategy. +//! +//! [`Batch::execute`] handles universal fast paths, then delegates to dense or valid-only +//! execution. + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::Batch; +use super::RowPolicy; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Constant; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::types::batch_constant; +use crate::validity::Validity; + +mod constant; +mod dense; +mod filter_scatter; +mod valid_only; + +mod output; +pub(crate) use output::finalize_kernel_output; + +impl Batch { + /// Apply constant folding and null handling around `kernel`. + /// + /// When the mask contains valid and invalid rows, `try_unfiltered` may avoid filtering. + /// `Ok(None)` filters the valid rows and scatters the output back. Every result is checked + /// against the planned shape and dtype. + pub(crate) fn execute( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.validity, Validity::AllInvalid) + || self.inputs.iter().any(|input| { + input + .as_opt::() + .is_some_and(|constant| constant.scalar().is_null()) + }) + { + return Ok(self.all_null()); + } + + // All inputs constant, and their conjoined validity proves every row non-null. This sees + // through extension and masked wrappers just like argument decoding does. + if self.row_count > 0 + && self.validity.definitely_no_nulls() + && self + .inputs + .iter() + .all(|input| batch_constant(input).is_some()) + { + return self.broadcast_one_row(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, false, ctx), + RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), + RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs new file mode 100644 index 00000000000..47fd8cdaf46 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; + +use super::super::Batch; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnId; + +impl Batch { + pub(super) fn all_null(&self) -> ArrayRef { + ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() + } + + pub(super) fn finalize_output( + &self, + values: ArrayRef, + expected_len: usize, + ) -> VortexResult { + reconcile_output(self.id, &self.result_dtype, expected_len, values) + } + + /// Validate the output from a row kernel before batch validity is attached. + pub(super) fn validate_kernel_output( + &self, + values: ArrayRef, + expected_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx) + } +} + +/// Validate the output produced directly by a row kernel. +/// +/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring +/// nullability, and every produced row **must** be valid. Batch execution owns strict null +/// propagation and attaches input-derived validity only after this boundary. +pub(crate) fn finalize_kernel_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + vortex_ensure!( + values.all_valid(ctx)?, + "the {id} row kernel must produce only valid rows, got at least one null row", + ); + + cast_output_nullability(result_dtype, values) +} + +/// Reconcile an output with the function's declared shape and nullability. +fn reconcile_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + + cast_output_nullability(result_dtype, values) +} + +/// Validate an output's shape and logical dtype without executing a nullability cast. +fn validate_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: &ArrayRef, +) -> VortexResult<()> { + vortex_ensure_eq!( + values.len(), + expected_len, + "the {id} kernel output must contain {expected_len} rows, got {}", + values.len(), + ); + vortex_ensure!( + values.dtype().eq_ignore_nullability(result_dtype), + "the {id} kernel output dtype must match {result_dtype} ignoring nullability, got {}", + values.dtype(), + ); + + Ok(()) +} + +/// Cast only the output nullability after its shape, dtype, and validity are accepted. +fn cast_output_nullability(result_dtype: &DType, values: ArrayRef) -> VortexResult { + if values.dtype() == result_dtype { + Ok(values) + } else { + values.cast(result_dtype.clone()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs new file mode 100644 index 00000000000..d14fed87370 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +/// The result of resolving batch validity. +enum ResolvedValidity { + /// The output for an all-valid or all-null batch. + Output(ArrayRef), + + /// A mask with both valid and invalid rows. + PartiallyValid(Mask), +} + +impl Batch { + /// Resolve validity, try unfiltered execution, then fall back to filtering. + pub(super) fn execute_valid_only( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedValidity::Output(output) => return Ok(output), + ResolvedValidity::PartiallyValid(valid) => valid, + }; + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Materialize validity and handle all-valid or all-null batches. + fn resolve_validity( + &self, + kernel: &impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Check all-true before all-false: an empty mask is both, and must not be treated as + // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). + if valid.all_true() { + let values = VortexResult::from(kernel( + self.execution_args(&self.inputs, self.row_count), + ctx, + )?)?; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = self.finalize_output(values, self.row_count)?; + + return Ok(ResolvedValidity::Output(values)); + } + + if valid.all_false() { + return Ok(ResolvedValidity::Output(self.all_null())); + } + + Ok(ResolvedValidity::PartiallyValid(valid)) + } + + /// Try execution against the original inputs, then mask a returned full-length result. + fn try_execute_unfiltered( + &self, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(execution) = try_unfiltered( + self.execution_args(&self.inputs, self.row_count), + valid, + ctx, + )? + else { + return Ok(None); + }; + let values = VortexResult::from(execution)?; + let values = self.validate_kernel_output(values, valid.len(), ctx)?; + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + self.finalize_output(values.mask(mask)?, valid.len()) + .map(Some) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs deleted file mode 100644 index 8a0e10d79df..00000000000 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs +++ /dev/null @@ -1,475 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Applies columnar semantics around typed row-kernel invocations. -//! -//! [`Batch`] owns strict null propagation, constant broadcasting, execution strategy selection, and -//! output validation. The row kernel therefore handles only decoded values and its selected output -//! capability. - -use smallvec::SmallVec; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_ensure_eq; -use vortex_mask::AllOr; -use vortex_mask::Mask; - -use super::BatchPlan; -use super::RowPolicy; -use super::args::BorrowedExecutionArgs; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrays::MaskedArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::scalar::Scalar; -use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ScalarFnId; -use crate::scalar_fn::unstable::row::execute::RowExecution; -use crate::scalar_fn::unstable::row::types::batch_constant; -use crate::validity::Validity; - -/// The result of resolving batch validity. -enum ResolvedValidity { - /// The output for an all-valid or all-null batch. - Output(ArrayRef), - - /// A mask with both valid and invalid rows. - PartiallyValid(Mask), -} - -/// One batch of inputs and the metadata needed before its row kernel runs. -pub(crate) struct Batch { - /// The function being executed, named in the errors this raises. - id: ScalarFnId, - - /// The number of rows in the original execution scope. - row_count: usize, - - /// The input columns, collected once: constant folding inspects them and the filter strategy - /// filters them. - inputs: SmallVec<[ArrayRef; 4]>, - - /// The input dtypes, collected with the columns and reused by both planning and execution. - arg_dtypes: SmallVec<[DType; 4]>, - - /// The conjoined input validity, so a row of the output is valid iff it is valid in every - /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. - validity: Validity, - - /// The dtype the function declares for these inputs, which the kernel's output is reconciled - /// against. Already widened to nullable if any input is nullable. - result_dtype: DType, - - /// The non-nullable dtype the dispatched output capability builds, computed while planning. - output_dtype: DType, - - /// How the concrete dispatch executes nullable rows. - policy: RowPolicy, -} - -impl Batch { - /// Collect the inputs and derive their dtypes, validity, and execution policy. - /// - /// **Not** for a nullary function: with no inputs there is no validity to propagate and no - /// per-row work to fold, and the all-constant check below would vacuously pass. - pub(crate) fn new( - id: ScalarFnId, - args: &dyn ExecutionArgs, - plan: impl FnOnce(&[DType]) -> VortexResult, - ) -> VortexResult { - let row_count = args.row_count(); - let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) - .map(|index| args.get(index)) - .collect::>()?; - - for (index, input) in inputs.iter().enumerate() { - vortex_ensure_eq!( - input.len(), - row_count, - "the {id} input {index} must have {row_count} rows, got {}", - input.len(), - ); - } - - let arg_dtypes: SmallVec<[DType; 4]> = - inputs.iter().map(|input| input.dtype().clone()).collect(); - let plan = plan(&arg_dtypes)?; - let result_dtype = plan.result_dtype(&arg_dtypes); - - let mut validity = Validity::NonNullable; - for input in &inputs { - validity = validity.and(input.validity()?)?; - } - - Ok(Self { - id, - row_count, - inputs, - arg_dtypes, - validity, - result_dtype, - output_dtype: plan.output_dtype, - policy: plan.policy, - }) - } - - /// Apply constant folding and null handling around `kernel`. - /// - /// When the mask contains valid and invalid rows, `try_unfiltered` may avoid filtering. - /// `Ok(None)` filters the valid rows and scatters the output back. Every result is checked - /// against the planned shape and dtype. - pub(crate) fn execute( - &self, - kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, - try_unfiltered: impl FnOnce( - BorrowedExecutionArgs<'_>, - &Mask, - &mut ExecutionCtx, - ) -> VortexResult>, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - // Strictness: an all-null batch has no observable row work. Keep the literal-constant - // check explicit alongside the conjoined validity invariant. - if matches!(self.validity, Validity::AllInvalid) - || self.inputs.iter().any(|input| { - input - .as_opt::() - .is_some_and(|constant| constant.scalar().is_null()) - }) - { - return Ok(self.all_null()); - } - - // All inputs constant, and their conjoined validity proves every row non-null. This sees - // through extension and masked wrappers just like argument decoding does. - if self.row_count > 0 - && self.validity.definitely_no_nulls() - && self - .inputs - .iter() - .all(|input| batch_constant(input).is_some()) - { - return self.broadcast_one_row(kernel, ctx); - } - - match self.policy { - RowPolicy::Dense => self.execute_dense(kernel, false, ctx), - RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), - RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx), - } - } - - /// Evaluate one row of constant inputs and broadcast the validated result. - fn broadcast_one_row( - &self, - kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let one_row: SmallVec<[ArrayRef; 4]> = self - .inputs - .iter() - .map(|input| input.slice(0..1)) - .collect::>()?; - - let result = VortexResult::from(kernel(self.execution_args(&one_row, 1), ctx)?)?; - let result = self.validate_kernel_output(result, 1, ctx)?; - let result = self.finalize_output(result, 1)?; - let scalar = result.execute_scalar(0, ctx)?; - - Ok(ConstantArray::new(scalar, self.row_count).into_array()) - } - - /// Run every stored payload, then attach the input validity without materializing its mask. - fn execute_dense( - &self, - kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, - retry_deferred_error: bool, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let values = match kernel(self.execution_args(&self.inputs, self.row_count), ctx)? { - RowExecution::Output(values) => values, - RowExecution::DeferredError(error) if retry_deferred_error => { - let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; - - // Unlike `resolve_validity`, all-true preserves the deferred error and all-false - // suppresses evidence that came entirely from null rows. An empty loop cannot - // produce deferred evidence, so the ambiguous empty mask cannot reach this arm. - if valid.all_true() { - return Err(error); - } - if valid.all_false() { - return Ok(self.all_null()); - } - - // Deferred retry receives only the dense kernel. Filter first so the retry cannot - // evaluate null rows again. - return self.filter_and_scatter(kernel, &valid, ctx); - } - RowExecution::DeferredError(error) => return Err(error), - }; - let values = self.validate_kernel_output(values, self.row_count, ctx)?; - - match self.validity.clone() { - Validity::NonNullable | Validity::AllValid => { - self.finalize_output(values, self.row_count) - } - Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), - // Handled by the guard above, before the kernel ran. - Validity::AllInvalid => Ok(self.all_null()), - } - } - - /// Materialize validity and handle all-valid or all-null batches. - fn resolve_validity( - &self, - kernel: &impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; - - // Check all-true before all-false: an empty mask is both, and must not be treated as - // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). - if valid.all_true() { - let values = VortexResult::from(kernel( - self.execution_args(&self.inputs, self.row_count), - ctx, - )?)?; - let values = self.validate_kernel_output(values, self.row_count, ctx)?; - let values = self.finalize_output(values, self.row_count)?; - - return Ok(ResolvedValidity::Output(values)); - } - - if valid.all_false() { - return Ok(ResolvedValidity::Output(self.all_null())); - } - - Ok(ResolvedValidity::PartiallyValid(valid)) - } - - /// Resolve validity, try unfiltered execution, then fall back to filtering. - fn execute_valid_only( - &self, - kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, - try_unfiltered: impl FnOnce( - BorrowedExecutionArgs<'_>, - &Mask, - &mut ExecutionCtx, - ) -> VortexResult>, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let valid = match self.resolve_validity(&kernel, ctx)? { - ResolvedValidity::Output(output) => return Ok(output), - ResolvedValidity::PartiallyValid(valid) => valid, - }; - - if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { - return Ok(result); - } - - self.filter_and_scatter(kernel, &valid, ctx) - } - - /// Try execution against the original inputs, then mask a returned full-length result. - fn try_execute_unfiltered( - &self, - try_unfiltered: impl FnOnce( - BorrowedExecutionArgs<'_>, - &Mask, - &mut ExecutionCtx, - ) -> VortexResult>, - valid: &Mask, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let Some(execution) = try_unfiltered( - self.execution_args(&self.inputs, self.row_count), - valid, - ctx, - )? - else { - return Ok(None); - }; - let values = VortexResult::from(execution)?; - let values = self.validate_kernel_output(values, valid.len(), ctx)?; - - let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); - self.finalize_output(values.mask(mask)?, valid.len()) - .map(Some) - } - - /// Filter to valid rows, run the kernel, then scatter into a null-padded output. - fn filter_and_scatter( - &self, - kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, - valid: &Mask, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let filtered: SmallVec<[ArrayRef; 4]> = self - .inputs - .iter() - .map(|input| input.filter(valid.clone())) - .collect::>()?; - - let values = VortexResult::from(kernel( - self.execution_args(&filtered, valid.true_count()), - ctx, - )?)?; - let values = self.validate_kernel_output(values, valid.true_count(), ctx)?; - - self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) - } - - fn all_null(&self) -> ArrayRef { - ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() - } - - /// Pair an input view with this batch's planning metadata. - fn execution_args<'b>( - &'b self, - arrays: &'b [ArrayRef], - row_count: usize, - ) -> BorrowedExecutionArgs<'b> { - BorrowedExecutionArgs::new( - arrays, - row_count, - &self.arg_dtypes, - &self.output_dtype, - self.policy, - ) - } - - fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { - reconcile_output(self.id, &self.result_dtype, expected_len, values) - } - - /// Validate the output from a row kernel before batch validity is attached. - fn validate_kernel_output( - &self, - values: ArrayRef, - expected_len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx) - } - - /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set - /// bits, producing an array of length `valid.len()` that is null at every unset position. - fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult { - vortex_ensure_eq!( - values.len(), - valid.true_count(), - "the {} kernel output must contain {} filtered rows, got {}", - self.id, - valid.true_count(), - values.len(), - ); - - let AllOr::Some(slices) = valid.slices() else { - // The caller handles the all-true and all-false masks. - vortex_bail!( - "scatter_valid requires valid and invalid rows, got an all-valid or all-invalid mask" - ); - }; - - // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index - // 0, and any in-bounds index would do since they are masked out below (values is non-empty - // here). - let mut indices = vec![0u64; valid.len()]; - let mut rank = 0u64; - for &(start, end) in slices { - for index in &mut indices[start..end] { - *index = rank; - rank += 1; - } - } - let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array(); - - let scattered = values.take(indices)?; - - // A nullable gathered array cannot be wrapped because a `Masked` child must be all valid. - // The general masking pass unions its nulls with the batch validity instead. - if scattered.dtype().is_nullable() { - let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); - return scattered.mask(mask); - } - - // The gathered values are all valid, so attaching validity is sufficient. - Ok(MaskedArray::try_new( - scattered, - Validity::from_mask(valid.clone(), Nullability::Nullable), - )? - .into_array()) - } -} - -/// Validate the output produced directly by a row kernel. -/// -/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring -/// nullability, and every produced row **must** be valid. Batch execution owns strict null -/// propagation and attaches input-derived validity only after this boundary. -pub(crate) fn finalize_kernel_output( - id: ScalarFnId, - result_dtype: &DType, - expected_len: usize, - values: ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { - validate_output(id, result_dtype, expected_len, &values)?; - vortex_ensure!( - values.all_valid(ctx)?, - "the {id} row kernel must produce only valid rows, got at least one null row", - ); - - cast_output_nullability(result_dtype, values) -} - -/// Reconcile an output with the function's declared shape and nullability. -fn reconcile_output( - id: ScalarFnId, - result_dtype: &DType, - expected_len: usize, - values: ArrayRef, -) -> VortexResult { - validate_output(id, result_dtype, expected_len, &values)?; - - cast_output_nullability(result_dtype, values) -} - -/// Validate an output's shape and logical dtype without executing a nullability cast. -fn validate_output( - id: ScalarFnId, - result_dtype: &DType, - expected_len: usize, - values: &ArrayRef, -) -> VortexResult<()> { - vortex_ensure_eq!( - values.len(), - expected_len, - "the {id} kernel output must contain {expected_len} rows, got {}", - values.len(), - ); - vortex_ensure!( - values.dtype().eq_ignore_nullability(result_dtype), - "the {id} kernel output dtype must match {result_dtype} ignoring nullability, got {}", - values.dtype(), - ); - - Ok(()) -} - -/// Cast only the output nullability after its shape, dtype, and validity are accepted. -fn cast_output_nullability(result_dtype: &DType, values: ArrayRef) -> VortexResult { - if values.dtype() == result_dtype { - Ok(values) - } else { - values.cast(result_dtype.clone()) - } -} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs index b501965fa61..a246ba3e30b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -11,15 +11,53 @@ //! applies that strategy, and [`BorrowedExecutionArgs`] pairs each kernel invocation with its //! planning metadata. +use smallvec::SmallVec; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; + mod args; pub(super) use args::BorrowedExecutionArgs; -mod execution; -pub(super) use execution::Batch; -pub(super) use execution::finalize_kernel_output; +mod execute; +pub(super) use execute::finalize_kernel_output; + +mod planning; pub(super) use super::visitor::BatchPlan; pub(super) use super::visitor::RowPolicy; +/// One batch of inputs and the metadata needed before its row kernel runs. +pub(crate) struct Batch { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The number of rows in the original execution scope. + row_count: usize, + + /// The input columns, collected once: constant folding inspects them and the filter strategy + /// filters them. + inputs: SmallVec<[ArrayRef; 4]>, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: SmallVec<[DType; 4]>, + + /// The conjoined input validity, so a row of the output is valid iff it is valid in every + /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. + validity: Validity, + + /// The dtype the function declares for these inputs, which the kernel's output is reconciled + /// against. Already widened to nullable if any input is nullable. + result_dtype: DType, + + /// The non-nullable dtype the dispatched output capability builds, computed while planning. + output_dtype: DType, + + /// How the concrete dispatch executes nullable rows. + policy: RowPolicy, +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs new file mode 100644 index 00000000000..24bb31b2709 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use super::Batch; +use super::BatchPlan; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; + +impl Batch { + /// Collect the inputs and derive their dtypes, validity, and execution policy. + /// + /// **Not** for a nullary function: with no inputs there is no validity to propagate and no + /// per-row work to fold, and the all-constant check would vacuously pass. + pub(crate) fn new( + id: ScalarFnId, + args: &dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let row_count = args.row_count(); + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>()?; + + for (index, input) in inputs.iter().enumerate() { + vortex_ensure_eq!( + input.len(), + row_count, + "the {id} input {index} must have {row_count} rows, got {}", + input.len(), + ); + } + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let result_dtype = plan.result_dtype(&arg_dtypes); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + row_count, + inputs, + arg_dtypes, + validity, + result_dtype, + output_dtype: plan.output_dtype, + policy: plan.policy, + }) + } + + /// Pair an input view with this batch's planning metadata. + pub(super) fn execution_args<'b>( + &'b self, + arrays: &'b [ArrayRef], + row_count: usize, + ) -> BorrowedExecutionArgs<'b> { + BorrowedExecutionArgs::new( + arrays, + row_count, + &self.arg_dtypes, + &self.output_dtype, + self.policy, + ) + } +} From 413796dd61aae2470d1139cf014bc8844d9ca792 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 12:50:10 -0400 Subject: [PATCH 153/160] Execute primitive numeric operators with RowFn Signed-off-by: Connor Tsui --- vortex-array/benches/binary_ops.rs | 58 +++ .../typed_view/primitive/numeric_operator.rs | 2 +- .../scalar_fn/fns/binary/numeric/checked.rs | 88 +---- .../src/scalar_fn/fns/binary/numeric/mod.rs | 12 +- .../scalar_fn/fns/binary/numeric/primitive.rs | 370 +++++------------- .../src/scalar_fn/fns/binary/numeric/row.rs | 135 +++++++ .../src/scalar_fn/fns/binary/numeric/tests.rs | 9 +- 7 files changed, 306 insertions(+), 368 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/numeric/row.rs diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 6a07d03f50b..ccd440ff939 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -38,10 +38,60 @@ static SESSION: LazyLock = LazyLock::new(array_session); const LEN: usize = 32_768; +const ROWFN_MATRIX_CASES: &[(usize, RowFnShape)] = &[ + (128, RowFnShape::PerRowPerRow), + (128, RowFnShape::PerRowConstant), + (128, RowFnShape::ConstantPerRow), + (128, RowFnShape::PerRowNullableConstant), + (LEN, RowFnShape::PerRowPerRow), + (LEN, RowFnShape::PerRowConstant), + (LEN, RowFnShape::ConstantPerRow), + (LEN, RowFnShape::PerRowNullableConstant), +]; + +#[derive(Clone, Copy, Debug)] +enum RowFnShape { + PerRowPerRow, + PerRowConstant, + ConstantPerRow, + PerRowNullableConstant, +} + /// Decimal Mul and Div cost far more per lane than Add, so they run over a shorter array to keep /// the instrumented CodSpeed runs quick. const DECIMAL_MUL_DIV_LEN: usize = 8_192; +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_add(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Add); +} + +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_subtract(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Sub); +} + +#[divan::bench(args = ROWFN_MATRIX_CASES)] +fn rowfn_multiply(bencher: Bencher, &(len, shape): &(usize, RowFnShape)) { + bench_rowfn_shape(bencher, len, shape, Operator::Mul); +} + +fn bench_rowfn_shape(bencher: Bencher, len: usize, shape: RowFnShape, operator: Operator) { + let per_row = + || PrimitiveArray::from_iter((0..len).map(|index| (index % 1_024) as i64 + 1)).into_array(); + let constant = || ConstantArray::new(17_i64, len).into_array(); + let nullable_constant = || ConstantArray::new(Some(17_i64), len).into_array(); + + let (lhs, rhs) = match shape { + RowFnShape::PerRowPerRow => (per_row(), per_row()), + RowFnShape::PerRowConstant => (per_row(), constant()), + RowFnShape::ConstantPerRow => (constant(), per_row()), + RowFnShape::PerRowNullableConstant => (per_row(), nullable_constant()), + }; + + bench_primitive(bencher, lhs, rhs, operator); +} + #[divan::bench] fn add_i64_nonnull(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); @@ -170,6 +220,14 @@ fn div_i64_nonnull(bencher: Bencher) { bench_primitive(bencher, lhs, rhs, Operator::Div); } +#[divan::bench] +fn div_i64_nullable(bencher: Bencher) { + let lhs = primitive_nullable(1_000_000, 7).into_array(); + let rhs = primitive_nullable(17, 5).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Div); +} + #[divan::bench] fn sub_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index 6b389a0554b..d7bcc8d0b4c 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_err; use crate::scalar_fn::fns::operators::Operator; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. pub enum NumericOperator { /// Binary element-wise addition of two arrays or of two scalars. diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index afb709abe1c..054846b7ef7 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Checked-lane execution for numeric kernels, driven by the shared +//! Checked-lane execution for the decimal kernels, driven by the shared //! `vortex-compute` lane kernels. - -use std::ops::BitOrAssign; +//! +//! The primitive widths do not come through here: they are computed one row at a time by +//! [`row`](super::row), which writes a value for every row and reduces failure evidence without +//! scanning the finished output. use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -13,33 +15,18 @@ use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_mask::AllOr; use vortex_mask::Mask; -/// Evidence that a lane failed, anything other than [`Default`] meaning failure. -/// -/// `bool` is the ordinary choice; [`map_checked_into`] explains why the wider members exist and -/// asserts the width bound that membership here does **not** imply. -/// -/// [`map_checked_into`]: IndexedSourceExt::map_checked_into -pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - -impl Failure for bool {} -impl Failure for u8 {} -impl Failure for u16 {} -impl Failure for u32 {} -impl Failure for u64 {} - /// Apply the fallible `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it returns `None` on a _valid_ lane. +/// `Err(first_failing_valid_lane)` only when it returns `None` on a valid lane. /// /// `apply` also runs on invalid lanes, whose failures are masked out and whose values are /// unspecified, so it must be total: no panics or side effects on any stored lane value. /// -/// This drives the one-pass early-exit kernels, which abort at the end of the enclosing 64-lane -/// chunk. Use it when per-lane failure handling is cheap relative to the operation, as in integer -/// division. Prefer [`checked_apply_lanes`] when the failure check itself vectorizes. +/// This drives the one-pass early-exit kernels: failures abort at the end of the enclosing +/// 64-lane chunk. It suits an operation whose per-lane failure handling is cheap relative to the +/// operation itself, which is what the decimal kernels and their per-lane casts are. /// -/// `#[inline]`: this must inline into the caller that builds the closure, so that a captured -/// constant operand flattens into a register rather than living behind a pointer the loop reloads -/// on every lane, which blocks vectorization. +/// Keep this wrapper inlineable so captured constants can become loop invariants in the caller. +/// The lane kernels retain their own inlining decisions. #[inline] pub(super) fn checked_lanes( source: S, @@ -61,7 +48,6 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, @@ -73,55 +59,3 @@ where Ok(values.freeze()) } - -/// Apply the split value/failure `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it flags a _valid_ lane. -/// -/// The hot pass writes every value unconditionally and OR-reduces one piece of evidence, leaving -/// the loop free of per-lane selects and per-chunk exit branches. Only if a lane flagged does a -/// cold second pass re-run `apply` through the early-exit kernels, which drop null-lane failures -/// and attribute the first valid one. -/// -/// Like [`checked_lanes`], `apply` runs on invalid lanes and must be total. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -pub(super) fn checked_apply_lanes( - source: S, - valid_rows: &Mask, - mut apply: Apply, -) -> Result, usize> -where - S: IndexedSource + Copy, - T: Copy + Default, - Fail: Failure, - Apply: FnMut(S::Item) -> (T, Fail), -{ - let len = source.len(); - debug_assert_eq!(len, valid_rows.len()); - - let valid_bits = match valid_rows.bit_buffer() { - AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), - AllOr::Some(valid_bits) => Some(valid_bits), - }; - - let mut values = BufferMut::::with_capacity(len); - - let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &mut apply) != Fail::default() { - let mut checked = |item: S::Item| { - let (value, failure) = apply(item); - (failure == Fail::default()).then_some(value) - }; - match valid_bits { - None => source.try_map_into(out, &mut checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?, - } - } - - // SAFETY: the kernels initialize every lane in `out`. - unsafe { values.set_len(len) }; - - Ok(values.freeze()) -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 6dc0de0fbea..a0c427b142e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,16 +4,19 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! +//! The primitive widths are computed by a +//! [`RowFn`](crate::scalar_fn::unstable::row::RowFn), which owns null handling, constants, and +//! validity for them; see [`row`]. Decimal keeps its own columnar implementation in [`decimal`]. +//! //! [`Binary`]: super::Binary mod checked; mod decimal; mod primitive; -#[cfg(test)] -mod tests; +mod row; use decimal::execute_numeric_decimal; -use primitive::execute_numeric_primitive; +use row::execute_numeric_primitive; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -81,3 +84,6 @@ fn build_empty_result( Ok(Canonical::empty(&result_dtype).into_array()) } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 8fd53d15216..7cbe88c7c8e 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -1,73 +1,52 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::checked::Failure; -use super::checked::checked_apply_lanes; -use super::checked::checked_lanes; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; +//! Checked arithmetic for one primitive row. + +use std::ops::BitOrAssign; + use crate::dtype::NativePType; -use crate::dtype::PType; use crate::dtype::half::f16; -use crate::match_each_native_ptype; -use crate::scalar::NumericOperator; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; -use crate::validity::Validity; - -struct CheckedAdd; -struct CheckedSub; +/// Checked addition, failing on integer overflow. +pub(super) struct CheckedAdd; -struct CheckedMul; +/// Checked subtraction, failing on integer overflow. +pub(super) struct CheckedSub; -struct CheckedDiv; +/// Checked multiplication, failing on integer overflow. +pub(super) struct CheckedMul; -trait CheckedPrimitiveOp: Sized { - /// Message for the error raised when this operation fails on a valid lane. - const ERROR: &'static str; +/// Checked division, failing on integer division by zero and on `MIN / -1`. +pub(super) struct CheckedDiv; - /// Whether to check inside the value loop and exit early, rather than reducing evidence over - /// the whole batch and re-running only once some lane has flagged. - const CHECKED_VALUE_LOOP: bool = false; +/// OR-reducible evidence that a row failed, with [`Default`] meaning success. +pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - /// How this operation reports a failing lane. See [`Failure`]. - type Failure: Failure; +impl Failure for bool {} +impl Failure for u8 {} +impl Failure for u16 {} +impl Failure for u32 {} +impl Failure for u64 {} - /// Compute a lane's value and its failure evidence together. - /// - /// Overflowing-style rather than `Option`, so the vectorizable kernels can write every - /// lane's value unconditionally and reduce the evidence separately. [`Self::checked`] is the - /// shape for the scalar and early-exit paths. - fn apply(lhs: T, rhs: T) -> (T, Self::Failure); +/// One arithmetic operator at one width, split into its value and failure evidence. +pub(super) trait CheckedPrimitiveOp: 'static + Sized { + /// The error reported for a batch in which some valid row failed. + const ERROR: &'static str; - /// [`Self::apply`] folded into the `Option` that the early-exit kernels take. - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - let (value, failed) = Self::apply(lhs, rhs); + /// How this operation reports a failing row. See [`Failure`]. + type Fail: Failure; - (failed == Self::Failure::default()).then_some(value) - } + /// The result of this operation, paired with evidence of whether the row failed. + fn apply(lhs: T, rhs: T) -> (T, Self::Fail); } impl CheckedPrimitiveOp for CheckedAdd { const ERROR: &'static str = "integer overflow in checked add"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.add_value(rhs), lhs.add_error(rhs)) } @@ -76,9 +55,9 @@ impl CheckedPrimitiveOp for CheckedAdd { impl CheckedPrimitiveOp for CheckedSub { const ERROR: &'static str = "integer overflow in checked sub"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.sub_value(rhs), lhs.sub_error(rhs)) } @@ -87,9 +66,9 @@ impl CheckedPrimitiveOp for CheckedSub { impl CheckedPrimitiveOp for CheckedMul { const ERROR: &'static str = "integer overflow in checked mul"; - type Failure = T::MulFailure; + type Fail = T::MulFailure; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, T::MulFailure) { (lhs.mul_value(rhs), lhs.mul_failure(rhs)) } @@ -97,16 +76,10 @@ impl CheckedPrimitiveOp for CheckedMul { impl CheckedPrimitiveOp for CheckedDiv { const ERROR: &'static str = "integer division by zero or overflow in checked div"; - // Integer division still lowers to scalar divides, so the split - // value/error-scan loop used to auto-vectorize add/sub/mul only adds a - // second full scan. Use the one-pass early-exit checked kernel for integer - // division, matching Arrow/Velox. Float division has no checked errors and - // stays on the split/vectorizable default path. - const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { let failed = lhs.div_error(rhs); let value = if failed { @@ -116,151 +89,16 @@ impl CheckedPrimitiveOp for CheckedDiv { }; (value, failed) } - - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - lhs.div_checked(rhs) - } } -/// Execute a numeric operation between two primitive-typed arrays. -pub(super) fn execute_numeric_primitive( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - - match_each_native_ptype!(ptype, |T| { - match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), - } - }) -} - -fn execute_checked_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType, - Op: CheckedPrimitiveOp, - Scalar: From, - Scalar: From>, -{ - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - let len = lhs.len(); - debug_assert_eq!(len, rhs.len()); - - let validity = lhs.validity().and(rhs.validity())?; - let valid_rows = validity.execute_mask(len, ctx)?; - - let values = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_op_lanes::<_, T, Op>( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - &valid_rows, - |(lhs, rhs)| (lhs, rhs), - ), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Capture the constant by value so it stays hoisted out of the lane loop. - let rhs = *rhs; - checked_op_lanes::<_, T, Op>(lhs.as_slice(), &valid_rows, move |lhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let lhs = *lhs; - checked_op_lanes::<_, T, Op>(rhs.as_slice(), &valid_rows, move |rhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let value = Op::checked(*lhs, *rhs) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - return Ok(constant_result_array(value, len, &result_dtype)); - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => Ok(Buffer::zeroed(len)), - } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - - primitive_result_array::(values, validity, &result_dtype) -} - -/// Run `Op` over the lanes of `source` in the loop shape it declares: the one-pass early-exit -/// kernel when `Op::CHECKED_VALUE_LOOP` is set, and the split value/failure kernel otherwise. +/// Per-width arithmetic used to compute values and failure evidence. /// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -fn checked_op_lanes( - source: S, - valid_rows: &Mask, - mut to_operands: impl FnMut(S::Item) -> (T, T), -) -> Result, usize> -where - S: IndexedSource + Copy, - T: NativePType, - Op: CheckedPrimitiveOp, -{ - if Op::CHECKED_VALUE_LOOP { - checked_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::checked(lhs, rhs) - }) - } else { - checked_apply_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::apply(lhs, rhs) - }) - } -} - -fn primitive_result_array( - values: Buffer, - validity: Validity, - dtype: &DType, -) -> VortexResult { - let array = PrimitiveArray::new(values, validity).into_array(); - if array.dtype() == dtype { - return Ok(array); - } - array.cast(dtype.clone()) -} - -fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef -where - T: NativePType, - Scalar: From + From>, -{ - if dtype.is_nullable() { - ConstantArray::new(Some(value), len).into_array() - } else { - ConstantArray::new(value, len).into_array() - } -} - -trait CheckedArithmetic: NativePType { - const DIV_CHECKS_IN_VALUE_LOOP: bool; - - /// How multiplication reports a failing lane, which is width-dependent in a way the other - /// operations are not. `impl_checked_unsigned!` and `impl_checked_signed!` say why each family - /// reports what it reports. +/// The add, subtract, and multiply value methods **must** be total over every stored lane value. +/// [`Self::div_value`] may assume that [`Self::div_error`] returned `false` for the same operands. +pub(super) trait CheckedArithmetic: NativePType { + /// How multiplication reports a failing row. + /// + /// This may be a word rather than `bool` when narrowing evidence would block vectorization. type MulFailure: Failure; fn add_value(self, rhs: Self) -> Self; @@ -269,18 +107,15 @@ trait CheckedArithmetic: NativePType { fn sub_error(self, rhs: Self) -> bool; fn mul_value(self, rhs: Self) -> Self; fn mul_failure(self, rhs: Self) -> Self::MulFailure; + + /// Divide operands that [`Self::div_error`] accepted. fn div_value(self, rhs: Self) -> Self; + + /// Return whether [`Self::div_value`] would trap for these operands. fn div_error(self, rhs: Self) -> bool; - fn div_checked(self, rhs: Self) -> Option; } -/// The integer arithmetic every width shares, given the two things that actually differ between -/// them: how multiplication reports a failing lane, and how add/sub/div detect one. -/// -/// Only `mul_failure` genuinely varies per width, so it is the macro's parameter and everything -/// else is written once. The `$mul_failure_ty` and body are supplied by the caller because the -/// choice between a widened high half and a `bool` is exactly the vectorization decision described -/// on [`Failure`]. +/// Generate the shared integer operations from their failure predicates. macro_rules! impl_checked_integer { ( $ty:ty, @@ -291,67 +126,57 @@ macro_rules! impl_checked_integer { = |$mf_lhs:ident, $mf_rhs:ident| $mul_failure:expr, ) => { impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = true; - type MulFailure = $mul_failure_ty; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self.wrapping_add(rhs) } - #[inline(always)] + #[inline] fn add_error(self, rhs: Self) -> bool { let ($add_lhs, $add_rhs) = (self, rhs); $add_error } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self.wrapping_sub(rhs) } - #[inline(always)] + #[inline] fn sub_error(self, rhs: Self) -> bool { let ($sub_lhs, $sub_rhs) = (self, rhs); $sub_error } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self.wrapping_mul(rhs) } - #[inline(always)] + #[inline] $(#[$mul_failure_attr])* fn mul_failure(self, rhs: Self) -> $mul_failure_ty { let ($mf_lhs, $mf_rhs) = (self, rhs); $mul_failure } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, rhs: Self) -> bool { let ($div_lhs, $div_rhs) = (self, rhs); $div_error } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - self.checked_div(rhs) - } } }; } -/// The unsigned widths. `add`, `sub` and `div` are written once here, and `widening_mul` covers -/// every width including the 64-bit one, which widens into `u128`: the discarded high half of the -/// widened product is the failure evidence, and costs none of the comparison LLVM folds into -/// `umul.with.overflow`. +/// Unsigned multiplication reports its discarded high half as failure evidence. macro_rules! impl_checked_unsigned { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_integer!( @@ -364,12 +189,7 @@ macro_rules! impl_checked_unsigned { }; } -/// The signed widths, on the same principle. `widening_mul` is the shorthand for the narrow widths, -/// whose two-sided range check over a wider product is not the shape LLVM folds into an overflow -/// intrinsic, so they vectorize while reporting a plain `bool`. The 64-bit width cannot: deriving a -/// `bool` there costs the comparison that scalarizes the loop, so `high_half_mul` hands back the -/// discarded high half as a word. Both arms derive every shift and bound from `$ty`, so -/// instantiating one at a new width cannot silently keep another width's constants. +/// Signed widths use a range check or discarded high-half evidence. macro_rules! impl_checked_signed { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_signed!($ty, mul_failure: bool = |lhs, rhs| { @@ -377,9 +197,6 @@ macro_rules! impl_checked_signed { product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide }); }; - // Zero exactly when the product fits: a signed multiply overflows iff the high half of the true - // product differs from the sign extension of the half that was kept, so the two XOR to zero on - // the lanes that fit. `tests::test_multiply_overflow_boundaries` pins the boundaries. ($ty:ty, high_half_mul: $wide:ty => $failure:ty) => { impl_checked_signed!($ty, mul_failure: #[expect( clippy::cast_possible_truncation, @@ -389,13 +206,17 @@ macro_rules! impl_checked_signed { let kept = wide as $ty; let discarded = (wide >> <$ty>::BITS) as $ty; + // A product fits exactly when its discarded half is the sign extension of the kept + // half. XOR reduces that comparison to zero evidence for success and nonzero evidence + // for overflow without converting the wide product to a branch. + (discarded ^ (kept >> (<$ty>::BITS - 1))) as $failure }); }; ( $ty:ty, mul_failure: $(#[$mul_failure_attr:meta])* $mul_failure_ty:ty - = |$l:ident, $r:ident| $mul_failure:expr + = |$lhs:ident, $rhs:ident| $mul_failure:expr ) => { impl_checked_integer!( $ty, @@ -408,7 +229,7 @@ macro_rules! impl_checked_signed { ((lhs ^ rhs) & (lhs ^ value)) < 0 }, div_error: |lhs, rhs| rhs == 0 || (lhs == <$ty>::MIN && rhs == -1), - mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$l, $r| $mul_failure, + mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$lhs, $rhs| $mul_failure, ); }; } @@ -417,54 +238,47 @@ macro_rules! impl_checked_float { ($($ty:ty),+ $(,)?) => { $( impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = false; - type MulFailure = bool; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self + rhs } - #[inline(always)] + #[inline] fn add_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self - rhs } - #[inline(always)] + #[inline] fn sub_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self * rhs } - #[inline(always)] + #[inline] fn mul_failure(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, _rhs: Self) -> bool { false } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - Some(self / rhs) - } } )+ }; @@ -484,30 +298,27 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; - /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change - /// pivots, and a spread of magnitudes that straddles the 64-bit split. + /// Values around zero, signed extrema, and 32- and 64-bit boundaries where the discarded + /// multiplication half or its sign extension changes. const PROBES: &[i64] = &[ - 0, // - 1, // - -1, // - 2, // - -2, // - 3, // - i64::MIN, // - i64::MIN + 1, // - i64::MAX, // - i64::MAX - 1, // - 1 << 31, // - 1 << 32, // - 1 << 62, // - -(1 << 62), // - 0x7FFF_FFFF, // - -0x8000_0000, // + 0, // Additive identity. + 1, // Smallest positive value. + -1, // All sign bits set. + 2, // Small positive power of two. + -2, // Small negative power of two. + 3, // Small non-power of two. + i64::MIN, // Minimum signed value. + i64::MIN + 1, // Minimum signed value's neighbor. + i64::MAX, // Maximum signed value. + i64::MAX - 1, // Maximum signed value's neighbor. + 1 << 31, // First positive value outside i32. + 1 << 32, // First value with bit 32 set. + 1 << 62, // Largest positive power of two in i64. + -(1 << 62), // Negative counterpart of the largest power of two. + 0x7FFF_FFFF, // Maximum i32 represented as i64. + -0x8000_0000, // Minimum i32 represented as i64. ]; - /// Every `mul_failure` impl is either a bit trick or a two-sided range check, so hold each - /// against the obvious reference: `reference` is the width's own `checked_mul`, whose `None` - /// _is_ the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { let failed = lhs.mul_failure(rhs) != ::default(); @@ -522,14 +333,11 @@ mod tests { assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); let (lhs, rhs) = (lhs as u64, rhs as u64); - assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); } } } - /// The 8-bit widths are cheap enough to check exhaustively, which pins the shift of the - /// unsigned formula and the range check of the signed one against every product that exists. #[test] fn mul_failure_agrees_with_checked_mul_exhaustively_at_8_bits() { for lhs in u8::MIN..=u8::MAX { diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs new file mode 100644 index 00000000000..b6f565d8ff4 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive arithmetic execution through [`RowFn`]. +//! +//! `Binary` keeps its registered contract; [`NumericBinary`] is only an execution helper. Decimal +//! arithmetic remains on its existing columnar path. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::primitive::CheckedAdd; +use super::primitive::CheckedArithmetic; +use super::primitive::CheckedDiv; +use super::primitive::CheckedMul; +use super::primitive::CheckedPrimitiveOp; +use super::primitive::CheckedSub; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::NumericOperator; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::unstable::row::InitializedElement; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::UninitElementSink; +use crate::scalar_fn::unstable::row::execute_rows; + +pub(super) fn execute_numeric_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + execute_rows(&NumericBinary, &op, &args, ctx) +} + +/// Internal row execution for the primitive arithmetic operators. +#[derive(Clone)] +struct NumericBinary; + +impl RowFn for NumericBinary { + type Options = NumericOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + // Fallibility is queried without input dtypes, so this conservatively covers integer widths. + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + // `NumericBinary` is a private implementation detail of `Binary`: it is never registered or + // serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. + ScalarFnVTable::id(&Binary) + } + + fn dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = PType::try_from( + args.first() + .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?, + )?; + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_div::(visitor), + } + }) + } +} + +fn visit_checked(visitor: V) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + V: RowVisitor, +{ + visitor.visit_deferred::<(T, T), T, Op::Fail>( + |(lhs, rhs)| Op::apply(lhs, rhs), + |failure| { + if failure != ::default() { + return Err(numeric_error(Op::ERROR)); + } + + Ok(()) + }, + ) +} + +fn visit_div(visitor: V) -> VortexResult +where + T: CheckedArithmetic, + V: RowVisitor, +{ + if T::PTYPE.is_float() { + return visit_checked::(visitor); + } + + // Integer division is scalar and expensive, so deferring its cheap failure check preserves no + // vectorization. Check each divide immediately and stop at the first failure. + // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when + // they need to skip invalid rows. + visitor.visit_into::<(T, T), UninitElementSink, _>(|(lhs, rhs), output| { + let (value, failed) = CheckedDiv::apply(lhs, rhs); + if failed { + return Err(numeric_error(>::ERROR)); + } + + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) + }) +} + +/// Keep rich error construction out of row closures so the closures remain inlineable. +#[cold] +#[inline(never)] +fn numeric_error(message: &'static str) -> VortexError { + vortex_err!(InvalidArgument: "{message}") +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 7ee98ae717e..3813c8612b3 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -201,8 +201,7 @@ fn test_integer_array_array_errors_on_valid_lanes() { assert!(result.is_err()); } -/// Multiply two non-nullable lanes of `lhs` by two of `rhs`, expecting `Some(product)` where the -/// product fits and `None` where the checked kernel must report overflow. +/// Assert one checked multiplication through the complete array execution path. #[track_caller] fn assert_multiply(lhs: T, rhs: T, expected: Option) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -297,13 +296,11 @@ fn test_multiply_overflow_on_null_lane_ignored( Ok(()) } -/// The hot pass OR-reduces evidence across whole 64-lane chunks before anything looks at it, so an -/// overflow in a late chunk must still be caught, and must still be suppressed when its lane is -/// null. Every other test here fits in a single chunk and cannot show either. +/// An overflow late in the batch must still be reported, unless its row is null. #[rstest] #[case::reported(true)] #[case::suppressed_by_null(false)] -fn test_multiply_overflow_survives_chunk_reduction( +fn test_multiply_overflow_survives_batch_reduction( #[case] lane_is_valid: bool, ) -> VortexResult<()> { const LEN: u32 = 1000; From c057706c943eacef3b930e9f86bc3a5ff53c8b9d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 12:52:33 -0400 Subject: [PATCH 154/160] Execute primitive comparisons with RowFn Signed-off-by: Connor Tsui --- encodings/runend/src/trace_tests.rs | 8 + vortex-array/benches/compare.rs | 115 ++++++++++ .../src/scalar_fn/fns/binary/compare/mod.rs | 8 +- .../scalar_fn/fns/binary/compare/primitive.rs | 200 ++++++++++-------- .../fns/binary/compare/primitive/columnar.rs | 122 +++++++++++ .../primitive/operand.rs} | 19 +- .../src/scalar_fn/fns/binary/compare/tests.rs | 141 ++++++++++++ vortex-array/src/scalar_fn/fns/binary/mod.rs | 1 - vortex-array/src/test_harness/trace/tests.rs | 8 + vortex-btrblocks/src/trace_tests.rs | 16 ++ 10 files changed, 546 insertions(+), 92 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs rename vortex-array/src/scalar_fn/fns/binary/{primitive_operand.rs => compare/primitive/operand.rs} (78%) diff --git a/encodings/runend/src/trace_tests.rs b/encodings/runend/src/trace_tests.rs index 96f2afff50a..be8fd0ba762 100644 --- a/encodings/runend/src/trace_tests.rs +++ b/encodings/runend/src/trace_tests.rs @@ -73,6 +73,14 @@ fn trace_compare_on_runend() -> VortexResult<()> { iter 0 current=vortex.runend(bool, len=9) builder_active=false execute_until target=AnyCanonical root=vortex.binary(bool, len=3) iter 0 current=vortex.binary(bool, len=3) builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 1 current=vortex.bool(bool, len=3) builder_active=false return output=vortex.bool(bool, len=3) diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 4a399760dc2..9e6dd3e4e5b 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -4,6 +4,7 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use divan::counter::ItemsCount; use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; @@ -38,6 +39,7 @@ const ARRAY_SIZE: usize = 65_536; fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); bencher + .counter(ItemsCount::new(ARRAY_SIZE)) .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input @@ -49,6 +51,31 @@ fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { }); } +fn u8_array(offset: u8) -> ArrayRef { + (0u8..=u8::MAX) + .cycle() + .take(ARRAY_SIZE) + .map(|value| value.wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn i32_array(offset: i32) -> ArrayRef { + (0i32..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + +fn u64_array(offset: u64) -> ArrayRef { + (0u64..) + .take(ARRAY_SIZE) + .map(|value| value.wrapping_mul(31).wrapping_add(offset)) + .collect::>() + .into_array() +} + fn bool_array(rng: &mut StdRng) -> ArrayRef { BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.5))).into_array() } @@ -87,6 +114,13 @@ fn float_array(rng: &mut StdRng) -> ArrayRef { .into_array() } +fn f32_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f32..1.0)) + .collect::>() + .into_array() +} + fn string_array(rng: &mut StdRng) -> ArrayRef { VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { let len = rng.random_range(1usize..24); @@ -153,6 +187,14 @@ fn compare_int_constant(bencher: Bencher) { bench_compare(bencher, arr, constant, Operator::Gte); } +#[divan::bench] +fn compare_int_constant_lhs(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let constant = ConstantArray::new(50_000_000i64, ARRAY_SIZE).into_array(); + let arr = int_array(&mut rng); + bench_compare(bencher, constant, arr, Operator::Gte); +} + #[divan::bench] fn compare_int_eq(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -161,6 +203,55 @@ fn compare_int_eq(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Eq); } +#[divan::bench] +fn compare_i32(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = i32_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_i32_constant(bencher: Bencher) { + let lhs = i32_array(1); + let rhs = ConstantArray::new(1_000_000i32, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = u8_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u8_constant(bencher: Bencher) { + let lhs = u8_array(1); + let rhs = ConstantArray::new(127u8, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_constant(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = ConstantArray::new(1_000_000u64, ARRAY_SIZE).into_array(); + bench_compare(bencher, lhs, rhs, Operator::Gte); +} + +#[divan::bench] +fn compare_u64_eq(bencher: Bencher) { + let lhs = u64_array(1); + let rhs = u64_array(17); + bench_compare(bencher, lhs, rhs, Operator::Eq); +} + #[divan::bench] fn compare_float(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -169,6 +260,30 @@ fn compare_float(bencher: Bencher) { bench_compare(bencher, arr1, arr2, Operator::Gte); } +#[divan::bench] +fn compare_float_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = float_array(&mut rng); + let arr2 = float_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_f32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_f32_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = f32_array(&mut rng); + let arr2 = f32_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + #[divan::bench] fn compare_decimal(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index d25a652ee57..a36d0a22bde 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -4,9 +4,9 @@ //! Native comparison kernels. //! //! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every -//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from -//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise -//! comparator for nested types. There is no Arrow fallback. +//! comparison directly over Vortex canonical arrays: bit buffers for booleans, row or fused lane +//! kernels for primitives, lane kernels for decimals, binary views for strings and bytes, and a +//! row-wise comparator for nested types. There is no Arrow fallback. //! //! Floating point values compare with Vortex's total ordering (`NaN` is the largest value, //! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics. @@ -211,7 +211,7 @@ fn compare_arrays( ) .into_array()), DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx), - DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), + DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, ctx), DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Map(..) => { diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 1247358dce1..93afcf538ed 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -1,28 +1,30 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Native comparison of primitive arrays via bit-packing lane kernels. +//! Primitive comparison execution through [`RowFn`]. + +mod columnar; +mod operand; -use vortex_buffer::BitBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_bail; +use vortex_error::vortex_err; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::BoolArray; -use crate::arrays::ConstantArray; +#[cfg(target_arch = "x86_64")] +use crate::arrays::Constant; use crate::dtype::DType; use crate::dtype::NativePType; -use crate::dtype::Nullability; use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::compare::collect_bits; -use crate::scalar_fn::fns::binary::compare::collect_zip_bits; -use crate::scalar_fn::fns::binary::compare::compare_validity; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; /// Compare two primitive arrays of the same [`PType`]. /// @@ -32,99 +34,125 @@ pub(super) fn compare_primitive( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - match_each_native_ptype!(ptype, |T| { - compare_primitive_typed::(lhs, rhs, op, nullability, ctx) - }) + compare_primitive_with_path(lhs, rhs, op, PrimitiveComparisonPath::Auto, ctx) +} + +/// Selects automatic production dispatch or a forced implementation in tests. +#[derive(Clone, Copy)] +pub(super) enum PrimitiveComparisonPath { + /// Use the architecture and operand-specific production policy. + Auto, + + /// Force row execution. + #[cfg(test)] + Row, + + /// Force fused columnar execution. + #[cfg(test)] + Columnar, } -fn compare_primitive_typed( +/// Compare primitives through the selected implementation. +pub(super) fn compare_primitive_with_path( lhs: &ArrayRef, rhs: &ArrayRef, op: CompareOperator, - nullability: Nullability, + path: PrimitiveComparisonPath, ctx: &mut ExecutionCtx, ) -> VortexResult { - let len = lhs.len(); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - if lhs.len() != rhs.len() { - vortex_bail!( - "compare operator requires equal lengths, got {} and {}", - lhs.len(), - rhs.len() - ); - } - - let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; - - let bits = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slices(lhs, rhs, op), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => compare_slice_constant(lhs, *rhs, op), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slice_constant(rhs, *lhs, op.swap()), - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Unreachable through `execute_compare` (constant-constant is folded there), but - // cheap to answer anyway. - BitBuffer::full(apply_op(*lhs, *rhs, op), len) - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { - return Ok( - ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) - .into_array(), - ); + let use_columnar = match path { + PrimitiveComparisonPath::Auto => { + #[cfg(target_arch = "x86_64")] + { + use_columnar_comparison(lhs, rhs, op)? + } + #[cfg(not(target_arch = "x86_64"))] + { + false + } } + #[cfg(test)] + PrimitiveComparisonPath::Row => false, + #[cfg(test)] + PrimitiveComparisonPath::Columnar => true, }; + if use_columnar { + return columnar::compare_primitive(lhs, rhs, op, ctx); + } - Ok(BoolArray::try_new(bits, validity)?.into_array()) + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + execute_rows(&PrimitiveCompare, &op, &args, ctx) } -#[inline(always)] -fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { - match op { - CompareOperator::Eq => lhs.is_eq(rhs), - CompareOperator::NotEq => !lhs.is_eq(rhs), - CompareOperator::Gt => lhs.is_gt(rhs), - CompareOperator::Gte => lhs.is_ge(rhs), - CompareOperator::Lt => lhs.is_lt(rhs), - CompareOperator::Lte => lhs.is_le(rhs), +/// Internal row execution for primitive comparison operators. +#[derive(Clone)] +struct PrimitiveCompare; + +impl RowFn for PrimitiveCompare { + type Options = CompareOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + // `PrimitiveCompare` is a private implementation detail of `Binary`: it is never registered + // or serialized independently. Reusing the public ID keeps execution errors attributed to + // `Binary`. If this type becomes registrable, it needs its own ID and persistence contract. + ScalarFnVTable::id(&Binary) } -} -fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { - // Dispatch the operator outside the lane loop so each instantiation vectorizes a single - // branch-free predicate. - match op { - CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)), - CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)), - CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), - CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), - CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), - CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + fn dispatch>( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = + PType::try_from(args.first().ok_or_else(|| { + vortex_err!("a comparison operator takes two operands, got none") + })?)?; + + match_each_native_ptype!(ptype, |T| { visit_compare::(*op, visitor) }) } } -fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { +#[cfg(target_arch = "x86_64")] +fn use_columnar_comparison( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, +) -> VortexResult { + let ptype = PType::try_from(lhs.dtype())?; + Ok(match (ptype, op) { + // Equality bit-packs efficiently for every type supported by the columnar path. + (PType::I64 | PType::U64 | PType::F64, CompareOperator::Eq | CompareOperator::NotEq) => { + true + } + // The fused comparison and bit-packing loop produces better x86 code for signed 64-bit + // integers and f64. The RowFn byte-output loop remains faster for narrower lanes. + (PType::I64 | PType::F64, _) => true, + // LLVM 22 vectorizes the mixed-constant RowFn loop at 16 CGUs without LTO. However, the + // fused comparison and bit-packing path is still about 38% faster in + // `compare_u64_constant`. Recheck that benchmark before changing this dispatch. + (PType::U64, _) => lhs.is::() || rhs.is::(), + _ => false, + }) +} + +fn visit_compare(op: CompareOperator, visitor: V) -> VortexResult +where + T: NativePType, + V: RowVisitor, +{ match op { - CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), - CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), - CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)), - CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)), - CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)), - CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), + CompareOperator::Eq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)), + CompareOperator::NotEq => visitor.visit::<(T, T), bool>(|(lhs, rhs)| !lhs.is_eq(rhs)), + CompareOperator::Gt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_gt(rhs)), + CompareOperator::Gte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_ge(rhs)), + CompareOperator::Lt => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_lt(rhs)), + CompareOperator::Lte => visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_le(rhs)), } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs new file mode 100644 index 00000000000..6728437e6a8 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/columnar.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fused comparison and bit-packing for wide primitive lanes. +//! +//! Production uses this implementation only for measured x86 paths. Keeping it portable lets the +//! semantic tests exercise the RowFn and fused paths on every target. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use super::operand::PrimitiveOperand; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// Compare primitive operands with one fused comparison and bit-packing loop. +pub(super) fn compare_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match PType::try_from(lhs.dtype())? { + PType::I64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::U64 => compare_primitive_typed::(lhs, rhs, op, ctx), + PType::F64 => compare_primitive_typed::(lhs, rhs, op, ctx), + ptype => vortex_bail!("columnar comparison is not selected for {ptype}"), + } +} + +fn compare_primitive_typed( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable()); + let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; + let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + let bits = match (&lhs, &rhs) { + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slices(lhs, rhs, op), + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => compare_slice_constant(lhs, *rhs, op), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slice_constant(rhs, *lhs, op.swap()), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => BitBuffer::full(apply_op(*lhs, *rhs, op), len), + (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) + .into_array(), + ); + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { + match op { + CompareOperator::Eq => lhs.is_eq(rhs), + CompareOperator::NotEq => !lhs.is_eq(rhs), + CompareOperator::Gt => lhs.is_gt(rhs), + CompareOperator::Gte => lhs.is_ge(rhs), + CompareOperator::Lt => lhs.is_lt(rhs), + CompareOperator::Lte => lhs.is_le(rhs), + } +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |lhs: T, rhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + } +} + +fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(lhs, |lhs: T| lhs.is_eq(rhs)), + CompareOperator::NotEq => collect_bits(lhs, |lhs: T| !lhs.is_eq(rhs)), + CompareOperator::Gt => collect_bits(lhs, |lhs: T| lhs.is_gt(rhs)), + CompareOperator::Gte => collect_bits(lhs, |lhs: T| lhs.is_ge(rhs)), + CompareOperator::Lt => collect_bits(lhs, |lhs: T| lhs.is_lt(rhs)), + CompareOperator::Lte => collect_bits(lhs, |lhs: T| lhs.is_le(rhs)), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs similarity index 78% rename from vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs rename to vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs index 71d1122fc79..1563b8e68e6 100644 --- a/vortex-array/src/scalar_fn/fns/binary/primitive_operand.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive/operand.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Decoding shared by primitive binary operators. +//! Operand decoding for the fused primitive comparison path. use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -15,19 +15,33 @@ use crate::validity::Validity; /// A materialized primitive column, a non-null constant, or an all-null constant. pub(super) enum PrimitiveOperand { + /// A per-row primitive column and its validity. Array { + /// The materialized values. values: Buffer, + + /// The validity of the values. validity: Validity, }, + + /// A non-null value repeated for every row. Constant { + /// The repeated value. value: T, + + /// The number of repeated rows. len: usize, + + /// The validity implied by the constant's dtype. validity: Validity, }, + + /// An all-null constant with this row count. Null(usize), } impl PrimitiveOperand { + /// Decode an operand once for the fused comparison loop. pub(super) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { if let Some(constant) = array.as_opt::() { return Ok( @@ -49,9 +63,11 @@ impl PrimitiveOperand { let array = array.clone().execute::(ctx)?; let validity = array.validity()?; let values = array.into_buffer::(); + Ok(Self::Array { values, validity }) } + /// Return the logical row count. pub(super) fn len(&self) -> usize { match self { Self::Array { values, .. } => values.len(), @@ -59,6 +75,7 @@ impl PrimitiveOperand { } } + /// Return the operand validity. pub(super) fn validity(&self) -> Validity { match self { Self::Array { validity, .. } => validity.clone(), diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index 9831a963354..d0f7a9b5e57 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -12,7 +12,9 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::VTable; use crate::array_session; +use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; @@ -21,6 +23,7 @@ use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::arrays::ScalarFn; use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; @@ -39,6 +42,8 @@ use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; use crate::scalar::DecimalValue; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::primitive::PrimitiveComparisonPath; +use crate::scalar_fn::fns::binary::compare::primitive::compare_primitive_with_path; use crate::scalar_fn::fns::binary::scalar_cmp; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -429,6 +434,142 @@ fn float_total_order() { ); } +#[rstest] +#[case::row_eq(PrimitiveComparisonPath::Row, CompareOperator::Eq)] +#[case::row_not_eq(PrimitiveComparisonPath::Row, CompareOperator::NotEq)] +#[case::row_lt(PrimitiveComparisonPath::Row, CompareOperator::Lt)] +#[case::columnar_eq(PrimitiveComparisonPath::Columnar, CompareOperator::Eq)] +#[case::columnar_not_eq(PrimitiveComparisonPath::Columnar, CompareOperator::NotEq)] +#[case::columnar_lt(PrimitiveComparisonPath::Columnar, CompareOperator::Lt)] +fn test_primitive_comparison_paths_preserve_semantics_and_encoding( + #[case] path: PrimitiveComparisonPath, + #[case] op: CompareOperator, +) -> VortexResult<()> { + let lhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::NAN, // Null on the left. + -0.0, // Signed zero ordering. + 1.0, // A finite value below NaN. + f64::NAN, // Null on the right. + ], + Validity::from_iter([ + true, // + false, // + true, // + true, // + true, // + ]), + ) + .into_array(); + let rhs = PrimitiveArray::new( + vec![ + f64::NAN, // Equal NaNs. + f64::INFINITY, // Null on the left. + 0.0, // Signed zero ordering. + f64::NAN, // A finite value below NaN. + 1.0, // Null on the right. + ], + Validity::from_iter([ + true, // + true, // + true, // + true, // + false, // + ]), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + + let actual = compare_primitive_with_path(&lhs, &rhs, op, path, &mut ctx)?; + let expected = match op { + CompareOperator::Eq => [ + Some(true), // Equal NaNs. + None, // Null on the left. + Some(false), // Distinct signed zeroes. + Some(false), // A finite value and NaN. + None, // Null on the right. + ], + CompareOperator::NotEq | CompareOperator::Lt => [ + Some(false), // Equal NaNs. + None, // Null on the left. + Some(true), // Distinct signed zeroes. + Some(true), // A finite value and NaN. + None, // Null on the right. + ], + _ => unreachable!(), + }; + let expected = BoolArray::from_iter(expected); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(actual.dtype(), &DType::Bool(Nullability::Nullable)); + + // This encoding difference is intentional: the fused path materializes bits and validity + // together, while the RowFn path keeps masking lazy. + match path { + PrimitiveComparisonPath::Columnar => assert_eq!(actual.encoding_id(), Bool.id()), + PrimitiveComparisonPath::Row => assert!(actual.as_opt::().is_some()), + PrimitiveComparisonPath::Auto => unreachable!(), + } + + Ok(()) +} + +#[cfg(target_arch = "x86_64")] +#[rstest] +#[case::i64_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::i64_not_eq( + buffer![1_i64, 2, 3].into_array(), + buffer![1_i64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::u64_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::u64_not_eq( + buffer![1_u64, 2, 3].into_array(), + buffer![1_u64, 4, 3].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +#[case::f64_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::Eq, + [true, false, true] +)] +#[case::f64_not_eq( + buffer![1_f64, 2.0, 3.0].into_array(), + buffer![1_f64, 4.0, 3.0].into_array(), + CompareOperator::NotEq, + [false, true, false] +)] +fn test_primitive_equality_auto_uses_columnar_for_supported_ptype( + #[case] lhs: ArrayRef, + #[case] rhs: ArrayRef, + #[case] op: CompareOperator, + #[case] expected: [bool; 3], +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let actual = + compare_primitive_with_path(&lhs, &rhs, op, PrimitiveComparisonPath::Auto, &mut ctx)?; + + assert_eq!(actual.encoding_id(), Bool.id()); + assert_arrays_eq!(actual, BoolArray::from_iter(expected), &mut ctx); + + Ok(()) +} + #[rstest] #[case(Operator::Eq, [true, false, true, true])] #[case(Operator::Lt, [false, true, false, false])] diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 80faff20e0a..a5b9fe70539 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -43,7 +43,6 @@ mod compare; pub use compare::*; mod numeric; pub(crate) use numeric::*; -mod primitive_operand; use crate::scalar::NumericOperator; use crate::scalar::Scalar; diff --git a/vortex-array/src/test_harness/trace/tests.rs b/vortex-array/src/test_harness/trace/tests.rs index 98043caec13..fd9685dcb35 100644 --- a/vortex-array/src/test_harness/trace/tests.rs +++ b/vortex-array/src/test_harness/trace/tests.rs @@ -685,6 +685,14 @@ fn trace_compare_on_dict() -> VortexResult<()> { iter 0 current=vortex.dict(bool, len=5) builder_active=false ExecuteSlot slot=1 parent=vortex.dict(bool, len=5) child=vortex.binary(bool, len=3) iter 1 current=vortex.binary(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=3) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=3) iter 2 current=vortex.bool(bool, len=3) stack_parent=vortex.dict(bool, len=5) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=5) diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index d779757bc77..bf2b00f563e 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -209,6 +209,14 @@ fn trace_scan_compare_on_compressed_shipdate() -> VortexResult<()> { Done array=vortex.primitive(i32, len=4096) iter 1 current=vortex.primitive(i32, len=4096) builder_active=false return output=vortex.primitive(i32, len=4096) + optimize root=vortex.slice(i32, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i32, len=1) child=vortex.constant(i32, len=4096) -> vortex.constant(i32, len=1) + done output=vortex.constant(i32, len=1) + execute_until target=AnyCanonical root=vortex.constant(i32, len=1) + iter 0 current=vortex.constant(i32, len=1) builder_active=false + Done array=vortex.primitive(i32, len=1) + iter 1 current=vortex.primitive(i32, len=1) builder_active=false + return output=vortex.primitive(i32, len=1) Done array=vortex.bool(bool, len=4096) iter 2 current=vortex.bool(bool, len=4096) builder_active=false return output=vortex.bool(bool, len=4096) @@ -267,6 +275,14 @@ fn trace_scan_compare_on_compressed_quantity() -> VortexResult<()> { Done array=vortex.primitive(i16, len=50) iter 1 current=vortex.primitive(i16, len=50) builder_active=false return output=vortex.primitive(i16, len=50) + optimize root=vortex.slice(i16, len=1) session=false + reduce_parent static:SliceReduceAdaptor(Constant) slot=0 parent=vortex.slice(i16, len=1) child=vortex.constant(i16, len=50) -> vortex.constant(i16, len=1) + done output=vortex.constant(i16, len=1) + execute_until target=AnyCanonical root=vortex.constant(i16, len=1) + iter 0 current=vortex.constant(i16, len=1) builder_active=false + Done array=vortex.primitive(i16, len=1) + iter 1 current=vortex.primitive(i16, len=1) builder_active=false + return output=vortex.primitive(i16, len=1) Done array=vortex.bool(bool, len=50) iter 6 current=vortex.bool(bool, len=50) stack_parent=vortex.dict(bool, len=4096) slot=1 builder_active=false pop_frame slot=1 output=vortex.dict(bool, len=4096) From b655d61815351d284cdf6f575cfa9c60d0a313af Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 12:57:15 -0400 Subject: [PATCH 155/160] Execute tensor L2 norm with RowFn Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 5 + .../unstable/row/batch/execute/mod.rs | 31 +- .../unstable/row/batch/execute/output.rs | 27 ++ .../unstable/row/batch/execute/valid_only.rs | 29 ++ .../src/scalar_fn/unstable/row/batch/tests.rs | 197 +++++++- .../src/scalar_fn/unstable/row/row_fn.rs | 39 +- .../src/scalar_fn/unstable/row/vtable.rs | 1 + vortex-tensor/Cargo.toml | 2 +- vortex-tensor/benches/l2_norm.rs | 36 +- .../src/scalar_fns/cosine_similarity.rs | 7 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 425 ++++-------------- vortex-tensor/src/scalar_fns/mod.rs | 4 + vortex-tensor/src/scalar_fns/row.rs | 166 +++++++ vortex-tensor/src/scalar_fns/tests/l2_norm.rs | 327 ++++++++++++++ vortex-tensor/src/scalar_fns/tests/mod.rs | 7 + vortex-tensor/src/scalar_fns/tests/row.rs | 97 ++++ vortex-tensor/src/utils.rs | 59 +++ 17 files changed, 1093 insertions(+), 366 deletions(-) create mode 100644 vortex-tensor/src/scalar_fns/row.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/l2_norm.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/mod.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/row.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index 781d15711be..922e3ee5095 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -55,6 +55,11 @@ impl<'a> BorrowedExecutionArgs<'a> { } } + /// Return the concrete arrays used by encoding-aware execution. + pub(crate) fn arrays(&self) -> &'a [ArrayRef] { + self.arrays + } + /// Return the original input dtypes used to select the row implementation. pub(crate) fn dtypes(&self) -> &'a [DType] { self.dtypes diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index f29fbcaad27..f2d4f605740 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -3,8 +3,8 @@ //! Selects a batch execution strategy. //! -//! [`Batch::execute`] handles universal fast paths, then delegates to dense or valid-only -//! execution. +//! [`Batch::execute`] handles universal fast paths and encoded reductions, then delegates to dense +//! or valid-only execution. use vortex_error::VortexResult; use vortex_mask::Mask; @@ -28,13 +28,18 @@ mod output; pub(crate) use output::finalize_kernel_output; impl Batch { - /// Apply constant folding and null handling around `kernel`. + /// Apply encoded reductions, constant folding, and null handling around `kernel`. /// - /// When the mask contains valid and invalid rows, `try_unfiltered` may avoid filtering. - /// `Ok(None)` filters the valid rows and scatters the output back. Every result is checked - /// against the planned shape and dtype. + /// `reduce` receives the original inputs before constant broadcasting. When the mask contains + /// valid and invalid rows, `try_unfiltered` may avoid filtering. `Ok(None)` filters the valid + /// rows and scatters the output back. Every result is checked against the planned shape and + /// dtype. pub(crate) fn execute( &self, + reduce: impl FnOnce( + BorrowedExecutionArgs<'_>, + &mut ExecutionCtx, + ) -> VortexResult>, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_unfiltered: impl FnOnce( BorrowedExecutionArgs<'_>, @@ -55,6 +60,20 @@ impl Batch { return Ok(self.all_null()); } + // An empty mask is both all-true and all-false, so deferred encoded evidence cannot be + // attributed to an observable row. Let the ordinary policy construct the typed empty + // output instead. + if self.row_count > 0 + && let Some(execution) = reduce(self.execution_args(&self.inputs, self.row_count), ctx)? + { + match execution { + RowExecution::Output(values) => return self.finalize_reduced(values, ctx), + RowExecution::DeferredError(error) => { + return self.resolve_reduced_error(error, kernel, try_unfiltered, ctx); + } + } + } + // All inputs constant, and their conjoined validity proves every row non-null. This sees // through extension and masked wrappers just like argument decoding does. if self.row_count > 0 diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs index 47fd8cdaf46..f4dfc44fae2 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -14,6 +14,7 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; impl Batch { pub(super) fn all_null(&self) -> ArrayRef { @@ -28,6 +29,32 @@ impl Batch { reconcile_output(self.id, &self.result_dtype, expected_len, values) } + /// Reconcile an encoding-aware result and apply the batch's strict input validity. + pub(super) fn finalize_reduced( + &self, + values: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + validate_output(self.id, &self.result_dtype, self.row_count, &values)?; + + let input_valid = self.validity.execute_mask(self.row_count, ctx)?; + let output_valid = values.validity()?.execute_mask(self.row_count, ctx)?; + vortex_ensure!( + input_valid.bitand_not(&output_valid).all_false(), + "the {} encoded reduction produced nulls for valid rows", + self.id, + ); + + let values = match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => values, + Validity::Array(valid) => values.mask(valid)?, + // Handled before the encoding-aware hook runs. + Validity::AllInvalid => return Ok(self.all_null()), + }; + + cast_output_nullability(&self.result_dtype, values) + } + /// Validate the output from a row kernel before batch validity is attached. pub(super) fn validate_kernel_output( &self, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs index d14fed87370..d68c2e57da8 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -24,6 +25,34 @@ enum ResolvedValidity { } impl Batch { + /// Resolve deferred evidence from the encoded path by executing only observable rows. + pub(super) fn resolve_reduced_error( + &self, + error: VortexError, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + /// Resolve validity, try unfiltered execution, then fall back to filtering. pub(super) fn execute_valid_only( &self, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 7ad4b69c4e7..ebf69580c54 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -32,7 +32,6 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::scalar_fn::EmptyOptions; -use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::InputElement; @@ -60,7 +59,13 @@ struct AddShort(ShortVisit); struct ShortDecodeI64; #[derive(Clone)] -struct Identity; +struct OriginalInputReducer; + +#[derive(Clone)] +struct InvalidEncodedReduction; + +#[derive(Clone)] +struct DeferredOriginalReducer; #[derive(Clone)] struct ValidOnlyIdentity; @@ -325,16 +330,103 @@ impl RowFn for RetryConstantAdd { }, ) } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 1 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(0u8, args[0].len()).into_array(), + ))); + } + + Ok(None) + } } -impl RowFn for Identity { +impl RowFn for OriginalInputReducer { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("test.identity"); + static ID: CachedId = CachedId::new("test.original_input_reducer"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 3 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(42_i64, 3).into_array(), + ))); + } + + Ok(None) + } +} + +impl RowFn for InvalidEncodedReduction { + type Options = usize; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_encoded_reduction"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + null_index: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::Output( + PrimitiveArray::new( + vec![10_i64, 20], + Validity::from_iter((0..2).map(|index| index != *null_index)), + ) + .into_array(), + ))) + } +} + +impl RowFn for DeferredOriginalReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.deferred_original_reducer"); *ID } @@ -346,6 +438,17 @@ impl RowFn for Identity { ) -> VortexResult { visitor.visit::<(i64,), i64>(|(value,)| value) } + + fn reduce_encoded( + &self, + _options: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::DeferredError(vortex_err!( + InvalidArgument: "encoded payload failed" + )))) + } } impl RowFn for ValidOnlyIdentity { @@ -550,7 +653,7 @@ fn test_short_constant_null_tolerant_decode_is_rejected() -> VortexResult<()> { } #[test] -fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> { +fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { let lhs = PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); let rhs = ConstantArray::new(1u8, 2).into_array(); @@ -584,13 +687,89 @@ fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { Ok(()) } +#[test] +fn test_reduce_encoded_defers_errors_behind_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_empty_batch_skips_deferred_encoded_error() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(Vec::::new()).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 0); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid(Validity::AllValid)] +#[case::mixed(Validity::from_iter([true, false]))] +fn test_reduce_encoded_rejects_nulls_on_valid_rows(#[case] validity: Validity) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![10_i64, 20], validity).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&InvalidEncodedReduction, &0, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an encoded reduction introduced a null on a valid row"), + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_encoded_reduction"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("encoded reduction produced nulls for valid rows"), + "the boundary error must identify invalid reduced output, got {error}", + ); + Ok(()) +} + +#[test] +fn test_reduce_encoded_preserves_input_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&InvalidEncodedReduction, &1, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + let expected = ConstantArray::new(42_i64, 3).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + #[test] fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { let input = ConstantArray::new(7_i64, 2).into_array(); let args = VecExecutionArgs::new(vec![input.clone()], 2); let mut ctx = array_session().create_execution_ctx(); - let actual = execute_rows(&Identity, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; assert_arrays_eq!(&actual, &input, &mut ctx); Ok(()) @@ -614,7 +793,8 @@ fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResul let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; @@ -642,7 +822,8 @@ fn test_valid_only_filters_and_scatters() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 831cc25f251..93da85db841 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -4,8 +4,8 @@ //! The [`RowFn`] contract for scalar functions whose natural kernel computes one row at a time. //! //! Implementations declare their arity and fallibility, then use [`RowFn::dispatch`] to select the -//! typed row signature for each supported dtype combination. Optional methods provide -//! serialization without putting persistence plumbing in the row kernel. +//! typed row signature for each supported dtype combination. Optional hooks provide serialization +//! and encoding-aware execution without putting columnar plumbing in the row kernel. use std::fmt::Debug; use std::fmt::Display; @@ -16,8 +16,11 @@ use vortex_error::vortex_bail; use vortex_session::VortexSession; use super::visitor::RowVisitor; +use crate::ArrayRef; +use crate::ExecutionCtx; use crate::dtype::DType; use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::unstable::row::RowExecution; /// A strict scalar function whose row kernel cannot produce null from valid inputs. /// @@ -40,12 +43,14 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// The arguments in display order. Its length is the function's exact arity. const ARG_NAMES: &'static [&'static str]; - /// Whether any dispatch can raise a semantic error. + /// Whether any dispatch or encoded reduction can raise a semantic error. /// /// See [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible) for a /// more detailed explanation of semantic errors. /// - /// The framework checks dispatched element and result types. A conservative `true` is allowed. + /// The framework checks dispatched element and result types, but cannot inspect + /// [`reduce_encoded`](Self::reduce_encoded). Set this to `true` when that hook can return a + /// semantic error or [`RowExecution::DeferredError`]. A conservative `true` is allowed. const FALLIBLE: bool; /// Returns the ID of the scalar function. @@ -76,4 +81,30 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { args: &[DType], visitor: V, ) -> VortexResult; + + /// Try an encoding-aware implementation before decoding the inputs into row elements. + /// + /// `None` continues to the row loop. [`Output`](RowExecution::Output) may remain encoded or + /// lazy. [`DeferredError`](RowExecution::DeferredError) retries only valid rows. Batch execution + /// calls this hook at most once with the original nonempty inputs. Nullary functions, empty + /// batches, slices, and compacted retries skip it. + /// + /// Like a dense row closure, this hook must be total over every stored payload, including + /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or + /// retried through the row layer. + /// + /// # Requirements + /// + /// - `output.len()` **must** equal `args[0].len()`. + /// - The output dtype **must** match the planned dtype when ignoring nullability. + /// - The output **must not** introduce a null where every input is valid. + fn reduce_encoded( + &self, + options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + _ = (options, args, ctx); + Ok(None) + } } diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 257ac3ca0b5..ac68ce46711 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -134,6 +134,7 @@ pub fn execute_rows( let batch = prepare_batch(function, options, args)?; batch.execute( + |args, ctx| function.reduce_encoded(options, args.arrays(), ctx), |args, ctx| execute_row_kernel(function, options, args, ctx), |args, valid, ctx| try_execute_rows_unfiltered(function, options, args, valid, ctx), ctx, diff --git a/vortex-tensor/Cargo.toml b/vortex-tensor/Cargo.toml index abdca676775..9706102f6d6 100644 --- a/vortex-tensor/Cargo.toml +++ b/vortex-tensor/Cargo.toml @@ -17,7 +17,7 @@ version = { workspace = true } workspace = true [dependencies] -vortex-array = { workspace = true } +vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } diff --git a/vortex-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index 6f597084113..bf8832f2520 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -15,11 +15,19 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::l2_norm::L2Norm; @@ -54,13 +62,25 @@ fn vectors(width: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { let session = vortex_array::array_session(); bencher .counter(ItemsCount::new(input.len())) .with_inputs(|| { ( - L2Norm::try_new_array(input.clone()).unwrap().into_array(), + L2Norm + .try_new_array(input.len(), EmptyOptions, [input.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -80,3 +100,17 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_l2_norm(bencher, input); } + +#[divan::bench(args = WIDTHS)] +fn constant(bencher: Bencher, width: usize) { + bench_l2_norm(bencher, constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let input = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_l2_norm(bencher, input); +} diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index ef6eed69e94..05a0078d3ae 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -10,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; @@ -142,8 +143,8 @@ impl ScalarFnVTable for CosineSimilarity { let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm::try_new_array(lhs_ref.clone())?; - let norm_rhs_arr = L2Norm::try_new_array(rhs_ref.clone())?; + let norm_lhs_arr = L2Norm.try_new_array(len, EmptyOptions, [lhs_ref.clone()])?; + let norm_rhs_arr = L2Norm.try_new_array(len, EmptyOptions, [rhs_ref.clone()])?; let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; // Execute to get the inner product and norms of the arrays. We only fully decompress @@ -286,7 +287,7 @@ impl CosineSimilarity { let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - let norm_arr = L2Norm::try_new_array(plain_ref.clone())?; + let norm_arr = L2Norm.try_new_array(len, EmptyOptions, [plain_ref.clone()])?; let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index a72b9da66db..78cabc961b0 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -3,51 +3,44 @@ //! L2 norm expression for tensor-like types. -use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::Normalized; -use crate::matcher::AnyTensor; -use crate::utils::extract_flat_elements; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::extract_normalized_children; -use crate::utils::reattach_validity; +use crate::utils::l2_norm_row; use crate::utils::validate_tensor_float_input; /// L2 norm (Euclidean norm) of a tensor or vector column. @@ -63,139 +56,112 @@ use crate::utils::validate_tensor_float_input; /// of the storage contract, not a separate lossy-compute mode. /// /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct L2Norm; impl L2Norm { /// Creates a new [`TypedScalarFnInstance`] wrapping the L2 norm operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(L2Norm, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the L2 norm over `child`. /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtype is + /// unsupported. pub fn try_new_array(child: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(L2Norm::new().erased(), vec![child]) + ScalarFnArray::try_new(Self::new().erased(), vec![child]) } } -impl ScalarFnVTable for L2Norm { +impl RowFn for L2Norm { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.l2_norm"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(1) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("input"), - _ => unreachable!("L2Norm must have exactly one child"), - } - } - - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let input_dtype = &arg_dtypes[0]; - let tensor_match = validate_tensor_float_input(input_dtype)?; - let ptype = tensor_match.element_ptype(); - - let nullability = Nullability::from(input_dtype.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn execute( + fn dispatch>( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let input_ref = args.get(0)?; - let row_count = args.row_count(); - - let ext = input_ref.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let tensor_flat_size = tensor_match.list_size() as usize; - let element_ptype = tensor_match.element_ptype(); - - let norm_dtype = DType::Primitive(element_ptype, ext.nullability()); - - // Stored norms are authoritative. Reattach the parent validity because the child is - // non-nullable. - if input_ref.is::() { - let (_, norms) = extract_normalized_children(&input_ref); - let norms = reattach_validity(norms, input_ref.validity()?)?; - vortex_ensure_eq!(norms.dtype(), &norm_dtype); - return Ok(norms); - } - - // Optimize for the constant array case. - if let Some(array) = input_ref.as_opt::() { - let scalar = array.scalar().as_extension().to_storage_scalar(); - - let Some(elements) = scalar.as_list().elements() else { - return Ok(ConstantArray::new(Scalar::null(norm_dtype), row_count).into_array()); - }; - - let norm_scalar = match_each_float_ptype!(element_ptype, |T| { - let values: Vec = elements - .iter() - .map(|s| { - s.as_primitive() - .as_::() - .vortex_expect("element was somehow not the correct float") - }) - .collect(); - let norm = l2_norm_row::(&values); - - Scalar::try_new(norm_dtype, Some(norm.into())) - })?; - - let norms = ConstantArray::new(norm_scalar, row_count).into_array(); - return Ok(norms); - } - - let input: ExtensionArray = input_ref.execute(ctx)?; - let validity = input.as_ref().validity()?; - - let storage = input.storage_array(); - let flat = extract_flat_elements(storage, tensor_flat_size, ctx)?; - - match_each_float_ptype!(flat.ptype(), |T| { - let buffer: Buffer = (0..row_count) - .map(|i| l2_norm_row(flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow,), UninitElementSink, _>(|(row,), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l2_norm_row(row)) } + }) }) } - fn validity( + /// `L2Norm` over a [`Normalized`]-encoded column is defined to read back the authoritative + /// stored norms. Callers of lossy encodings opt into that storage semantics instead of forcing + /// a decode-and-recompute path here. + fn reduce_encoded( &self, _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if the input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let input = &args[0]; + if input.is::() { + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let (_, norms) = extract_normalized_children(input); + vortex_ensure!( + norms.dtype().is_primitive(), + "normalized norms must be primitive, got {}", + norms.dtype(), + ); + vortex_ensure_eq!(norms.dtype().as_ptype(), element_ptype); + return Ok(Some(RowExecution::Output(norms))); + } - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + let Some(constant) = input.as_opt::() else { + return Ok(None); + }; + let element_ptype = validate_tensor_float_input(input.dtype())?.element_ptype(); + let norm_dtype = + DType::Primitive(element_ptype, input.dtype().as_extension().nullability()); + let storage = constant.scalar().as_extension().to_storage_scalar(); + + let Some(elements) = storage.as_list().elements() else { + let output = ConstantArray::new(Scalar::null(norm_dtype), input.len()); + return Ok(Some(RowExecution::Output(output.into_array()))); + }; + + let norm = match_each_float_ptype!(element_ptype, |T| { + let values: Vec = elements + .iter() + .map(|element| { + element + .as_primitive() + .as_::() + .vortex_expect("tensor element must match its declared ptype") + }) + .collect(); + Scalar::try_new(norm_dtype, Some(l2_norm_row::(&values).into())) + })?; + let output = ConstantArray::new(norm, input.len()); + Ok(Some(RowExecution::Output(output.into_array()))) } } @@ -241,230 +207,3 @@ impl ScalarFnArrayVTable for L2Norm { }) } } - -/// Computes the L2 norm (Euclidean norm) of a float slice. -/// -/// Returns `sqrt(sum(v_i^2))`. A zero-length or all-zero input produces `0.0`. -fn l2_norm_row(v: &[T]) -> T { - let mut sum_sq = T::zero(); - for &x in v { - sum_sq = sum_sq + x * x; - } - sum_sq.sqrt() -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::EmptyMetadata; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Constant; - use vortex_array::arrays::ConstantArray; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::dtype::extension::ExtDType; - use vortex_array::scalar::Scalar; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::l2_norm::L2Norm; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::literal_vector_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. - fn eval_l2_norm(input: ArrayRef) -> VortexResult> { - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[rstest] - #[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] - #[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] - #[case::single_element(&[1], &[7.0], &[7.0])] - #[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] - fn known_norms( - #[case] shape: &[usize], - #[case] elements: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let arr = tensor_array(shape, elements)?; - assert_close(&eval_l2_norm(arr)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let arr = tensor_array( - &[3], - &[ - 3.0, 4.0, 0.0, // norm = 5.0 - 0.0, 0.0, 0.0, // norm = 0.0 - 1.0, 1.0, 1.0, // norm = sqrt(3) - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); - Ok(()) - } - - #[test] - fn vector_multiple_rows() -> VortexResult<()> { - let arr = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // norm = 1.0 - 3.0, 4.0, 0.0, // norm = 5.0 - ], - )?; - assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 is masked as null. - let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; - let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: norm = 5.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is a non-null tensor should short-circuit to a - /// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so - /// execution stops at the [`Constant`] encoding instead of canonicalizing into a - /// [`PrimitiveArray`]. - #[test] - fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { - let input = literal_vector_array(&[3.0f64, 4.0], 4); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let output = result.execute_until::(&mut ctx)?; - - let constant = output - .as_opt::() - .expect("L2Norm over a constant input must produce a constant output"); - assert_eq!(constant.len(), 4); - let norm = constant - .scalar() - .as_primitive() - .as_::() - .expect("norm scalar must be a non-null primitive"); - assert_close(&[norm], &[5.0]); - Ok(()) - } - - /// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of - /// the correct primitive dtype and length. - #[test] - fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { - let storage_dtype = DType::FixedSizeList( - DType::Primitive(PType::F64, Nullability::NonNullable).into(), - 2, - Nullability::Nullable, - ); - let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); - let null_scalar = Scalar::null(DType::Extension(ext_dtype)); - let input = ConstantArray::new(null_scalar, 3).into_array(); - - let scalar_fn = L2Norm::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let output = result.execute_until::(&mut ctx)?; - - let constant = output - .as_opt::() - .expect("null constant input must produce a constant output"); - assert_eq!(constant.len(), 3); - assert!(constant.scalar().is_null()); - assert_eq!( - constant.dtype(), - &DType::Primitive(PType::F64, Nullability::Nullable) - ); - Ok(()) - } - - #[test] - fn reads_through_a_nullable_normalized_column() -> VortexResult<()> { - let normalized = vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?; - let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - - let mut ctx = SESSION.create_execution_ctx(); - let validity = Validity::from_iter([true, false]); - let input = Normalized::try_new(normalized, norms, validity, &mut ctx)?.into_array(); - - let result = ScalarFnArray::try_new(L2Norm::new().erased(), vec![input])?.into_array(); - let prim: PrimitiveArray = result.execute(&mut ctx)?; - - assert_eq!( - prim.dtype(), - &DType::Primitive(PType::F64, Nullability::Nullable) - ); - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[5.0]); - - Ok(()) - } - - #[rstest] - #[case::fixed_shape_tensor(l2_norm_tensor_child())] - #[case::vector(l2_norm_vector_child())] - fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { - let original = L2Norm::try_new_array(child.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(L2Norm); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("L2Norm serialize must produce metadata"); - - let children = vec![child]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn l2_norm_tensor_child() -> ArrayRef { - tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") - } - - fn l2_norm_vector_child() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } -} diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 68f10ca6b01..da9b8950e7a 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -6,3 +6,7 @@ pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; +pub(crate) mod row; + +#[cfg(test)] +mod tests; diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs new file mode 100644 index 00000000000..a99ad76ea6e --- /dev/null +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the tensor scalar functions add to the row-function machinery: an element type that reads a +//! tensor row and the width rule they share. + +use std::marker::PhantomData; + +use num_traits::Float; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::Masked; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::masked::MaskedArraySlotsExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::scalar_fn::unstable::row::InputElement; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::utils::extract_flat_elements; +use crate::utils::validate_tensor_float_input; +use crate::utils::validate_tensor_float_inputs; + +/// The width rule the tensor scalar functions share: every argument is the same float tensor dtype, +/// and the width is its element ptype. +pub(crate) fn tensor_element_ptype(args: &[DType]) -> VortexResult { + Ok(validate_tensor_float_inputs(args)?.element_ptype()) +} + +/// Marker for tensor-valued input elements: accepts any tensor-like extension column whose +/// elements are `T`, and presents each row as its flat elements, `&[T]`. +pub struct TensorRow(PhantomData); + +/// The decoded form of a [`TensorRow`] column: one flat typed buffer plus the stride to read it at. +/// +/// Typed at decode time rather than per row. `FlatElements::row` re-derives its typed slice on every +/// call, which costs a ptype check and a buffer downcast per row; a row loop reads every row, so it +/// pays that once here instead. +pub struct TensorRows { + /// Every row's elements, back to back. + elements: Buffer, + + /// Number of logical tensor rows, stored so zero-width tensors retain their length. + rows: usize, + + /// Elements per row, the length of each row slice. + list_size: usize, + + /// `list_size` for a full column and `0` for constant-backed storage, so `index * stride` pins a + /// constant to its single materialized row without a branch in the loop. + stride: usize, +} + +// SAFETY: `TensorRows` records the row count validated during decode, and both checked and +// unchecked access use the same stride and row width. +unsafe impl InputElement for TensorRow { + type Column = TensorRows; + type View<'a> = &'a TensorRows; + type Elem<'a> = &'a [T]; + + // Tensor storage is a fully materialized non-nullable primitive buffer, so the elements behind + // a null row are arbitrary values rather than an unresolvable reference. + const DENSE_SAFE: bool = true; + // Tensor storage is a primitive buffer; reading it cannot fail on account of its values. + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + let tensor_match = validate_tensor_float_input(dtype)?; + let expected = T::PTYPE; + vortex_ensure_eq!( + tensor_match.element_ptype(), + expected, + "expected a tensor of {expected} elements, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // Dense batch execution owns the mask and restores it on the result. Decode the values + // directly so a nullable tensor does not rebuild its extension storage under that mask. + let array = match array.as_opt::() { + Some(masked) => masked.child().clone(), + None => array, + }; + + let rows = array.len(); + let list_size = validate_tensor_float_input(array.dtype())?.list_size() as usize; + let ext: ExtensionArray = array.execute(ctx)?; + let flat = extract_flat_elements(ext.storage_array(), list_size, ctx)?; + let list_size = flat.list_size(); + let stride = flat.row_stride(); + let elements = flat.into_buffer::(); + + let expected_elements = if stride == 0 { + list_size + } else { + vortex_ensure_eq!( + stride, + list_size, + "per-row tensor stride must equal its width, got {stride}", + ); + let Some(expected_elements) = rows.checked_mul(stride) else { + vortex_bail!( + "tensor row storage length must fit usize, got {rows} rows of width {stride}", + ); + }; + + expected_elements + }; + vortex_ensure_eq!( + elements.len(), + expected_elements, + "tensor row storage must contain {expected_elements} elements, got {}", + elements.len(), + ); + + Ok(TensorRows { + elements, + rows, + list_size, + stride, + }) + } + + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + fn get(column: &Self::Column, index: usize) -> &[T] { + let start = index * column.stride; + &column.elements.as_slice()[start..start + column.list_size] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.rows + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + Self::get(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a [T] + where + Self: 'a, + { + let start = index * view.stride; + + // SAFETY: decode established one complete stored row for stride 0, or `rows` contiguous + // `list_size`-element rows otherwise. The caller guarantees `index < rows`. + unsafe { + std::slice::from_raw_parts(view.elements.as_slice().as_ptr().add(start), view.list_size) + } + } +} diff --git a/vortex-tensor/src/scalar_fns/tests/l2_norm.rs b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs new file mode 100644 index 00000000000..38ce74d8c65 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/l2_norm.rs @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::l2_norm::L2Norm; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. +fn eval_l2_norm(input: ArrayRef) -> VortexResult> { + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = L2Norm::new(); + let array = L2Norm::try_new_array(tensor_array(&[1], &[3.0])?)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + assert_close( + &eval_l2_norm(zero_width_vector_array::(3)?)?, + &[0.0, 0.0, 0.0], + ); + assert!(eval_l2_norm(vector_array(2, &[] as &[f64])?)?.is_empty()); + + let constant = Vector::constant_array::(&[], 3)?; + assert_close(&eval_l2_norm(constant)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +#[rstest] +#[case::three_four_five(&[2], &[3.0, 4.0], &[5.0])] +#[case::zero_vector(&[3], &[0.0, 0.0, 0.0], &[0.0])] +#[case::single_element(&[1], &[7.0], &[7.0])] +#[case::negative_elements(&[2], &[-3.0, -4.0], &[5.0])] +fn known_norms( + #[case] shape: &[usize], + #[case] elements: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let arr = tensor_array(shape, elements)?; + assert_close(&eval_l2_norm(arr)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let arr = tensor_array( + &[3], + &[ + 3.0, 4.0, 0.0, // norm = 5.0 + 0.0, 0.0, 0.0, // norm = 0.0 + 1.0, 1.0, 1.0, // norm = sqrt(3) + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]); + Ok(()) +} + +#[test] +fn vector_multiple_rows() -> VortexResult<()> { + let arr = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // norm = 1.0 + 3.0, 4.0, 0.0, // norm = 5.0 + ], + )?; + assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 is masked as null. + let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![arr])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: norm = 5.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is a non-null tensor should short-circuit to a +/// [`ConstantArray`] output whose scalar is the precomputed norm. Uses [`execute_until`] so +/// execution stops at the [`Constant`] encoding instead of canonicalizing into a +/// [`PrimitiveArray`]. +#[test] +fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { + let input = literal_vector_array(&[3.0f64, 4.0], 4); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over a constant input must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// An extension array over constant storage is folded just like a top-level constant instead of +/// recomputing the same norm once per row. +#[test] +fn extension_backed_constant_yields_constant_output() -> VortexResult<()> { + let input = Vector::constant_array(&[3.0f64, 4.0], 4)?; + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("L2Norm over constant-backed extension storage must produce a constant output"); + assert_eq!(constant.len(), 4); + let norm = constant + .scalar() + .as_primitive() + .as_::() + .expect("norm scalar must be a non-null primitive"); + assert_close(&[norm], &[5.0]); + Ok(()) +} + +/// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of +/// the correct primitive dtype and length. +#[test] +fn constant_null_input_yields_null_constant_output() -> VortexResult<()> { + let storage_dtype = DType::FixedSizeList( + DType::Primitive(PType::F64, Nullability::NonNullable).into(), + 2, + Nullability::Nullable, + ); + let ext_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?.erased(); + let null_scalar = Scalar::null(DType::Extension(ext_dtype)); + let input = ConstantArray::new(null_scalar, 3).into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![input])?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let output = result.execute_until::(&mut ctx)?; + + let constant = output + .as_opt::() + .expect("null constant input must produce a constant output"); + assert_eq!(constant.len(), 3); + assert!(constant.scalar().is_null()); + assert_eq!( + constant.dtype(), + &DType::Primitive(PType::F64, Nullability::Nullable) + ); + Ok(()) +} + +/// An `f32` column must dispatch at `f32` and produce an `f32` result, which is the property that +/// makes width polymorphism load-bearing rather than decorative. +#[rstest] +#[case::f32(&[3.0f32, 4.0], PType::F32)] +#[case::f64(&[3.0f64, 4.0], PType::F64)] +fn dispatches_at_input_width( + #[case] elements: &[T], + #[case] expected: PType, +) -> VortexResult<()> { + let arr = tensor_array(&[2], elements)?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L2Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .execute(&mut ctx)?; + assert_eq!(prim.ptype(), expected); + Ok(()) +} + +/// `L2Norm(Normalized(normalized, norms))` reads back the authoritative stored norms rather than +/// recomputing over decoded coordinates. The normalized child here is deliberately *not* +/// unit-norm, mimicking lossy storage, so readthrough and recompute disagree: row 0 decodes to +/// `[6, 8]` (norm `10`) and row 1 to `[6, 0]` (norm `6`), while the stored norms are `5` and `2`. +#[test] +fn normalized_readthrough_returns_stored_norms() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 2.0]).into_array(); + // SAFETY: A focused test of the lossy storage contract: the stored norms are authoritative + // even though this normalized child violates the unit-norm invariant. + let denorm = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::NonNullable) }.into_array(); + + assert_close(&eval_l2_norm(denorm)?, &[5.0, 2.0]); + Ok(()) +} + +/// The readthrough must survive a partially-null column. +/// +/// This pins the dense policy the row contract derives. Filtering could hand `reduce_encoded` a +/// filtered input, which is no longer an `ExactScalarFn`, silently falling back to +/// decode-and-recompute. For a lossy child that changes the answer: row 0 below would come back as +/// `10` (recomputed from `[6, 8]`) instead of the authoritative stored `5`. +#[test] +fn normalized_readthrough_survives_null_rows() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[1.2, 1.6, 3.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 2.0]).into_array(); + // SAFETY: Intentionally lossy, as in `normalized_readthrough_returns_stored_norms`, so that + // a recompute fallback is observable. + let denorm = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::from_iter([true, false])) } + .into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +/// The readthrough must still propagate validity carried by the `Normalized` parent. +#[test] +fn normalized_readthrough_propagates_parent_validity() -> VortexResult<()> { + let normalized = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let denorm = Normalized::try_new( + normalized, + norms, + Validity::from_iter([true, false]), + &mut ctx, + )? + .into_array(); + + let scalar_fn = L2Norm.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![denorm])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + Ok(()) +} + +#[rstest] +#[case::fixed_shape_tensor(l2_norm_tensor_child())] +#[case::vector(l2_norm_vector_child())] +fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { + let original = L2Norm.try_new_array(child.len(), EmptyOptions, [child.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(L2Norm); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("L2Norm serialize must produce metadata"); + + let children = vec![child]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn l2_norm_tensor_child() -> ArrayRef { + tensor_array(&[3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid tensor array") +} + +fn l2_norm_vector_child() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/mod.rs b/vortex-tensor/src/scalar_fns/tests/mod.rs new file mode 100644 index 00000000000..5447772adda --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/mod.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests for the tensor scalar functions. + +mod l2_norm; +mod row; diff --git a/vortex-tensor/src/scalar_fns/tests/row.rs b/vortex-tensor/src/scalar_fns/tests/row.rs new file mode 100644 index 00000000000..8a7fa23edd4 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/row.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::Float; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::match_each_float_ptype; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::tensor_array; + +/// The marginal cost of a new tensor scalar function is this entire definition. Everything else +/// (null propagation, constants, validity, f16/f32/f64 dispatch, dtype checks, and constructors) is +/// derived. +#[derive(Clone, Debug, Default)] +struct L1Norm; + +impl RowFn for L1Norm { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.test.l1_norm"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow,), UninitElementSink, _>(|(row,), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, l1_norm_row(row)) } + }) + }) + } +} + +fn l1_norm_row(row: &[T]) -> T { + row.iter().fold(T::zero(), |acc, &x| acc + x.abs()) +} + +#[test] +fn derived_fn_executes_with_nulls() -> VortexResult<()> { + let arr = tensor_array(&[2], &[3.0, -4.0, 1.0, 1.0])?; + let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); + + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + let prim: PrimitiveArray = L1Norm + .try_new_array(arr.len(), EmptyOptions, [arr])? + .execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[7.0]); + Ok(()) +} + +/// A kernel written once serves every float width. +#[test] +fn derived_fn_dispatches_at_input_width() -> VortexResult<()> { + let mut ctx = crate::tests::SESSION.create_execution_ctx(); + + let f32_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f32, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f32_result.ptype(), PType::F32); + assert_eq!(f32_result.as_slice::(), &[7.0f32]); + + let f64_result: PrimitiveArray = L1Norm + .try_new_array(1, EmptyOptions, [tensor_array(&[2], &[3.0f64, -4.0])?])? + .execute(&mut ctx)?; + assert_eq!(f64_result.ptype(), PType::F64); + Ok(()) +} diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 3e33fe20db9..c7339787b73 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Shared helpers for tensor scalar functions. + use half::f16; +use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -22,6 +25,7 @@ use vortex_array::dtype::PType; use vortex_array::dtype::proto::dtype as pb; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::validity::Validity; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -60,6 +64,16 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { SAFETY_FACTOR as f64 * machine_epsilon * dimensions_root } +/// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. +pub(crate) fn l2_norm_row(row: &[T]) -> T { + let mut sum_squared = T::zero(); + for &element in row { + sum_squared = sum_squared + element * element; + } + + sum_squared.sqrt() +} + /// Extracts the `(normalized, norms)` children of a [`Normalized`]-encoded array. /// /// # Panics @@ -120,6 +134,22 @@ pub fn validate_binary_tensor_float_inputs<'a>( validate_tensor_float_input(lhs) } +/// Validates that every argument has the same float tensor dtype, ignoring nullability. +pub fn validate_tensor_float_inputs(args: &[DType]) -> VortexResult> { + let (first, rest) = args + .split_first() + .ok_or_else(|| vortex_err!("tensor expression expects at least one input"))?; + + for arg in rest { + vortex_ensure!( + first.eq_ignore_nullability(arg), + "tensor expression expects inputs to have the same dtype, got {first} and {arg}" + ); + } + + validate_tensor_float_input(first) +} + /// The flat primitive elements of a tensor storage array, with typed row access. /// /// This struct hides the stride detail that arises from the [`ConstantArray`] optimization: a @@ -148,6 +178,23 @@ impl FlatElements { let slice = self.elems.as_slice::(); &slice[row_idx * self.list_size..][..self.list_size] } + + /// Returns the number of elements in each row. + #[must_use] + pub fn list_size(&self) -> usize { + self.list_size + } + + /// Returns the physical distance between rows, or zero when every row uses one stored value. + #[must_use] + pub fn row_stride(&self) -> usize { + if self.is_constant { 0 } else { self.list_size } + } + + /// Returns the elements as a typed buffer, performing the ptype check once for the batch. + pub fn into_buffer(self) -> Buffer { + self.elems.into_buffer::() + } } /// Extracts the flat primitive elements from a tensor storage array (FixedSizeList). @@ -353,6 +400,18 @@ pub mod test_helpers { Vector::try_new_vector_array(flat_fsl(elements, dim)) } + /// Builds `rows` zero-width vectors over an empty typed element buffer. + pub fn zero_width_vector_array(rows: usize) -> VortexResult { + let storage = FixedSizeListArray::new( + Buffer::::empty().into_array(), + 0, + Validity::NonNullable, + rows, + ) + .into_array(); + Vector::try_new_vector_array(storage) + } + /// Builds a [`FixedShapeTensor`] extension array whose storage is a [`ConstantArray`], /// representing a single query tensor broadcast to `len` rows. pub fn constant_tensor_array>( From 2f55638bcdcc2876eab193c85464d4aff74313a0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 12:58:51 -0400 Subject: [PATCH 156/160] Execute tensor product functions with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/benches/cosine_similarity.rs | 35 +- vortex-tensor/benches/inner_product.rs | 43 +- vortex-tensor/src/encodings/normalized/mod.rs | 1 - .../src/scalar_fns/cosine_similarity.rs | 924 +++++------------- vortex-tensor/src/scalar_fns/inner_product.rs | 496 ++-------- vortex-tensor/src/scalar_fns/row.rs | 17 + .../src/scalar_fns/tests/cosine_similarity.rs | 610 ++++++++++++ .../src/scalar_fns/tests/inner_product.rs | 282 ++++++ vortex-tensor/src/scalar_fns/tests/mod.rs | 2 + vortex-tensor/src/utils.rs | 18 +- vortex-tensor/src/vector_search.rs | 4 +- 11 files changed, 1298 insertions(+), 1134 deletions(-) create mode 100644 vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/inner_product.rs diff --git a/vortex-tensor/benches/cosine_similarity.rs b/vortex-tensor/benches/cosine_similarity.rs index 6cc5eb867ef..49551fbf701 100644 --- a/vortex-tensor/benches/cosine_similarity.rs +++ b/vortex-tensor/benches/cosine_similarity.rs @@ -21,11 +21,14 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity; @@ -41,15 +44,11 @@ fn main() { } /// Total `f64` elements per operand, held constant across widths: the row count is -/// `ELEMENTS / width`. This budget is a quarter of the one the other tensor benches use, because -/// the constant arms recompute the broadcast vector's norm per row and cost roughly ten times the -/// column arms per element. It is what keeps every arm inside the 1 ms per-iteration limit from -/// `docs/developer-guide/benchmarking.md`, measured against CodSpeed's CPU simulation. +/// `ELEMENTS / width`. The smaller budget keeps the wider cosine kernels inside the 1 ms +/// per-iteration limit from `docs/developer-guide/benchmarking.md` under CodSpeed simulation. const ELEMENTS: usize = 2_048; -/// Widths chosen to separate the two costs, as in `l2_norm.rs`: the redundant norm pass is -/// `O(rows * width)`, one third of the closure's arithmetic, so wide tensors show the hoist -/// while a narrow one is dominated by per-row framework costs. +/// Widths that expose both fixed row-framework costs and the `O(width)` kernel work. const WIDTHS: &[usize] = &[2, 32, 256]; /// `ELEMENTS / width` vectors of `width` `f64` elements, non-nullable. `seed` offsets the values so @@ -85,9 +84,9 @@ fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { bencher .with_inputs(|| { ( - CosineSimilarity::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + CosineSimilarity + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -106,6 +105,22 @@ fn column_x_constant(bencher: Bencher, width: usize) { bench_cosine(bencher, vectors(width, 0), constant_vector(width)); } +/// The lhs is a broadcast query vector, whose norm is the same in every row. +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_cosine(bencher, constant_vector(width), vectors(width, 31)); +} + +/// A nullable broadcast rhs exercises constant preparation and output validity together. +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_cosine(bencher, vectors(width, 0), rhs); +} + /// One query vector represented as an extension array over constant storage. fn extension_constant_vector(width: usize) -> ArrayRef { let ext_dtype = vectors(width, 0).dtype().as_extension().clone(); diff --git a/vortex-tensor/benches/inner_product.rs b/vortex-tensor/benches/inner_product.rs index 796e9b648d6..5c4adf1c7ec 100644 --- a/vortex-tensor/benches/inner_product.rs +++ b/vortex-tensor/benches/inner_product.rs @@ -15,11 +15,19 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::inner_product::InnerProduct; @@ -56,15 +64,25 @@ fn vectors(width: usize, seed: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_inner_product(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { let session = vortex_array::array_session(); bencher .counter(ItemsCount::new(lhs.len())) .with_inputs(|| { ( - InnerProduct::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + InnerProduct + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -84,3 +102,22 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_inner_product(bencher, lhs, vectors(width, 31)); } + +#[divan::bench(args = WIDTHS)] +fn column_x_constant(bencher: Bencher, width: usize) { + bench_inner_product(bencher, vectors(width, 0), constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_inner_product(bencher, constant_vector(width), vectors(width, 31)); +} + +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_inner_product(bencher, vectors(width, 0), rhs); +} diff --git a/vortex-tensor/src/encodings/normalized/mod.rs b/vortex-tensor/src/encodings/normalized/mod.rs index 0f79ce97bee..ed840a5771f 100644 --- a/vortex-tensor/src/encodings/normalized/mod.rs +++ b/vortex-tensor/src/encodings/normalized/mod.rs @@ -21,7 +21,6 @@ pub use array::NormalizedSlots; mod compress; pub use compress::NormalizedScheme; pub use compress::normalize; -pub(crate) use compress::try_build_constant_normalized; mod execute; diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 05a0078d3ae..1582ef9c7d0 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Cosine similarity expression for tensor-like types. +//! Cosine similarity between two tensor columns. +use num_traits::Float; use num_traits::Zero; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -14,36 +15,39 @@ use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; +use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::encodings::normalized::try_build_constant_normalized; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_norm::L2Norm; +use crate::scalar_fns::row::TensorRow; +#[cfg(test)] +use crate::scalar_fns::row::probe; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; +use crate::utils::l2_norm_row; /// Cosine similarity between two columns. /// /// Computes `dot(a, b) / (||a|| * ||b||)` over the flat backing buffer of each tensor or vector. /// The shape and permutation do not affect the result because cosine similarity only depends on the -/// element values, not their logical arrangement. +/// element values, not their logical arrangement. A zero norm on either side yields `0.0`. /// /// Both inputs must be tensor-like extension arrays ([`FixedShapeTensor`] or [`Vector`]) with the /// same dtype and a float element type. The output is a float column of the same float type. @@ -56,13 +60,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct CosineSimilarity; impl CosineSimilarity { /// Creates a new [`TypedScalarFnInstance`] wrapping the cosine similarity operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(CosineSimilarity, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the cosine similarity between `lhs` and @@ -70,127 +74,90 @@ impl CosineSimilarity { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(CosineSimilarity::new().erased(), vec![lhs, rhs]) + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) } } -impl ScalarFnVTable for CosineSimilarity { +impl RowFn for CosineSimilarity { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.cosine_similarity"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("CosineSimilarity must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_prepared_into::<(TensorRow, TensorRow), UninitElementSink, _, _>( + |(lhs, rhs)| { + #[cfg(test)] + probe::record(lhs.is_some(), rhs.is_some()); + ConstantNorms { + lhs: lhs.map(l2_norm_row), + rhs: rhs.map(l2_norm_row), + } + }, + |norms, (lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { + InitializedElement::write( + output, + cosine_similarity_row_prepared(norms, lhs, rhs), + ) + } + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands make the _stored_ norms and normalized children + /// authoritative: `cos(D(x, s), D(y, t)) = dot(x, y)` and `cos(D(x, s), y) = dot(x, y) / + /// ||y||`, in both cases forced to `0.0` on rows where any authoritative norm is `0.0` (even + /// for lossy children whose decoded coordinates are nonzero). + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, + args: &[ArrayRef], ctx: &mut ExecutionCtx, - ) -> VortexResult { - let mut lhs_ref = args.get(0)?; - let mut rhs_ref = args.get(1)?; - let len = args.row_count(); - - // Normalize extension-level constants so the encoded fast path can use them. - if let Some(normalized_array) = try_build_constant_normalized(&lhs_ref, len, ctx)? { - lhs_ref = normalized_array.into_array(); - } - if let Some(normalized_array) = try_build_constant_normalized(&rhs_ref, len, ctx)? { - rhs_ref = normalized_array.into_array(); - } + ) -> VortexResult> { + let lhs = args[0].clone(); + let rhs = args[1].clone(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { - NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); - } + match NormalizedOrientation::classify(&lhs, &rhs) { + NormalizedOrientation::Both { lhs, rhs } => cosine_both_normalized(lhs, rhs, ctx) + .map(|output| Some(RowExecution::Output(output))), NormalizedOrientation::One { normalized_array, plain, - } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); - } - NormalizedOrientation::Neither => {} + } => cosine_one_normalized(normalized_array, plain, ctx) + .map(|output| Some(RowExecution::Output(output))), + NormalizedOrientation::Neither => Ok(None), } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm.try_new_array(len, EmptyOptions, [lhs_ref.clone()])?; - let norm_rhs_arr = L2Norm.try_new_array(len, EmptyOptions, [rhs_ref.clone()])?; - let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; - - // Execute to get the inner product and norms of the arrays. We only fully decompress - // because we need to perform special logic (guard against 0) during division. - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - let norm_l: PrimitiveArray = norm_lhs_arr.into_array().execute(ctx)?; - let norm_r: PrimitiveArray = norm_rhs_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norm_l.as_slice::(); - let norms_r = norm_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - let denom = norms_l[i] * norms_r[i]; - - if denom == T::zero() { - T::zero() - } else { - dots[i] / denom - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false } } @@ -220,584 +187,177 @@ impl ScalarFnArrayVTable for CosineSimilarity { } } -impl CosineSimilarity { - /// Both sides are [`Normalized`]-encoded: treat the normalized children as authoritative, so - /// `cosine_similarity = dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - // `Normalized` makes the normalized children authoritative, so their dot product is the - // cosine similarity even for lossy storage wrappers, except that a zero stored norm still - // represents a zero vector. - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norms_l.as_slice::(); - let norms_r = norms_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if norms_l[i] == T::zero() || norms_r[i] == T::zero() { - T::zero() - } else { - dots[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: treat the normalized child as authoritative, so - /// `cosine_similarity = dot(n, b) / ||b||`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, normalized_norms) = extract_normalized_children(normalized_ref); - - let dot_arr = InnerProduct::try_new_array(normalized, plain_ref.clone())?; - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - - let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - - let norm_arr = L2Norm.try_new_array(len, EmptyOptions, [plain_ref.clone()])?; - let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let normalized_norms = normalized_norms.as_slice::(); - let plain_norms = plain_norm.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if normalized_norms[i] == T::zero() || plain_norms[i] == T::zero() { - T::zero() - } else { - dots[i] / plain_norms[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } +/// Per-batch state for the cosine row kernel: the L2 norm of each operand that is constant for +/// the batch. +/// +/// A broadcast query vector holds the same elements in every row, so its norm is the same in +/// every row too. Computing it in the prepare step hoists an `O(width)` pass and a `sqrt` per row +/// out of the row loop. `None` marks an operand that varies by row, whose norm the row closure +/// computes exactly as it did before the hoist. +struct ConstantNorms { + /// The norm of the lhs when it is batch-constant. + lhs: Option, + + /// The norm of the rhs when it is batch-constant. + rhs: Option, } -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::cosine_similarity::CosineSimilarity; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::constant_tensor_array; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. - fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[test] - fn unit_vectors_1d() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 0.0, 1.0, 0.0, // Tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 1.0, 0.0, 0.0, // Tensor 2 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - /// Single-row cosine similarity for various vector pairs. - #[rstest] - // Antiparallel -> -1.0. - #[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] - // dot=24, both magnitudes=5 -> 24/25 = 0.96. - #[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] - // Zero vector -> guarded to 0.0. - #[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); - Ok(()) - } - - /// Self-similarity across various tensor shapes should always produce 1.0. - #[rstest] - // 2x3 matrix, flattened to 6 elements. - #[case::matrix_2d( - &[2, 3], - &[ - 1.0, 0.0, 0.0, // row 0 - 0.0, 0.0, 0.0, // row 1 - ], - )] - // 2x2x2 tensor, 8 elements. - #[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] - fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { - let lhs = tensor_array(shape, elements)?; - let rhs = tensor_array(shape, elements)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn scalar_0d() -> VortexResult<()> { - // 0-dimensional tensor: each "tensor" is a single scalar value. - let lhs = tensor_array(&[], &[5.0, 3.0])?; - let rhs = tensor_array(&[], &[5.0, -3.0])?; - - // Same sign -> 1.0, opposite sign -> -1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); - Ok(()) - } - - #[test] - fn many_rows() -> VortexResult<()> { - // 5 tensors of shape [4] compared against themselves -> all 1.0. - let lhs = tensor_array( - &[4], - &[ - 1.0, 2.0, 3.0, 4.0, // tensor 0 - 0.0, 1.0, 0.0, 0.0, // tensor 1 - 5.0, 0.0, 5.0, 0.0, // tensor 2 - 1.0, 1.0, 1.0, 1.0, // tensor 3 - 0.0, 0.0, 0.0, 7.0, // tensor 4 - ], - )?; - let rhs = lhs.clone(); - - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0, 1.0, 1.0, 1.0, 1.0], - ); - Ok(()) - } - - #[test] - fn constant_query_tensor() -> VortexResult<()> { - // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. - let data = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 0.0, 1.0, 0.0, // tensor 1 - 0.0, 0.0, 1.0, // tensor 2 - 1.0, 0.0, 0.0, // tensor 3 - ], - )?; - let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn vector_unit_vectors() -> VortexResult<()> { - let lhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 1.0, 0.0, 0.0, // vector 1 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn vector_constant_query() -> VortexResult<()> { - let data = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - 0.0, 0.0, 1.0, // vector 2 - 1.0, 0.0, 0.0, // vector 3 - ], - )?; - let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. - let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; - let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; - let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: self-similarity = 1.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_self_similarity() -> VortexResult<()> { - // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. - // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Self-similarity should always be 1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); - Ok(()) - } - - #[test] - fn both_normalized_orthogonal() -> VortexResult<()> { - // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. - // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn both_normalized_zero_norm() -> VortexResult<()> { - // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS is plain [3.0, 4.0]. - // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[3.0, 4.0])?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 0.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); - Ok(()) - } - - #[test] - fn both_normalized_null_rows() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_r = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - let validity = Validity::from_iter([true, false]); - let rhs = Normalized::try_new(normalized_r, norms_r, validity, &mut ctx)?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine - // similarity for that row must be `0.0` even though the dot product of the normalized - // children is nonzero. - let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; - let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally violates the unit-norm invariant by pairing a nonzero normalized row - // with a stored norm of `0.0`, mimicking lossy storage. - // SAFETY: The children are structurally valid. - let lhs = - unsafe { Normalized::new_unchecked(normalized_l, norms_l, Validity::NonNullable) } - .into_array(); - - let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; - let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Same as above for the rhs operand. - // SAFETY: The children are structurally valid. - let rhs = - unsafe { Normalized::new_unchecked(normalized_r, norms_r, Validity::NonNullable) } - .into_array(); - - // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both - // `0.0`, so cosine similarity must be `0.0`. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. The plain side is a normal nonzero - // tensor with positive norm. cosine similarity must still be `0.0` because the - // authoritative stored norm on the normalized_array side is `0.0`. - let normalized = tensor_array(&[2], &[0.6, 0.8])?; - let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally pairs a nonzero normalized row with a stored norm of `0.0`, mimicking - // lossy storage where the stored norm is authoritative. - // SAFETY: The children are structurally valid. - let normalized_array = - unsafe { Normalized::new_unchecked(normalized, norms, Validity::NonNullable) } - .into_array(); - - let plain = tensor_array(&[2], &[1.0, 0.0])?; - - // Normalized encoding on the lhs: `One { normalized_array: lhs, plain: rhs }`. - assert_close( - &eval_cosine_similarity(normalized_array.clone(), plain.clone())?, - &[0.0], - ); - - // Normalized encoding on the rhs: `One { normalized_array: rhs, plain: lhs }`. The same - // zero-norm guard must fire regardless of operand order. - assert_close(&eval_cosine_similarity(plain, normalized_array)?, &[0.0]); - Ok(()) - } - - #[test] - fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { - // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. - // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. - let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 - 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 - 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 - 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { - // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn both_constant_tensors() -> VortexResult<()> { - // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). - let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; - let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; - let expected = 1.0 / 2.0_f64.sqrt(); - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[expected, expected, expected], - ); - Ok(()) - } - - #[test] - fn constant_zero_norm_query() -> VortexResult<()> { - // A zero-norm constant query must produce `0.0` for every row via the zero-norm guard in - // `execute_one_normalized` and `execute_both_normalized`. - let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 2.0, 3.0, // - 4.0, 5.0, 6.0, // - 7.0, 8.0, 9.0, // - ], - )?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); - Ok(()) - } - - #[test] - fn constant_self_similarity_nonunit() -> VortexResult<()> { - // A non-unit constant query compared to itself must produce `1.0`. This exercises the - // helper's division: after normalization, both sides must be exactly unit so the - // Normalized fast path's inner product yields 1. - let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); - Ok(()) - } - - #[test] - fn vector_constant_matches_plain() -> VortexResult<()> { - // Exercise the `Vector` extension variant through the new pre-pass. - let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[rstest] - #[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] - #[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = CosineSimilarity::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("CosineSimilarity serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) +/// Computes the cosine similarity of one row, taking any hoisted norm from `norms` and computing +/// the rest exactly as [`cosine_similarity_row`] does. +/// +/// Each arm accumulates the same values in the same order as [`cosine_similarity_row`], and the +/// denominator keeps its lhs-times-rhs order, so the result is bit-identical whether a norm was +/// hoisted or not. The match costs one predictable branch per row: the arm is the same for the +/// whole batch. +fn cosine_similarity_row_prepared( + norms: &ConstantNorms, + lhs: &[T], + rhs: &[T], +) -> T { + match (norms.lhs, norms.rhs) { + (None, None) => cosine_similarity_row(lhs, rhs), + (Some(lhs_norm), None) => { + let mut dot = T::zero(); + let mut rhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + rhs_norm_squared = rhs_norm_squared + rhs_element * rhs_element; + } + cosine_from_parts(dot, lhs_norm * rhs_norm_squared.sqrt()) + } + (None, Some(rhs_norm)) => { + let mut dot = T::zero(); + let mut lhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + lhs_norm_squared = lhs_norm_squared + lhs_element * lhs_element; + } + cosine_from_parts(dot, lhs_norm_squared.sqrt() * rhs_norm) + } + (Some(lhs_norm), Some(rhs_norm)) => { + let mut dot = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + } + cosine_from_parts(dot, lhs_norm * rhs_norm) + } } +} - fn cosine_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") - } +/// Computes the cosine similarity of two equal-length float slices. +/// +/// Returns `dot(a, b) / (||a|| * ||b||)`, or `0.0` when either norm is zero. +fn cosine_similarity_row(lhs: &[T], rhs: &[T]) -> T { + let mut dot = T::zero(); + let mut lhs_norm_squared = T::zero(); + let mut rhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + lhs_norm_squared = lhs_norm_squared + lhs_element * lhs_element; + rhs_norm_squared = rhs_norm_squared + rhs_element * rhs_element; + } + + cosine_from_parts(dot, lhs_norm_squared.sqrt() * rhs_norm_squared.sqrt()) +} - fn cosine_vector_rhs() -> ArrayRef { - vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +/// The shared tail of every cosine arm: `dot / denominator`, guarded to `0.0` when it is +/// zero. +fn cosine_from_parts(dot: T, denominator: T) -> T { + if denominator == T::zero() { + T::zero() + } else { + dot / denominator } +} - fn cosine_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") - } +/// Both sides are [`Normalized`]-encoded: the normalized children are authoritative, so their dot +/// product is the cosine similarity, except that a row with a zero _stored_ norm is a zero vector. +/// +/// Unlike [`InnerProduct::reduce_encoded`], which composes lazy `Mul` arrays over the norm columns, +/// this executes and materializes. The zero-norm guard is a conditional per row rather than an +/// arithmetic factor, so there is no lazy array that expresses it; the norm columns are one value +/// per row rather than one per coordinate, so materializing them is cheap next to the decode this +/// avoids. +/// +/// [`InnerProduct::reduce_encoded`]: InnerProduct::reduce_encoded +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_both_normalized( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized_l, normalized_r])? + .execute(ctx)?; + let norms_l: PrimitiveArray = norms_l.execute(ctx)?; + let norms_r: PrimitiveArray = norms_r.execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let norms_l = norms_l.as_slice::(); + let norms_r = norms_r.as_slice::(); + // Zipped rather than indexed by `0..len`: one bounds check per iterator instead of three + // per row. A length disagreement between the children shortens the result, which the + // lifting reports against the batch row count rather than panicking mid-loop. + let buffer: Buffer = dots + .iter() + .zip(norms_l) + .zip(norms_r) + .map(|((&dot, &norm_l), &norm_r)| { + if norm_l.is_zero() || norm_r.is_zero() { + T::zero() + } else { + dot + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) +} - fn cosine_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") - } +/// One side is [`Normalized`]-encoded: `cos = dot(normalized, plain) / ||plain||`, forced to `0.0` +/// on rows where the stored norm or the plain norm is `0.0`. +/// +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_one_normalized( + normalized_array: &ArrayRef, + plain: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = normalized_array.len(); + let (normalized, normalized_norms) = extract_normalized_children(normalized_array); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized, plain.clone()])? + .execute(ctx)?; + let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; + let plain_norm: PrimitiveArray = L2Norm + .try_new_array(len, EmptyOptions, [plain.clone()])? + .execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let normalized_norms = normalized_norms.as_slice::(); + let plain_norms = plain_norm.as_slice::(); + // Zipped for the same reason as [`cosine_both_normalized`]. + let buffer: Buffer = dots + .iter() + .zip(normalized_norms) + .zip(plain_norms) + .map(|((&dot, &stored_norm), &plain_norm)| { + if stored_norm.is_zero() || plain_norm.is_zero() { + T::zero() + } else { + dot / plain_norm + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) } diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 74eb184045f..a31cce5ddda 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -6,40 +6,34 @@ use num_traits::Float; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::matcher::AnyTensor; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; -use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; /// Inner product (dot product) between two columns. /// @@ -52,13 +46,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct InnerProduct; impl InnerProduct { /// Creates a new [`TypedScalarFnInstance`] wrapping the inner product operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(InnerProduct, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the inner product between `lhs` and @@ -66,116 +60,89 @@ impl InnerProduct { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(InnerProduct::new().erased(), vec![lhs, rhs]) + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) } } -impl ScalarFnVTable for InnerProduct { +impl RowFn for InnerProduct { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.inner_product"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("InnerProduct must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - // TODO(connor): relax the float-only gate once integer tensors are supported. - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow, TensorRow), UninitElementSink, _>( + |(lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, inner_product_row(lhs, rhs)) } + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands factor through their stored norms: with `D(x, s)` denoting + /// `x * s` rowwise, `dot(D(x, s), D(y, t)) = s * t * dot(x, y)` and + /// `dot(D(x, s), y) = s * dot(x, y)`. The rewrite is expressed with lazy [`Operator::Mul`] + /// arrays over the (much smaller) norm columns, so no denormalized coordinates are decoded. + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let lhs_ref = args.get(0)?; - let rhs_ref = args.get(1)?; - let len = args.row_count(); + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let len = args[0].len(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { + Ok(match NormalizedOrientation::classify(&args[0], &args[1]) { NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized_l, normalized_r])?; + Some( + dot.binary(norms_l, Operator::Mul)? + .binary(norms_r, Operator::Mul)?, + ) } NormalizedOrientation::One { normalized_array, plain, } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); + let (normalized, norms) = extract_normalized_children(normalized_array); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized, plain.clone()])?; + Some(dot.binary(norms, Operator::Mul)?) } - NormalizedOrientation::Neither => {} + NormalizedOrientation::Neither => None, } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Canonicalize so we can perform the math directly. - let lhs: ExtensionArray = lhs_ref.execute(ctx)?; - let rhs: ExtensionArray = rhs_ref.execute(ctx)?; - - // We validated that both inputs have the same type. - let ext = lhs.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let dimensions = tensor_match.list_size() as usize; - - // Extract the storage array from each extension input. We pass the storage (FSL) rather - // than the extension array to avoid canonicalizing the extension wrapper. - let lhs_storage = lhs.storage_array(); - let rhs_storage = rhs.storage_array(); - - let lhs_flat = extract_flat_elements(lhs_storage, dimensions, ctx)?; - let rhs_flat = extract_flat_elements(rhs_storage, dimensions, ctx)?; - - match_each_float_ptype!(lhs_flat.ptype(), |T| { - let buffer: Buffer = (0..len) - .map(|i| inner_product_row(lhs_flat.row::(i), rhs_flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + .map(RowExecution::Output)) } } @@ -205,329 +172,12 @@ impl ScalarFnArrayVTable for InnerProduct { } } -impl InnerProduct { - /// Both sides are [`Normalized`]-encoded: `inner_product = s_l * s_r * dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let nl = norms_l.as_slice::(); - let nr = norms_r.as_slice::(); - let buffer: Buffer = (0..len).map(|i| nl[i] * nr[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: `inner_product = s * dot(n, other)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, norms) = extract_normalized_children(normalized_ref); - let normalized_norms: PrimitiveArray = norms.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized, plain_ref.clone())? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let ns = normalized_norms.as_slice::(); - let buffer: Buffer = (0..len).map(|i| ns[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } -} - /// Computes the inner product (dot product) of two equal-length float slices. /// /// Returns `sum(a_i * b_i)`. -fn inner_product_row(a: &[T], b: &[T]) -> T { - a.iter() - .zip(b.iter()) - .map(|(&x, &y)| x * y) - .fold(T::zero(), |acc, v| acc + v) -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::inner_product::InnerProduct; - use crate::tests::SESSION; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates inner product between two tensor arrays and returns the result as `Vec`. - fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - /// Single-row inner product for various vector pairs. - #[rstest] - // Orthogonal: [1, 0] . [0, 1] = 0. - #[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] - // Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. - #[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] - // Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. - #[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] - // Scaled: [2, 0] . [3, 0] = 6. - #[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_inner_product(lhs, rhs)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 3.0, 4.0, 0.0, // tensor 1 - 1.0, 1.0, 1.0, // tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 0.0, 1.0, 0.0, // tensor 0: dot = 0 - 3.0, 4.0, 0.0, // tensor 1: dot = 25 - 2.0, 2.0, 2.0, // tensor 2: dot = 6 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); - Ok(()) - } - - #[test] - fn vector_inner_product() -> VortexResult<()> { - let lhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0 - 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0: dot = 25 - 0.0, 1.0, // vector 1: dot = 0 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. - let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; - let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; - let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert!(prim.is_valid(2, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[23.0]); - assert_close(&[prim.as_slice::()[2]], &[127.0]); - Ok(()) - } - - #[test] - fn rejects_non_extension_dtype() { - let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); - let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - } - - #[test] - fn rejects_mismatched_dtypes() -> VortexResult<()> { - let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; - let rhs = vector_array(2, &[3.0_f64, 4.0])?; - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn both_normalized() -> VortexResult<()> { - // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). - // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). - // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; - - // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. - assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); - Ok(()) - } - - #[test] - fn both_normalized_multiple_rows() -> VortexResult<()> { - // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. - // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS: plain [1.0, 2.0]. - // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[1.0, 2.0])?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS: plain [1.0, 2.0]. - // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 2.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn both_normalized_null_rows() -> VortexResult<()> { - let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_l = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - let validity = Validity::from_iter([true, false]); - let lhs = Normalized::try_new(normalized_l, norms_l, validity, &mut ctx)?.into_array(); - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[25.0]); - Ok(()) - } - - #[rstest] - #[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] - #[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = InnerProduct::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(InnerProduct); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("InnerProduct serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn inner_product_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } - - fn inner_product_vector_rhs() -> ArrayRef { - vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") - } - - fn inner_product_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") - } - - fn inner_product_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") - } +fn inner_product_row(lhs: &[T], rhs: &[T]) -> T { + lhs.iter() + .zip(rhs) + .map(|(&lhs_element, &rhs_element)| lhs_element * rhs_element) + .fold(T::zero(), |sum, product| sum + product) } diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index a99ad76ea6e..0b3cf3634ee 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -164,3 +164,20 @@ unsafe impl InputElement for TensorRow { } } } + +/// Records which operands a test's prepare step received as batch constants. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + thread_local! { + /// Bit 0 records the lhs and bit 1 records the rhs. Thread-local storage prevents + /// concurrent tests from racing; row execution remains on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Records whether each operand was constant for the current batch. + pub(crate) fn record(lhs_constant: bool, rhs_constant: bool) { + SEEN_CONSTANTS.set(u8::from(lhs_constant) | (u8::from(rhs_constant) << 1)); + } +} diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs new file mode 100644 index 00000000000..13eed7b0224 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -0,0 +1,610 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::assert_arrays_eq; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::cosine_similarity::CosineSimilarity; +use crate::scalar_fns::row::probe; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::constant_tensor_array; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. +fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = CosineSimilarity::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = CosineSimilarity::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_cosine_similarity(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +/// Like [`eval_cosine_similarity`], but returns the executed array for exact array comparisons. +fn eval_cosine_similarity_array( + lhs: ArrayRef, + rhs: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + Ok(result + .into_array() + .execute::(ctx)? + .into_array()) +} + +#[test] +fn unit_vectors_1d() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 0.0, 1.0, 0.0, // Tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 1.0, 0.0, 0.0, // Tensor 2 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +/// Single-row cosine similarity for various vector pairs. +#[rstest] +// Antiparallel -> -1.0. +#[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] +// dot=24, both magnitudes=5 -> 24/25 = 0.96. +#[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] +// Zero vector -> guarded to 0.0. +#[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); + Ok(()) +} + +/// Self-similarity across various tensor shapes should always produce 1.0. +#[rstest] +// 2x3 matrix, flattened to 6 elements. +#[case::matrix_2d( + &[2, 3], + &[ + 1.0, 0.0, 0.0, // row 0 + 0.0, 0.0, 0.0, // row 1 + ], +)] +// 2x2x2 tensor, 8 elements. +#[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] +fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { + let lhs = tensor_array(shape, elements)?; + let rhs = tensor_array(shape, elements)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn scalar_0d() -> VortexResult<()> { + // 0-dimensional tensor: each "tensor" is a single scalar value. + let lhs = tensor_array(&[], &[5.0, 3.0])?; + let rhs = tensor_array(&[], &[5.0, -3.0])?; + + // Same sign -> 1.0, opposite sign -> -1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); + Ok(()) +} + +#[test] +fn many_rows() -> VortexResult<()> { + // 5 tensors of shape [4] compared against themselves -> all 1.0. + let lhs = tensor_array( + &[4], + &[ + 1.0, 2.0, 3.0, 4.0, // tensor 0 + 0.0, 1.0, 0.0, 0.0, // tensor 1 + 5.0, 0.0, 5.0, 0.0, // tensor 2 + 1.0, 1.0, 1.0, 1.0, // tensor 3 + 0.0, 0.0, 0.0, 7.0, // tensor 4 + ], + )?; + let rhs = lhs.clone(); + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0, 1.0, 1.0, 1.0, 1.0], + ); + Ok(()) +} + +#[test] +fn constant_query_tensor() -> VortexResult<()> { + // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. + let data = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 0.0, 1.0, 0.0, // tensor 1 + 0.0, 0.0, 1.0, // tensor 2 + 1.0, 0.0, 0.0, // tensor 3 + ], + )?; + let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn vector_unit_vectors() -> VortexResult<()> { + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 1.0, 0.0, 0.0, // vector 1 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn vector_constant_query() -> VortexResult<()> { + let data = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + 0.0, 0.0, 1.0, // vector 2 + 1.0, 0.0, 0.0, // vector 3 + ], + )?; + let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. + let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; + let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; + let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: self-similarity = 1.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_self_similarity() -> VortexResult<()> { + // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. + // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Self-similarity should always be 1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); + Ok(()) +} + +#[test] +fn both_normalized_orthogonal() -> VortexResult<()> { + // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. + // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn both_normalized_zero_norm() -> VortexResult<()> { + // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS is plain [3.0, 4.0]. + // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[3.0, 4.0])?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 0.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); + Ok(()) +} + +#[test] +fn both_normalized_null_rows() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_r = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + let validity = Validity::from_iter([true, false]); + let rhs = Normalized::try_new(normalized_r, norms_r, validity, &mut ctx)?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine + // similarity for that row must be `0.0` even though the dot product of the normalized + // children is nonzero. + let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; + let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally violates the unit-norm invariant by + // pairing a nonzero normalized row with a stored norm of `0.0`, mimicking lossy storage. + let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l, Validity::NonNullable) } + .into_array(); + + let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; + let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: Same as above for the rhs operand. + let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r, Validity::NonNullable) } + .into_array(); + + // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both + // `0.0`, so cosine similarity must be `0.0`. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. The plain side is a normal nonzero + // tensor with positive norm. cosine similarity must still be `0.0` because the + // authoritative stored norm on the denorm side is `0.0`. + let normalized = tensor_array(&[2], &[0.6, 0.8])?; + let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally pairs a nonzero normalized row with a + // stored norm of `0.0`, mimicking lossy storage where the stored norm is authoritative. + let denorm = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::NonNullable) }.into_array(); + + let plain = tensor_array(&[2], &[1.0, 0.0])?; + + // Denorm on the lhs: `One { denorm: lhs, plain: rhs }`. + assert_close( + &eval_cosine_similarity(denorm.clone(), plain.clone())?, + &[0.0], + ); + + // Denorm on the rhs: `One { denorm: rhs, plain: lhs }`. The same zero-norm guard must + // fire regardless of operand order. + assert_close(&eval_cosine_similarity(plain, denorm)?, &[0.0]); + Ok(()) +} + +#[test] +fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { + // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. + // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. + let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 + 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 + 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 + 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 + ], + )?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { + // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn both_constant_tensors() -> VortexResult<()> { + // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). + let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; + let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[expected, expected, expected], + ); + Ok(()) +} + +#[test] +fn constant_zero_norm_query() -> VortexResult<()> { + // A zero-norm constant query must produce `0.0` through the prepared row kernel's + // zero-denominator guard. + let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 2.0, 3.0, // + 4.0, 5.0, 6.0, // + 7.0, 8.0, 9.0, // + ], + )?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +#[test] +fn constant_self_similarity_nonunit() -> VortexResult<()> { + // The prepared path hoists both norms and computes the same dot product for every row. + let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); + Ok(()) +} + +/// An extension array over constant storage (what [`Vector::constant_array`] builds) is a batch +/// constant like any other. The row layer sees through the wrapper, so `prepare` hoists its norm. +#[test] +fn vector_constant_matches_plain() -> VortexResult<()> { + let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + assert_eq!( + probe::SEEN_CONSTANTS.get(), + 0b01, + "the extension-over-constant lhs must reach prepare as a batch constant", + ); + Ok(()) +} + +/// Both literal and extension-wrapped constant storage reach the prepared row path. The probe +/// ensures that the literal query remains a batch constant instead of becoming a per-row column. +/// +/// [`ConstantArray`]: vortex_array::arrays::ConstantArray +#[test] +fn literal_constant_rhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(lhs.clone(), literal_vector_array(&query, 4), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(lhs, vector_array(3, &query.repeat(4))?, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// The mirror of [`literal_constant_rhs_matches_expanded_column`], exercising the hoisted-lhs arm +/// of the prepared kernel. +#[test] +fn literal_constant_lhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(literal_vector_array(&query, 4), rhs.clone(), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(vector_array(3, &query.repeat(4))?, rhs, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// A zero-norm literal constant query must be guarded to `0.0` on every row by the prepared row +/// kernel, exactly as the unprepared kernel guards it. +#[test] +fn literal_constant_zero_norm_query_yields_zero() -> VortexResult<()> { + let lhs = vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = literal_vector_array(&[0.0f64, 0.0, 0.0], 2); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0]); + Ok(()) +} + +/// Two literal constants are folded to a single-row execution by the row lifting, and that row +/// still runs the prepared kernel with both norms hoisted. +#[test] +fn both_literal_constants() -> VortexResult<()> { + let lhs = literal_vector_array(&[1.0f64, 0.0, 0.0], 3); + let rhs = literal_vector_array(&[1.0f64, 1.0, 0.0], 3); + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[expected; 3]); + Ok(()) +} + +#[rstest] +#[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] +#[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("CosineSimilarity serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn cosine_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_vector_rhs() -> ArrayRef { + vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn cosine_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/inner_product.rs b/vortex-tensor/src/scalar_fns/tests/inner_product.rs new file mode 100644 index 00000000000..03d23aaacad --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/inner_product.rs @@ -0,0 +1,282 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::inner_product::InnerProduct; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates inner product between two tensor arrays and returns the result as `Vec`. +fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = InnerProduct::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = InnerProduct::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_inner_product(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +/// Single-row inner product for various vector pairs. +#[rstest] +// Orthogonal: [1, 0] . [0, 1] = 0. +#[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] +// Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. +#[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] +// Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. +#[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] +// Scaled: [2, 0] . [3, 0] = 6. +#[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_inner_product(lhs, rhs)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 3.0, 4.0, 0.0, // tensor 1 + 1.0, 1.0, 1.0, // tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 0.0, 1.0, 0.0, // tensor 0: dot = 0 + 3.0, 4.0, 0.0, // tensor 1: dot = 25 + 2.0, 2.0, 2.0, // tensor 2: dot = 6 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); + Ok(()) +} + +#[test] +fn vector_inner_product() -> VortexResult<()> { + let lhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0 + 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0: dot = 25 + 0.0, 1.0, // vector 1: dot = 0 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. + let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; + let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert!(prim.is_valid(2, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[23.0]); + assert_close(&[prim.as_slice::()[2]], &[127.0]); + Ok(()) +} + +#[test] +fn rejects_non_extension_dtype() { + let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); + let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); +} + +#[test] +fn rejects_mismatched_dtypes() -> VortexResult<()> { + let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; + let rhs = vector_array(2, &[3.0_f64, 4.0])?; + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn both_normalized() -> VortexResult<()> { + // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). + // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). + // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; + + // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. + assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); + Ok(()) +} + +#[test] +fn both_normalized_multiple_rows() -> VortexResult<()> { + // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. + // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS: plain [1.0, 2.0]. + // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[1.0, 2.0])?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS: plain [1.0, 2.0]. + // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 2.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn both_normalized_null_rows() -> VortexResult<()> { + let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_l = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let validity = Validity::from_iter([true, false]); + let lhs = Normalized::try_new(normalized_l, norms_l, validity, &mut ctx)?.into_array(); + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[25.0]); + Ok(()) +} + +#[rstest] +#[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] +#[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(InnerProduct); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("InnerProduct serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn inner_product_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} + +fn inner_product_vector_rhs() -> ArrayRef { + vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") +} + +fn inner_product_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn inner_product_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/mod.rs b/vortex-tensor/src/scalar_fns/tests/mod.rs index 5447772adda..bb3726e9329 100644 --- a/vortex-tensor/src/scalar_fns/tests/mod.rs +++ b/vortex-tensor/src/scalar_fns/tests/mod.rs @@ -3,5 +3,7 @@ //! Tests for the tensor scalar functions. +mod cosine_similarity; +mod inner_product; mod l2_norm; mod row; diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index c7339787b73..73d98730092 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -65,6 +65,9 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { } /// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. +/// +/// L2 norm and cosine similarity share this implementation so prepared constant norms use the +/// same accumulation order as rows from per-row inputs. pub(crate) fn l2_norm_row(row: &[T]) -> T { let mut sum_squared = T::zero(); for &element in row { @@ -121,19 +124,6 @@ pub fn validate_tensor_float_input(input_dtype: &DType) -> VortexResult( - lhs: &'a DType, - rhs: &DType, -) -> VortexResult> { - vortex_ensure!( - lhs.eq_ignore_nullability(rhs), - "binary tensor expression expects inputs to have the same dtype, got {lhs} and {rhs}" - ); - validate_tensor_float_input(lhs) -} - /// Validates that every argument has the same float tensor dtype, ignoring nullability. pub fn validate_tensor_float_inputs(args: &[DType]) -> VortexResult> { let (first, rest) = args @@ -332,7 +322,7 @@ impl BinaryTensorOpMetadata { let lhs_dtype = DType::from_proto(lhs_pb, session)?; let rhs_dtype = DType::from_proto(rhs_pb, session)?; - validate_binary_tensor_float_inputs(&lhs_dtype, &rhs_dtype)?; + validate_tensor_float_inputs(&[lhs_dtype.clone(), rhs_dtype.clone()])?; let lhs = children.get(0, &lhs_dtype, len)?; let rhs = children.get(1, &rhs_dtype, len)?; diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index ad3b96d1bff..492bc837b89 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -35,11 +35,13 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; @@ -79,7 +81,7 @@ pub fn build_similarity_search_tree>( let num_rows = data.len(); let query_vec = Vector::constant_array(query, num_rows)?; - let cosine = CosineSimilarity::try_new_array(data, query_vec)?.into_array(); + let cosine = CosineSimilarity.try_new_array(num_rows, EmptyOptions, [data, query_vec])?; let threshold_scalar = Scalar::primitive(threshold, Nullability::NonNullable); let threshold_array = ConstantArray::new(threshold_scalar, num_rows).into_array(); From 2a2d8ae9da281aa39480fdcc61d6cbad430b6924 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 12:59:39 -0400 Subject: [PATCH 157/160] Execute spatial distance with RowFn Signed-off-by: Connor Tsui --- vortex-spatial/Cargo.toml | 2 +- vortex-spatial/src/extension/mod.rs | 57 ++++++++++++ vortex-spatial/src/extension/point.rs | 18 ++++ vortex-spatial/src/extension/polygon.rs | 18 ++++ vortex-spatial/src/scalar_fn/distance.rs | 105 +++++++---------------- vortex-spatial/src/scalar_fn/mod.rs | 1 + vortex-spatial/src/scalar_fn/row.rs | 94 ++++++++++++++++++++ 7 files changed, 218 insertions(+), 77 deletions(-) create mode 100644 vortex-spatial/src/scalar_fn/row.rs diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 3be2b2d9d66..cd306089325 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -22,7 +22,7 @@ geo-types = { workspace = true } geoarrow = { workspace = true } geoarrow-cast = { workspace = true } prost = { workspace = true } -vortex-array = { workspace = true } +vortex-array = { workspace = true, features = ["unstable_row_fns"] } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } vortex-edition = { workspace = true } diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index f24f31f02aa..8a670bb0e12 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -178,6 +178,63 @@ pub(crate) fn geometries( } } +/// The geometry a null row decodes to under [`geometries_null_tolerant`]. Arbitrary: the caller +/// guarantees null rows are never read. +pub(crate) fn placeholder_geometry() -> Geometry { + Geometry::Point(geo_types::Point::new(0.0, 0.0)) +} + +/// Whether [`geometries_null_tolerant`] supports this array without filtering null rows first. +pub(crate) fn can_decode_geometries_null_tolerant(array: &ArrayRef) -> VortexResult { + if array.validity()?.definitely_no_nulls() { + return Ok(true); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + + Ok(ext.is::() || ext.is::()) +} + +/// Decode a native geometry column that may contain null rows, writing [`placeholder_geometry`] +/// into their slots. The caller guarantees null rows are never read. +/// +/// `Ok(None)` means this geometry type has no null-tolerant decode yet (`Point` and `Polygon` are +/// covered), and the caller falls back to the filter strategy, which never decodes a null row. A +/// column with definitely no nulls delegates to the ordinary [`geometries`] for any type. +pub(crate) fn geometries_null_tolerant( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>>> { + if array.validity()?.definitely_no_nulls() { + return geometries(array, ctx).map(Some); + } + + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + + if ext.is::() { + point_geometries_null_tolerant(&storage, ctx).map(Some) + } else if ext.is::() { + polygon_geometries_null_tolerant(&storage, ctx).map(Some) + } else { + Ok(None) + } +} + /// Decode a constant operand scalar to one geometry, a constant of any /// supported geometry type is decoded exactly like a column. pub(crate) fn single_geometry( diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index e6a00fe8fea..b774f624c4c 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -50,6 +50,7 @@ use super::coordinate::coordinate_from_struct; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A single location: `geoarrow.point`, stored as `Struct` of non-nullable `f64`. @@ -150,6 +151,23 @@ pub(crate) fn point_geometries( .collect() } +/// Like [`point_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn point_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + point_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + impl ArrowExportVTable for Point { fn arrow_ext_id(&self) -> Id { *ARROW_POINT diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index a4c88b07b22..bce33efe6e5 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -53,6 +53,7 @@ use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; use super::geoarrow_metadata; use super::geoarrow_to_wkb; +use super::placeholder_geometry; use super::spatial_metadata_from_arrow; /// A polygon: `geoarrow.polygon`, stored as `List>>` (rings of vertices). @@ -153,6 +154,23 @@ pub(crate) fn polygon_geometries( .collect() } +/// Like [`polygon_geometries`], but a null row decodes to the placeholder geometry instead of +/// failing. The caller guarantees null rows are never read. +pub(crate) fn polygon_geometries_null_tolerant( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + polygon_array(storage, ctx)? + .iter() + .map(|geometry| match geometry { + None => Ok(placeholder_geometry()), + Some(geometry) => Ok(geometry + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()), + }) + .collect() +} + /// Build a geoarrow `PolygonArray` from a `Polygon`'s `List>` storage. fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let polygon_type = polygon_type( diff --git a/vortex-spatial/src/scalar_fn/distance.rs b/vortex-spatial/src/scalar_fn/distance.rs index dfd3d09ed23..e41999338a6 100644 --- a/vortex-spatial/src/scalar_fn/distance.rs +++ b/vortex-spatial/src/scalar_fn/distance.rs @@ -6,43 +6,20 @@ use geo::Distance; use geo::Euclidean; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Distance`. -fn validate_distance_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: distance requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: distance operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -60,66 +37,41 @@ impl SpatialDistance { } } -impl ScalarFnVTable for SpatialDistance { +impl RowFn for SpatialDistance { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.distance"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { - Ok(EmptyOptions) - } - - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("distance has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_distance_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Primitive(PType::F64, nullability)) - } - - fn execute( + fn deserialize( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Distance is a value, not a verdict: no bounding-rect test can decide it. - execute_binary_geo_types(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx) + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn validity( + fn dispatch>( &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } - - fn is_strict(&self, _: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _: &Self::Options) -> bool { - false + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(GeometryRow, GeometryRow), UninitElementSink, _>( + |(a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, Euclidean.distance(a, b)) } + }, + ) } } @@ -196,8 +148,9 @@ mod tests { Ok(()) } - /// Distance passes no bounding-rect rejection: a point far outside a constant polygon's - /// bounding rect still gets its true distance, alongside an inside point at distance zero. + /// Distance is a value rather than a verdict, so no bounding-rect rejection may fire for it: a + /// point far outside a constant polygon's rect still gets its true distance. Carried over from + /// #9076, which added the rejection to the predicates but deliberately not to this function. #[test] fn distance_to_constant_polygon_is_exact() -> VortexResult<()> { let session = vortex_array::array_session(); diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index 99fe5d28528..6291075246a 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -13,3 +13,4 @@ mod execute; pub mod intersects; pub mod length; pub mod make_line; +pub(crate) mod row; diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs new file mode 100644 index 00000000000..b7353e10c1c --- /dev/null +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What the geo scalar functions add to the row-function machinery: an element type that decodes a +//! native geometry column into `geo_types` geometries. + +use geo_types::Geometry; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::scalar_fn::unstable::row::InputElement; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::extension::can_decode_geometries_null_tolerant; +use crate::extension::geometries; +use crate::extension::geometries_null_tolerant; +use crate::extension::is_native_geometry; + +/// Marker for native geometry input elements: accepts any native geometry column and presents each +/// row as a decoded `geo_types` geometry. +/// +/// The two operands of a binary geo function need not share a geometry type, since distance, +/// containment and intersection across types are all meaningful, so this validates only that the +/// column is _some_ native geometry. +pub(crate) struct GeometryRow; + +// SAFETY: [`view`](InputElement::view) returns the decoded geometry slice and +// [`view_len`](InputElement::view_len) reports that slice's exact length. +unsafe impl InputElement for GeometryRow { + type Column = Vec>; + type View<'a> = &'a [Geometry]; + type Elem<'a> = &'a Geometry; + + // A geometry row is decoded from its coordinate storage, which behind a null row holds arbitrary + // coordinates that need not describe a well-formed geometry. + const DENSE_SAFE: bool = false; + // Decoding builds a geometry from stored coordinates, and a malformed one in a *valid* row is a + // domain error rather than an infrastructural failure. + const DECODE_FALLIBLE: bool = true; + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + is_native_geometry(dtype), + "spatial: operand {dtype} is not a native geometry type" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + geometries(&array, ctx) + } + + fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { + can_decode_geometries_null_tolerant(array) + } + + fn get(column: &Self::Column, index: usize) -> &Geometry { + &column[index] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column.as_slice() + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.len() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry + where + Self: 'a, + { + &view[index] + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> &'a Geometry + where + Self: 'a, + { + // SAFETY: The caller established that `index` is below the slice length returned by + // `view_len` for this exact view. + unsafe { view.get_unchecked(index) } + } + + /// Null rows decode to a placeholder geometry that the branch-and-skip row loop never reads. + /// `Point` and `Polygon` columns are covered; other geometry types return `Ok(None)` and the + /// batch falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + geometries_null_tolerant(&array, ctx) + } +} From b37103bfe10defe9f80cc8bc4d3e1d92f78a0700 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 13:00:11 -0400 Subject: [PATCH 158/160] Execute spatial predicates with RowFn Signed-off-by: Connor Tsui --- Cargo.toml | 8 +- vortex-spatial/benches/binary_predicates.rs | 26 +- vortex-spatial/src/scalar_fn/contains.rs | 646 +++++++++++++++--- vortex-spatial/src/scalar_fn/execute.rs | 7 +- .../src/scalar_fn/execute/binary.rs | 253 +------ .../src/scalar_fn/execute/geo_types.rs | 24 - vortex-spatial/src/scalar_fn/intersects.rs | 244 +++++-- vortex-spatial/src/scalar_fn/row.rs | 94 ++- 8 files changed, 854 insertions(+), 448 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 65452f18300..fc38587579e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -162,7 +162,13 @@ flatbuffers = "25.2.10" fsst-rs = "0.6.0" futures = { version = "0.3.31", default-features = false } fuzzy-matcher = "0.3" -geo = "0.31.0" +# `vortex-spatial`'s `contains_route` transcribes geo's `impl_contains_from_relate!` dispatch +# table, so any bump that moves a row silently changes containment verdicts. The tests stay green +# wherever relate and the direct algorithm agree. Pinned exactly so that taking any new geo, +# patch releases included, is a deliberate edit of this line that re-verifies the table; a caret +# requirement would let `cargo update` (or automated lockfile maintenance) take 0.31.x with no diff +# to review. See `vortex-spatial/src/scalar_fn/contains.rs`. +geo = "=0.31.0" geo-traits = "0.3.0" geo-types = "0.7.19" geoarrow = "0.8.0" diff --git a/vortex-spatial/benches/binary_predicates.rs b/vortex-spatial/benches/binary_predicates.rs index b84ab67b17c..281578e4e16 100644 --- a/vortex-spatial/benches/binary_predicates.rs +++ b/vortex-spatial/benches/binary_predicates.rs @@ -12,11 +12,6 @@ //! column-x-column arms are the control: no operand is constant, so a prepared path has nothing to //! hoist and must not regress them. //! -//! `contains` has no all-overlapping arm. One `contains(query polygon, contained square)` row -//! builds a topology graph over the constant's 128 edges, which CodSpeed's CPU simulation charges -//! around 120 µs, so no row count both fits the per-iteration budget and exercises the row loop. -//! [`intersects::polygons_overlapping_x_constant`] covers the never-rejects case instead. -//! //! Run with `cargo bench -p vortex-spatial --bench binary_predicates`. #![expect(clippy::unwrap_used)] @@ -64,6 +59,10 @@ const ROWS: usize = 1 << 7; /// pairwise predicate. It needs a smaller fixture than [`ROWS`] to stay inside the same budget. const OVERLAPPING_POLYGON_ROWS: usize = 1 << 5; +/// Containment builds a topology graph for each polygon pair. Four rows fit the benchmark budget +/// while exercising construction followed by reuse of the prepared constant geometry. +const CONTAINED_POLYGON_ROWS: usize = 4; + /// Deterministic pseudo-random value in `[0, 1)`. fn unit(i: usize) -> f64 { ((i.wrapping_mul(2654435761) >> 8) % 10_000) as f64 / 10_000.0 @@ -225,6 +224,23 @@ mod contains { }); } + /// Constant container against contained polygons: every bbox check passes, the first row + /// prepares the constant geometry, and the remaining rows reuse it for the full predicate. + #[divan::bench] + fn constant_x_polygons_overlapping(bencher: Bencher) { + let mut ctx = SESSION.create_execution_ctx(); + let query = query_constant(&mut ctx, CONTAINED_POLYGON_ROWS); + let polygons = squares_mostly_overlapping(CONTAINED_POLYGON_ROWS); + bencher + .counter(ItemsCount::new(CONTAINED_POLYGON_ROWS)) + .bench_local(|| { + execute( + SpatialContains::try_new_array(query.clone(), polygons.clone()), + &mut ctx, + ) + }); + } + /// Constant container against a point column with one null row in eight. #[divan::bench] fn constant_x_nullable_points(bencher: Bencher) { diff --git a/vortex-spatial/src/scalar_fn/contains.rs b/vortex-spatial/src/scalar_fn/contains.rs index 8850e59f751..9b22a471204 100644 --- a/vortex-spatial/src/scalar_fn/contains.rs +++ b/vortex-spatial/src/scalar_fn/contains.rs @@ -3,44 +3,31 @@ //! `ST_Contains`: OGC containment test between two native geometries. +use std::cell::OnceCell; + +use geo::BoundingRect; use geo::Contains; +use geo::PreparedGeometry; +use geo::Relate; +use geo_types::Geometry; +use geo_types::Rect; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Contains`. -fn validate_contains_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: contains requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: contains operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; +#[cfg(test)] +use crate::scalar_fn::row::probe; /// OGC `ST_Contains` between two native geometry operands, each a column or a constant /// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone @@ -59,83 +46,314 @@ impl SpatialContains { } } -impl ScalarFnVTable for SpatialContains { +impl RowFn for SpatialContains { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.contains"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { Ok(EmptyOptions) } - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) + /// Containment is not symmetric, so `a` is always the container and `b` the contained. + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), UninitElementSink, _, _>( + |(a, b)| { + #[cfg(test)] + probe::record(a.is_some(), b.is_some()); + ConstOperands { + a: a.map(PreparedOperand::new), + b: b.map(PreparedOperand::new), + } + }, + |operands, (a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, contains_row_prepared(operands, a, b)) } + }, + ) } +} + +/// Per-batch state for the contains row kernel: the prepared form of whichever operand is +/// constant for the batch. `None` marks an operand that varies by row. +struct ConstOperands { + /// Operand `a` (the container) when it is batch-constant. + a: Option, + + /// Operand `b` (the contained) when it is batch-constant. + b: Option, +} + +/// One batch-constant operand: its bounding rectangle and the [`PreparedGeometry`] built on the +/// first row whose pairing routes through relate. +/// +/// The build is lazy because preparation (self-noding the topology graph plus an R*-tree over the +/// edges) costs `O(edges log edges)` and pays off only on relate-routed pairings; a batch of +/// point rows against a constant polygon never touches it, and preparing a large constant eagerly +/// would charge such a batch for nothing. +struct PreparedOperand { + /// The constant's bounding rectangle, folded once for conservative row rejection. + bbox: Option>, + + /// The constant's prepared form, initialized only when a relate route needs it. + prepared: OnceCell, f64>>, +} - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("contains has exactly two children"), +impl PreparedOperand { + fn new(geometry: &Geometry) -> Self { + Self { + bbox: finite_bounding_rect(geometry), + prepared: OnceCell::new(), } } - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_contains_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Bool(nullability)) + /// Return the prepared geometry, cloning the decoded constant only on first use. + /// + /// `geometry` **must** be the constant represented by this state. The row kernel maintains + /// that relationship by passing the operand from the same decoded constant column that + /// produced this [`PreparedOperand`]. + fn get(&self, geometry: &Geometry) -> &PreparedGeometry<'static, Geometry, f64> { + self.prepared + .get_or_init(|| PreparedGeometry::from(geometry.clone())) } +} - fn execute( - &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Containment is not symmetric: `a` is always the container and `b` the contained. A - // container's rect must cover the contained's rect (`Rect::contains` is the closed - // test), so a contained rect poking outside proves the row false. - execute_binary_geo_types( - &a, - &b, - |a, b| a.contains(b), - Some(|ra, rb| (!ra.contains(rb)).then_some(false)), - ctx, - ) - } +/// Returns a bounding rectangle only when ordered comparisons can conservatively reject a row. +/// +/// Geo permits non-finite coordinates. A rectangle containing NaN cannot prove non-containment, +/// because its ordered comparisons can return false even when the exact algorithm accepts the +/// geometry. +fn finite_bounding_rect(geometry: &Geometry) -> Option> { + let bbox = geometry.bounding_rect()?; + let min = bbox.min(); + let max = bbox.max(); + + [min.x, min.y, max.x, max.y] + .into_iter() + .all(f64::is_finite) + .then_some(bbox) +} - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) +/// How geo's `a.contains(b)` computes its verdict for a pairing. +enum ContainsRoute { + /// `a.relate(b).is_contains()`. + ForwardRelate, + + /// `b.relate(a).is_within()`, how geo phrases relate for `MultiPolygon` containers. + ReversedRelate, + + /// A direct algorithm (coordinate position, point arithmetic); nothing to prepare. + Direct, +} + +/// The route geo 0.31's `Contains` dispatch takes for `a.contains(b)`. +/// +/// The prepared substitution in [`contains_row_prepared`] **must** run relate exactly where geo +/// runs relate, with the same argument order, because geo's direct algorithms are not everywhere +/// bit-identical to a relate matrix query (they resolve degenerate and boundary cases with +/// different arithmetic). The relate rows below transcribe geo's `impl_contains_from_relate!` +/// lists per container type; everything else, notably every `Point`/`MultiPoint` contained side +/// and every `Point` container, is direct. +/// +/// **This table is coupled to the geo version.** It transcribes a dispatch that geo is free to +/// reshuffle in any release, and a wrong row is a silently wrong verdict rather than a build error. +/// The workspace therefore pins `geo = "=0.31.0"`: taking any new geo, patch releases included, is +/// a deliberate edit of that line, and the edit must re-verify this table against +/// `impl_contains_from_relate!`. +/// +/// `constant_operands_agree_with_columns` is the mechanical check, and it is **not** complete: it +/// compares the prepared route against plain `a.contains(b)` only for the container types it has +/// cases for. `routes_agree_with_geo_for_every_container` covers the rest, one representative +/// pairing per container variant, and is the one to extend when geo grows a geometry type. Both +/// stay green wherever relate and the direct algorithm agree, so neither replaces the pin. +fn contains_route(a: &Geometry, b: &Geometry) -> ContainsRoute { + use Geometry as G; + + match (a, b) { + // Line contains [Polygon, MultiLineString, MultiPolygon, GeometryCollection, Rect, + // Triangle]. + ( + G::Line(_), + G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // LineString contains [Polygon, MultiPoint, MultiLineString, MultiPolygon, + // GeometryCollection, Rect, Triangle]. + | ( + G::LineString(_), + G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // MultiLineString contains everything except Point. + | ( + G::MultiLineString(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // MultiPoint contains [Line, LineString, Polygon, MultiLineString, MultiPolygon, + // GeometryCollection, Rect, Triangle]. + | ( + G::MultiPoint(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // Polygon contains everything except Point and MultiPoint. + | ( + G::Polygon(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) + // Rect contains [Line, LineString, MultiPoint, MultiLineString, MultiPolygon, + // GeometryCollection, Triangle]; Rect contains Rect and Polygon are direct. + | ( + G::Rect(_), + G::Line(_) + | G::LineString(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Triangle(_), + ) + // Triangle and GeometryCollection contain everything except Point. + | ( + G::Triangle(_) | G::GeometryCollection(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiPoint(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) => ContainsRoute::ForwardRelate, + + // MultiPolygon contains everything except Point and MultiPoint, phrased reversed. + ( + G::MultiPolygon(_), + G::Line(_) + | G::LineString(_) + | G::Polygon(_) + | G::MultiLineString(_) + | G::MultiPolygon(_) + | G::GeometryCollection(_) + | G::Rect(_) + | G::Triangle(_), + ) => ContainsRoute::ReversedRelate, + + _ => ContainsRoute::Direct, } +} - fn is_strict(&self, _: &Self::Options) -> bool { - true +/// Computes one row of contains, substituting a prepared graph for a constant operand on the +/// pairings geo itself answers through relate. +/// +/// [`PreparedGeometry`] carries the operand's self-noded topology graph and edge R*-tree, so a +/// relate against it skips rebuilding both and reads its bounding rect from cache; geo asserts +/// the cached graph equal to a freshly built one (its `swap_arg_index` test), which is what makes +/// the substitution result-preserving. Before dispatch, a disjoint constant-side bounding rect +/// conservatively rejects the row, matching the columnar implementation's #9076 optimization. +/// All other rows delegate to the same direct or relate route as `a.contains(b)`. +fn contains_row_prepared(operands: &ConstOperands, a: &Geometry, b: &Geometry) -> bool { + let rejected = match (&operands.a, &operands.b) { + (None, None) => false, + (Some(const_a), Some(const_b)) => const_a + .bbox + .zip(const_b.bbox) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + (Some(const_a), None) => const_a + .bbox + .zip(finite_bounding_rect(b)) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + (None, Some(const_b)) => finite_bounding_rect(a) + .zip(const_b.bbox) + .is_some_and(|(bbox_a, bbox_b)| !bbox_a.contains(&bbox_b)), + }; + + if rejected { + return false; } - fn is_fallible(&self, _: &Self::Options) -> bool { - false + match contains_route(a, b) { + ContainsRoute::Direct => a.contains(b), + ContainsRoute::ForwardRelate => match (&operands.a, &operands.b) { + (Some(const_a), Some(const_b)) => const_a.get(a).relate(const_b.get(b)).is_contains(), + (Some(const_a), None) => const_a.get(a).relate(b).is_contains(), + (None, Some(const_b)) => a.relate(const_b.get(b)).is_contains(), + (None, None) => a.contains(b), + }, + ContainsRoute::ReversedRelate => match (&operands.a, &operands.b) { + (Some(const_a), Some(const_b)) => const_b.get(b).relate(const_a.get(a)).is_within(), + (Some(const_a), None) => b.relate(const_a.get(a)).is_within(), + (None, Some(const_b)) => const_b.get(b).relate(a).is_within(), + (None, None) => a.contains(b), + }, } } #[cfg(test)] mod tests { + use geo::Contains; + use geo_types::Coord; use geo_types::Geometry; + use geo_types::GeometryCollection; + use geo_types::Line; use geo_types::LineString; + use geo_types::MultiLineString; + use geo_types::MultiPoint; + use geo_types::MultiPolygon; use geo_types::Point; use geo_types::Polygon; + use geo_types::Rect; + use geo_types::Triangle; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -144,6 +362,7 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::MaskedArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -158,10 +377,15 @@ mod tests { use vortex_error::vortex_err; use wkb::writer::WriteOptions; + use super::ConstOperands; + use super::PreparedOperand; use super::SpatialContains; + use super::contains_row_prepared; + use crate::scalar_fn::row::probe::assert_prepared_agrees_with_columns; use crate::test_harness::linestring_column; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; + use crate::test_harness::polygon_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { @@ -218,6 +442,21 @@ mod tests { assert_contains(container, other, [expected; 3]) } + /// A non-finite bounding rectangle cannot reject a containment that the exact geometry + /// algorithm accepts. + #[test] + fn nan_bounding_rect_does_not_reject_containment() { + let container = multipoint(vec![(f64::NAN, f64::NAN), (1.0, 1.0)]); + let contained = point(1.0, 1.0); + let operands = ConstOperands { + a: Some(PreparedOperand::new(&container)), + b: Some(PreparedOperand::new(&contained)), + }; + + assert!(container.contains(&contained)); + assert!(contains_row_prepared(&operands, &container, &contained)); + } + /// Partially overlapping polygons contain each other in neither direction. #[test] fn overlapping_polygons_contain_neither_way() -> VortexResult<()> { @@ -246,6 +485,20 @@ mod tests { assert_contains(container, points, [true, false, false]) } + /// Constant container vs a linestring column: a row whose bounding rect pokes outside the + /// container's is not contained, while one wholly inside is. Carried over from the columnar + /// bounding-rect rejection in #9076, since it constrains the verdict rather than the mechanism. + #[test] + fn constant_container_vs_row_rect_poking_outside() -> VortexResult<()> { + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let lines = linestring_column(vec![ + vec![(1.0, 1.0), (3.0, 3.0)], + vec![(1.0, 1.0), (9.0, 1.0)], + vec![(5.0, 5.0), (9.0, 9.0)], + ])?; + assert_contains(container, lines, [true, false, false]) + } + /// Polygon column vs constant point: only the polygon around the point contains it. #[test] fn polygon_column_vs_constant_point() -> VortexResult<()> { @@ -266,20 +519,6 @@ mod tests { assert_contains(away, point, [false; 2]) } - /// Constant container vs a linestring column: a row whose bounding rect pokes outside the - /// container's rect is proven false by the rect pre-check alone; a fully inside row still - /// needs (and passes) the exact test. - #[test] - fn constant_container_vs_row_rect_poking_outside() -> VortexResult<()> { - let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; - let lines = linestring_column(vec![ - vec![(1.0, 1.0), (3.0, 3.0)], - vec![(1.0, 1.0), (9.0, 1.0)], - vec![(5.0, 5.0), (9.0, 9.0)], - ])?; - assert_contains(container, lines, [true, false, false]) - } - /// Column vs column pairs rows: each polygon row is tested against the point row at the /// same position. #[test] @@ -410,6 +649,83 @@ mod tests { Ok(()) } + /// A nullable polygon column: unit squares at `centers`, the rows where `nulls` is true + /// masked out, spelled as `Masked` over non-nullable storage. + fn nullable_squares(centers: &[(f64, f64)], nulls: &[bool]) -> VortexResult { + let squares = centers + .iter() + .map(|&(x, y)| { + vec![vec![ + (x - 1.0, y - 1.0), + (x + 1.0, y - 1.0), + (x + 1.0, y + 1.0), + (x - 1.0, y + 1.0), + (x - 1.0, y - 1.0), + ]] + }) + .collect(); + let polygons = polygon_column(squares)?; + + Ok( + MaskedArray::try_new(polygons, Validity::from_iter(nulls.iter().map(|n| !n)))? + .into_array(), + ) + } + + /// Nullable geometry operands conjoin their validity before computing containment. + #[test] + fn contains_nullable_geometries_conjoins_validity() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let centers = [(0.0, 0.0), (5.0, 5.0), (0.5, -0.2), (9.0, 9.0), (0.0, 1.0)]; + let nulls = [false, true, false, false, true]; + let polygons = nullable_squares(¢ers, &nulls)?; + let points = nullable_point_column(vec![ + Some((0.0, 0.0)), + Some((5.0, 5.0)), + None, + Some((0.0, 0.0)), + Some((0.0, 1.0)), + ])?; + + let actual = SpatialContains::try_new_array(polygons, points)? + .into_array() + .execute::(&mut ctx)? + .into_array(); + let expected = BoolArray::from_iter([Some(true), None, None, Some(false), None]); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// Geometry types without a null-tolerant decode fall back to filtering valid rows. + #[test] + fn contains_unsupported_geometry_falls_back_to_filter() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let validity = Validity::from_iter([true, false, true, true]); + let lines = linestring_column(vec![ + vec![(0.0, 0.0), (4.0, 4.0)], + vec![(0.0, 0.0), (1.0, 1.0)], + vec![(2.0, 2.0), (3.0, 3.0)], + vec![(0.0, 4.0), (4.0, 0.0)], + ])?; + let nullable_lines = MaskedArray::try_new(lines.clone(), validity.clone())?.into_array(); + let point = geometry_constant(&Geometry::Point(Point::new(2.0, 2.0)), 4)?; + + let expected = SpatialContains::try_new_array(lines, point.clone())?.into_array(); + let expected = MaskedArray::try_new(expected, validity)?.into_array(); + let actual = SpatialContains::try_new_array(nullable_lines, point)? + .into_array() + .execute::(&mut ctx)? + .into_array(); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + /// A non-geometry operand dtype is rejected up front, before execution. #[test] fn non_geometry_operand_is_rejected() -> VortexResult<()> { @@ -419,4 +735,166 @@ mod tests { assert!(result.is_err()); Ok(()) } + + // The prepared-vs-expanded agreement grid: every constant arrangement of a pairing must + // return exactly what the fully expanded columns return. + + /// A point geometry. + fn point(x: f64, y: f64) -> Geometry { + Geometry::Point(Point::new(x, y)) + } + + /// A linestring geometry through `coords`. + fn line(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::LineString(LineString::from(coords)) + } + + /// A multipoint geometry over `coords`. + fn multipoint(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::MultiPoint(MultiPoint::from(coords)) + } + + /// A two-point line segment geometry, the `Line` container variant. + fn line_geometry(start: (f64, f64), end: (f64, f64)) -> Geometry { + Geometry::Line(Line::new( + Coord { + x: start.0, + y: start.1, + }, + Coord { x: end.0, y: end.1 }, + )) + } + + /// A multilinestring geometry over one linestring per entry of `parts`. + fn multilinestring(parts: Vec>) -> Geometry { + Geometry::MultiLineString(MultiLineString::new( + parts.into_iter().map(LineString::from).collect(), + )) + } + + /// A geometry collection wrapping `parts`. + fn collection(parts: Vec) -> Geometry { + Geometry::GeometryCollection(GeometryCollection::from(parts)) + } + + /// An axis-aligned rectangle geometry, the `Rect` container variant. + fn rect_geometry(x0: f64, y0: f64, x1: f64, y1: f64) -> Geometry { + Geometry::Rect(Rect::new(Coord { x: x0, y: y0 }, Coord { x: x1, y: y1 })) + } + + /// A triangle geometry large enough to contain the small test polygons. + fn triangle_geometry() -> Geometry { + Geometry::Triangle(Triangle::new( + Coord { x: 0.0, y: 0.0 }, + Coord { x: 8.0, y: 0.0 }, + Coord { x: 0.0, y: 8.0 }, + )) + } + + /// A two-part multipolygon: `4x4` squares at the origin and at `(10, 10)`. + fn two_part_multipolygon() -> Geometry { + Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 4.0, 4.0), + rect_polygon(10.0, 10.0, 14.0, 14.0), + ])) + } + + /// Every container variant `contains_route` distinguishes, checked against plain + /// `a.contains(b)` in all four constant arrangements. + /// + /// Every case is a containment geo answers `true`, which the test asserts: a pairing that is + /// false regardless of route (a lower-dimensional container, say) also agrees regardless of + /// route, and pins nothing. A true case fails when the prepared substitution diverges from + /// geo: a table row whose relate phrasing disagrees with geo's dispatch on this input, or a + /// bounding-rect prescreen that wrongly rejects a contained row. It is **not** a version + /// tripwire: a geo release that reshuffles its dispatch stays green wherever relate and the + /// direct algorithm agree, which is why the workspace pins `geo` exactly. + /// + /// This is the table's own regression, and the one to extend when geo grows a geometry type: + /// `constant_operands_agree_with_columns` below goes through real arrays and so is the better + /// end-to-end check, but it only covers the container types it has cases for, and WKB decoding + /// limits which types those can be. The MultiPoint and Line containers route relate only for + /// contained types a MultiPoint or Line can rarely contain, so their true cases lean on + /// `GeometryCollection` membership and collinear `MultiLineString` parts respectively. + #[rstest] + #[case::point(point(1.0, 1.0), point(1.0, 1.0))] + #[case::line(line_geometry((0.0, 0.0), (4.0, 4.0)), point(2.0, 2.0))] + #[case::line_x_multilinestring(line_geometry((0.0, 0.0), (4.0, 4.0)), multilinestring(vec![vec![(1.0, 1.0), (2.0, 2.0)]]))] + #[case::linestring(line(vec![(0.0, 0.0), (4.0, 4.0)]), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::multipoint(multipoint(vec![(0.0, 0.0), (2.0, 2.0), (4.0, 4.0)]), collection(vec![point(2.0, 2.0)]))] + #[case::multilinestring(multilinestring(vec![vec![(0.0, 0.0), (4.0, 4.0)]]), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipolygon(two_part_multipolygon(), rect_polygon(1.0, 1.0, 3.0, 3.0).into())] + #[case::geometrycollection(collection(vec![rect_polygon(0.0, 0.0, 8.0, 8.0).into()]), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::rect(rect_geometry(0.0, 0.0, 8.0, 8.0), line(vec![(2.0, 2.0), (4.0, 4.0)]))] + #[case::triangle(triangle_geometry(), rect_polygon(1.0, 1.0, 2.0, 2.0).into())] + fn routes_agree_with_geo_for_every_container(#[case] a: Geometry, #[case] b: Geometry) { + let expected = a.contains(&b); + assert!( + expected, + "route cases must be containments geo answers true, or every route agrees vacuously", + ); + + let arrangements = [ + (None, None), + (Some(PreparedOperand::new(&a)), None), + (None, Some(PreparedOperand::new(&b))), + ( + Some(PreparedOperand::new(&a)), + Some(PreparedOperand::new(&b)), + ), + ]; + + for (index, (const_a, const_b)) in arrangements.into_iter().enumerate() { + let operands = ConstOperands { + a: const_a, + b: const_b, + }; + assert_eq!( + contains_row_prepared(&operands, &a, &b), + expected, + "arrangement {index} disagrees with geo's own contains", + ); + } + } + + /// Constant arrangements agree with expanded columns across the routes the prepared kernel + /// distinguishes: forward relate (polygon, linestring and multipoint containers), reversed + /// relate (multipolygon containers), and the direct pairings (a point on either side, + /// multipoint over multipoint, polygon over multipoint), including boundary contact, + /// crossing, disjoint and empty cases. + #[rstest] + #[case::polygon_nested_polygon(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::polygon_touching_from_inside(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(0.0, 2.0, 2.0, 4.0).into())] + #[case::polygon_overlapping_polygon(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(2.0, 2.0, 6.0, 6.0).into())] + #[case::polygon_disjoint_polygon(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::polygon_x_point_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(2.0, 2.0))] + #[case::polygon_x_point_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(0.0, 2.0))] + #[case::polygon_x_point_outside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(20.0, 20.0))] + #[case::polygon_x_nan_point(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), point(f64::NAN, 2.0))] + #[case::polygon_x_linestring_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon_x_linestring_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(0.0, 1.0), (0.0, 3.0)]))] + #[case::polygon_x_linestring_crossing(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![(-2.0, 2.0), (2.0, 2.0)]))] + #[case::polygon_x_empty_linestring(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), line(vec![]))] + #[case::polygon_x_multipoint_inside(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::polygon_x_multipoint_on_boundary(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), multipoint(vec![(0.0, 1.0), (0.0, 3.0)]))] + #[case::linestring_x_multipoint_on_line(line(vec![(0.0, 0.0), (4.0, 4.0)]), multipoint(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipoint_x_multipoint_subset(multipoint(vec![(0.0, 0.0), (2.0, 2.0), (4.0, 4.0)]), multipoint(vec![(2.0, 2.0)]))] + #[case::multipoint_x_linestring_between_points(multipoint(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(1.0, 1.0), (2.0, 2.0)]))] + #[case::multipolygon_x_polygon_in_one_part(two_part_multipolygon(), rect_polygon(1.0, 1.0, 3.0, 3.0).into())] + #[case::multipolygon_x_polygon_straddling(two_part_multipolygon(), rect_polygon(3.0, 3.0, 11.0, 11.0).into())] + #[case::multipolygon_x_polygon_disjoint(two_part_multipolygon(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::multipolygon_x_point_inside(two_part_multipolygon(), point(11.0, 11.0))] + #[case::point_x_point_equal(point(1.0, 1.0), point(1.0, 1.0))] + #[case::point_x_polygon(point(2.0, 2.0), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + fn constant_operands_agree_with_columns( + #[case] a: Geometry, + #[case] b: Geometry, + ) -> VortexResult<()> { + assert_prepared_agrees_with_columns( + SpatialContains::try_new_array, + geometry_constant(&a, 3)?, + geometry_constant(&b, 3)?, + ) + } } diff --git a/vortex-spatial/src/scalar_fn/execute.rs b/vortex-spatial/src/scalar_fn/execute.rs index 3a7494bcb39..2acdce76e36 100644 --- a/vortex-spatial/src/scalar_fn/execute.rs +++ b/vortex-spatial/src/scalar_fn/execute.rs @@ -7,17 +7,14 @@ //! propagation without prescribing how a kernel represents geometries or builds its output. //! Native columnar kernels such as `ST_MakeLine` use these dispatchers directly. //! -//! [`execute_unary_geo_types`] and [`execute_binary_geo_types`] are convenience adapters for -//! row-oriented algorithms from the `geo` ecosystem. They decode valid inputs into -//! `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`], such as an `f64` or -//! boolean array. +//! [`execute_unary_geo_types`] adapts row-oriented algorithms from the `geo` ecosystem. It decodes +//! valid inputs into `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`]. mod binary; mod geo_types; mod unary; pub(crate) use binary::dispatch_binary; -pub(crate) use binary::execute_binary_geo_types; pub(crate) use unary::dispatch_unary; pub(crate) use unary::execute_unary_geo_types; use vortex_array::ArrayRef; diff --git a/vortex-spatial/src/scalar_fn/execute/binary.rs b/vortex-spatial/src/scalar_fn/execute/binary.rs index f2c03bd1beb..5cf639461b0 100644 --- a/vortex-spatial/src/scalar_fn/execute/binary.rs +++ b/vortex-spatial/src/scalar_fn/execute/binary.rs @@ -1,28 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Binary pairwise dispatch, plus an adapter for row-oriented `geo_types` kernels. +//! Binary constant-and-column operand dispatch. -use geo::BoundingRect; -use geo_types::Geometry; -use geo_types::Rect; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_mask::Mask; use super::Execution; use super::Operand; -use super::geo_types::GeoTypesOutput; -use super::geo_types::eval_column; -use super::geo_types::eval_column_pair; -use crate::extension::single_geometry; /// Dispatch a binary strict geometry kernel over constants and columns. /// @@ -80,6 +72,7 @@ where if len != 0 && valid.all_false() { return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); } + kernel( Execution { operands: [left, right], @@ -90,245 +83,3 @@ where ctx, ) } - -/// A bounding-rectangle pre-check for [`execute_binary_geo_types`]'s one-constant paths. -/// -/// Called per row with rectangles in operand order, it returns `Some(result)` when they prove the -/// result and `None` when the exact kernel must run. -pub(crate) type BboxPrecheck = fn(&Rect, &Rect) -> Option; - -/// Run a binary row-oriented kernel whose inputs are decoded to `geo_types::Geometry`. -/// -/// The `geo_types` name describes the values passed to `compute`, not the output. `T` is converted -/// into a Vortex array before this function returns. Nulls propagate from either operand. With -/// exactly one constant operand, `bbox_precheck` may prove a result from the fixed constant -/// bounding rectangle and the current row's rectangle before the exact kernel runs. -pub(crate) fn execute_binary_geo_types( - left: &ArrayRef, - right: &ArrayRef, - compute: F, - bbox_precheck: Option>, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T + Copy, -{ - let nullability = Nullability::from(left.dtype().is_nullable() || right.dtype().is_nullable()); - dispatch_binary( - left, - right, - T::dtype(nullability), - |execution, ctx| match execution.operands { - [Operand::Constant(left), Operand::Constant(right)] => { - let left = single_geometry(&left, ctx)?; - let right = single_geometry(&right, ctx)?; - Ok(ConstantArray::new( - compute(&left, &right).into_scalar(execution.nullability), - execution.len, - ) - .into_array()) - } - [Operand::Constant(left), Operand::Column(right)] => { - let left = single_geometry(&left, ctx)?; - let prescreen = bbox_precheck.zip(left.bounding_rect()); - eval_column( - &right, - &execution.valid, - |right| { - prescreen - .and_then(|(precheck, fixed)| precheck(&fixed, &right.bounding_rect()?)) - .unwrap_or_else(|| compute(&left, right)) - }, - execution.nullability, - ctx, - ) - } - [Operand::Column(left), Operand::Constant(right)] => { - let right = single_geometry(&right, ctx)?; - let prescreen = bbox_precheck.zip(right.bounding_rect()); - eval_column( - &left, - &execution.valid, - |left| { - prescreen - .and_then(|(precheck, fixed)| precheck(&left.bounding_rect()?, &fixed)) - .unwrap_or_else(|| compute(left, &right)) - }, - execution.nullability, - ctx, - ) - } - [Operand::Column(left), Operand::Column(right)] => eval_column_pair( - &left, - &right, - &execution.valid, - compute, - execution.nullability, - ctx, - ), - }, - ctx, - ) -} - -#[cfg(test)] -mod tests { - use std::cell::Cell; - - use geo::Contains; - use geo::Intersects; - use geo_types::Geometry; - use vortex_array::ArrayRef; - use vortex_array::ExecutionCtx; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::ConstantArray; - use vortex_array::assert_arrays_eq; - use vortex_array::validity::Validity; - use vortex_buffer::BitBuffer; - use vortex_error::VortexResult; - - use super::BboxPrecheck; - use super::execute_binary_geo_types; - use crate::test_harness::linestring_column; - use crate::test_harness::nullable_point_column; - use crate::test_harness::point_column; - use crate::test_harness::polygon_column; - - const DISJOINT_PRECHECK: BboxPrecheck = - |left, right| (!left.intersects(right)).then_some(false); - - fn triangle_constant(len: usize, ctx: &mut ExecutionCtx) -> VortexResult { - let ring = vec![(0.0, 0.0), (10.0, 0.0), (0.0, 10.0), (0.0, 0.0)]; - let scalar = polygon_column(vec![vec![ring]])?.execute_scalar(0, ctx)?; - Ok(ConstantArray::new(scalar, len).into_array()) - } - - fn counting_intersects( - counter: &Cell, - ) -> impl Fn(&Geometry, &Geometry) -> bool + Copy { - move |left, right| { - counter.set(counter.get() + 1); - left.intersects(right) - } - } - - #[test] - fn bbox_precheck_skips_exact_test() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(3, &mut ctx)?; - let probes = point_column(vec![50.0, 8.0, 2.0], vec![50.0, 8.0, 2.0])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - #[test] - fn bbox_precheck_leaves_nulls_alone() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(3, &mut ctx)?; - let probes = nullable_point_column(vec![Some((50.0, 50.0)), None, Some((2.0, 2.0))])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - let expected = BoolArray::new( - BitBuffer::from_iter([false, false, true]), - Validity::from_iter([true, false, true]), - ) - .into_array(); - - assert_arrays_eq!(result, expected, &mut ctx); - assert_eq!(exact_runs.get(), 1); - Ok(()) - } - - #[test] - fn bbox_precheck_sees_rects_in_operand_order() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let triangle = triangle_constant(2, &mut ctx)?; - let exact_runs = Cell::new(0); - let counted = |left: &Geometry, right: &Geometry| { - exact_runs.set(exact_runs.get() + 1); - left.contains(right) - }; - - let result = execute_binary_geo_types( - &probes, - &triangle, - counted, - Some(|left, right| (!left.contains(right)).then_some(false)), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 0); - Ok(()) - } - - #[test] - fn empty_constant_falls_through_to_exact() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let scalar = linestring_column(vec![vec![]])?.execute_scalar(0, &mut ctx)?; - let empty = ConstantArray::new(scalar, 2).into_array(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let exact_runs = Cell::new(0); - - let result = execute_binary_geo_types( - &empty, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_PRECHECK), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - #[test] - fn bbox_precheck_matches_exact_results() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - let triangle = triangle_constant(6, &mut ctx)?; - let probes = nullable_point_column(vec![ - Some((50.0, 50.0)), - Some((8.0, 8.0)), - Some((2.0, 2.0)), - None, - Some((0.0, 0.0)), - Some((10.0, 0.0)), - ])?; - let exact = |left: &Geometry, right: &Geometry| left.intersects(right); - - let with_precheck = - execute_binary_geo_types(&triangle, &probes, exact, Some(DISJOINT_PRECHECK), &mut ctx)?; - let exact_only = execute_binary_geo_types(&triangle, &probes, exact, None, &mut ctx)?; - - assert_arrays_eq!(with_precheck, exact_only, &mut ctx); - Ok(()) - } -} diff --git a/vortex-spatial/src/scalar_fn/execute/geo_types.rs b/vortex-spatial/src/scalar_fn/execute/geo_types.rs index 038aca46502..7007f02cfc6 100644 --- a/vortex-spatial/src/scalar_fn/execute/geo_types.rs +++ b/vortex-spatial/src/scalar_fn/execute/geo_types.rs @@ -118,27 +118,3 @@ where let values = decoded.iter().map(compute).collect(); Ok(T::build_array(len, valid, values, nullability)) } - -/// Evaluate a decoded kernel over rows where both geometry columns are valid. -pub(super) fn eval_column_pair( - left: &ArrayRef, - right: &ArrayRef, - valid: &Mask, - compute: F, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T, -{ - let len = left.len(); - let left = geometries(&left.filter(valid.clone())?, ctx)?; - let right = geometries(&right.filter(valid.clone())?, ctx)?; - let values = left - .iter() - .zip(&right) - .map(|(left, right)| compute(left, right)) - .collect(); - Ok(T::build_array(len, valid, values, nullability)) -} diff --git a/vortex-spatial/src/scalar_fn/intersects.rs b/vortex-spatial/src/scalar_fn/intersects.rs index 77d33886ff3..9a3a198e838 100644 --- a/vortex-spatial/src/scalar_fn/intersects.rs +++ b/vortex-spatial/src/scalar_fn/intersects.rs @@ -3,44 +3,27 @@ //! `ST_Intersects`: OGC intersection test between two native geometries. +use geo::BoundingRect; use geo::Intersects; +use geo_types::Geometry; +use geo_types::Rect; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::is_native_geometry; -use crate::scalar_fn::execute::execute_binary_geo_types; - -/// Validate the two native geometry operands accepted by `ST_Intersects`. -fn validate_intersects_operands(dtypes: &[DType]) -> VortexResult<()> { - vortex_ensure!( - dtypes.len() == 2, - "spatial: intersects requires exactly two geometry operands, got {}", - dtypes.len() - ); - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "spatial: intersects operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} +use crate::scalar_fn::row::GeometryRow; +#[cfg(test)] +use crate::scalar_fn::row::probe; /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry /// operands, each a column or a constant literal. @@ -58,74 +41,100 @@ impl SpatialIntersects { } } -impl ScalarFnVTable for SpatialIntersects { +impl RowFn for SpatialIntersects { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["a", "b"]; + const FALLIBLE: bool = true; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.st.intersects"); *ID } - fn serialize(&self, _: &Self::Options) -> VortexResult>> { + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { Ok(Some(vec![])) } - fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { Ok(EmptyOptions) } - fn arity(&self, _: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("a"), - 1 => ChildName::from("b"), - _ => unreachable!("intersects has exactly two children"), - } - } - - fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_intersects_operands(dtypes)?; - let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); - Ok(DType::Bool(nullability)) - } - - fn execute( + fn dispatch>( &self, - _: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let a = args.get(0)?; - let b = args.get(1)?; - // Disjoint bounding rects prove the geometries disjoint; rect contact (closed test) - // falls through to the exact test. - execute_binary_geo_types( - &a, - &b, - |x, y| x.intersects(y), - Some(|ra, rb| (!ra.intersects(rb)).then_some(false)), - ctx, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_prepared_into::<(GeometryRow, GeometryRow), UninitElementSink, _, _>( + |(a, b)| { + #[cfg(test)] + probe::record(a.is_some(), b.is_some()); + ConstBboxes::new(a, b) + }, + |bboxes, (a, b), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, intersects_row_prepared(bboxes, a, b)) } + }, ) } +} - fn validity( - &self, - _: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } +/// Per-batch state for the intersects row kernel: the bounding rect of each operand that is +/// constant for the batch. +/// +/// geo opens many intersects pairings with `has_disjoint_bboxes`, an early-out that folds +/// [`bounding_rect`] over both operands. For a batch-constant operand that fold recomputes the +/// same rect every row, so it is hoisted here and [`intersects_row_prepared`] replays the +/// comparison with the hoisted value. `None` marks an operand that varies by row or has no +/// bounding rect (an empty geometry); both mean no early-out, exactly as `has_disjoint_bboxes` +/// treats a missing rect. +/// +/// [`bounding_rect`]: BoundingRect::bounding_rect +struct ConstBboxes { + /// The bounding rect of operand `a` when it is batch-constant. + a: Option>, + + /// The bounding rect of operand `b` when it is batch-constant. + b: Option>, +} - fn is_strict(&self, _: &Self::Options) -> bool { - true +impl ConstBboxes { + fn new(a: Option<&Geometry>, b: Option<&Geometry>) -> Self { + Self { + a: a.and_then(BoundingRect::bounding_rect), + b: b.and_then(BoundingRect::bounding_rect), + } } +} - fn is_fallible(&self, _: &Self::Options) -> bool { - false - } +/// Computes one row of intersects, spending any bounding rect hoisted into `bboxes`. +/// +/// Disjoint bounding rectangles conservatively prove that the geometries do not intersect. The +/// fall-through delegates to the unchanged `a.intersects(b)`, which refolds both rects internally, +/// so a batch where every row overlaps pays one extra `bounding_rect` fold over the row operand; +/// the win concentrates where most rows are disjoint, the usual spatial-filter shape. +fn intersects_row_prepared(bboxes: &ConstBboxes, a: &Geometry, b: &Geometry) -> bool { + let disjoint = match (bboxes.a, bboxes.b) { + (None, None) => false, + (Some(bbox_a), Some(bbox_b)) => !bbox_a.intersects(&bbox_b), + (Some(bbox_a), None) => b + .bounding_rect() + .is_some_and(|bbox_b| !bbox_a.intersects(&bbox_b)), + (None, Some(bbox_b)) => a + .bounding_rect() + .is_some_and(|bbox_a| !bbox_a.intersects(&bbox_b)), + }; + + if disjoint { + return false; + } + + a.intersects(b) } #[cfg(test)] @@ -133,7 +142,9 @@ mod tests { use geo_types::Coord; use geo_types::Geometry; use geo_types::LineString; + use geo_types::MultiPoint; use geo_types::MultiPolygon; + use geo_types::Point; use geo_types::Polygon; use rstest::rstest; use vortex_array::ArrayRef; @@ -158,8 +169,10 @@ mod tests { use wkb::writer::WriteOptions; use super::SpatialIntersects; + use crate::scalar_fn::row::probe::assert_prepared_agrees_with_columns; use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; + use crate::test_harness::rect_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. fn rect_polygon(x0: f64, y0: f64, x1: f64, y1: f64) -> Polygon { @@ -441,4 +454,85 @@ mod tests { assert!(result.is_err()); Ok(()) } + + // The prepared-vs-expanded agreement grid: every constant arrangement of a pairing must + // return exactly what the fully expanded columns return. + + /// A point geometry. + fn point(x: f64, y: f64) -> Geometry { + Geometry::Point(Point::new(x, y)) + } + + /// A linestring geometry through `coords`. + fn line(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::LineString(LineString::from(coords)) + } + + /// A multipoint geometry over `coords`. + fn multipoint(coords: Vec<(f64, f64)>) -> Geometry { + Geometry::MultiPoint(MultiPoint::from(coords)) + } + + /// Constant arrangements agree with expanded columns across the pairing classes the prepared + /// kernel treats differently: bbox-prechecked pairs (polygon x polygon, linestring x + /// anything, multipolygon blankets), direct pairs (points), the excluded `MultiPoint` route, + /// and an empty geometry whose bounding rect does not exist. + #[rstest] + #[case::polygons_overlapping(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(2.0, 2.0, 6.0, 6.0).into())] + #[case::polygons_touching_edge(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(4.0, 0.0, 8.0, 4.0).into())] + #[case::polygons_touching_corner(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(4.0, 4.0, 8.0, 8.0).into())] + #[case::polygons_disjoint(rect_polygon(0.0, 0.0, 4.0, 4.0).into(), rect_polygon(20.0, 20.0, 24.0, 24.0).into())] + #[case::polygons_nested(rect_polygon(0.0, 0.0, 8.0, 8.0).into(), rect_polygon(2.0, 2.0, 4.0, 4.0).into())] + #[case::polygon_x_point_inside(donut(), point(2.0, 2.0))] + #[case::polygon_x_point_on_boundary(donut(), point(0.0, 5.0))] + #[case::polygon_x_point_in_hole(donut(), point(5.0, 5.0))] + #[case::point_outside_x_polygon(point(20.0, 20.0), donut())] + #[case::nan_point_x_polygon(point(f64::NAN, 2.0), donut())] + #[case::polygon_x_nan_point(donut(), point(f64::NAN, 2.0))] + #[case::points_equal(point(1.0, 1.0), point(1.0, 1.0))] + #[case::points_distinct(point(1.0, 1.0), point(2.0, 1.0))] + #[case::linestring_crossing_polygon(line(vec![(-2.0, -2.0), (2.0, 2.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::linestring_disjoint_polygon(line(vec![(-2.0, -2.0), (-6.0, -6.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::linestring_in_polygon_hole(line(vec![(4.5, 4.5), (5.5, 5.5)]), donut())] + #[case::linestrings_crossing(line(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(0.0, 4.0), (4.0, 0.0)]))] + #[case::linestrings_disjoint(line(vec![(0.0, 0.0), (4.0, 4.0)]), line(vec![(10.0, 10.0), (14.0, 14.0)]))] + #[case::empty_linestring_x_polygon(line(vec![]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipoint_straddling_polygon(multipoint(vec![(2.0, 2.0), (20.0, 20.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipoint_outside_polygon(multipoint(vec![(20.0, 20.0), (30.0, 30.0)]), rect_polygon(0.0, 0.0, 4.0, 4.0).into())] + #[case::multipolygon_disjoint_polygon( + Geometry::MultiPolygon(MultiPolygon::new(vec![ + rect_polygon(0.0, 0.0, 2.0, 2.0), + rect_polygon(10.0, 10.0, 12.0, 12.0), + ])), + rect_polygon(20.0, 20.0, 24.0, 24.0).into() + )] + fn constant_operands_agree_with_columns( + #[case] a: Geometry, + #[case] b: Geometry, + ) -> VortexResult<()> { + assert_prepared_agrees_with_columns( + SpatialIntersects::try_new_array, + geometry_constant(&a, 3)?, + geometry_constant(&b, 3)?, + ) + } + + /// `Rect` has no WKB form, so its constant comes from a one-row rect column; its conservative + /// bbox early-out and exact fall-through must agree with the expanded form like the rest. + #[test] + fn rect_operand_agrees_with_columns() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let rect_scalar = rect_column(vec![(0.0, 0.0, 4.0, 4.0)])?.execute_scalar(0, &mut ctx)?; + let rect_constant = ConstantArray::new(rect_scalar, 3).into_array(); + let polygon_constant = + geometry_constant(&Geometry::Polygon(rect_polygon(2.0, 2.0, 6.0, 6.0)), 3)?; + + assert_prepared_agrees_with_columns( + SpatialIntersects::try_new_array, + rect_constant, + polygon_constant, + ) + } } diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index b7353e10c1c..750497e5f32 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -32,10 +32,10 @@ unsafe impl InputElement for GeometryRow { type View<'a> = &'a [Geometry]; type Elem<'a> = &'a Geometry; - // A geometry row is decoded from its coordinate storage, which behind a null row holds arbitrary - // coordinates that need not describe a well-formed geometry. + // A geometry row is decoded from its coordinate storage, which behind a null row holds + // arbitrary coordinates that need not describe a well-formed geometry. const DENSE_SAFE: bool = false; - // Decoding builds a geometry from stored coordinates, and a malformed one in a *valid* row is a + // Decoding builds a geometry from stored coordinates, and a malformed one in a _valid_ row is a // domain error rather than an infrastructural failure. const DECODE_FALLIBLE: bool = true; fn validate(dtype: &DType) -> VortexResult<()> { @@ -92,3 +92,91 @@ unsafe impl InputElement for GeometryRow { geometries_null_tolerant(&array, ctx) } } + +/// Test-only support for the prepared geo row kernels: a probe recording which operands a +/// `prepare` step saw as batch-constant, and the shared prepared-vs-expanded agreement check +/// built on it. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::ScalarFnArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + + thread_local! { + /// Which operands the last `prepare` saw as constant, as a bitmask (bit 0 for `a`, bit 1 + /// for `b`). Thread-local rather than a process global so concurrent tests in one process + /// (plain `cargo test`) cannot race it; execution runs on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Record which operands `prepare` saw as constant. + pub(crate) fn record(a_constant: bool, b_constant: bool) { + SEEN_CONSTANTS.set(u8::from(a_constant) | (u8::from(b_constant) << 1)); + } + + /// Execute `build(a, b)` and assert that `prepare` saw exactly `expect_seen` as its constant + /// operands, so the test knows which decode path the inputs took. + fn run_probed( + build: &impl Fn(ArrayRef, ArrayRef) -> VortexResult, + a: ArrayRef, + b: ArrayRef, + expect_seen: u8, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + SEEN_CONSTANTS.set(u8::MAX); + let result = build(a, b)? + .into_array() + .execute::(ctx)? + .into_array(); + + assert_eq!( + SEEN_CONSTANTS.get(), + expect_seen, + "prepare saw the wrong constant operands", + ); + Ok(result) + } + + /// Assert that every constant-operand arrangement of `build(a, b)` returns exactly what the + /// fully expanded columns return, and that each arrangement's constness really reached + /// `prepare` (so the constants exercised the stride-0 path rather than a decoded column). + /// + /// Arrangements: `a` constant, `b` constant, and both constant with `a` masked. A plain + /// constant pair folds to a single-row execution before the row loop, so masking one side is + /// what drives the both-hoisted arm across rows; that run is compared against the same mask + /// over the expanded column. + pub(crate) fn assert_prepared_agrees_with_columns( + build: impl Fn(ArrayRef, ArrayRef) -> VortexResult, + const_a: ArrayRef, + const_b: ArrayRef, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let col_a = const_a.clone().execute::(&mut ctx)?.into_array(); + let col_b = const_b.clone().execute::(&mut ctx)?.into_array(); + + let baseline = run_probed(&build, col_a.clone(), col_b.clone(), 0b00, &mut ctx)?; + let a_hoisted = run_probed(&build, const_a.clone(), col_b.clone(), 0b01, &mut ctx)?; + let b_hoisted = run_probed(&build, col_a.clone(), const_b.clone(), 0b10, &mut ctx)?; + assert_arrays_eq!(a_hoisted, baseline, &mut ctx); + assert_arrays_eq!(b_hoisted, baseline, &mut ctx); + + let validity = Validity::from_iter((0..col_a.len()).map(|row| row != 1)); + let masked_const_a = MaskedArray::try_new(const_a, validity.clone())?.into_array(); + let masked_col_a = MaskedArray::try_new(col_a, validity)?.into_array(); + let both_hoisted = run_probed(&build, masked_const_a, const_b, 0b11, &mut ctx)?; + let masked_baseline = run_probed(&build, masked_col_a, col_b, 0b00, &mut ctx)?; + assert_arrays_eq!(both_hoisted, masked_baseline, &mut ctx); + + Ok(()) + } +} From 97b5a990abb4c9828dbc711c23a0a0f065140d78 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 13:00:49 -0400 Subject: [PATCH 159/160] Automate RowFn benchmark comparisons Signed-off-by: Connor Tsui --- scripts/benchmark-rowfn.sh | 488 ++++++++++++++++++++++++ scripts/rowfn_benchmark.py | 482 +++++++++++++++++++++++ scripts/tests/test_rowfn_benchmark.py | 197 ++++++++++ vortex-array/Cargo.toml | 8 + vortex-array/benches/row_fn_executor.rs | 332 ++++++++++++++++ vortex-array/benches/strict_validity.rs | 216 +++++++++++ 6 files changed, 1723 insertions(+) create mode 100755 scripts/benchmark-rowfn.sh create mode 100755 scripts/rowfn_benchmark.py create mode 100644 scripts/tests/test_rowfn_benchmark.py create mode 100644 vortex-array/benches/row_fn_executor.rs create mode 100644 vortex-array/benches/strict_validity.rs diff --git a/scripts/benchmark-rowfn.sh b/scripts/benchmark-rowfn.sh new file mode 100755 index 00000000000..7536e37ec87 --- /dev/null +++ b/scripts/benchmark-rowfn.sh @@ -0,0 +1,488 @@ +#!/usr/bin/env bash + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +set -Eeu -o pipefail + +script_directory=$(dirname "$(realpath "${BASH_SOURCE[0]}")") + +usage() { + cat >&2 <<'EOF' +Usage: benchmark-rowfn.sh [OPTIONS] + +Options: + --suite NAME Select a preset or benchmark label. Repeatable; defaults to full. + --filter PATTERN Pass a Divan benchmark filter. Repeatable. + --build-only Build and record benchmark executables without measuring. + --measure-only Measure previously recorded benchmark executables without building. + --config NAME repository (16 CGUs/no LTO, default) or primary (1 CGU/fat LTO). + --target-root PATH Parent for reusable baseline and candidate Cargo targets. + --baseline-target PATH Reusable Cargo target for the baseline revision. + --candidate-target PATH + Reusable Cargo target for the candidate revision. + --codegen-units N Override the selected configuration. + --lto VALUE Override LTO with false, thin, or fat. + --rustflags FLAGS Override RUSTFLAGS; defaults to -C target-cpu=native. + --build-jobs N Jobs per concurrent revision build; defaults to 8 and cannot exceed 8. + --bench-cpu N Logical CPU used for every timed process; defaults to 4. + --warm-runs N Warm runs per revision; defaults to 2. + --measured-pairs N Alternating measured pairs; defaults to 7. + --sample-count N Divan sample count; defaults to 100. + --min-time SECONDS Divan minimum time; defaults to 0.25. + --max-time SECONDS Divan maximum time; defaults to 0.5. + --lock-file PATH Global timed-run lock; defaults to /tmp/vortex-rowfn-benchmark.lock. + --list-suites Print presets and benchmark labels, then exit. +EOF +} + +suite_catalog=( + "array-binary_ops|vortex-array|binary_ops|array,numeric,design-a-matrix,full" + "array-compare|vortex-array|compare|array,compare,full" + "array-row_fn_executor|vortex-array|row_fn_executor|array,framework,full" + "array-strict_validity|vortex-array|strict_validity|array,framework,full" + "array-like|vortex-array|like|array,full" + "array-take_filter|vortex-array|take_filter|array,full" + "array-varbinview_compact|vortex-array|varbinview_compact|array,full" + "tensor-l2_norm|vortex-tensor|l2_norm|tensor,full" + "tensor-inner_product|vortex-tensor|inner_product|tensor,full" + "tensor-cosine_similarity|vortex-tensor|cosine_similarity|tensor,full" + "tensor-normalized|vortex-tensor|normalized|tensor,full" + "spatial-binary_predicates|vortex-spatial|binary_predicates|spatial,full" + "spatial-distance|vortex-spatial|distance|spatial,full" + "spatial-envelope|vortex-spatial|envelope|spatial,full" + "spatial-predicate_bbox|vortex-spatial|predicate_bbox|spatial,full" +) + +requested_suites=() +filters=() +run_build=true +run_measure=true +configuration=repository +target_root= +baseline_target_override= +candidate_target_override= +codegen_units_override= +lto_override= +rustflags_override= +build_jobs=8 +bench_cpu=4 +warm_runs=2 +measured_pairs=7 +sample_count=100 +min_time=0.25 +max_time=0.5 +lock_file=/tmp/vortex-rowfn-benchmark.lock + +while [[ $# -gt 0 ]]; do + case $1 in + --suite) requested_suites+=("$2"); shift 2 ;; + --filter) filters+=("$2"); shift 2 ;; + --build-only) + if [[ $run_build == false ]]; then + echo "--build-only and --measure-only are mutually exclusive." >&2 + exit 1 + fi + run_measure=false + shift + ;; + --measure-only) + if [[ $run_measure == false ]]; then + echo "--build-only and --measure-only are mutually exclusive." >&2 + exit 1 + fi + run_build=false + shift + ;; + --config) configuration=$2; shift 2 ;; + --target-root) target_root=$2; shift 2 ;; + --baseline-target) baseline_target_override=$2; shift 2 ;; + --candidate-target) candidate_target_override=$2; shift 2 ;; + --codegen-units) codegen_units_override=$2; shift 2 ;; + --lto) lto_override=$2; shift 2 ;; + --rustflags) rustflags_override=$2; shift 2 ;; + --build-jobs) build_jobs=$2; shift 2 ;; + --bench-cpu) bench_cpu=$2; shift 2 ;; + --warm-runs) warm_runs=$2; shift 2 ;; + --measured-pairs) measured_pairs=$2; shift 2 ;; + --sample-count) sample_count=$2; shift 2 ;; + --min-time) min_time=$2; shift 2 ;; + --max-time) max_time=$2; shift 2 ;; + --lock-file) lock_file=$2; shift 2 ;; + --list-suites) + echo "Presets: full array framework numeric design-a-matrix compare tensor spatial" + printf '%s\n' "${suite_catalog[@]}" | cut -d '|' -f 1 + exit 0 + ;; + -h|--help) usage; exit 0 ;; + --*) echo "Unknown option: $1" >&2; usage; exit 1 ;; + *) break ;; + esac +done + +if [[ $# -ne 3 ]]; then + usage + exit 1 +fi +if [[ $(uname -m) != x86_64 ]]; then + echo "RowFn native performance decisions require an x86_64 host." >&2 + exit 1 +fi +if ((build_jobs < 1 || build_jobs > 8)); then + echo "--build-jobs must be between 1 and 8 so two builds cannot exceed 16 jobs." >&2 + exit 1 +fi +if [[ $run_measure == true ]]; then + command -v flock >/dev/null || { echo "benchmark-rowfn.sh requires flock." >&2; exit 1; } +fi + +baseline=$(realpath "$1") +candidate=$(realpath "$2") +output=$(realpath -m "$3") +if [[ -e $output ]]; then + echo "Output path already exists: $output" >&2 + exit 1 +fi + +case $configuration in + primary) codegen_units=1; lto=fat ;; + repository) codegen_units=16; lto=false ;; + *) echo "Unknown configuration: $configuration" >&2; exit 1 ;; +esac +codegen_units=${codegen_units_override:-$codegen_units} +lto=${lto_override:-$lto} +rustflags=${rustflags_override:--C target-cpu=native} + +if ((${#requested_suites[@]} == 0)); then + requested_suites=(full) +fi +selected_suites=() +declare -A selected_labels=() +for request in "${requested_suites[@]}"; do + matched=false + for entry in "${suite_catalog[@]}"; do + IFS='|' read -r label _ _ groups <<<"$entry" + if [[ $request == "$label" || ,$groups, == *,$request,* ]]; then + matched=true + if [[ -z ${selected_labels[$label]:-} ]]; then + selected_suites+=("$entry") + selected_labels[$label]=1 + fi + fi + done + if [[ $matched == false ]]; then + echo "Unknown suite or benchmark label: $request" >&2 + exit 1 + fi +done + +common_suites=() +skipped_suites=() +for entry in "${selected_suites[@]}"; do + IFS='|' read -r label package bench _ <<<"$entry" + baseline_source="$baseline/$package/benches/$bench.rs" + candidate_source="$candidate/$package/benches/$bench.rs" + + if [[ -f $baseline_source && -f $candidate_source ]]; then + common_suites+=("$entry") + elif [[ -f $baseline_source ]]; then + skipped_suites+=("$label (baseline only)") + elif [[ -f $candidate_source ]]; then + skipped_suites+=("$label (candidate only)") + else + skipped_suites+=("$label (missing from both revisions)") + fi +done +if ((${#common_suites[@]} == 0)); then + echo "No requested benchmark targets exist in both revisions; no comparison is possible." >&2 + printf 'Skipped: %s\n' "${skipped_suites[@]}" >&2 + exit 1 +fi +selected_suites=("${common_suites[@]}") +if ((${#skipped_suites[@]} != 0)); then + printf 'Skipping one-sided benchmark target: %s\n' "${skipped_suites[@]}" >&2 +fi + +common_git_dir=$(git -C "$candidate" rev-parse --path-format=absolute --git-common-dir) +repository_root=$(dirname "$common_git_dir") +if [[ -n $target_root && (-n $baseline_target_override || -n $candidate_target_override) ]]; then + echo "--target-root cannot be combined with revision-specific target paths." >&2 + exit 1 +fi +if [[ -z $target_root ]]; then + target_root="$repository_root/target/rowfn-benchmark/$(basename "$output")" +fi +target_root=$(realpath -m "$target_root") +baseline_target=$(realpath -m "${baseline_target_override:-$target_root/baseline}") +candidate_target=$(realpath -m "${candidate_target_override:-$target_root/candidate}") +if [[ $baseline_target == "$candidate_target" ]]; then + echo "Baseline and candidate must use different Cargo target directories." >&2 + exit 1 +fi + +mkdir -p "$output" +if [[ $run_build == true ]]; then + mkdir -p "$output/build" "$baseline_target" "$candidate_target" +fi +if [[ $run_measure == true ]]; then + mkdir -p "$output/warm" "$output/measured" +fi +parser="$script_directory/rowfn_benchmark.py" + +{ + echo "RowFn benchmark machine record" + echo "Date: $(date --iso-8601=seconds)" + echo "Host: $(hostname)" + echo "Kernel: $(uname -srvmo)" + echo "Benchmark CPU: $bench_cpu" + echo "Configuration: $configuration" + echo "Cargo profile: bench, $codegen_units codegen units, LTO $lto" + echo "RUSTFLAGS: $rustflags" + echo "Warm runs: $warm_runs" + echo "Measured pairs: $measured_pairs" + echo "Divan: TSC timer, $sample_count samples, min $min_time s, max $max_time s" + if ((${#skipped_suites[@]} == 0)); then + echo "Skipped one-sided benchmark targets: none" + else + printf 'Skipped one-sided benchmark target: %s\n' "${skipped_suites[@]}" + fi + echo + echo "Baseline toolchain:" + (cd "$baseline" && rustc -vV && cargo -V) + echo + echo "Candidate toolchain:" + (cd "$candidate" && rustc -vV && cargo -V) + echo + lscpu + echo + rg -m1 '^microcode' /proc/cpuinfo || true + for path in \ + /sys/devices/system/cpu/cpu"$bench_cpu"/cpufreq/scaling_governor \ + /sys/devices/system/cpu/cpu"$bench_cpu"/cpufreq/energy_performance_preference \ + /sys/devices/system/cpu/cpufreq/boost; do + [[ -r $path ]] && echo "$path: $(<"$path")" + done +} >"$output/machine.txt" + +build_revision() { + local worktree=$1 + local target=$2 + local log=$3 + + ( + cd "$worktree" + export CARGO_TARGET_DIR=$target + export CARGO_PROFILE_BENCH_CODEGEN_UNITS=$codegen_units + export CARGO_PROFILE_BENCH_LTO=$lto + export RUSTFLAGS=$rustflags + for package in vortex-array vortex-tensor vortex-spatial; do + local command=(cargo bench --no-run -j "$build_jobs" -p "$package") + local has_bench=false + for entry in "${selected_suites[@]}"; do + IFS='|' read -r _ suite_package bench _ <<<"$entry" + if [[ $suite_package == "$package" ]]; then + command+=(--bench "$bench") + has_bench=true + fi + done + if [[ $has_bench == true ]]; then + "${command[@]}" + fi + done + ) >"$log" 2>&1 +} + +if [[ $run_build == true ]]; then + echo "Building baseline and candidate with $build_jobs jobs each." + build_revision "$baseline" "$baseline_target" "$output/build/baseline.txt" & + baseline_pid=$! + build_revision "$candidate" "$candidate_target" "$output/build/candidate.txt" & + candidate_pid=$! + baseline_status=0 + candidate_status=0 + wait "$baseline_pid" || baseline_status=$? + wait "$candidate_pid" || candidate_status=$? + if [[ $baseline_status -ne 0 || $candidate_status -ne 0 ]]; then + echo "Benchmark build failed; see $output/build/." >&2 + exit 1 + fi +fi + +find_benchmark() { + local target=$1 + local name=$2 + local binary + + binary=$(find "$target/release/deps" -maxdepth 1 -type f -executable -name "$name-*" \ + -printf '%T@ %p\n' | sort -nr | head -n 1 | cut -d ' ' -f 2-) + [[ -n $binary ]] || { echo "Cannot find benchmark $name under $target." >&2; exit 1; } + echo "$binary" +} + +declare -A baseline_binaries=() +declare -A candidate_binaries=() +build_settings=( + --setting "configuration=$configuration" + --setting "codegen_units=$codegen_units" + --setting "lto=$lto" + --setting "rustflags=$rustflags" +) + +record_build() { + local revision=$1 + local worktree=$2 + local target=$3 + local metadata="$target/rowfn-benchmark-build.json" + local arguments=( + record-build + --output "$metadata" + --worktree "$worktree" + --target "$target" + "${build_settings[@]}" + ) + + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ bench _ <<<"$entry" + local binary + binary=$(find_benchmark "$target" "$bench") + arguments+=(--binary "$label=$binary") + done + python3 "$parser" "${arguments[@]}" + echo "Recorded $revision build metadata: $metadata" +} + +if [[ $run_build == true ]]; then + record_build baseline "$baseline" "$baseline_target" + record_build candidate "$candidate" "$candidate_target" +fi + +load_binaries() { + local worktree=$1 + local target=$2 + local metadata="$target/rowfn-benchmark-build.json" + local arguments=( + validate-build + --metadata "$metadata" + --worktree "$worktree" + --target "$target" + "${build_settings[@]}" + ) + + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + arguments+=(--suite "$label") + done + python3 "$parser" "${arguments[@]}" +} + +baseline_binary_output=$(load_binaries "$baseline" "$baseline_target") +candidate_binary_output=$(load_binaries "$candidate" "$candidate_target") +mapfile -t baseline_binary_records <<<"$baseline_binary_output" +mapfile -t candidate_binary_records <<<"$candidate_binary_output" +for record in "${baseline_binary_records[@]}"; do + label=${record%%=*} + baseline_binaries[$label]=${record#*=} +done +for record in "${candidate_binary_records[@]}"; do + label=${record%%=*} + candidate_binaries[$label]=${record#*=} +done + +manifest_args=( + manifest + --output "$output/manifest.json" + --machine-record "$output/machine.txt" + --baseline-worktree "$baseline" + --candidate-worktree "$candidate" + --baseline-target "$baseline_target" + --candidate-target "$candidate_target" + --setting "configuration=$configuration" + --setting "codegen_units=$codegen_units" + --setting "lto=$lto" + --setting "rustflags=$rustflags" + --setting "bench_cpu=$bench_cpu" + --setting "warm_runs=$warm_runs" + --setting "measured_pairs=$measured_pairs" + --setting "sample_count=$sample_count" + --setting "min_time=$min_time" + --setting "max_time=$max_time" +) +for filter in "${filters[@]}"; do + manifest_args+=(--filter "$filter") +done +for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + manifest_args+=( + --suite "$label" + --baseline-binary "$label=${baseline_binaries[$label]}" + --candidate-binary "$label=${candidate_binaries[$label]}" + ) +done +python3 "$parser" "${manifest_args[@]}" + +if [[ $run_measure == false ]]; then + echo "Build evidence: $output" + echo "Baseline target: $baseline_target" + echo "Candidate target: $candidate_target" + exit 0 +fi + +run_suite() { + local revision=$1 + local label=$2 + local destination=$3 + local binary + local command + + if [[ $revision == baseline ]]; then + binary=${baseline_binaries[$label]} + else + binary=${candidate_binaries[$label]} + fi + command=( + taskset -c "$bench_cpu" "$binary" + --bench --timer tsc --sample-count "$sample_count" + --min-time "$min_time" --max-time "$max_time" --color never + "${filters[@]}" + ) + echo "Running $label ($revision) -> $destination" + "${command[@]}" >"$destination" 2>&1 +} + +echo "Waiting for the global timed benchmark lock: $lock_file" +exec {benchmark_lock}>"$lock_file" +flock "$benchmark_lock" +if pgrep -x cargo >/dev/null || pgrep -x rustc >/dev/null; then + echo "Cargo or rustc is active after acquiring the benchmark lock; refusing to measure." >&2 + exit 1 +fi + +for ((round = 1; round <= warm_runs; round++)); do + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + if ((round % 2 == 1)); then + run_suite baseline "$label" "$output/warm/$label-baseline-$round.txt" + run_suite candidate "$label" "$output/warm/$label-candidate-$round.txt" + else + run_suite candidate "$label" "$output/warm/$label-candidate-$round.txt" + run_suite baseline "$label" "$output/warm/$label-baseline-$round.txt" + fi + done +done + +for ((pair = 1; pair <= measured_pairs; pair++)); do + for entry in "${selected_suites[@]}"; do + IFS='|' read -r label _ _ _ <<<"$entry" + if ((pair % 2 == 1)); then + run_suite baseline "$label" "$output/measured/$label-baseline-$pair.txt" + run_suite candidate "$label" "$output/measured/$label-candidate-$pair.txt" + else + run_suite candidate "$label" "$output/measured/$label-candidate-$pair.txt" + run_suite baseline "$label" "$output/measured/$label-baseline-$pair.txt" + fi + done +done + +python3 "$parser" summarize "$output" +echo "Raw results: $output" +echo "Summary: $output/summary.md" diff --git a/scripts/rowfn_benchmark.py b/scripts/rowfn_benchmark.py new file mode 100755 index 00000000000..35c08918d75 --- /dev/null +++ b/scripts/rowfn_benchmark.py @@ -0,0 +1,482 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Capture and summarize evidence from ``benchmark-rowfn.sh`` runs.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import re +import statistics +import subprocess +from collections.abc import Iterable +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from pathlib import Path + +RESULT_FILE = re.compile(r"^(?P.+)-(?Pbaseline|candidate)-(?P\d+)\.txt$") +TREE_ROW = re.compile(r"^(?P(?:│ | )*)(?:├─ |╰─ )(?P.*)$") +TIMING = re.compile(r"(?P\d+(?:\.\d+)?)\s*(?Pps|ns|µs|us|ms|s)\s*$") +UNIT_TO_NS = { + "ps": 0.001, + "ns": 1.0, + "µs": 1_000.0, + "us": 1_000.0, + "ms": 1_000_000.0, + "s": 1_000_000_000.0, +} + + +@dataclass(frozen=True) +class BenchmarkSummary: + suite: str + benchmark: str + pairs: int + baseline_median_ns: float + candidate_median_ns: float + median_ratio: float + minimum_ratio: float + maximum_ratio: float + ratio_mad: float + + +def run_git(worktree: Path, *args: str, binary: bool = False) -> str | bytes: + """Run one read-only Git command in ``worktree``.""" + + result = subprocess.run( + ["git", "-C", str(worktree), *args], + check=True, + capture_output=True, + text=not binary, + ) + return result.stdout if binary else result.stdout.strip() + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + + return digest.hexdigest() + + +def toolchain_record(worktree: Path) -> dict[str, str]: + """Capture the tools selected from a revision's working directory.""" + + def version(*command: str) -> str: + result = subprocess.run( + command, + cwd=worktree, + check=True, + capture_output=True, + text=True, + ) + return result.stdout.strip() + + return {"rustc": version("rustc", "-vV"), "cargo": version("cargo", "-V")} + + +def revision_record(worktree: Path, target: Path, binaries: Iterable[str]) -> dict[str, object]: + """Describe the exact revision, dirty patch, targets, and benchmark executables.""" + + status = str(run_git(worktree, "status", "--short")).splitlines() + diff = run_git(worktree, "diff", "--binary", "HEAD", binary=True) + assert isinstance(diff, bytes) + + untracked = run_git(worktree, "ls-files", "--others", "--exclude-standard", "-z", binary=True) + assert isinstance(untracked, bytes) + dirty_digest = hashlib.sha256(diff) + dirty_digest.update(untracked) + for relative_path in filter(None, untracked.decode().split("\0")): + path = worktree / relative_path + if path.is_file(): + dirty_digest.update(relative_path.encode()) + dirty_digest.update(bytes.fromhex(sha256_file(path))) + + executable_records: dict[str, object] = {} + for entry in binaries: + label, separator, raw_path = entry.partition("=") + if not separator: + raise ValueError(f"expected LABEL=PATH for benchmark binary, got {entry!r}") + path = Path(raw_path).resolve() + executable_records[label] = { + "path": str(path), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + + return { + "worktree": str(worktree.resolve()), + "head": run_git(worktree, "rev-parse", "HEAD"), + "changed_paths": status, + "tracked_diff_sha256": hashlib.sha256(diff).hexdigest(), + "dirty_state_sha256": dirty_digest.hexdigest(), + "target": str(target.resolve()), + "binaries": executable_records, + } + + +def build_identity(worktree: Path, target: Path, settings: dict[str, str]) -> dict[str, object]: + revision = revision_record(worktree, target, []) + revision.pop("binaries") + return { + "settings": settings, + "toolchain": toolchain_record(worktree), + "revision": revision, + } + + +def binary_records(entries: Iterable[str]) -> dict[str, object]: + records: dict[str, object] = {} + for entry in entries: + label, separator, raw_path = entry.partition("=") + if not separator: + raise ValueError(f"expected LABEL=PATH for benchmark binary, got {entry!r}") + path = Path(raw_path).resolve() + records[label] = { + "path": str(path), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + return records + + +def write_build_record(args: argparse.Namespace) -> None: + output = Path(args.output) + worktree = Path(args.worktree) + target = Path(args.target) + settings = dict(setting.split("=", 1) for setting in args.setting) + identity = build_identity(worktree, target, settings) + binaries = binary_records(args.binary) + + if output.exists(): + previous = json.loads(output.read_text(encoding="utf-8")) + previous_identity = {key: previous.get(key) for key in identity} + if previous_identity == identity: + binaries = {**previous.get("binaries", {}), **binaries} + + record = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + **identity, + "binaries": binaries, + } + output.write_text(f"{json.dumps(record, indent=2, sort_keys=True)}\n", encoding="utf-8") + + +def validated_build_binaries(args: argparse.Namespace) -> dict[str, str]: + metadata = Path(args.metadata) + if not metadata.is_file(): + raise ValueError(f"build metadata does not exist: {metadata}") + + record = json.loads(metadata.read_text(encoding="utf-8")) + if record.get("schema_version") != 1: + raise ValueError(f"unsupported build metadata schema in {metadata}") + + settings = dict(setting.split("=", 1) for setting in args.setting) + identity = build_identity(Path(args.worktree), Path(args.target), settings) + mismatches = [key for key in identity if record.get(key) != identity[key]] + if mismatches: + fields = ", ".join(mismatches) + raise ValueError(f"stale benchmark build metadata ({fields} changed): {metadata}") + + binaries = record.get("binaries", {}) + resolved: dict[str, str] = {} + for suite in args.suite: + stored = binaries.get(suite) + if stored is None: + raise ValueError(f"benchmark suite {suite!r} was not recorded in {metadata}") + path = Path(stored["path"]) + if not path.is_file(): + raise ValueError(f"recorded benchmark binary does not exist: {path}") + current = { + "path": str(path.resolve()), + "sha256": sha256_file(path), + "size": path.stat().st_size, + } + if current != stored: + raise ValueError(f"recorded benchmark binary changed: {path}") + resolved[suite] = str(path.resolve()) + + return resolved + + +def validate_build_record(args: argparse.Namespace) -> None: + for suite, path in validated_build_binaries(args).items(): + print(f"{suite}={path}") + + +def write_manifest(args: argparse.Namespace) -> None: + settings = dict(setting.split("=", 1) for setting in args.setting) + manifest = { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + "settings": settings, + "suites": args.suite, + "filters": args.filter, + "machine_record": str(Path(args.machine_record).resolve()), + "baseline": revision_record( + Path(args.baseline_worktree), + Path(args.baseline_target), + args.baseline_binary, + ), + "candidate": revision_record( + Path(args.candidate_worktree), + Path(args.candidate_target), + args.candidate_binary, + ), + } + output = Path(args.output) + output.write_text(f"{json.dumps(manifest, indent=2, sort_keys=True)}\n", encoding="utf-8") + + +def timing_ns(field: str) -> float: + match = TIMING.search(field.strip()) + if match is None: + raise ValueError(f"cannot parse Divan timing from {field!r}") + + return float(match.group("value")) * UNIT_TO_NS[match.group("unit")] + + +def parse_divan(path: Path) -> dict[str, float]: + """Return benchmark paths and median nanoseconds from one Divan table.""" + + parents: dict[int, str] = {} + timings: dict[str, float] = {} + for line in path.read_text(encoding="utf-8").splitlines(): + fields = re.split(r"\s+│\s+", line) + tree_match = TREE_ROW.match(fields[0]) + if tree_match is None: + continue + + depth = len(tree_match.group("prefix")) // 3 + body = tree_match.group("body").rstrip() + timing_match = TIMING.search(body) + name = body[: timing_match.start()].rstrip() if timing_match else body.strip() + parents = {level: parent for level, parent in parents.items() if level < depth} + + if timing_match is None: + parents[depth] = name + continue + if len(fields) < 3: + raise ValueError(f"timed Divan row has no median column in {path}: {line}") + + components = [parents[level] for level in sorted(parents) if level < depth] + benchmark = "/".join([*components, name]) + if benchmark in timings: + raise ValueError(f"duplicate benchmark {benchmark!r} in {path}") + timings[benchmark] = timing_ns(fields[2]) + + if not timings: + raise ValueError(f"no Divan benchmark timings found in {path}") + + return timings + + +def read_measurements(directory: Path) -> dict[tuple[str, str, int, str], float]: + measurements: dict[tuple[str, str, int, str], float] = {} + for path in sorted(directory.glob("*.txt")): + match = RESULT_FILE.match(path.name) + if match is None: + continue + suite = match.group("suite") + revision = match.group("revision") + pair = int(match.group("pair")) + for benchmark, median_ns in parse_divan(path).items(): + measurements[suite, revision, pair, benchmark] = median_ns + + if not measurements: + raise ValueError(f"no measured result files found in {directory}") + + return measurements + + +def summarize(measurements: dict[tuple[str, str, int, str], float]) -> list[BenchmarkSummary]: + inventories: dict[tuple[str, str], set[str]] = {} + for suite, revision, _, benchmark in measurements: + inventories.setdefault((suite, revision), set()).add(benchmark) + + suites = {suite for suite, _ in inventories} + comparable = { + (suite, benchmark) + for suite in suites + for benchmark in inventories.get((suite, "baseline"), set()) & inventories.get((suite, "candidate"), set()) + } + groups = { + (suite, pair, benchmark) for suite, _, pair, benchmark in measurements if (suite, benchmark) in comparable + } + incomplete = [ + group + for group in groups + if (group[0], "baseline", group[1], group[2]) not in measurements + or (group[0], "candidate", group[1], group[2]) not in measurements + ] + if incomplete: + raise ValueError(f"unpaired benchmark measurements: {sorted(incomplete)!r}") + if not groups: + raise ValueError("unpaired benchmark measurements: no comparable benchmarks") + + by_benchmark: dict[tuple[str, str], list[tuple[float, float]]] = {} + for suite, pair, benchmark in sorted(groups): + baseline = measurements[suite, "baseline", pair, benchmark] + candidate = measurements[suite, "candidate", pair, benchmark] + by_benchmark.setdefault((suite, benchmark), []).append((baseline, candidate)) + + summaries = [] + for (suite, benchmark), pairs in sorted(by_benchmark.items()): + baseline_values = [baseline for baseline, _ in pairs] + candidate_values = [candidate for _, candidate in pairs] + ratios = [candidate / baseline for baseline, candidate in pairs] + median_ratio = statistics.median(ratios) + summaries.append( + BenchmarkSummary( + suite=suite, + benchmark=benchmark, + pairs=len(pairs), + baseline_median_ns=statistics.median(baseline_values), + candidate_median_ns=statistics.median(candidate_values), + median_ratio=median_ratio, + minimum_ratio=min(ratios), + maximum_ratio=max(ratios), + ratio_mad=statistics.median(abs(ratio - median_ratio) for ratio in ratios), + ) + ) + + return summaries + + +def inventory_differences( + measurements: dict[tuple[str, str, int, str], float], +) -> list[tuple[str, str, str]]: + """Return benchmarks that exist in only one revision.""" + + inventories: dict[tuple[str, str], set[str]] = {} + for suite, revision, _, benchmark in measurements: + inventories.setdefault((suite, revision), set()).add(benchmark) + + differences = [] + for suite in sorted({suite for suite, _ in inventories}): + baseline = inventories.get((suite, "baseline"), set()) + candidate = inventories.get((suite, "candidate"), set()) + differences.extend((suite, "baseline only", benchmark) for benchmark in baseline - candidate) + differences.extend((suite, "candidate only", benchmark) for benchmark in candidate - baseline) + + return sorted(differences) + + +def format_ns(value: float) -> str: + for divisor, unit in ((1_000_000_000, "s"), (1_000_000, "ms"), (1_000, "µs")): + if value >= divisor: + return f"{value / divisor:.3f} {unit}" + + return f"{value:.3f} ns" + + +def write_summary( + output_directory: Path, + summaries: list[BenchmarkSummary], + differences: Iterable[tuple[str, str, str]] = (), +) -> None: + csv_path = output_directory / "ratios.csv" + with csv_path.open("w", encoding="utf-8", newline="") as file: + writer = csv.DictWriter(file, fieldnames=list(asdict(summaries[0]))) + writer.writeheader() + writer.writerows(asdict(summary) for summary in summaries) + + markdown = [ + "# RowFn benchmark comparison", + "", + "Ratios are paired candidate/baseline medians. Lower is faster.", + "", + "| Suite | Benchmark | Pairs | Baseline | Candidate | Ratio | Change | MAD |", + "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for summary in sorted(summaries, key=lambda result: result.median_ratio, reverse=True): + change = (summary.median_ratio - 1.0) * 100.0 + markdown.append( + f"| {summary.suite} | `{summary.benchmark}` | {summary.pairs} " + f"| {format_ns(summary.baseline_median_ns)} " + f"| {format_ns(summary.candidate_median_ns)} " + f"| {summary.median_ratio:.6f} | {change:+.2f}% | {summary.ratio_mad:.6f} |" + ) + differences = list(differences) + if differences: + markdown.extend( + [ + "", + "## Unpaired benchmark inventory", + "", + "These benchmarks were recorded for only one revision and are excluded from ratios.", + "", + ] + ) + markdown.extend(f"- `{suite}/{benchmark}`: {revision}." for suite, revision, benchmark in differences) + markdown.append("") + (output_directory / "summary.md").write_text("\n".join(markdown), encoding="utf-8") + + +def summarize_directory(args: argparse.Namespace) -> None: + output_directory = Path(args.output_directory) + measurements = read_measurements(output_directory / "measured") + summaries = summarize(measurements) + write_summary(output_directory, summaries, inventory_differences(measurements)) + + +def argument_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(required=True) + + manifest = subparsers.add_parser("manifest", help="capture revisions and executable hashes") + manifest.add_argument("--output", required=True) + manifest.add_argument("--machine-record", required=True) + manifest.add_argument("--baseline-worktree", required=True) + manifest.add_argument("--candidate-worktree", required=True) + manifest.add_argument("--baseline-target", required=True) + manifest.add_argument("--candidate-target", required=True) + manifest.add_argument("--setting", action="append", default=[]) + manifest.add_argument("--suite", action="append", default=[]) + manifest.add_argument("--filter", action="append", default=[]) + manifest.add_argument("--baseline-binary", action="append", default=[]) + manifest.add_argument("--candidate-binary", action="append", default=[]) + manifest.set_defaults(function=write_manifest) + + record_build = subparsers.add_parser("record-build", help="record reusable benchmark binaries") + record_build.add_argument("--output", required=True) + record_build.add_argument("--worktree", required=True) + record_build.add_argument("--target", required=True) + record_build.add_argument("--setting", action="append", default=[]) + record_build.add_argument("--binary", action="append", default=[]) + record_build.set_defaults(function=write_build_record) + + validate_build = subparsers.add_parser("validate-build", help="validate a reusable build") + validate_build.add_argument("--metadata", required=True) + validate_build.add_argument("--worktree", required=True) + validate_build.add_argument("--target", required=True) + validate_build.add_argument("--setting", action="append", default=[]) + validate_build.add_argument("--suite", action="append", default=[]) + validate_build.set_defaults(function=validate_build_record) + + summary = subparsers.add_parser("summarize", help="write ratios.csv and summary.md") + summary.add_argument("output_directory") + summary.set_defaults(function=summarize_directory) + + return parser + + +def main() -> None: + parser = argument_parser() + args = parser.parse_args() + try: + args.function(args) + except (OSError, subprocess.CalledProcessError, ValueError) as error: + parser.error(str(error)) + + +if __name__ == "__main__": + main() diff --git a/scripts/tests/test_rowfn_benchmark.py b/scripts/tests/test_rowfn_benchmark.py new file mode 100644 index 00000000000..ace0797c6a3 --- /dev/null +++ b/scripts/tests/test_rowfn_benchmark.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +import importlib.util +import sys +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +REPO_ROOT = Path(__file__).resolve().parents[2] +SCRIPT = REPO_ROOT / "scripts" / "rowfn_benchmark.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("rowfn_benchmark", SCRIPT) + assert spec is not None + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def write_divan(path: Path, rows: list[str]) -> None: + path.write_text( + "\n".join( + [ + "Timer precision: 20 ns", + "bench fastest │ slowest │ median │ mean │ samples │ iters", + *rows, + "", + ] + ), + encoding="utf-8", + ) + + +class RowFnBenchmarkTest(unittest.TestCase): + def setUp(self) -> None: + self.module = load_module() + self.temporary_directory = tempfile.TemporaryDirectory() + self.directory = Path(self.temporary_directory.name) + + def tearDown(self) -> None: + self.temporary_directory.cleanup() + + def test_parse_divan_preserves_nested_benchmark_names_and_converts_units(self) -> None: + output = self.directory / "result.txt" + write_divan( + output, + [ + "├─ non_nullable │ │ │ │ │", + "│ ├─ 2 17.18 µs │ 18 µs │ 17.33 µs │ 17.4 µs │ 100 │ 100", + "│ ╰─ 32 6.709 µs │ 8 µs │ 6.829 µs │ 7 µs │ 100 │ 100", + "╰─ nullable │ │ │ │ │", + " ╰─ 2 799.7 ns │ 1 µs │ 979.7 ns │ 986 ns │ 100 │ 100", + ], + ) + + self.assertEqual( + self.module.parse_divan(output), + { + "non_nullable/2": 17_330.0, + "non_nullable/32": 6_829.0, + "nullable/2": 979.7, + }, + ) + + def test_summarize_writes_paired_ratios_and_slowest_first_markdown(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + results = { + "numeric-baseline-1.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-1.txt": ["├─ add 12 ns │ 12 ns │ 12 ns │ 12 ns │ 100 │ 100"], + "numeric-baseline-2.txt": ["├─ add 20 ns │ 20 ns │ 20 ns │ 20 ns │ 100 │ 100"], + "numeric-candidate-2.txt": ["├─ add 18 ns │ 18 ns │ 18 ns │ 18 ns │ 100 │ 100"], + "numeric-baseline-3.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-3.txt": ["├─ add 11 ns │ 11 ns │ 11 ns │ 11 ns │ 100 │ 100"], + "numeric-baseline-4.txt": ["├─ mul 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-4.txt": ["├─ mul 9 ns │ 9 ns │ 9 ns │ 9 ns │ 100 │ 100"], + } + for name, rows in results.items(): + write_divan(measured / name, rows) + + summaries = self.module.summarize(self.module.read_measurements(measured)) + self.module.write_summary(self.directory, summaries) + + add = next(summary for summary in summaries if summary.benchmark == "add") + self.assertEqual(add.pairs, 3) + self.assertAlmostEqual(add.median_ratio, 1.1) + self.assertAlmostEqual(add.ratio_mad, 0.1) + + csv_output = (self.directory / "ratios.csv").read_text(encoding="utf-8") + markdown = (self.directory / "summary.md").read_text(encoding="utf-8") + self.assertIn("suite,benchmark,pairs", csv_output) + self.assertLess(markdown.index("`add`"), markdown.index("`mul`")) + + def test_summarize_rejects_unpaired_measurements(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + write_divan( + measured / "numeric-baseline-1.txt", + ["╰─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + ) + + with self.assertRaisesRegex(ValueError, "unpaired benchmark measurements"): + self.module.summarize(self.module.read_measurements(measured)) + + def test_summarize_excludes_and_reports_revision_only_benchmarks(self) -> None: + measured = self.directory / "measured" + measured.mkdir() + results = { + "numeric-baseline-1.txt": ["├─ add 10 ns │ 10 ns │ 10 ns │ 10 ns │ 100 │ 100"], + "numeric-candidate-1.txt": [ + "├─ add 11 ns │ 11 ns │ 11 ns │ 11 ns │ 100 │ 100", + "╰─ candidate 5 ns │ 5 ns │ 5 ns │ 5 ns │ 100 │ 100", + ], + } + for name, rows in results.items(): + write_divan(measured / name, rows) + + measurements = self.module.read_measurements(measured) + summaries = self.module.summarize(measurements) + differences = self.module.inventory_differences(measurements) + self.module.write_summary(self.directory, summaries, differences) + + self.assertEqual([summary.benchmark for summary in summaries], ["add"]) + self.assertEqual(differences, [("numeric", "candidate only", "candidate")]) + markdown = (self.directory / "summary.md").read_text(encoding="utf-8") + self.assertIn("`numeric/candidate`: candidate only.", markdown) + + def test_build_record_validates_identity_and_executable(self) -> None: + target = self.directory / "target" + target.mkdir() + binary = target / "binary_ops-123" + binary.write_bytes(b"first binary") + metadata = target / "rowfn-benchmark-build.json" + identity = { + "settings": {"codegen_units": "1", "lto": "fat"}, + "toolchain": {"rustc": "rustc 1.97.1", "cargo": "cargo 1.97.1"}, + "revision": {"head": "abc123", "dirty_state_sha256": "clean"}, + } + arguments = SimpleNamespace( + output=str(metadata), + worktree=str(self.directory), + target=str(target), + setting=["codegen_units=1", "lto=fat"], + binary=[f"numeric={binary}"], + ) + + with mock.patch.object(self.module, "build_identity", return_value=identity): + self.module.write_build_record(arguments) + + validation = SimpleNamespace( + metadata=str(metadata), + worktree=str(self.directory), + target=str(target), + setting=["codegen_units=1", "lto=fat"], + suite=["numeric"], + ) + with mock.patch.object(self.module, "build_identity", return_value=identity): + self.assertEqual( + self.module.validated_build_binaries(validation), + {"numeric": str(binary.resolve())}, + ) + + changed_identities = { + "settings": {**identity, "settings": {"codegen_units": "16", "lto": "false"}}, + "toolchain": { + **identity, + "toolchain": {"rustc": "rustc 1.98.0", "cargo": "cargo 1.98.0"}, + }, + "revision": { + **identity, + "revision": {"head": "def456", "dirty_state_sha256": "changed"}, + }, + } + for field, changed_identity in changed_identities.items(): + with ( + self.subTest(field=field), + mock.patch.object(self.module, "build_identity", return_value=changed_identity), + self.assertRaisesRegex(ValueError, f"{field} changed"), + ): + self.module.validated_build_binaries(validation) + + binary.write_bytes(b"second binary") + with ( + mock.patch.object(self.module, "build_identity", return_value=identity), + self.assertRaisesRegex(ValueError, "binary changed"), + ): + self.module.validated_build_binaries(validation) + + +if __name__ == "__main__": + unittest.main() diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 2af2eacf238..68bd189ef11 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -134,6 +134,10 @@ harness = false name = "binary_ops" harness = false +[[bench]] +name = "row_fn_executor" +harness = false + [[bench]] name = "kleene_bool" harness = false @@ -213,6 +217,10 @@ harness = false name = "validity_is_valid" harness = false +[[bench]] +name = "strict_validity" +harness = false + [[bench]] name = "dict_unreferenced_mask" harness = false diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs new file mode 100644 index 00000000000..84a2e029412 --- /dev/null +++ b/vortex-array/benches/row_fn_executor.rs @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares owned-output, sink-writing, and hand-written primitive row loops. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::OutputSink; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const ROWS: usize = 65_536; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +#[derive(Clone)] +struct RowWrappingAdd; + +impl RowFn for RowWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_wrapping_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64), i64>(|(lhs, rhs)| lhs.wrapping_add(rhs)) + } +} + +#[derive(Clone)] +struct RowCheckedAdd; + +impl RowFn for RowCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_checked_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(i64, i64), i64, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(checked_add_error()); + } + Ok(()) + }, + ) + } +} + +/// Keep error construction out of the benchmarked success path. +#[cold] +#[inline(never)] +fn checked_add_error() -> VortexError { + vortex_err!("integer overflow in row checked add") +} + +/// A benchmark sink that writes one `i64` per row. +struct I64Sink( + /// The output values written by the row loop. + BufferMut, +); + +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +#[derive(Clone)] +struct RowSinkWrappingAdd; + +impl RowFn for RowSinkWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_wrapping_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64, i64), I64Sink, _>(|(lhs, rhs), out| { + *out = lhs.wrapping_add(rhs); + }) + } +} + +#[derive(Clone)] +struct RowSinkCheckedAdd; + +impl RowFn for RowSinkCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_checked_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64, i64), UninitElementSink, _>( + |(lhs, rhs), output| -> VortexResult { + let value = lhs.checked_add(rhs).ok_or_else(checked_add_error)?; + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + Ok(unsafe { InitializedElement::write(output, value) }) + }, + ) + } +} + +fn inputs() -> (ArrayRef, ArrayRef) { + let lhs = (0..ROWS) + .map(|index| index as i64) + .collect::>() + .into_array(); + let rhs = (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>() + .into_array(); + (lhs, rhs) +} + +fn constant_inputs() -> (ArrayRef, ArrayRef) { + let (lhs, _) = inputs(); + let rhs = ConstantArray::new(Scalar::from(7i64), ROWS).into_array(); + (lhs, rhs) +} + +fn nullable_inputs() -> (ArrayRef, ArrayRef) { + let lhs = PrimitiveArray::new( + (0..ROWS).map(|index| index as i64).collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(5))), + ) + .into_array(); + let rhs = PrimitiveArray::new( + (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(7))), + ) + .into_array(); + (lhs, rhs) +} + +fn bench_row_fn(bencher: Bencher, function: F, make_inputs: fn() -> (ArrayRef, ArrayRef)) +where + F: RowFn, +{ + bencher + .with_inputs(make_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + function + .clone() + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, constant_inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, nullable_inputs); +} + +#[divan::bench] +fn handrolled_sink_wrapping_add(bencher: Bencher) { + bencher + .with_inputs(inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = lhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let rhs = rhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let mut output = BufferMut::zeroed(ROWS); + for ((out, lhs), rhs) in output + .as_mut_slice() + .iter_mut() + .zip(lhs.as_slice()) + .zip(rhs.as_slice()) + { + *out = lhs.wrapping_add(*rhs); + } + PrimitiveArray::new(output.freeze(), Validity::NonNullable).into_array() + }); +} + +#[divan::bench] +fn row_checked_add(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, inputs); +} + +#[divan::bench] +fn row_wrapping_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, nullable_inputs); +} + +#[divan::bench] +fn row_sink_checked_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowSinkCheckedAdd, nullable_inputs); +} + +#[divan::bench] +fn row_wrapping_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, nullable_inputs); +} diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs new file mode 100644 index 00000000000..9a713220fab --- /dev/null +++ b/vortex-array/benches/strict_validity.rs @@ -0,0 +1,216 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks how the row lifting applies validity to a dense kernel. +//! +//! Both arms run the same kernel over the same nullable input and differ only in how the output's +//! nulls are restored: +//! +//! - `lazy` is what the lifting behind [`RowFn`] does: conjoin the input validities (itself lazy) +//! and hand the resulting boolean array straight to `mask`. +//! - `eager` is the strategy it replaced: materialize the conjunction into a `Mask`, copy it into a +//! `BitBuffer`, wrap that in a `BoolArray`, and mask with it. +//! +//! `chain` composes three of the same function, which is where a per-call materialization +//! compounds. + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const SIZES: &[usize] = &[65_536, 1 << 20]; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +/// The shared kernel: double every lane, ignoring validity. +fn doubled(input: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let input = input.clone().execute::(ctx)?; + let values: Buffer = input + .as_slice::() + .iter() + .map(|value| value.wrapping_mul(2)) + .collect(); + Ok(PrimitiveArray::new(values, Validity::NonNullable).into_array()) +} + +/// Dense row function: the lifting decides how validity is applied. +/// +/// Its [`reduce_encoded`](RowFn::reduce_encoded) answers with [`doubled`] before the row loop is +/// reached, so this arm runs exactly the kernel [`EagerDouble`] runs and the two differ only in +/// validity. The row closure below is what makes it a [`RowFn`] at all, and never executes. +#[derive(Clone)] +struct LazyDouble; + +impl RowFn for LazyDouble { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.lazy_double"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i32,), i32>(|(value,)| value.wrapping_mul(2)) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + doubled(&args[0], ctx).map(|output| Some(RowExecution::Output(output))) + } +} + +/// The same function, applying validity the way the adapter used to: materialize a mask first. +#[derive(Clone)] +struct EagerDouble; + +impl ScalarFnVTable for EagerDouble { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.eager_double"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { + ChildName::from("input") + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + Ok(DType::Primitive(PType::I32, args[0].nullability())) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let valid = input.validity()?.execute_mask(args.row_count(), ctx)?; + let values = doubled(&input, ctx)?; + + if valid.all_true() { + return values.cast(DType::Primitive(PType::I32, input.dtype().nullability())); + } + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + values.mask(mask) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// A nullable i32 column with array-backed validity (~10% nulls). +fn nullable_input(len: usize) -> ArrayRef { + PrimitiveArray::new( + (0..len as i32).collect::>(), + Validity::from_iter((0..len).map(|index| !index.is_multiple_of(10))), + ) + .into_array() +} + +fn bench_depth(bencher: Bencher, function: F, len: usize, depth: usize) +where + F: ScalarFnVTable + Clone, +{ + bencher + .with_inputs(|| nullable_input(len)) + .bench_local_values(|input| { + let mut ctx = SESSION.create_execution_ctx(); + let mut array = input; + for _ in 0..depth { + array = function + .clone() + .try_new_array(len, EmptyOptions, [array]) + .unwrap() + .into_array(); + } + array.execute::(&mut ctx).unwrap() + }); +} + +#[divan::bench(args = SIZES)] +fn lazy(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn eager(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn lazy_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 3); +} + +#[divan::bench(args = SIZES)] +fn eager_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 3); +} From 5c302bce65ec54b1aff5cd9fbff57b406765681d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 14 Aug 2026 18:34:36 -0400 Subject: [PATCH 160/160] Record the final RowFn framework investigation Signed-off-by: Connor Tsui --- research/rowfn-review-followup/README.md | 12 ++ .../2026-08-14-framework-refinement.md | 193 ++++++++++++++++++ 2 files changed, 205 insertions(+) create mode 100644 research/rowfn-review-followup/README.md create mode 100644 research/rowfn-review-followup/codegen/2026-08-14-framework-refinement.md diff --git a/research/rowfn-review-followup/README.md b/research/rowfn-review-followup/README.md new file mode 100644 index 00000000000..66161b98aab --- /dev/null +++ b/research/rowfn-review-followup/README.md @@ -0,0 +1,12 @@ + + + +# RowFn review follow-up + +This directory preserves investigations that informed the focused RowFn pull-request stack. The +umbrella branch merges the focused stack with its tree unchanged, so these notes live on the +umbrella's first-parent history. + +- [`codegen/2026-08-14-framework-refinement.md`](codegen/2026-08-14-framework-refinement.md) + records the final framework review, compiler-output experiments, benchmark results, and stack + publication procedure. diff --git a/research/rowfn-review-followup/codegen/2026-08-14-framework-refinement.md b/research/rowfn-review-followup/codegen/2026-08-14-framework-refinement.md new file mode 100644 index 00000000000..cdceee0b93a --- /dev/null +++ b/research/rowfn-review-followup/codegen/2026-08-14-framework-refinement.md @@ -0,0 +1,193 @@ + + + +# RowFn framework refinement and stack gate + +Date: 2026-08-14. + +Platform: Apple M4 Max (`aarch64-apple-darwin`), rustc 1.97.1, LLVM 22.1.6. The compiler-output +gate used the bench profile, 16 codegen units, no LTO, `-C target-cpu=native`, loop-vectorizer +remarks, and line-table debug information. The timing evidence is local Apple Silicon evidence and +does not establish x86 performance. + +## Final design decisions + +- The framework is stricter than an ordinary strict scalar function. Strict input validity is + propagated by batch execution, while a row kernel must return a valid value for every valid input + row. A function that can create null from valid inputs cannot use this row-kernel contract. +- `skip-invalid` means that the executor initializes invalid output rows, invokes the row kernel only + for valid rows, then publishes the completed sink. `filter-and-scatter` means that batch execution + filters every input to the valid rows, runs the ordinary row kernel, then scatters those results + into their original positions with nulls elsewhere. +- Documentation uses _partially valid_ for a mask with both valid and invalid rows. It uses + _batch-constant_ for a decoded input with one addressable value and _non-constant_ for the + opposite. Vague uses of _mixed_ and _varying_ were removed. +- `Args::views_no_constants` names the fast path that exists only when no input is batch-constant. + The partial-constant path keeps representation dispatch visible to LLVM so loop unswitching can + specialize the runtime constant orientation and broadcast the constant operand. +- Decoded batch constants are validated when `ArgColumn` is constructed. Both ordinary and + null-tolerant decoding call `ArgColumn::try_from_constant`, so a malformed zero-length decode is + rejected before row zero can be read. The private constant representation therefore has exactly + one addressable row. +- `RowExecution` remains public and re-exported because the tensor L2 branch consumes it through + `RowFn::reduce_encoded`. Encoding-aware reductions remain on the first branch that consumes them + rather than expanding the framework PR in advance. +- Masked-array cleanup was kept out of the framework change. It belongs with the separate work + tracked by issue #9403. + +## Source structure + +`Batch` now lives in `batch/mod.rs`. Planning lives in `batch/planning.rs`. `batch/execute/mod.rs` +owns the high-level router and universal fast paths, while named leaves own constant, dense, +valid-only, filter-and-scatter, and output behavior. This keeps the central state and reader-facing +control flow in their parent modules without rebuilding a monolithic execution file. + +Owned and sink executors separate algorithm comments from compiler-sensitive source constraints. +The local constraints point to the relevant accessors and describe the observed non-LTO behavior +without embedding machine-specific timing numbers in production code. Safety comments are attached +to the exact unsafe operation and use named intermediate values when that makes the proof readable. + +`BitBuffer::try_for_each_set_index` generalizes fallible set-bit traversal in `vortex-buffer`. It +retains the word-oriented and all-ones traversal paths and returns immediately after the callback +fails. `execute_sink_valid_rows` uses it instead of storing a deferred row error and continuing the +scan. + +## Compiler-output experiments + +The final owning-library build was generated with: + +```text +CARGO_TARGET_DIR=/private/tmp/rowfn-ir-batch-split-stacked-candidate \ +CARGO_PROFILE_BENCH_CODEGEN_UNITS=16 \ +CARGO_PROFILE_BENCH_LTO=false \ +RUSTFLAGS='-C target-cpu=native -C remark=loop-vectorize -C debuginfo=line-tables-only' \ +cargo rustc -p vortex-array --lib --profile bench \ + --features _test-harness,table-display,unstable_row_fns -- \ + --emit=llvm-ir,asm,link +``` + +The module split was compared with the same executable source before the split. NumericBinary +functions emitted out of line fell from 147 to 94 because 36 owned, eight dense-sink, and eight +valid-row-sink helpers inlined into their visitors. End-to-end NumericBinary AArch64 instructions +fell from 44,986 to 44,179. Dense and owned paths fell from 30,190 to 29,300 instructions; sparse +paths fell from 10,082 to 9,995. + +Hot-loop structure was preserved: + +- Owned paths retained 152 vector-body references, 164 splat references, and the same vector-load + widths. +- Sparse paths retained 16 vector-body references, 48 splat references, and the same widths. +- Dense paths had no vector-body references in either build. +- Planning retained 19 functions and 1,191 AArch64 instructions. + +Splitting `Args::indexed_source` into a named local produced equivalent optimized LLVM IR and +identical AArch64 instructions for all 60 inspected owned NumericBinary monomorphizations. Removing +the owned lexical scope likewise changed only IR numbering and debug placement; the machine +instructions and vector, broadcast, and bounds-check metrics were identical. + +Moving bindings at the start of owned execution preserved every vectorized and unswitched hot path. +It also reduced the inspected bounds-failure references from six to four per monomorphization and +reduced representative executor instruction counts. This remains target- and compiler-specific +evidence, not a language-level guarantee. + +The first sink source rewrite exposed an important measurement trap. Bounds checks reported inside +each outlined sparse callback were not newly created checks. The baseline called a shared +`ElementTuple::get` that contained the same checks, while the candidate inlined that getter into +each callback. The confirmed difference was CGU placement, inlining, and code duplication. Reports +must compare the same logical work across callers and callees rather than count references inside +one symbol. + +## Partial-constant execution + +The final valid-row setup hoists the no-constant versus partial-constant dispatch outside set-bit +traversal. This duplicates both optimized word traversals in each specialization, growing the +static sparse assembly from 6,275 to 8,778 instructions and the full valid-row symbols from 8,000 +to 10,576. Static size alone was not treated as a regression because the relevant question was hot +execution and instruction-cache behavior. + +Five warmed, alternating A/B runs with 5,000 iterations per case compared the hoisted dispatch with +the earlier single traversal: + +| Case | Change in median | +| --- | ---: | +| Dense, left batch-constant | -57% | +| Dense, right batch-constant | -23% | +| Sparse, left batch-constant | -31% | +| Sparse, right batch-constant | -24% | +| Dense, no constants | -14.5% | +| Sparse, no constants | +1.7% | +| Owned control | -1.9% | + +The partial-constant improvements are material. The sparse no-constant and owned controls are within +run noise. A standalone dense-division control was noisy at +7.8% by its median, but focused pairs +overlapped and its hot instruction sequence was identical, so it was not defensible regression +evidence. + +Moving sink allocation after decoding, validation, and preparation was isolated with five +alternating 100,000-iteration runs. The retained ordering had a 22.28 microsecond median with a +0.28 microsecond median absolute deviation. Allocation before decoding had a 23.82 microsecond +median with a 1.17 microsecond deviation, 6.9% slower and noisier in this experiment. + +## Rejected alternatives + +- `set_indices().try_for_each` made the sparse code smaller but lost the vectorized remainder paths + and the all-ones fast path. +- New unchecked `ArgColumn` and `ElementTuple` access removed 16 panic sites but recovered only 398 + instructions from the 8,778-instruction sparse result. The unsafe API was not justified. +- A generic `#[inline(never)]` traversal worsened the full valid-row code to 12,982 instructions. +- One shared dynamic-callback traversal reduced static code to 6,079 instructions plus a shared + 435-instruction traversal, but introduced two indirect `blr` calls in the set-bit loops. That + per-valid-row dispatch was rejected. +- Moving constant-length error formatting to a cold, no-inline helper preserved the hot paths but + did not reliably reduce whole-binary text across consumers. The simpler construction-time check + was retained. + +## Stack integration + +The focused stack order is framework, primitive numeric operators, primitive comparisons, tensor +L2, tensor products, spatial distance, spatial predicates, and benchmark tools. The already merged +RowFn types PR is below `develop` and was not rewritten. + +| Branch | Final commit | +| --- | --- | +| `ct/row-fn-framework` | `8d7edf34dd445e152fc91be8c9b73d027d0131ad` | +| `ct/row-fn-numeric-operators` | `413796dd61aae2470d1139cf014bc8844d9ca792` | +| `ct/row-fn-primitive-comparisons` | `c057706c943eacef3b930e9f86bc3a5ff53c8b9d` | +| `ct/row-fn-tensor-l2` | `b655d61815351d284cdf6f575cfa9c60d0a313af` | +| `ct/row-fn-tensor-products` | `2f55638bcdcc2876eab193c85464d4aff74313a0` | +| `ct/row-fn-spatial-distance` | `2a2d8ae9da281aa39480fdcc61d6cbad430b6924` | +| `ct/row-fn-spatial-predicates` | `b37103bfe10defe9f80cc8bc4d3e1d92f78a0700` | +| `ct/row-fn-benchmark-tools` | `97b5a990abb4c9828dbc711c23a0a0f065140d78` | + +The framework split required one conflict resolution in tensor L2. The encoded-reduction hook moved +from the deleted monolithic execution file into the new high-level router, output reconciliation, +and valid-only modules. Every other dependent commit replayed without a semantic adaptation. Each +rewritten commit retained its DCO trailer, and every layer was checked with `git range-diff`. + +The final focused checks passed: + +```text +cargo test -p vortex-buffer set_index + 11 passed + +cargo test -p vortex-array scalar_fn::unstable::row + 40 passed + +cargo test -p vortex-tensor scalar_fns + 73 passed + +cargo test -p vortex-spatial scalar_fn + 178 passed + +cargo +nightly fmt --all -- --check + passed + +PYO3_PYTHON=/Users/connor/spiral/vortex-data/vortex1/.venv/bin/python3 \ + cargo clippy -p vortex-buffer -p vortex-array -p vortex-tensor -p vortex-spatial \ + --all-targets --all-features + passed +``` + +The historical `ct/row-fn` branch is synchronized with a two-parent merge. Its first parent keeps +the umbrella investigation history and this report. Its second parent is the final focused stack +tip, and the merge tree exactly equals that second parent.