rebase

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs

@@ -2,8 +2,9 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_error::VortexResult;
-use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct;
+
 use crate::ExecutionCtx;
+use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct;
 use crate::arrays::VariantArray;
 use crate::arrays::variant::VariantArrayExt;
 

vortex-array/src/aggregate_fn/fns/count/grouped.rs

@@ -91,13 +91,14 @@ mod tests {
     use vortex_error::VortexResult;
 
     use crate::ArrayRef;
+    use crate::ExecutionCtx;
     use crate::IntoArray;
-    use crate::LEGACY_SESSION;
     use crate::VortexSessionExecute;
     use crate::aggregate_fn::DynGroupedAccumulator;
     use crate::aggregate_fn::GroupedAccumulator;
     use crate::aggregate_fn::NumericalAggregateOpts;
     use crate::aggregate_fn::fns::count::Count;
+    use crate::array_session;
     use crate::arrays::FixedSizeListArray;
     use crate::arrays::ListViewArray;
     use crate::arrays::PrimitiveArray;
@@ -110,13 +111,17 @@ mod tests {
     use crate::validity::Validity;
 
     /// Run a grouped count through the accumulator.
-    fn grouped_count_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult<ArrayRef> {
+    fn grouped_count_actual(
+        groups: &ArrayRef,
+        elem_dtype: &DType,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
         let mut acc = GroupedAccumulator::try_new(
             Count,
             NumericalAggregateOpts::default(),
             elem_dtype.clone(),
         )?;
-        acc.accumulate_list(groups, &mut LEGACY_SESSION.create_execution_ctx())?;
+        acc.accumulate_list(groups, ctx)?;
         acc.finish()
     }
 
@@ -125,7 +130,7 @@ mod tests {
         elements: &ArrayRef,
         ranges: &[(usize, usize)],
     ) -> VortexResult<ArrayRef> {
-        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let mut ctx = array_session().create_execution_ctx();
         let counts: Buffer<u64> = ranges
             .iter()
             .map(|&(offset, size)| {
@@ -151,49 +156,46 @@ mod tests {
 
     #[test]
     fn listview_counts_all_valid() -> VortexResult<()> {
-        let assertion_session = crate::array_session();
-        let mut assertion_ctx = assertion_session.create_execution_ctx();
+        let mut ctx = array_session().create_execution_ctx();
         let elements =
             PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array();
         let elem_dtype = DType::Primitive(PType::I32, NonNullable);
         let ranges = [(0, 2), (2, 1), (3, 3), (6, 0)];
 
         let groups = listview(elements.clone(), &ranges)?;
-        let actual = grouped_count_actual(&groups, &elem_dtype)?;
+        let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?;
         let expected = grouped_count_reference(&elements, &ranges)?;
 
         let direct =
             PrimitiveArray::new(buffer![2u64, 1, 3, 0], Validity::NonNullable).into_array();
-        assert_arrays_eq!(&actual, &direct, &mut assertion_ctx);
-        assert_arrays_eq!(&actual, &expected, &mut assertion_ctx);
+        assert_arrays_eq!(&actual, &direct, &mut ctx);
+        assert_arrays_eq!(&actual, &expected, &mut ctx);
         Ok(())
     }
 
     #[test]
     fn listview_counts_with_nulls() -> VortexResult<()> {
-        let assertion_session = crate::array_session();
-        let mut assertion_ctx = assertion_session.create_execution_ctx();
+        let mut ctx = array_session().create_execution_ctx();
         let elements =
             PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)])
                 .into_array();
         let elem_dtype = DType::Primitive(PType::I32, Nullable);
         let ranges = [(0, 3), (3, 2), (5, 1)];
 
         let groups = listview(elements.clone(), &ranges)?;
-        let actual = grouped_count_actual(&groups, &elem_dtype)?;
+        let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?;
         let expected = grouped_count_reference(&elements, &ranges)?;
 
         // Group 0: {1, null, 3} -> 2. Group 1: {null, null} -> 0. Group 2: {9} -> 1.
         let direct = PrimitiveArray::new(buffer![2u64, 0, 1], Validity::NonNullable).into_array();
-        assert_arrays_eq!(&actual, &direct, &mut assertion_ctx);
-        assert_arrays_eq!(&actual, &expected, &mut assertion_ctx);
+        assert_arrays_eq!(&actual, &direct, &mut ctx);
+        assert_arrays_eq!(&actual, &expected, &mut ctx);
         Ok(())
     }
 
     #[test]
     fn listview_counts_varbinview_with_nulls() -> VortexResult<()> {
-        let assertion_session = crate::array_session();
-        let mut assertion_ctx = assertion_session.create_execution_ctx();
+        let mut ctx = array_session().create_execution_ctx();
         let elements = VarBinViewArray::from_iter_nullable_str([
             Some("a"),
             None,
@@ -206,17 +208,18 @@ mod tests {
         let ranges = [(0, 2), (2, 2), (4, 1)];
 
         let groups = listview(elements.clone(), &ranges)?;
-        let actual = grouped_count_actual(&groups, &elem_dtype)?;
+        let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?;
         let expected = grouped_count_reference(&elements, &ranges)?;
 
         let direct = PrimitiveArray::new(buffer![1u64, 1, 1], Validity::NonNullable).into_array();
-        assert_arrays_eq!(&actual, &direct, &mut assertion_ctx);
-        assert_arrays_eq!(&actual, &expected, &mut assertion_ctx);
+        assert_arrays_eq!(&actual, &direct, &mut ctx);
+        assert_arrays_eq!(&actual, &expected, &mut ctx);
         Ok(())
     }
 
     #[test]
     fn fixed_size_counts_float_nans() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
         let elements =
             PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(2.0)])
                 .into_array();
@@ -225,32 +228,31 @@ mod tests {
             FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into_array();
 
         // NaNs are excluded by default and counted otherwise.
-        let actual = grouped_count_actual(&groups, &elem_dtype)?;
+        let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?;
         let expected = PrimitiveArray::new(buffer![1u64, 1], Validity::NonNullable).into_array();
-        assert_arrays_eq!(&actual, &expected);
+        assert_arrays_eq!(&actual, &expected, &mut ctx);
 
         let mut acc =
             GroupedAccumulator::try_new(Count, NumericalAggregateOpts::include_nans(), elem_dtype)?;
-        acc.accumulate_list(&groups, &mut LEGACY_SESSION.create_execution_ctx())?;
+        acc.accumulate_list(&groups, &mut ctx)?;
         let actual = acc.finish()?;
         let expected = PrimitiveArray::new(buffer![2u64, 1], Validity::NonNullable).into_array();
-        assert_arrays_eq!(&actual, &expected);
+        assert_arrays_eq!(&actual, &expected, &mut ctx);
         Ok(())
     }
 
     #[test]
     fn fixed_size_counts_with_nulls() -> VortexResult<()> {
-        let assertion_session = crate::array_session();
-        let mut assertion_ctx = assertion_session.create_execution_ctx();
+        let mut ctx = array_session().create_execution_ctx();
         let elements =
             PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array();
         let elem_dtype = DType::Primitive(PType::I32, Nullable);
         let groups =
             FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into_array();
 
-        let actual = grouped_count_actual(&groups, &elem_dtype)?;
+        let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?;
         let direct = PrimitiveArray::new(buffer![1u64, 2], Validity::NonNullable).into_array();
-        assert_arrays_eq!(&actual, &direct, &mut assertion_ctx);
+        assert_arrays_eq!(&actual, &direct, &mut ctx);
         Ok(())
     }
 }