IsSortedKernel

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-array/src/aggregate_fn/accumulator.rs

@@ -146,7 +146,8 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
     }
 
     fn flush(&mut self) -> VortexResult<Scalar> {
-        let partial = self.vtable.flush(&mut self.partial)?;
+        let partial = self.vtable.to_scalar(&self.partial)?;
+        self.vtable.reset(&mut self.partial);
 
         #[cfg(debug_assertions)]
         {
@@ -162,8 +163,8 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
     }
 
     fn finish(&mut self) -> VortexResult<Scalar> {
-        let partial = self.flush()?;
-        let result = self.vtable.finalize_scalar(partial)?;
+        let result = self.vtable.finalize_scalar(&self.partial)?;
+        self.vtable.reset(&mut self.partial);
 
         vortex_ensure!(
             result.dtype() == &self.return_dtype,

vortex-array/src/aggregate_fn/fns/is_constant/mod.rs

@@ -234,9 +234,10 @@ impl IsConstantPartial {
     }
 }
 
+static NAMES: std::sync::LazyLock<FieldNames> =
+    std::sync::LazyLock::new(|| FieldNames::from(["is_constant", "value"]));
+
 pub fn make_is_constant_partial_dtype(element_dtype: &DType) -> DType {
-    static NAMES: std::sync::LazyLock<FieldNames> =
-        std::sync::LazyLock::new(|| FieldNames::from(["is_constant", "value"]));
     DType::Struct(
         StructFields::new(
             NAMES.clone(),
@@ -313,9 +314,9 @@ impl AggregateFnVTable for IsConstant {
         Ok(())
     }
 
-    fn flush(&self, partial: &mut Self::Partial) -> VortexResult<Scalar> {
+    fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
         let dtype = make_is_constant_partial_dtype(&partial.element_dtype);
-        let result = match &partial.first_value {
+        Ok(match &partial.first_value {
             None => {
                 // Empty accumulator — return null struct.
                 Scalar::null(dtype)
@@ -329,13 +330,12 @@ impl AggregateFnVTable for IsConstant {
                         .cast(&partial.element_dtype.as_nullable())?,
                 ],
             ),
-        };
+        })
+    }
 
-        // Reset state.
+    fn reset(&self, partial: &mut Self::Partial) {
         partial.is_constant = true;
         partial.first_value = None;
-
-        Ok(result)
     }
 
     #[inline]
@@ -409,21 +409,15 @@ impl AggregateFnVTable for IsConstant {
     }
 
     fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
-        // TODO: extract is_constant field from struct array
-        Ok(partials)
+        partials.get_item(NAMES.get(0))
     }
 
-    fn finalize_scalar(&self, partial: Scalar) -> VortexResult<Scalar> {
-        if partial.is_null() {
+    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
+        if partial.first_value.is_none() {
             // Empty accumulator → return false.
             return Ok(Scalar::bool(false, Nullability::NonNullable));
         }
-        let is_constant_val = partial
-            .as_struct()
-            .field_by_idx(0)
-            .map(|s| s.as_bool().value().unwrap_or(false))
-            .unwrap_or(false);
-        Ok(Scalar::bool(is_constant_val, Nullability::NonNullable))
+        Ok(Scalar::bool(partial.is_constant, Nullability::NonNullable))
     }
 }
 

vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs

@@ -30,6 +30,7 @@ use crate::aggregate_fn::AggregateFnVTable;
 use crate::aggregate_fn::DynAccumulator;
 use crate::arrays::Constant;
 use crate::arrays::Null;
+use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::dtype::FieldNames;
 use crate::dtype::Nullability;
@@ -200,10 +201,11 @@ pub struct IsSortedPartial {
     element_dtype: DType,
 }
 
+static NAMES: std::sync::LazyLock<FieldNames> = std::sync::LazyLock::new(|| {
+    FieldNames::from(["is_sorted", "strict", "first_value", "last_value"])
+});
+
 pub fn make_is_sorted_partial_dtype(element_dtype: &DType) -> DType {
-    static NAMES: std::sync::LazyLock<FieldNames> = std::sync::LazyLock::new(|| {
-        FieldNames::from(["is_sorted", "strict", "first_value", "last_value"])
-    });
     DType::Struct(
         StructFields::new(
             NAMES.clone(),
@@ -315,56 +317,48 @@ impl AggregateFnVTable for IsSorted {
         Ok(())
     }
 
-    fn flush(&self, partial: &mut Self::Partial) -> VortexResult<Scalar> {
+    fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
         let dtype = make_is_sorted_partial_dtype(&partial.element_dtype);
-
-        // Take ownership of the values to avoid cloning.
-        let first = partial.first_value.take();
-        let last = partial.last_value.take();
-        let is_sorted = partial.is_sorted;
-        let strict = partial.strict;
-
-        // Reset state.
-        partial.is_sorted = true;
-
-        let result = match (first, last) {
+        Ok(match (&partial.first_value, &partial.last_value) {
             (None, _) => {
                 // Empty accumulator — return null struct.
                 Scalar::null(dtype)
             }
             (Some(first_value), Some(last_value)) => {
-                // Values are already nullable from into_nullable_unchecked in accumulate.
                 // SAFETY: We constructed partial_dtype and the children match its field dtypes.
                 unsafe {
                     Scalar::struct_unchecked(
                         dtype,
                         [
-                            Scalar::bool(is_sorted, Nullability::NonNullable),
-                            Scalar::bool(strict, Nullability::NonNullable),
-                            first_value,
-                            last_value,
+                            Scalar::bool(partial.is_sorted, Nullability::NonNullable),
+                            Scalar::bool(partial.strict, Nullability::NonNullable),
+                            first_value.clone(),
+                            last_value.clone(),
                         ],
                     )
                 }
             }
             (Some(first_value), None) => {
-                let cloned = first_value.clone();
                 // SAFETY: We constructed partial_dtype and the children match its field dtypes.
                 unsafe {
                     Scalar::struct_unchecked(
                         dtype,
                         [
-                            Scalar::bool(is_sorted, Nullability::NonNullable),
-                            Scalar::bool(strict, Nullability::NonNullable),
-                            first_value,
-                            cloned,
+                            Scalar::bool(partial.is_sorted, Nullability::NonNullable),
+                            Scalar::bool(partial.strict, Nullability::NonNullable),
+                            first_value.clone(),
+                            first_value.clone(),
                         ],
                     )
                 }
             }
-        };
+        })
+    }
 
-        Ok(result)
+    fn reset(&self, partial: &mut Self::Partial) {
+        partial.is_sorted = true;
+        partial.first_value = None;
+        partial.last_value = None;
     }
 
     #[inline]
@@ -479,21 +473,15 @@ impl AggregateFnVTable for IsSorted {
     }
 
     fn finalize(&self, partials: ArrayRef) -> VortexResult<ArrayRef> {
-        // TODO: extract is_sorted field from struct array
-        Ok(partials)
+        partials.get_item(NAMES.get(0))
     }
 
-    fn finalize_scalar(&self, partial: Scalar) -> VortexResult<Scalar> {
-        if partial.is_null() {
+    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
+        if partial.first_value.is_none() {
             // Empty accumulator → vacuously sorted.
             return Ok(Scalar::bool(true, Nullability::NonNullable));
         }
-        let is_sorted_val = partial
-            .as_struct()
-            .field_by_idx(0)
-            .map(|s| s.as_bool().value().unwrap_or(false))
-            .unwrap_or(false);
-        Ok(Scalar::bool(is_sorted_val, Nullability::NonNullable))
+        Ok(Scalar::bool(partial.is_sorted, Nullability::NonNullable))
     }
 }
 

vortex-array/src/aggregate_fn/fns/min_max/mod.rs

@@ -210,13 +210,17 @@ impl AggregateFnVTable for MinMax {
         Ok(())
     }
 
-    fn flush(&self, partial: &mut Self::Partial) -> VortexResult<Scalar> {
+    fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
         let dtype = make_minmax_dtype(&partial.element_dtype);
-        let result = match (partial.min.take(), partial.max.take()) {
-            (Some(min), Some(max)) => Scalar::struct_(dtype, vec![min, max]),
+        Ok(match (&partial.min, &partial.max) {
+            (Some(min), Some(max)) => Scalar::struct_(dtype, vec![min.clone(), max.clone()]),
             _ => Scalar::null(dtype),
-        };
-        Ok(result)
+        })
+    }
+
+    fn reset(&self, partial: &mut Self::Partial) {
+        partial.min = None;
+        partial.max = None;
     }
 
     #[inline]
@@ -266,8 +270,8 @@ impl AggregateFnVTable for MinMax {
         Ok(partials)
     }
 
-    fn finalize_scalar(&self, partial: Scalar) -> VortexResult<Scalar> {
-        Ok(partial)
+    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
+        self.to_scalar(partial)
     }
 }
 
@@ -450,7 +454,7 @@ mod tests {
         let scalar2 = Scalar::struct_(struct_dtype, vec![Scalar::from(2i32), Scalar::from(10i32)]);
         MinMax.combine_partials(&mut state, scalar2)?;
 
-        let result = MinMaxResult::from_scalar(MinMax.flush(&mut state)?)?
+        let result = MinMaxResult::from_scalar(MinMax.to_scalar(&state)?)?
             .vortex_expect("should have result");
         assert_eq!(result.min, Scalar::from(2i32));
         assert_eq!(result.max, Scalar::from(15i32));

vortex-array/src/aggregate_fn/fns/nan_count/mod.rs

@@ -115,10 +115,12 @@ impl AggregateFnVTable for NanCount {
         Ok(())
     }
 
-    fn flush(&self, partial: &mut Self::Partial) -> VortexResult<Scalar> {
-        let result = Scalar::primitive(*partial, NonNullable);
+    fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
+        Ok(Scalar::primitive(*partial, NonNullable))
+    }
+
+    fn reset(&self, partial: &mut Self::Partial) {
         *partial = 0;
-        Ok(result)
     }
 
     #[inline]
@@ -157,8 +159,8 @@ impl AggregateFnVTable for NanCount {
         Ok(partials)
     }
 
-    fn finalize_scalar(&self, partial: Scalar) -> VortexResult<Scalar> {
-        Ok(partial)
+    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
+        self.to_scalar(partial)
     }
 }
 
@@ -236,7 +238,8 @@ mod tests {
         let scalar2 = Scalar::primitive(3u64, Nullability::NonNullable);
         NanCount.combine_partials(&mut state, scalar2)?;
 
-        let result = NanCount.flush(&mut state)?;
+        let result = NanCount.to_scalar(&state)?;
+        NanCount.reset(&mut state);
         assert_eq!(result.as_primitive().typed_value::<u64>(), Some(8));
         Ok(())
     }

vortex-array/src/aggregate_fn/fns/sum/decimal.rs

@@ -320,7 +320,7 @@ mod tests {
         let small = Scalar::decimal(DecimalValue::from(9i64), DecimalDType::new(14, 0), Nullable);
         Sum.combine_partials(&mut state, small)?;
 
-        let result = Sum.flush(&mut state)?;
+        let result = Sum.to_scalar(&state)?;
         assert!(!result.is_null());
         assert_eq!(
             result.as_decimal().decimal_value(),
@@ -351,7 +351,7 @@ mod tests {
             Scalar::decimal(DecimalValue::from(1i64), DecimalDType::new(14, 0), Nullable);
         Sum.combine_partials(&mut state, one_more)?;
 
-        let result = Sum.flush(&mut state)?;
+        let result = Sum.to_scalar(&state)?;
         assert!(result.is_null());
         assert_eq!(
             result.dtype(),
@@ -380,7 +380,7 @@ mod tests {
         );
         Sum.combine_partials(&mut state, one_more)?;
 
-        let result = Sum.flush(&mut state)?;
+        let result = Sum.to_scalar(&state)?;
         assert!(result.is_null());
         Ok(())
     }
@@ -413,7 +413,7 @@ mod tests {
         let mut ctx = LEGACY_SESSION.create_execution_ctx();
         Sum.accumulate(&mut state, &columnar, &mut ctx)?;
 
-        let result = Sum.flush(&mut state)?;
+        let result = Sum.to_scalar(&state)?;
         assert!(result.is_null());
         Ok(())
     }

vortex-array/src/aggregate_fn/fns/sum/mod.rs

@@ -180,8 +180,8 @@ impl AggregateFnVTable for Sum {
         Ok(())
     }
 
-    fn flush(&self, partial: &mut Self::Partial) -> VortexResult<Scalar> {
-        let result = match &partial.current {
+    fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
+        Ok(match &partial.current {
             None => Scalar::null(partial.return_dtype.as_nullable()),
             Some(SumState::Unsigned(v)) => Scalar::primitive(*v, Nullability::Nullable),
             Some(SumState::Signed(v)) => Scalar::primitive(*v, Nullability::Nullable),
@@ -193,12 +193,11 @@ impl AggregateFnVTable for Sum {
                     .vortex_expect("return dtype must be decimal");
                 Scalar::decimal(*value, decimal_dtype, Nullability::Nullable)
             }
-        };
+        })
+    }
 
-        // Reset the state
+    fn reset(&self, partial: &mut Self::Partial) {
         partial.current = Some(make_zero_state(&partial.return_dtype));
-
-        Ok(result)
     }
 
     #[inline]
@@ -250,8 +249,8 @@ impl AggregateFnVTable for Sum {
         Ok(partials)
     }
 
-    fn finalize_scalar(&self, partial: Scalar) -> VortexResult<Scalar> {
-        Ok(partial)
+    fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar> {
+        self.to_scalar(partial)
     }
 }
 
@@ -461,7 +460,8 @@ mod tests {
         let scalar2 = Scalar::primitive(50i64, Nullable);
         Sum.combine_partials(&mut state, scalar2)?;
 
-        let result = Sum.flush(&mut state)?;
+        let result = Sum.to_scalar(&state)?;
+        Sum.reset(&mut state);
         assert_eq!(result.as_primitive().typed_value::<i64>(), Some(150));
         Ok(())
     }