assert_arrays_eq and assert_nth_scalar require ExecutionCtx (#8509)

Part of migration to populate execution context everywhere

Author: robert3005

encodings/alp/src/alp/array.rs

@@ -471,7 +471,6 @@ mod tests {
     use rstest::rstest;
     use vortex_array::Canonical;
     use vortex_array::IntoArray;
-    use vortex_array::LEGACY_SESSION;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::assert_arrays_eq;
@@ -482,7 +481,11 @@ mod tests {
     use crate::alp_encode;
     use crate::decompress_into_array;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = vortex_array::array_session();
+        crate::initialize(&session);
+        session
+    });
 
     #[rstest]
     #[case(0)]
@@ -507,10 +510,9 @@ mod tests {
                 .unwrap()
         };
         // Compare against the traditional array-based decompress path
-        let expected =
-            decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap();
+        let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
 
-        assert_arrays_eq!(result_canonical.into_array(), expected);
+        assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx);
     }
 
     #[rstest]
@@ -525,26 +527,19 @@ mod tests {
     #[case(2049)]
     fn test_execute_f64(#[case] size: usize) {
         let values = PrimitiveArray::from_iter((0..size).map(|i| i as f64));
-        let encoded = alp_encode(
-            values.as_view(),
-            None,
-            &mut LEGACY_SESSION.create_execution_ctx(),
-        )
-        .unwrap();
+        let encoded =
+            alp_encode(values.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
 
-        let result_canonical = {
-            let mut ctx = SESSION.create_execution_ctx();
-            encoded
-                .clone()
-                .into_array()
-                .execute::<Canonical>(&mut ctx)
-                .unwrap()
-        };
+        let mut ctx = SESSION.create_execution_ctx();
+        let result_canonical = encoded
+            .clone()
+            .into_array()
+            .execute::<Canonical>(&mut ctx)
+            .unwrap();
         // Compare against the traditional array-based decompress path
-        let expected =
-            decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap();
+        let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
 
-        assert_arrays_eq!(result_canonical.into_array(), expected);
+        assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx);
     }
 
     #[rstest]
@@ -564,27 +559,20 @@ mod tests {
             .collect();
 
         let array = PrimitiveArray::from_iter(values);
-        let encoded = alp_encode(
-            array.as_view(),
-            None,
-            &mut LEGACY_SESSION.create_execution_ctx(),
-        )
-        .unwrap();
+        let encoded =
+            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
         assert!(encoded.patches().unwrap().array_len() > 0);
 
-        let result_canonical = {
-            let mut ctx = SESSION.create_execution_ctx();
-            encoded
-                .clone()
-                .into_array()
-                .execute::<Canonical>(&mut ctx)
-                .unwrap()
-        };
+        let mut ctx = SESSION.create_execution_ctx();
+        let result_canonical = encoded
+            .clone()
+            .into_array()
+            .execute::<Canonical>(&mut ctx)
+            .unwrap();
         // Compare against the traditional array-based decompress path
-        let expected =
-            decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap();
+        let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
 
-        assert_arrays_eq!(result_canonical.into_array(), expected);
+        assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx);
     }
 
     #[rstest]
@@ -603,26 +591,19 @@ mod tests {
             .collect();
 
         let array = PrimitiveArray::from_option_iter(values);
-        let encoded = alp_encode(
-            array.as_view(),
-            None,
-            &mut LEGACY_SESSION.create_execution_ctx(),
-        )
-        .unwrap();
+        let encoded =
+            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
 
-        let result_canonical = {
-            let mut ctx = SESSION.create_execution_ctx();
-            encoded
-                .clone()
-                .into_array()
-                .execute::<Canonical>(&mut ctx)
-                .unwrap()
-        };
+        let mut ctx = SESSION.create_execution_ctx();
+        let result_canonical = encoded
+            .clone()
+            .into_array()
+            .execute::<Canonical>(&mut ctx)
+            .unwrap();
         // Compare against the traditional array-based decompress path
-        let expected =
-            decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap();
+        let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
 
-        assert_arrays_eq!(result_canonical.into_array(), expected);
+        assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx);
     }
 
     #[rstest]
@@ -643,27 +624,20 @@ mod tests {
             .collect();
 
         let array = PrimitiveArray::from_option_iter(values);
-        let encoded = alp_encode(
-            array.as_view(),
-            None,
-            &mut LEGACY_SESSION.create_execution_ctx(),
-        )
-        .unwrap();
+        let encoded =
+            alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap();
         assert!(encoded.patches().unwrap().array_len() > 0);
 
-        let result_canonical = {
-            let mut ctx = SESSION.create_execution_ctx();
-            encoded
-                .clone()
-                .into_array()
-                .execute::<Canonical>(&mut ctx)
-                .unwrap()
-        };
+        let mut ctx = SESSION.create_execution_ctx();
+        let result_canonical = encoded
+            .clone()
+            .into_array()
+            .execute::<Canonical>(&mut ctx)
+            .unwrap();
         // Compare against the traditional array-based decompress path
-        let expected =
-            decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap();
+        let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap();
 
-        assert_arrays_eq!(result_canonical.into_array(), expected);
+        assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx);
     }
 
     #[rstest]
@@ -720,7 +694,7 @@ mod tests {
     #[case(1000, 200)]
     #[case(2048, 512)]
     fn test_sliced_to_primitive(#[case] size: usize, #[case] slice_start: usize) {
-        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let mut ctx = SESSION.create_execution_ctx();
         let values: Vec<Option<f64>> = (0..size)
             .map(|i| {
                 if i % 5 == 0 {
@@ -784,7 +758,7 @@ mod tests {
         let normally_encoded = alp_encode(
             original.as_view(),
             None,
-            &mut LEGACY_SESSION.create_execution_ctx(),
+            &mut SESSION.create_execution_ctx(),
         )
         .unwrap();
         assert!(
@@ -818,7 +792,7 @@ mod tests {
         // The legacy decompress_into_array path should work correctly.
         let result_legacy = decompress_into_array(
             alp_without_chunk_offsets.clone(),
-            &mut LEGACY_SESSION.create_execution_ctx(),
+            &mut SESSION.create_execution_ctx(),
         )
         .unwrap();
         let legacy_slice = result_legacy.as_slice::<f64>();