add extract constant flat row

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

Author: connortsui20

vortex-tensor/public-api.lock

@@ -142,7 +142,7 @@ impl vortex_tensor::fixed_shape::FixedShapeTensorMatcherMetadata<'_>
 
 pub fn vortex_tensor::fixed_shape::FixedShapeTensorMatcherMetadata<'_>::element_ptype(&self) -> vortex_array::dtype::ptype::PType
 
-pub fn vortex_tensor::fixed_shape::FixedShapeTensorMatcherMetadata<'_>::list_size(&self) -> usize
+pub fn vortex_tensor::fixed_shape::FixedShapeTensorMatcherMetadata<'_>::flat_list_size(&self) -> u32
 
 pub fn vortex_tensor::fixed_shape::FixedShapeTensorMatcherMetadata<'_>::metadata(&self) -> &vortex_tensor::fixed_shape::FixedShapeTensorMetadata
 
@@ -222,7 +222,7 @@ impl vortex_tensor::matcher::TensorMatch<'_>
 
 pub fn vortex_tensor::matcher::TensorMatch<'_>::element_ptype(self) -> vortex_array::dtype::ptype::PType
 
-pub fn vortex_tensor::matcher::TensorMatch<'_>::list_size(self) -> usize
+pub fn vortex_tensor::matcher::TensorMatch<'_>::list_size(self) -> u32
 
 impl<'a> core::clone::Clone for vortex_tensor::matcher::TensorMatch<'a>
 
