Clarify docs in session accessors

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

Author: gatesn

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

@@ -3,10 +3,11 @@
 
 mod array;
 pub use array::StructArray;
-
 mod compute;
+mod rules;
 
 mod vtable;
+pub use rules::*;
 pub use vtable::StructVTable;
 
 #[cfg(test)]

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

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::arrays::ConstantArray;
+use crate::arrays::ExactScalarFn;
+use crate::arrays::ScalarFnArrayExt;
+use crate::arrays::ScalarFnArrayView;
+use crate::arrays::StructArray;
+use crate::arrays::StructVTable;
+use crate::expr::EmptyOptions;
+use crate::expr::GetItem;
+use crate::expr::Mask;
+use crate::optimizer::rules::ArrayParentReduceRule;
+use crate::optimizer::rules::Exact;
+use crate::validity::Validity;
+use crate::vtable::ValidityHelper;
+use crate::Array;
+use crate::ArrayRef;
+use crate::IntoArray;
+
+#[derive(Debug)]
+pub struct StructGetItemRule;
+impl ArrayParentReduceRule<Exact<StructVTable>, ExactScalarFn<GetItem>> for StructGetItemRule {
+    fn child(&self) -> Exact<StructVTable> {
+        Exact::from(&StructVTable)
+    }
+
+    fn parent(&self) -> ExactScalarFn<GetItem> {
+        ExactScalarFn::from(&GetItem)
+    }
+
+    fn reduce_parent(
+        &self,
+        child: &StructArray,
+        parent: ScalarFnArrayView<'_, GetItem>,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let field_name = parent.options;
+        let Some(field) = child.field_by_name_opt(field_name) else {
+            return Ok(None);
+        };
+
+        match child.validity() {
+            Validity::NonNullable | Validity::AllValid => {
+                // If the struct is non-nullable or all valid, the field's validity is unchanged
+                Ok(Some(field.clone()))
+            }
+            Validity::AllInvalid => {
+                // If everything is invalid, the field is also all invalid
+                Ok(Some(
+                    ConstantArray::new(
+                        vortex_scalar::Scalar::null(field.dtype().clone()),
+                        field.len(),
+                    )
+                    .into_array(),
+                ))
+            }
+            Validity::Array(mask) => {
+                // If the validity is an array, we need to combine it with the field's validity
+                Mask.try_new_array(field.len(), EmptyOptions, [field.clone(), mask.clone()])
+                    .map(Some)
+            }
+        }
+    }
+}

vortex-array/src/executor.rs

@@ -1,18 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
+use vortex_error::VortexResult;
 use vortex_session::VortexSession;
 use vortex_vector::Datum;
 use vortex_vector::Vector;
 use vortex_vector::VectorOps;
 
-use crate::Array;
-use crate::ArrayRef;
 use crate::arrays::ConstantVTable;
 use crate::kernel::BindCtx;
 use crate::session::ArraySessionExt;
+use crate::Array;
+use crate::ArrayRef;
 
 /// Executor for exporting a Vortex [`Vector`] or [`Datum`] from an [`ArrayRef`].
 pub trait VectorExecutor {
@@ -36,6 +36,8 @@ impl VectorExecutor for ArrayRef {
     }
 
     fn execute_datum(&self, session: &VortexSession) -> VortexResult<Datum> {
+        log::error!("Executing array: {}", self.display_tree());
+
         // Attempt to short-circuit constant arrays.
         if let Some(constant) = self.as_opt::<ConstantVTable>() {
             return Ok(Datum::Scalar(constant.scalar().to_vector_scalar()));

vortex-array/src/expr/vtable.rs

@@ -433,6 +433,15 @@ impl<V: VTable> DynExprVTable for VTableAdapter<V> {
         {
             use vortex_error::vortex_ensure;
             use vortex_vector::datum_matches_dtype;
+
+            if !datum_matches_dtype(&result, &expected_dtype) {
+                vortex_bail!(
+                    "Expression execution returned datum of invalid dtype. Expected {}, got {:?}",
+                    expected_dtype,
+                    result
+                );
+            }
+
             vortex_ensure!(
                 datum_matches_dtype(&result, &expected_dtype),
                 "Expression execution invalid for dtype {}",

vortex-array/src/session/mod.rs

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_session::registry::Registry;
 use vortex_session::Ref;
 use vortex_session::SessionExt;
-use vortex_session::registry::Registry;
 
 use crate::arrays::BoolMaskedValidityRule;
 use crate::arrays::BoolVTable;
@@ -19,6 +19,7 @@ use crate::arrays::MaskedVTable;
 use crate::arrays::NullVTable;
 use crate::arrays::PrimitiveMaskedValidityRule;
 use crate::arrays::PrimitiveVTable;
+use crate::arrays::StructGetItemRule;
 use crate::arrays::StructVTable;
 use crate::arrays::VarBinVTable;
 use crate::arrays::VarBinViewVTable;
@@ -96,6 +97,7 @@ impl Default for ArraySession {
         optimizer.register_parent_rule(BoolMaskedValidityRule);
         optimizer.register_parent_rule(PrimitiveMaskedValidityRule);
         optimizer.register_parent_rule(DecimalMaskedValidityRule);
+        optimizer.register_parent_rule(StructGetItemRule);
 
         session
     }