Type Coercion (#6960)

This allows Vortex expressions to be typed in contexts that currently
don't have their own logical planner, for example expressions passed
into PyVortex.

---------

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

Author: gatesn

vortex-array/public-api.lock

@@ -75,6 +75,11 @@ impl AggregateFnRef {
         AggregateFnOptions { inner: &*self.0 }
     }
 
+    /// Coerce the input type for this aggregate function.
+    pub fn coerce_args(&self, input_dtype: &DType) -> VortexResult<DType> {
+        self.0.coerce_args(input_dtype)
+    }
+
     /// Compute the return [`DType`] per group given the input element type.
     pub fn return_dtype(&self, input_dtype: &DType) -> VortexResult<DType> {
         self.0.return_dtype(input_dtype)

vortex-array/src/aggregate_fn/erased.rs

@@ -39,6 +39,7 @@ pub(super) trait DynAggregateFn: 'static + Send + Sync + super::sealed::Sealed {
     fn id(&self) -> AggregateFnId;
     fn options_any(&self) -> &dyn Any;
 
+    fn coerce_args(&self, input_dtype: &DType) -> VortexResult<DType>;
     fn return_dtype(&self, input_dtype: &DType) -> VortexResult<DType>;
     fn state_dtype(&self, input_dtype: &DType) -> VortexResult<DType>;
     fn accumulator(
@@ -84,6 +85,10 @@ impl<V: AggregateFnVTable> DynAggregateFn for AggregateFnInner<V> {
         &self.options
     }
 
+    fn coerce_args(&self, input_dtype: &DType) -> VortexResult<DType> {
+        V::coerce_args(&self.vtable, &self.options, input_dtype)
+    }
+
     fn return_dtype(&self, input_dtype: &DType) -> VortexResult<DType> {
         V::return_dtype(&self.vtable, &self.options, input_dtype)
     }

vortex-array/src/aggregate_fn/typed.rs

@@ -60,6 +60,15 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync {
         vortex_bail!("Aggregate function {} is not deserializable", self.id());
     }
 
+    /// Coerce the input type for this aggregate function.
+    ///
+    /// This is optionally used by Vortex users when performing type coercion over a Vortex
+    /// expression. The default implementation returns the input type unchanged.
+    fn coerce_args(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult<DType> {
+        let _ = options;
+        Ok(input_dtype.clone())
+    }
+
     /// The return [`DType`] of the aggregate.
     fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult<DType>;