@@ -382,7 +382,7 @@ pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::validity(&self, _options:
 
 pub fn vortex_tensor::scalar_fns::l2_denorm::normalize_as_l2_denorm(input: vortex_array::array::erased::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::vtable::ScalarFnArray>
 
-pub fn vortex_tensor::scalar_fns::l2_denorm::validate_l2_normalized_rows(input: &vortex_array::array::erased::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<()>
+pub fn vortex_tensor::scalar_fns::l2_denorm::validate_l2_normalized_rows_against_norms(normalized: &vortex_array::array::erased::ArrayRef, norms: core::option::Option<&vortex_array::array::erased::ArrayRef>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<()>
 
 pub mod vortex_tensor::scalar_fns::l2_norm
 
@@ -570,12 +570,10 @@ pub struct vortex_tensor::vector::VectorMatcherMetadata
 
 impl vortex_tensor::vector::VectorMatcherMetadata
 
-pub fn vortex_tensor::vector::VectorMatcherMetadata::dimensions(&self) -> usize
+pub fn vortex_tensor::vector::VectorMatcherMetadata::dimensions(&self) -> u32
 
 pub fn vortex_tensor::vector::VectorMatcherMetadata::element_ptype(&self) -> vortex_array::dtype::ptype::PType
 
-pub fn vortex_tensor::vector::VectorMatcherMetadata::list_size(&self) -> u32
-
 pub fn vortex_tensor::vector::VectorMatcherMetadata::try_new(element_ptype: vortex_array::dtype::ptype::PType, dimensions: u32) -> vortex_error::VortexResult<Self>
 
 impl core::clone::Clone for vortex_tensor::vector::VectorMatcherMetadata

vortex-tensor/src/scalar_fns/inner_product.rs

@@ -55,6 +55,7 @@ use crate::matcher::AnyTensor;
 use crate::scalar_fns::l2_denorm::DenormOrientation;
 use crate::scalar_fns::sorf_transform::SorfMatrix;
 use crate::scalar_fns::sorf_transform::SorfTransform;
+use crate::utils::extract_constant_flat_row;
 use crate::utils::extract_flat_elements;
 use crate::utils::extract_l2_denorm_children;
 use crate::utils::validate_binary_tensor_float_inputs;
@@ -425,15 +426,15 @@ impl InnerProduct {
         let seed = sorf_view.options.seed;
         let padded_dim = dim.next_power_of_two();
 
-        // Extract the single stored row of the constant via the `is_constant` short-circuit.
-        let flat = extract_flat_elements(&const_storage, dim, ctx)?;
+        // Extract the single stored row of the constant.
+        let flat = extract_constant_flat_row(&const_storage, ctx)?;
         if flat.ptype() != PType::F32 {
             return Ok(None);
         }
 
         // Zero-pad the query from `dim` to `padded_dim` and forward-rotate.
         let mut padded_query = vec![0.0f32; padded_dim];
-        padded_query[..dim].copy_from_slice(flat.row::<f32>(0));
+        padded_query[..dim].copy_from_slice(flat.as_slice::<f32>());
 
         let rotation = SorfMatrix::try_new(seed, dim, num_rounds)?;
         let mut rotated_query = vec![0.0f32; padded_dim];
@@ -533,7 +534,7 @@ impl InnerProduct {
 
         let padded_dim = usize::try_from(fsl.list_size()).vortex_expect("fsl list_size fits usize");
 
-        let flat = extract_flat_elements(&const_storage, padded_dim, ctx)?;
+        let flat = extract_constant_flat_row(&const_storage, ctx)?;
         if flat.ptype() != PType::F32 {
             // TODO(connor): case 2 is f32-only. For f16/f64 we fall through to the standard
             // path, which computes the inner product with the correct element type.
@@ -552,7 +553,7 @@ impl InnerProduct {
             return Ok(Some(empty.into_array()));
         }
 
-        let q: &[f32] = flat.row::<f32>(0);
+        let q: &[f32] = flat.as_slice::<f32>();
         debug_assert_eq!(q.len(), padded_dim);
         let codes: &[u8] = codes_prim.as_slice::<u8>();
         let values: &[f32] = values_prim.as_slice::<f32>();

vortex-tensor/src/scalar_fns/l2_denorm.rs

@@ -60,6 +60,7 @@ use vortex_session::VortexSession;
 
 use crate::matcher::AnyTensor;
 use crate::scalar_fns::l2_norm::L2Norm;
+use crate::utils::extract_constant_flat_row;
 use crate::utils::extract_flat_elements;
 use crate::utils::validate_tensor_float_input;
 
@@ -525,12 +526,12 @@ pub(crate) fn try_build_constant_l2_denorm(
     let ext_dtype = input.dtype().as_extension().clone();
     let storage_fsl_nullability = storage.dtype().nullability();
 
-    // `extract_flat_elements` takes the `is_constant` single-row path for `Constant` storage, so
-    // this is cheap and does not expand the constant to the full column length.
-    let flat = extract_flat_elements(storage, list_size, ctx)?;
+    // Materialize just the single stored row; this does not expand the constant to the full
+    // column length.
+    let flat = extract_constant_flat_row(storage, ctx)?;
 
     let (normalized_fsl_scalar, norms_scalar) = match_each_float_ptype!(flat.ptype(), |T| {
-        let row = flat.row::<T>(0);
+        let row = flat.as_slice::<T>();
 
         let mut sum_sq = T::zero();
         for &x in row {
@@ -605,19 +606,14 @@ fn unit_norm_tolerance(element_ptype: PType) -> f64 {
     }
 }
 
-/// Validates that every valid row of `input` is already L2-normalized (either length 1 or 0).
-pub fn validate_l2_normalized_rows(input: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
-    validate_l2_normalized_rows_against_norms(input, None, ctx)
-}
-
 /// Validates that `normalized` and (when supplied) the matching `norms` jointly satisfy the
 /// [`L2Denorm`] invariants:
 ///
 /// - Every valid row of `normalized` has L2 norm `1.0` or `0.0` (within element-precision
 ///   tolerance).
 /// - When `norms` is supplied, every stored norm is non-negative and any row whose stored norm is
 ///   `0.0` is exactly the zero vector in `normalized`.
-fn validate_l2_normalized_rows_against_norms(
+pub fn validate_l2_normalized_rows_against_norms(
     normalized: &ArrayRef,
     norms: Option<&ArrayRef>,
     ctx: &mut ExecutionCtx,
@@ -771,7 +767,7 @@ mod tests {
 
     use crate::scalar_fns::l2_denorm::L2Denorm;
     use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm;
-    use crate::scalar_fns::l2_denorm::validate_l2_normalized_rows;
+    use crate::scalar_fns::l2_denorm::validate_l2_normalized_rows_against_norms;
     use crate::tests::SESSION;
     use crate::utils::test_helpers::assert_close;
     use crate::utils::test_helpers::constant_tensor_array;
@@ -906,15 +902,15 @@ mod tests {
         let input = vector_array(2, &[3.0f32, 4.0, 0.0, 0.0].map(half::f16::from_f32))?;
         let mut ctx = SESSION.create_execution_ctx();
         let roundtrip = normalize_as_l2_denorm(input, &mut ctx)?;
-        validate_l2_normalized_rows(&roundtrip.child_at(0).clone(), &mut ctx)?;
+        validate_l2_normalized_rows_against_norms(&roundtrip.child_at(0).clone(), None, &mut ctx)?;
         Ok(())
     }
 
     #[test]
     fn validate_l2_normalized_rows_rejects_unnormalized_input() -> VortexResult<()> {
         let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0])?;
         let mut ctx = SESSION.create_execution_ctx();
-        let result = validate_l2_normalized_rows(&input, &mut ctx);
+        let result = validate_l2_normalized_rows_against_norms(&input, None, &mut ctx);
         assert!(result.is_err());
         Ok(())
     }

vortex-tensor/src/utils.rs

@@ -9,6 +9,7 @@ use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
+use vortex_array::arrays::primitive::PrimitiveArrayExt;
 use vortex_array::arrays::scalar_fn::ExactScalarFn;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::NativePType;
@@ -124,12 +125,11 @@ impl FlatElements {
     }
 }
 
-// TODO(connor): Usage of this function is sometimes incorrect / not performant.
-// Make sure to fix them.
 /// Extracts the flat primitive elements from a tensor storage array (FixedSizeList).
 ///
 /// When the input is a [`ConstantArray`] (e.g., a literal query vector), only a single row is
-/// materialized to avoid expanding it to the full column length.
+/// materialized to avoid expanding it to the full column length. Callers that have already
+/// confirmed the storage is constant-backed should prefer [`extract_constant_flat_row`].
 pub fn extract_flat_elements(
     storage: &ArrayRef,
     list_size: usize,
@@ -146,13 +146,69 @@ pub fn extract_flat_elements(
 
     let fsl: FixedSizeListArray = source.execute(ctx)?;
     let elems: PrimitiveArray = fsl.elements().clone().execute(ctx)?;
+    vortex_ensure!(
+        !elems.nullability().is_nullable(),
+        "tensor storage elements must be non-nullable, got {}",
+        elems.dtype(),
+    );
     Ok(FlatElements {
         elems,
         list_size,
         is_constant,
     })
 }
 
+/// The single stored row of a constant-backed tensor storage array.
+///
+/// Contrast with [`FlatElements`], which exposes arbitrary row indices: a `FlatRow` statically
+/// encodes "there is exactly one row available," so call sites that have gated on a constant input
+/// read the row via [`Self::as_slice`] instead of `row(0)`.
+pub struct FlatRow {
+    elems: PrimitiveArray,
+}
+
+impl FlatRow {
+    /// Returns the [`PType`] of the underlying elements.
+    #[must_use]
+    pub fn ptype(&self) -> PType {
+        self.elems.ptype()
+    }
+
+    /// Returns the stored row as a typed slice. Its length equals the storage scalar's
+    /// fixed-size-list size.
+    #[must_use]
+    pub fn as_slice<T: NativePType>(&self) -> &[T] {
+        self.elems.as_slice::<T>()
+    }
+}
+
+/// Extracts the single stored row from a [`Constant`]-backed tensor storage array.
+///
+/// The caller must have confirmed that `storage` is a [`Constant`] encoding whose scalar is a
+/// non-null fixed-size list. This is the fast path for constant query vectors: exactly one row is
+/// materialized regardless of the column length.
+///
+/// # Panics
+///
+/// Panics if `storage` is not a [`Constant`] encoding.
+pub fn extract_constant_flat_row(
+    storage: &ArrayRef,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<FlatRow> {
+    let constant = storage
+        .as_opt::<Constant>()
+        .vortex_expect("extract_constant_flat_row requires Constant-backed storage");
+    let single = ConstantArray::new(constant.scalar().clone(), 1).into_array();
+    let fsl: FixedSizeListArray = single.execute(ctx)?;
+    let elems: PrimitiveArray = fsl.elements().clone().execute(ctx)?;
+    vortex_ensure!(
+        !elems.nullability().is_nullable(),
+        "tensor storage elements must be non-nullable, got {}",
+        elems.dtype(),
+    );
+    Ok(FlatRow { elems })
+}
+
 /// Extracts the `(normalized, norms)` children from an [`L2Denorm`] scalar function array.
 ///
 /// [`L2Denorm`]: crate::scalar_fns::l2_denorm::L2Denorm