Execute tensor L2 norm with RowFn

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

Author: connortsui20

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::<Vector>(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);
+}

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.