Clarify docs in session accessors

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

Cargo.lock

@@ -25,6 +25,7 @@ vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }
 vortex-mask = { workspace = true }
 vortex-scalar = { workspace = true }
+vortex-session = { workspace = true }
 
 [lints]
 workspace = true

encodings/runend/Cargo.toml

@@ -11,6 +11,7 @@ pub mod compress;
 mod compute;
 mod iter;
 mod ops;
+mod rules;
 
 #[doc(hidden)]
 pub mod _benchmarking {
@@ -23,11 +24,17 @@ pub mod _benchmarking {
 use vortex_array::ArrayBufferVisitor;
 use vortex_array::ArrayChildVisitor;
 use vortex_array::Canonical;
+use vortex_array::session::ArraySession;
+use vortex_array::session::ArraySessionExt;
+use vortex_array::vtable::ArrayVTableExt;
 use vortex_array::vtable::EncodeVTable;
 use vortex_array::vtable::VisitorVTable;
 use vortex_error::VortexResult;
+use vortex_session::SessionExt;
+use vortex_session::VortexSession;
 
 use crate::compress::runend_encode;
+use crate::rules::RunEndScalarFnRule;
 
 impl EncodeVTable<RunEndVTable> for RunEndVTable {
     fn encode(
@@ -59,6 +66,15 @@ impl VisitorVTable<RunEndVTable> for RunEndVTable {
     }
 }
 
+/// Initialize run-end encoding in the given session.
+pub fn initialize(session: &mut VortexSession) {
+    session.arrays().register(RunEndVTable.as_vtable());
+    session
+        .get_mut::<ArraySession>()
+        .optimizer_mut()
+        .register_parent_rule(RunEndScalarFnRule);
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_array::ProstMetadata;

encodings/runend/src/lib.rs

@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::arrays::AnyScalarFn;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::ConstantVTable;
+use vortex_array::arrays::ScalarFnArray;
+use vortex_array::optimizer::rules::ArrayParentReduceRule;
+use vortex_array::optimizer::rules::Exact;
+use vortex_dtype::DType;
+use vortex_error::VortexResult;
+
+use crate::RunEndArray;
+use crate::RunEndVTable;
+
+/// A rule to push down scalar functions through run-end encoding into the values array.
+///
+/// This only works if all other children of the scalar function array are constants.
+#[derive(Debug)]
+pub(crate) struct RunEndScalarFnRule;
+
+impl ArrayParentReduceRule<Exact<RunEndVTable>, AnyScalarFn> for RunEndScalarFnRule {
+    fn child(&self) -> Exact<RunEndVTable> {
+        Exact::from(&RunEndVTable)
+    }
+
+    fn parent(&self) -> AnyScalarFn {
+        AnyScalarFn
+    }
+
+    fn reduce_parent(
+        &self,
+        run_end: &RunEndArray,
+        parent: &ScalarFnArray,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        for (idx, child) in parent.children().iter().enumerate() {
+            if idx == child_idx {
+                // Skip ourselves
+                continue;
+            }
+
+            if !child.is::<ConstantVTable>() {
+                // We can only push down if all other children are constants
+                return Ok(None);
+            }
+        }
+
+        // TODO(ngates): relax this constraint and implement run-end decoding for all vector types.
+        if !matches!(parent.dtype(), DType::Bool(_) | DType::Primitive(..)) {
+            return Ok(None);
+        }
+
+        let values_len = run_end.values().len();
+        let mut new_children = parent.children();
+        for (idx, child) in new_children.iter_mut().enumerate() {
+            if idx == child_idx {
+                // Replace ourselves with run end values
+                *child = run_end.values().clone();
+                continue;
+            }
+
+            // Replace other children with their constant scalar value with length adjusted
+            // to the length of the run end values.
+            let constant = child.as_::<ConstantVTable>();
+            *child = ConstantArray::new(constant.scalar().clone(), values_len).into_array();
+        }
+
+        let new_values =
+            ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children, values_len)?
+                .into_array();
+
+        Ok(Some(
+            RunEndArray::try_new_offset_length(
+                run_end.ends().clone(),
+                new_values,
+                run_end.offset(),
+                run_end.len(),
+            )?
+            .into_array(),
+        ))
+    }
+}

encodings/runend/src/rules.rs

@@ -44,4 +44,9 @@ impl ScalarFnArray {
             stats: Default::default(),
         })
     }
+
+    /// Get the scalar function bound to this array.
+    pub fn scalar_fn(&self) -> &ScalarFn {
+        &self.scalar_fn
+    }
 }

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

