remove all to_canonical and replace with execute

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

Author: connortsui20

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

@@ -56,9 +56,12 @@ impl Default for TurboQuantConfig {
 ///
 /// All quantization (rotation, centroid lookup) happens in f32. f16 is upcast; f64 is truncated.
 #[allow(clippy::cast_possible_truncation)]
-fn extract_f32_elements(fsl: &FixedSizeListArray) -> VortexResult<PrimitiveArray> {
+fn extract_f32_elements(
+    fsl: &FixedSizeListArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<PrimitiveArray> {
     let elements = fsl.elements();
-    let primitive = elements.to_canonical()?.into_primitive();
+    let primitive = elements.clone().execute::<PrimitiveArray>(ctx)?;
     let ptype = primitive.ptype();
 
     match ptype {
@@ -128,7 +131,7 @@ fn turboquant_quantize_core(
     let rotation = RotationMatrix::try_new(seed, dimension)?;
     let padded_dim = rotation.padded_dim();
 
-    let f32_elements = extract_f32_elements(fsl)?;
+    let f32_elements = extract_f32_elements(fsl, ctx)?;
 
     let centroids = get_centroids(padded_dim as u32, bit_width)?;
     let boundaries = compute_centroid_boundaries(&centroids);
@@ -220,7 +223,7 @@ pub fn turboquant_encode(
 ) -> VortexResult<ArrayRef> {
     let ext_dtype = ext.dtype().clone();
     let storage = ext.storage_array();
-    let fsl = storage.to_canonical()?.into_fixed_size_list();
+    let fsl = storage.clone().execute::<FixedSizeListArray>(ctx)?;
 
     vortex_ensure!(
         config.bit_width >= 1 && config.bit_width <= TurboQuant::MAX_BIT_WIDTH,

vortex-tensor/src/encodings/turboquant/decompress.rs

@@ -73,7 +73,10 @@ pub fn execute_decompress(
 
     // Unpack codes from FixedSizeListArray -> flat u8 elements.
     let codes_fsl = array.codes().clone().execute::<FixedSizeListArray>(ctx)?;
-    let codes_prim = codes_fsl.elements().to_canonical()?.into_primitive();
+    let codes_prim = codes_fsl
+        .elements()
+        .clone()
+        .execute::<PrimitiveArray>(ctx)?;
     let indices = codes_prim.as_slice::<u8>();
 
     // Read norms in their native precision. Norms carry the validity of the array.

vortex-tensor/src/encodings/turboquant/tests.rs

@@ -125,26 +125,24 @@ fn encode_decode(
     fsl: &FixedSizeListArray,
     config: &TurboQuantConfig,
 ) -> VortexResult<(Vec<f32>, Vec<f32>)> {
+    let mut ctx = SESSION.create_execution_ctx();
     let original: Vec<f32> = {
-        let prim = fsl.elements().to_canonical().unwrap().into_primitive();
+        let prim = fsl.elements().clone().execute::<PrimitiveArray>(&mut ctx)?;
         prim.as_slice::<f32>().to_vec()
     };
     let ext = make_vector_ext(fsl);
     let config = config.clone();
-    let mut ctx = SESSION.create_execution_ctx();
     let encoded = turboquant_encode(ext.as_view(), &config, &mut ctx)?;
     let decoded_ext = encoded.execute::<ExtensionArray>(&mut ctx)?;
     let decoded_fsl = decoded_ext
         .storage_array()
-        .to_canonical()
-        .unwrap()
-        .into_fixed_size_list();
+        .clone()
+        .execute::<FixedSizeListArray>(&mut ctx)?;
     let decoded_elements: Vec<f32> = {
         let prim = decoded_fsl
             .elements()
-            .to_canonical()
-            .unwrap()
-            .into_primitive();
+            .clone()
+            .execute::<PrimitiveArray>(&mut ctx)?;
         prim.as_slice::<f32>().to_vec()
     };
     Ok((original, decoded_elements))
@@ -409,8 +407,8 @@ fn f16_input_encodes_successfully() -> VortexResult<()> {
     let decoded_ext = encoded.execute::<ExtensionArray>(&mut ctx)?;
     let decoded_fsl = decoded_ext
         .storage_array()
-        .to_canonical()?
-        .into_fixed_size_list();
+        .clone()
+        .execute::<FixedSizeListArray>(&mut ctx)?;
     assert_eq!(decoded_fsl.len(), num_rows);
     Ok(())
 }
@@ -470,9 +468,12 @@ fn stored_rotation_signs_produce_correct_decode() -> VortexResult<()> {
         .execute::<ExtensionArray>(&mut ctx)?;
     let decoded_fsl = decoded_ext
         .storage_array()
-        .to_canonical()?
-        .into_fixed_size_list();
-    let decoded = decoded_fsl.elements().to_canonical()?.into_primitive();
+        .clone()
+        .execute::<FixedSizeListArray>(&mut ctx)?;
+    let decoded = decoded_fsl
+        .elements()
+        .clone()
+        .execute::<PrimitiveArray>(&mut ctx)?;
     let decoded_slice = decoded.as_slice::<f32>();
 
     // Verify stored signs match seed-derived signs.
@@ -514,20 +515,26 @@ fn slice_preserves_data() -> VortexResult<()> {
     let full_decoded = encoded.clone().execute::<ExtensionArray>(&mut ctx)?;
     let full_fsl = full_decoded
         .storage_array()
-        .to_canonical()?
-        .into_fixed_size_list();
+        .clone()
+        .execute::<FixedSizeListArray>(&mut ctx)?;
     let expected = full_fsl.slice(5..10)?;
-    let expected_prim = expected.to_canonical()?.into_fixed_size_list();
-    let expected_elements = expected_prim.elements().to_canonical()?.into_primitive();
+    let expected_fsl = expected.execute::<FixedSizeListArray>(&mut ctx)?;
+    let expected_elements = expected_fsl
+        .elements()
+        .clone()
+        .execute::<PrimitiveArray>(&mut ctx)?;
 
     // Slice then decompress.
     let sliced = encoded.slice(5..10)?;
     let sliced_decoded = sliced.execute::<ExtensionArray>(&mut ctx)?;
     let sliced_fsl = sliced_decoded
         .storage_array()
-        .to_canonical()?
-        .into_fixed_size_list();
-    let actual_elements = sliced_fsl.elements().to_canonical()?.into_primitive();
+        .clone()
+        .execute::<FixedSizeListArray>(&mut ctx)?;
+    let actual_elements = sliced_fsl
+        .elements()
+        .clone()
+        .execute::<PrimitiveArray>(&mut ctx)?;
 
     assert_eq!(
         expected_elements.as_slice::<f32>(),
@@ -570,10 +577,10 @@ fn l2_norm_readthrough() -> VortexResult<()> {
     let tq = encoded.as_opt::<TurboQuant>().unwrap();
 
     // Stored norms should match the actual L2 norms of the input.
-    let norms_prim = tq.norms().to_canonical()?.into_primitive();
+    let norms_prim = tq.norms().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let stored_norms = norms_prim.as_slice::<f32>();
 
-    let input_prim = fsl.elements().to_canonical()?.into_primitive();
+    let input_prim = fsl.elements().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let input_f32 = input_prim.as_slice::<f32>();
     for row in 0..10 {
         let vec = &input_f32[row * 128..(row + 1) * 128];
@@ -601,7 +608,7 @@ fn cosine_similarity_quantized_accuracy() -> VortexResult<()> {
     let tq = encoded.as_opt::<TurboQuant>().unwrap();
 
     // Compute exact cosine similarity from original data.
-    let input_prim = fsl.elements().to_canonical()?.into_primitive();
+    let input_prim = fsl.elements().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let input_f32 = input_prim.as_slice::<f32>();
 
     // Read quantized codes, norms, and centroids for approximate computation.
@@ -610,7 +617,10 @@ fn cosine_similarity_quantized_accuracy() -> VortexResult<()> {
     let norms_prim = tq.norms().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let norms = norms_prim.as_slice::<f32>();
     let codes_fsl = tq.codes().clone().execute::<FixedSizeListArray>(&mut ctx)?;
-    let codes_prim = codes_fsl.elements().to_canonical()?.into_primitive();
+    let codes_prim = codes_fsl
+        .elements()
+        .clone()
+        .execute::<PrimitiveArray>(&mut ctx)?;
     let all_codes = codes_prim.as_slice::<u8>();
     let centroids_prim = tq.centroids().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let centroid_vals = centroids_prim.as_slice::<f32>();
@@ -674,15 +684,18 @@ fn dot_product_quantized_accuracy() -> VortexResult<()> {
     let encoded = turboquant_encode(ext.as_view(), &config, &mut ctx)?;
     let tq = encoded.as_opt::<TurboQuant>().unwrap();
 
-    let input_prim = fsl.elements().to_canonical()?.into_primitive();
+    let input_prim = fsl.elements().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let input_f32 = input_prim.as_slice::<f32>();
 
     let mut ctx = SESSION.create_execution_ctx();
     let pd = tq.padded_dim() as usize;
     let norms_prim = tq.norms().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let norms = norms_prim.as_slice::<f32>();
     let codes_fsl = tq.codes().clone().execute::<FixedSizeListArray>(&mut ctx)?;
-    let codes_prim = codes_fsl.elements().to_canonical()?.into_primitive();
+    let codes_prim = codes_fsl
+        .elements()
+        .clone()
+        .execute::<PrimitiveArray>(&mut ctx)?;
     let all_codes = codes_prim.as_slice::<u8>();
     let centroids_prim = tq.centroids().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let centroid_vals = centroids_prim.as_slice::<f32>();
@@ -815,13 +828,16 @@ fn nullable_vectors_roundtrip() -> VortexResult<()> {
 
     let decoded_fsl = decoded_ext
         .storage_array()
-        .to_canonical()?
-        .into_fixed_size_list();
-    let decoded_prim = decoded_fsl.elements().to_canonical()?.into_primitive();
+        .clone()
+        .execute::<FixedSizeListArray>(&mut ctx)?;
+    let decoded_prim = decoded_fsl
+        .elements()
+        .clone()
+        .execute::<PrimitiveArray>(&mut ctx)?;
     let decoded_f32 = decoded_prim.as_slice::<f32>();
 
     // Original f32 elements for non-null row comparison.
-    let orig_prim = fsl.elements().to_canonical()?.into_primitive();
+    let orig_prim = fsl.elements().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let orig_f32 = orig_prim.as_slice::<f32>();
 
     // Non-null rows should have reasonable reconstruction (within MSE bounds).
@@ -889,7 +905,7 @@ fn nullable_l2_norm_readthrough() -> VortexResult<()> {
     let norms: PrimitiveArray = norm_sfn.into_array().execute(&mut ctx)?;
 
     // Null rows should have null norms, valid rows should have correct norms.
-    let orig_prim = fsl.elements().to_canonical()?.into_primitive();
+    let orig_prim = fsl.elements().clone().execute::<PrimitiveArray>(&mut ctx)?;
     let orig_f32 = orig_prim.as_slice::<f32>();
     for row in 0..5 {
         if row % 2 == 0 {

vortex-tensor/src/scalar_fns/cosine_similarity.rs

@@ -206,15 +206,20 @@ impl ScalarFnVTable for CosineSimilarity {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::LazyLock;
+
     use rstest::rstest;
     use vortex_array::ArrayRef;
     use vortex_array::IntoArray;
-    use vortex_array::ToCanonical;
+    use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::MaskedArray;
+    use vortex_array::arrays::PrimitiveArray;
     use vortex_array::arrays::ScalarFnArray;
     use vortex_array::scalar_fn::ScalarFn;
+    use vortex_array::session::ArraySession;
     use vortex_array::validity::Validity;
     use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
 
     use crate::scalar_fns::ApproxOptions;
     use crate::scalar_fns::cosine_similarity::CosineSimilarity;
@@ -224,11 +229,15 @@ mod tests {
     use crate::utils::test_helpers::tensor_array;
     use crate::utils::test_helpers::vector_array;
 
+    static SESSION: LazyLock<VortexSession> =
+        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
+
     /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec<f64>`.
     fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef, len: usize) -> VortexResult<Vec<f64>> {
         let scalar_fn = ScalarFn::new(CosineSimilarity, ApproxOptions::Exact).erased();
         let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], len)?;
-        let prim = result.as_array().to_primitive();
+        let mut ctx = SESSION.create_execution_ctx();
+        let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?;
         Ok(prim.as_slice::<f64>().to_vec())
     }
 
@@ -398,7 +407,8 @@ mod tests {
 
         let scalar_fn = ScalarFn::new(CosineSimilarity, ApproxOptions::Exact).erased();
         let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 2)?;
-        let prim = result.as_array().to_primitive();
+        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)?);

vortex-tensor/src/scalar_fns/inner_product.rs

@@ -213,28 +213,36 @@ fn inner_product_row<T: Float + NativePType>(a: &[T], b: &[T]) -> T {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::LazyLock;
+
     use rstest::rstest;
     use vortex_array::ArrayRef;
     use vortex_array::IntoArray;
-    use vortex_array::ToCanonical;
+    use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::MaskedArray;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::arrays::ScalarFnArray;
     use vortex_array::scalar_fn::ScalarFn;
+    use vortex_array::session::ArraySession;
     use vortex_array::validity::Validity;
     use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
 
     use crate::scalar_fns::ApproxOptions;
     use crate::scalar_fns::inner_product::InnerProduct;
     use crate::utils::test_helpers::assert_close;
     use crate::utils::test_helpers::tensor_array;
     use crate::utils::test_helpers::vector_array;
 
+    static SESSION: LazyLock<VortexSession> =
+        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
+
     /// Evaluates inner product between two tensor arrays and returns the result as `Vec<f64>`.
     fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef, len: usize) -> VortexResult<Vec<f64>> {
         let scalar_fn = ScalarFn::new(InnerProduct, ApproxOptions::Exact).erased();
         let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], len)?;
-        let prim = result.as_array().to_primitive();
+        let mut ctx = SESSION.create_execution_ctx();
+        let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?;
         Ok(prim.as_slice::<f64>().to_vec())
     }
 
@@ -311,7 +319,8 @@ mod tests {
 
         let scalar_fn = ScalarFn::new(InnerProduct, ApproxOptions::Exact).erased();
         let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs], 3)?;
-        let prim = result.as_array().to_primitive();
+        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)?);

vortex-tensor/src/scalar_fns/l2_norm.rs

@@ -196,27 +196,36 @@ fn l2_norm_row<T: Float + NativePType>(v: &[T]) -> T {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::LazyLock;
+
     use rstest::rstest;
     use vortex_array::ArrayRef;
     use vortex_array::IntoArray;
-    use vortex_array::ToCanonical;
+    use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::MaskedArray;
+    use vortex_array::arrays::PrimitiveArray;
     use vortex_array::arrays::ScalarFnArray;
     use vortex_array::scalar_fn::ScalarFn;
+    use vortex_array::session::ArraySession;
     use vortex_array::validity::Validity;
     use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
 
     use crate::scalar_fns::ApproxOptions;
     use crate::scalar_fns::l2_norm::L2Norm;
     use crate::utils::test_helpers::assert_close;
     use crate::utils::test_helpers::tensor_array;
     use crate::utils::test_helpers::vector_array;
 
+    static SESSION: LazyLock<VortexSession> =
+        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
+
     /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec<f64>`.
     fn eval_l2_norm(input: ArrayRef, len: usize) -> VortexResult<Vec<f64>> {
         let scalar_fn = ScalarFn::new(L2Norm, ApproxOptions::Exact).erased();
         let result = ScalarFnArray::try_new(scalar_fn, vec![input], len)?;
-        let prim = result.as_array().to_primitive();
+        let mut ctx = SESSION.create_execution_ctx();
+        let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?;
         Ok(prim.as_slice::<f64>().to_vec())
     }
 
@@ -270,7 +279,8 @@ mod tests {
 
         let scalar_fn = ScalarFn::new(L2Norm, ApproxOptions::Exact).erased();
         let result = ScalarFnArray::try_new(scalar_fn, vec![arr], 2)?;
-        let prim = result.as_array().to_primitive();
+        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)?);