Faster ScalarFnArray creation, lazy ArrayStats (#10329)

ScalarFnArray creation now uses 1 allocation instead of 3.
ArrayStats are now lazy and aren't constructed unless read or written

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

vortex-array/src/array/mod.rs

@@ -425,7 +425,7 @@ impl<V: VTable> DynArrayData for ArrayData<V> {
             slots,
         );
         let encoding_id = self.vtable.id();
-        let stats = this.statistics().to_array_stats();
+        let stats = this.statistics().share_existing();
         let uninit = Arc::new_uninit();
 
         // SAFETY: `uninit` is new.
@@ -492,7 +492,7 @@ impl<V: VTable> DynArrayData for ArrayData<V> {
     fn execute(&self, this: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
         let len = this.len();
         let dtype = this.dtype().clone();
-        let stats = this.statistics().to_array_stats();
+        let stats = this.statistics().to_owned();
         let result = unsafe { self.execute_unchecked(this, ctx)? };
 
         if matches!(result.step(), ExecutionStep::Done) {
@@ -511,10 +511,7 @@ impl<V: VTable> DynArrayData for ArrayData<V> {
                 );
             }
 
-            result
-                .array()
-                .statistics()
-                .set_iter(crate::stats::StatsSet::from(stats).into_iter());
+            result.array().statistics().set_iter(stats.into_iter());
         }
 
         Ok(result)

vortex-array/src/arrays/scalar_fn/array.rs

@@ -4,6 +4,7 @@
 use std::fmt::Display;
 use std::fmt::Formatter;
 
+use smallvec::SmallVec;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -15,6 +16,7 @@ use crate::array::Array;
 use crate::array::ArrayParts;
 use crate::array::TypedArrayRef;
 use crate::arrays::ScalarFn;
+use crate::dtype::DType;
 use crate::scalar_fn::ScalarFnRef;
 
 // ScalarFnArray has a variable number of slots (one per child)
@@ -74,9 +76,13 @@ impl<T: TypedArrayRef<ScalarFn>> ScalarFnArrayExt for T {}
 
 impl Array<ScalarFn> {
     /// Create a new ScalarFnArray from a scalar function and its children.
-    pub fn try_new(scalar_fn: ScalarFnRef, children: Vec<ArrayRef>) -> VortexResult<Self> {
-        let len = Self::infer_len(&children)?;
-        Self::try_new_with_len(scalar_fn, children, len)
+    pub fn try_new(
+        scalar_fn: ScalarFnRef,
+        children: impl IntoIterator<Item = ArrayRef>,
+    ) -> VortexResult<Self> {
+        let slots: ArraySlots = children.into_iter().map(Some).collect();
+        let len = Self::infer_len(&slots)?;
+        Self::try_new_from_slots(scalar_fn, slots, len)
     }
 
     /// Create a new ScalarFnArray from a scalar function, children, and an explicit length.
@@ -85,37 +91,61 @@ impl Array<ScalarFn> {
     /// child array to infer the length from.
     pub fn try_new_with_len(
         scalar_fn: ScalarFnRef,
-        children: Vec<ArrayRef>,
+        children: impl IntoIterator<Item = ArrayRef>,
+        len: usize,
+    ) -> VortexResult<Self> {
+        let slots: ArraySlots = children.into_iter().map(Some).collect();
+        Self::try_new_from_slots(scalar_fn, slots, len)
+    }
+
+    #[inline]
+    fn try_new_from_slots(
+        scalar_fn: ScalarFnRef,
+        slots: ArraySlots,
         len: usize,
     ) -> VortexResult<Self> {
-        Self::validate_arity(&scalar_fn, children.len())?;
-        Self::validate_children_len(&children, len)?;
+        Self::validate_arity(&scalar_fn, slots.len())?;
+        Self::validate_children_len(&slots, len)?;
 
-        let arg_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
+        let arg_dtypes = Self::arg_dtypes(&slots);
         let dtype = scalar_fn.return_dtype(&arg_dtypes)?;
-        let data = ScalarFnData {
-            scalar_fn: scalar_fn.clone(),
-        };
+
+        Ok(unsafe { Self::new_from_validated_slots(scalar_fn, dtype, slots, len) })
+    }
+
+    /// # Safety
+    /// The caller must ensure arity and child lengths were validated and
+    /// "dtype" is scalar function's return dtype.
+    #[inline]
+    unsafe fn new_from_validated_slots(
+        scalar_fn: ScalarFnRef,
+        dtype: DType,
+        slots: ArraySlots,
+        len: usize,
+    ) -> Self {
         let vtable = ScalarFn { id: scalar_fn.id() };
+        let data = ScalarFnData { scalar_fn };
+        unsafe { Array::from_parts_unchecked(ArrayParts::new(vtable, dtype, len, data, slots)) }
+    }
 
-        Ok(unsafe {
-            Array::from_parts_unchecked(ArrayParts::new(
-                vtable,
-                dtype,
-                len,
-                data,
-                children.into_iter().map(Some).collect::<ArraySlots>(),
-            ))
-        })
+    #[inline]
+    fn arg_dtypes(slots: &[Option<ArrayRef>]) -> SmallVec<[DType; 4]> {
+        let mut arg_dtypes = SmallVec::with_capacity(slots.len());
+        for child in slots.iter().flatten() {
+            arg_dtypes.push(child.dtype().clone());
+        }
+        arg_dtypes
     }
 
-    fn infer_len(children: &[ArrayRef]) -> VortexResult<usize> {
-        let Some(child) = children.first() else {
+    #[inline]
+    fn infer_len(slots: &[Option<ArrayRef>]) -> VortexResult<usize> {
+        let Some(child) = slots.first().and_then(Option::as_ref) else {
             vortex_bail!("ScalarFnArray length cannot be inferred without children");
         };
         Ok(child.len())
     }
 
+    #[inline]
     fn validate_arity(scalar_fn: &ScalarFnRef, child_count: usize) -> VortexResult<()> {
         let arity = scalar_fn.signature().arity();
         vortex_ensure!(
@@ -125,11 +155,14 @@ impl Array<ScalarFn> {
         Ok(())
     }
 
-    fn validate_children_len(children: &[ArrayRef], len: usize) -> VortexResult<()> {
-        vortex_ensure!(
-            children.iter().all(|c| c.len() == len),
-            "ScalarFnArray must have children equal to the array length"
-        );
+    #[inline]
+    fn validate_children_len(slots: &[Option<ArrayRef>], len: usize) -> VortexResult<()> {
+        for child in slots.iter().flatten() {
+            vortex_ensure!(
+                child.len() == len,
+                "ScalarFnArray must have children equal to the array length"
+            );
+        }
         Ok(())
     }
 }

vortex-array/src/arrays/scalar_fn/rules.rs

@@ -190,7 +190,7 @@ mod tests {
                     ConstantArray::new(7i32, 100).into_array()
                 }
             })
-            .collect();
+            .collect::<Vec<_>>();
         let array = ScalarFnArray::try_new(
             TypedScalarFnInstance::new(Binary, Operator::Add).erased(),
             children,

vortex-array/src/executor.rs

@@ -45,7 +45,6 @@ use crate::optimizer::ArrayOptimizer;
 use crate::optimizer::kernels::ArrayKernelsExt;
 use crate::optimizer::kernels::ParentExecutionKernels;
 use crate::optimizer::kernels::execute_parent_key;
-use crate::stats::ArrayStats;
 use crate::stats::StatsSet;
 use crate::trace_op;
 
@@ -268,7 +267,7 @@ impl ArrayRef {
 
             let expected_len = current_array.len();
             let expected_dtype = current_array.dtype().clone();
-            let stats = current_array.statistics().to_array_stats();
+            let stats = current_array.statistics().to_owned();
             let encoding_id = current_array.encoding_id();
             trace_op!(record_execute_encoding(&current_array));
             let result = current_array.execute_encoding_unchecked(ctx)?;
@@ -592,7 +591,7 @@ fn finalize_done(
     mut builder: Option<Box<dyn ArrayBuilder>>,
     expected_len: usize,
     expected_dtype: DType,
-    stats: ArrayStats,
+    stats: StatsSet,
     encoding_id: ArrayId,
 ) -> VortexResult<(ArrayRef, Option<Box<dyn ArrayBuilder>>)> {
     let output = if let Some(mut builder) = builder.take() {
@@ -616,9 +615,7 @@ fn finalize_done(
         );
     }
 
-    output
-        .statistics()
-        .set_iter(StatsSet::from(stats).into_iter());
+    output.statistics().set_iter(stats.into_iter());
     Ok((output, None))
 }
 

vortex-array/src/expression.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use itertools::Itertools;
+use smallvec::SmallVec;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
@@ -30,13 +30,13 @@ impl ArrayRef {
             return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array());
         }
 
-        let children: Vec<_> = children
+        let len = self.len();
+        let children = children
             .iter()
             .map(|child| self.clone().apply_bound(child))
-            .try_collect()?;
+            .collect::<VortexResult<SmallVec<[ArrayRef; 4]>>>()?;
 
-        let array =
-            ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array();
+        let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array();
 
         array.optimize()
     }
@@ -54,18 +54,18 @@ impl ArrayRef {
         }
 
         // Otherwise, collect the child arrays.
-        let children: Vec<_> = expr
+        let len = self.len();
+        let children = expr
             .children()
             .iter()
             .map(|e| self.clone().apply(e))
-            .try_collect()?;
+            .collect::<VortexResult<SmallVec<[ArrayRef; 4]>>>()?;
 
         // And wrap the scalar function up in an array.
         let scalar_fn = expr
             .as_scalar()
             .vortex_expect("root and literal were handled above, so this is a scalar node");
-        let array =
-            ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array();
+        let array = ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, len)?.into_array();
 
         // Optimize the resulting array's root.
         array.optimize()

vortex-array/src/scalar_fn/fns/between/mod.rs

@@ -197,7 +197,7 @@ impl Between {
         upper: ArrayRef,
         options: BetweenOptions,
     ) -> VortexResult<ScalarFnArray> {
-        ScalarFnArray::try_new(Between.bind(options), vec![array, lower, upper])
+        ScalarFnArray::try_new(Between.bind(options), [array, lower, upper])
     }
 }
 

vortex-array/src/scalar_fn/fns/binary/mod.rs

@@ -63,7 +63,7 @@ impl Binary {
         rhs: ArrayRef,
         operator: Operator,
     ) -> VortexResult<ScalarFnArray> {
-        ScalarFnArray::try_new(Binary.bind(operator), vec![lhs, rhs])
+        ScalarFnArray::try_new(Binary.bind(operator), [lhs, rhs])
     }
 }
 

vortex-array/src/scalar_fn/fns/cast/mod.rs

@@ -58,7 +58,7 @@ impl Cast {
     /// Creates a lazy cast of `input` to `target_dtype`.
     #[expect(clippy::new_ret_no_self, reason = "constructs the lazy result array")]
     pub fn new(input: ArrayRef, target_dtype: DType) -> ScalarFnArray {
-        ScalarFnArray::try_new(Cast.bind(target_dtype), vec![input])
+        ScalarFnArray::try_new(Cast.bind(target_dtype), [input])
             .vortex_expect("Cast has one child and an infallible return dtype")
     }
 }

vortex-array/src/scalar_fn/fns/fill_null/mod.rs

@@ -46,7 +46,7 @@ impl FillNull {
     ///
     /// Returns an error if the children have different lengths or incompatible dtypes.
     pub fn try_new(input: ArrayRef, fill_value: ArrayRef) -> VortexResult<ScalarFnArray> {
-        ScalarFnArray::try_new(FillNull.bind(EmptyOptions), vec![input, fill_value])
+        ScalarFnArray::try_new(FillNull.bind(EmptyOptions), [input, fill_value])
     }
 }
 

vortex-array/src/scalar_fn/fns/get_item.rs

@@ -50,7 +50,7 @@ impl GetItem {
         input: ArrayRef,
         field_name: impl Into<FieldName>,
     ) -> VortexResult<ScalarFnArray> {
-        ScalarFnArray::try_new(GetItem.bind(field_name.into()), vec![input])
+        ScalarFnArray::try_new(GetItem.bind(field_name.into()), [input])
     }
 }