From 75c47073241fc24f86cf876dac7194bccaff3f2d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 6 Aug 2026 02:06:01 +0000 Subject: [PATCH] Apply the review findings on the row scalar function branch Cleanup, corrected documentation, and test coverage from a review of the RowFn API, its execution logic, and the scalar function implementations. The author-facing API is unchanged and every visit_prepared_into monomorph emits identical IR to the pre-review commit. Signed-off-by: Connor Tsui --- Cargo.toml | 8 +- SCALAR_FN_HANDOFF.md | 146 ++++++++++++++++++ .../src/scalar_fn/fns/binary/numeric/row.rs | 30 +++- vortex-array/src/scalar_fn/mod.rs | 11 +- .../src/scalar_fn/row/element/bool.rs | 7 - vortex-array/src/scalar_fn/row/element/mod.rs | 10 +- .../src/scalar_fn/row/element/primitive.rs | 7 - .../src/scalar_fn/row/element/tuple.rs | 6 +- vortex-array/src/scalar_fn/row/lift.rs | 109 +++++++++---- vortex-array/src/scalar_fn/row/row_fn.rs | 3 +- .../scalar_fn/row/tests/null_strategies.rs | 78 ++++++++++ vortex-array/src/scalar_fn/row/vtable.rs | 2 +- vortex-compute/src/lane_kernels/map_into.rs | 66 -------- vortex-geo/src/scalar_fn/contains.rs | 127 ++++++++++++++- vortex-geo/src/test_harness.rs | 2 +- .../src/scalar_fns/cosine_similarity.rs | 44 +++--- vortex-tensor/src/scalar_fns/l2_norm.rs | 14 +- vortex-tensor/src/scalar_fns/row.rs | 7 - vortex-tensor/src/utils.rs | 15 ++ 19 files changed, 521 insertions(+), 171 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 1484b75de11..3d65e572140 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-geo`'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-geo/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/SCALAR_FN_HANDOFF.md b/SCALAR_FN_HANDOFF.md index 96f2addf3d8..f99f4cb0eaa 100644 --- a/SCALAR_FN_HANDOFF.md +++ b/SCALAR_FN_HANDOFF.md @@ -278,6 +278,152 @@ The checks recorded for the final API state are: 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 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 ddb194c38d5..aab36bf196b 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -121,12 +121,23 @@ fn operand_ptype(args: &[DType]) -> VortexResult { } /// 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), @@ -145,9 +156,21 @@ where /// 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: a -/// deferred-error kernel runs densely and retries valid rows only on its cold error path. +/// 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, @@ -160,8 +183,7 @@ struct CheckedSink> { op: PhantomData, } -/// The uninitialized output slots of a [`CheckedSink`], borrowed once for the row loop, together -/// with the batch-wide failure reduction they contribute to. +/// 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, diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 1edf7b6afeb..11bcefe0325 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -13,9 +13,10 @@ //! 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: `vortex.byte_length`, `vortex.tensor.l2_norm`, -//! `vortex.tensor.inner_product`, `vortex.geo.distance`. Name the element types and write the row -//! closure, and the rest is derived, including which rows get visited. +//! 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 @@ -63,6 +64,10 @@ //! - **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; diff --git a/vortex-array/src/scalar_fn/row/element/bool.rs b/vortex-array/src/scalar_fn/row/element/bool.rs index 1f78417094b..d4fc51c769c 100644 --- a/vortex-array/src/scalar_fn/row/element/bool.rs +++ b/vortex-array/src/scalar_fn/row/element/bool.rs @@ -36,13 +36,6 @@ impl InputElement for bool { Ok(array.execute::(ctx)?.into_bit_buffer()) } - fn decode_null_tolerant( - array: ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Self::decode(array, ctx).map(Some) - } - fn get(column: &Self::Column, index: usize) -> bool { column.value(index) } diff --git a/vortex-array/src/scalar_fn/row/element/mod.rs b/vortex-array/src/scalar_fn/row/element/mod.rs index 15a3472b327..fc88690b6b0 100644 --- a/vortex-array/src/scalar_fn/row/element/mod.rs +++ b/vortex-array/src/scalar_fn/row/element/mod.rs @@ -96,10 +96,12 @@ pub 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 may use the default, because its ordinary - /// decode already tolerates null payloads. Other elements may override this by writing 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. + /// 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. diff --git a/vortex-array/src/scalar_fn/row/element/primitive.rs b/vortex-array/src/scalar_fn/row/element/primitive.rs index 1afc3e4c8b0..d54c922bf2d 100644 --- a/vortex-array/src/scalar_fn/row/element/primitive.rs +++ b/vortex-array/src/scalar_fn/row/element/primitive.rs @@ -43,13 +43,6 @@ impl InputElement for T { Ok(array.execute::(ctx)?.into_buffer::()) } - fn decode_null_tolerant( - array: ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Self::decode(array, ctx).map(Some) - } - fn get(column: &Self::Column, index: usize) -> T { column[index] } diff --git a/vortex-array/src/scalar_fn/row/element/tuple.rs b/vortex-array/src/scalar_fn/row/element/tuple.rs index b5ed2bfe272..a3c31cfeb2e 100644 --- a/vortex-array/src/scalar_fn/row/element/tuple.rs +++ b/vortex-array/src/scalar_fn/row/element/tuple.rs @@ -143,6 +143,10 @@ pub(in crate::scalar_fn::row) fn batch_constant(array: &ArrayRef) -> Option Option>; diff --git a/vortex-array/src/scalar_fn/row/lift.rs b/vortex-array/src/scalar_fn/row/lift.rs index 5916c38e2bc..a842ef5deef 100644 --- a/vortex-array/src/scalar_fn/row/lift.rs +++ b/vortex-array/src/scalar_fn/row/lift.rs @@ -118,6 +118,18 @@ pub(super) enum RowPolicy { } 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 { @@ -133,6 +145,15 @@ impl RowPolicy { } } +/// 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. @@ -313,6 +334,11 @@ impl<'a> Batch<'a> { .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); } @@ -320,6 +346,13 @@ impl<'a> Batch<'a> { 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), @@ -337,6 +370,40 @@ impl<'a> Batch<'a> { } } + /// 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. /// @@ -363,23 +430,10 @@ impl<'a> Batch<'a> { filtered_decode_cost: usize, 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(), - ); - } - - if valid.all_false() { - return Ok(self.all_null()); - } + 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)? @@ -621,23 +675,10 @@ impl Batch<'_> { strategy: NullStrategy, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let valid = self - .validity - .clone() - .execute_mask(self.args.row_count(), ctx)?; - - 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(Some); - } - - if valid.all_false() { - return Ok(Some(self.all_null())); - } + 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), diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs index a89aa19e668..f2a63636b1d 100644 --- a/vortex-array/src/scalar_fn/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -81,7 +81,8 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// 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. + /// 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 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 9b269846c02..d8a7fe18816 100644 --- a/vortex-array/src/scalar_fn/row/tests/null_strategies.rs +++ b/vortex-array/src/scalar_fn/row/tests/null_strategies.rs @@ -143,6 +143,7 @@ fn branch_propagates_real_errors() { /// 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; @@ -390,6 +391,83 @@ mod selection { 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)] diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs index af67be035c8..cff4831aa8f 100644 --- a/vortex-array/src/scalar_fn/row/vtable.rs +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -267,7 +267,7 @@ pub(super) fn row_policy( .map(|plan| plan.policy) } -/// Every [`RowFn`] is a [`ScalarFnVTable`], the row loop lifted by [`Batch`]. +/// 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 diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs index f5de42a5ce9..258913e9fc7 100644 --- a/vortex-compute/src/lane_kernels/map_into.rs +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -217,52 +217,6 @@ pub trait IndexedSourceExt: IndexedSource + Sized { } } - /// Split value/error map with **no validity awareness at all**: apply - /// `f(value) -> (R, bool)` to every lane, write the value unconditionally, and - /// OR-reduce the per-lane error flags into the returned `bool`. - /// - /// This is the highest-throughput checked shape for operations whose error - /// check auto-vectorizes (e.g. checked integer add/sub/mul): the loop carries a - /// single flag reduction with no per-lane select and no per-chunk exit branch, - /// so it runs at the speed of the unchecked [`map_into`]. The trade-off is that - /// it reports only *whether* some lane failed, not which one, and it never - /// exits early. Callers that need attribution or null-lane filtering should - /// re-run the (now known cold) input through [`try_map_into`] or - /// [`try_map_masked_into`] when this returns `true`. - /// - /// The flag is deliberately reduced inside the kernel rather than through a - /// closure-captured `&mut bool`: an escaped flag becomes a loop-carried memory - /// dependence that blocks auto-vectorization of the whole loop. - /// - /// [`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 f: F) -> bool - where - R: Copy, - F: FnMut(Self::Item) -> (R, bool), - { - let values = self; - let len = values.len(); - assert_eq!(out.len(), len, "out must have the same length as values"); - - let mut failed = false; - 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, error) = f(val); - failed |= error; - // 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. /// @@ -592,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(); diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs index 4acf16624e8..0e892b4d075 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -154,6 +154,18 @@ enum ContainsRoute { /// 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; @@ -311,12 +323,19 @@ fn contains_row_prepared(operands: &ConstOperands, a: &Geometry, b: &Geomet #[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; @@ -341,7 +360,10 @@ mod tests { use vortex_error::vortex_err; use wkb::writer::WriteOptions; + use super::ConstOperands; use super::GeoContains; + use super::PreparedOperand; + 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; @@ -732,6 +754,43 @@ mod tests { 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![ @@ -740,10 +799,70 @@ mod tests { ])) } + /// 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 and linestring containers), reversed relate - /// (multipolygon containers), and the direct pairings (a point on either side, polygon over - /// multipoint), including boundary contact, crossing, disjoint and empty cases. + /// 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())] @@ -760,6 +879,8 @@ mod tests { #[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())] diff --git a/vortex-geo/src/test_harness.rs b/vortex-geo/src/test_harness.rs index 7bec0f90ee9..c2dff4f85b6 100644 --- a/vortex-geo/src/test_harness.rs +++ b/vortex-geo/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() { diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 674d32f0cd5..e786d49b47c 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -37,6 +37,7 @@ 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::l2_norm_row; /// Cosine similarity between two columns. /// @@ -172,16 +173,6 @@ struct ConstNorms { rhs: Option, } -/// The L2 norm of one row, accumulating in the same order as [`cosine_similarity_row`] so a -/// hoisted norm is bit-identical to the one computed per row. -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() -} - /// Computes the cosine similarity of one row, taking any hoisted norm from `norms` and computing /// the rest exactly as [`cosine_similarity_row`] does. /// @@ -253,6 +244,13 @@ fn cosine_from_parts(dot: T, denom: T) -> T { /// 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, @@ -273,12 +271,18 @@ fn cosine_both_normalized( 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() { + // 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 { - dots[i] + dot } }) .collect(); @@ -311,12 +315,16 @@ fn cosine_one_normalized( 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() { + // 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 { - dots[i] / plain_norms[i] + dot / plain_norm } }) .collect(); diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 6eb5242d843..433a6527636 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -3,7 +3,6 @@ //! L2 norm expression for tensor-like types. -use num_traits::Float; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -13,7 +12,6 @@ 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::proto::dtype as pb; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::ElementSink; @@ -32,6 +30,7 @@ use crate::encodings::normalized::Normalized; 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. @@ -148,14 +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() -} diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index f0f6c8ea95f..3c02a1d2615 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -89,13 +89,6 @@ impl InputElement for TensorRow { }) } - fn decode_null_tolerant( - array: ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - Self::decode(array, ctx).map(Some) - } - fn get(column: &Self::Column, index: usize) -> &[T] { let start = index * column.stride; &column.elements.as_slice()[start..start + column.list_size] diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 0cd06e4c762..460dde82ea7 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -4,6 +4,7 @@ //! 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; @@ -63,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