@@ -17,6 +17,6 @@ pub struct ScalarFnMetadata {
 // Array tree display wrongly uses debug...
 impl Debug for ScalarFnMetadata {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.scalar_fn)
+        write!(f, "{}", self.scalar_fn.options())
     }
 }

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

@@ -4,7 +4,7 @@
 mod array;
 mod kernel;
 mod metadata;
-mod rules;
+pub(crate) mod rules;
 mod vtable;
 
 pub use array::*;

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

@@ -1,10 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
 use vortex_vector::Datum;
+use vortex_vector::VectorOps;
 use vortex_vector::scalar_matches_dtype;
 
 use crate::Array;
@@ -37,17 +37,23 @@ impl ArrayReduceRule<AnyScalarFn> for ScalarFnConstantRule {
             .collect();
         let input_dtypes = array.children.iter().map(|c| c.dtype().clone()).collect();
 
-        let result = array
-            .scalar_fn
-            .execute(ExecutionArgs {
-                datums: input_datums,
-                dtypes: input_dtypes,
-                row_count: array.len,
-                return_dtype: array.dtype.clone(),
-            })?
-            .into_scalar()
-            .vortex_expect("Scalar inputs should produce scalar output");
+        let result = array.scalar_fn.execute(ExecutionArgs {
+            datums: input_datums,
+            dtypes: input_dtypes,
+            row_count: array.len,
+            return_dtype: array.dtype.clone(),
+        })?;
 
+        let result = match result {
+            Datum::Scalar(s) => s,
+            Datum::Vector(v) => {
+                log::warn!(
+                    "Scalar function {} returned vector from execution over all scalar inputs",
+                    array.scalar_fn
+                );
+                v.scalar_at(0)
+            }
+        };
         assert!(scalar_matches_dtype(&result, &array.dtype));
 
         Ok(Some(

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

@@ -136,6 +136,7 @@ impl VTable for Binary {
 
         match op {
             Operator::And => {
+                // FIXME(ngates): implement logical compute over datums
                 let lhs = lhs.ensure_vector(args.row_count).into_bool().into_arrow()?;
                 let rhs = rhs.ensure_vector(args.row_count).into_bool().into_arrow()?;
                 return Ok(Datum::Vector(
@@ -145,6 +146,7 @@ impl VTable for Binary {
                 ));
             }
             Operator::Or => {
+                // FIXME(ngates): implement logical compute over datums
                 let lhs = lhs.ensure_vector(args.row_count).into_bool().into_arrow()?;
                 let rhs = rhs.ensure_vector(args.row_count).into_bool().into_arrow()?;
                 return Ok(Datum::Vector(

vortex-array/src/expr/exprs/binary.rs

@@ -23,6 +23,7 @@ use crate::arrays::StructGetItemRule;
 use crate::arrays::StructVTable;
 use crate::arrays::VarBinVTable;
 use crate::arrays::VarBinViewVTable;
+use crate::arrays::rules::ScalarFnConstantRule;
 use crate::optimizer::ArrayOptimizer;
 use crate::vtable::ArrayVTable;
 use crate::vtable::ArrayVTableExt;
@@ -94,6 +95,9 @@ impl Default for ArraySession {
         };
 
         let optimizer = session.optimizer_mut();
+
+        optimizer.register_reduce_rule(ScalarFnConstantRule);
+
         optimizer.register_parent_rule(BoolMaskedValidityRule);
         optimizer.register_parent_rule(PrimitiveMaskedValidityRule);
         optimizer.register_parent_rule(DecimalMaskedValidityRule);