Fix Arrow struct export (#5731)

And...
* Adds unary filter push-down rule
* Fixes ordering of filter + project in layout reader

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

Author: gatesn

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

@@ -31,7 +31,8 @@ impl FilterArray {
         }
     }
 
-    pub fn mask(&self) -> &Mask {
+    /// The mask used to filter the child array.
+    pub fn filter_mask(&self) -> &Mask {
         &self.mask
     }
 }

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

@@ -19,21 +19,26 @@ use crate::ArrayVisitor;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
 use crate::arrays::ConstantVTable;
+use crate::arrays::FilterArray;
+use crate::arrays::FilterVTable;
 use crate::arrays::ScalarFnArray;
 use crate::arrays::ScalarFnVTable;
 use crate::expr::ExecutionArgs;
 use crate::expr::ReduceCtx;
 use crate::expr::ReduceNode;
 use crate::expr::ReduceNodeRef;
 use crate::expr::ScalarFn;
+use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::ArrayReduceRule;
+use crate::optimizer::rules::Exact;
 use crate::optimizer::rules::ParentRuleSet;
 use crate::optimizer::rules::ReduceRuleSet;
 
 pub(super) const RULES: ReduceRuleSet<ScalarFnVTable> =
     ReduceRuleSet::new(&[&ScalarFnConstantRule, &ScalarFnAbstractReduceRule]);
 
-pub(super) const PARENT_RULES: ParentRuleSet<ScalarFnVTable> = ParentRuleSet::new(&[]);
+pub(super) const PARENT_RULES: ParentRuleSet<ScalarFnVTable> =
+    ParentRuleSet::new(&[ParentRuleSet::lift(&ScalarFnUnaryFilterPushDownRule)]);
 
 #[derive(Debug)]
 struct ScalarFnConstantRule;
@@ -152,3 +157,50 @@ impl ReduceCtx for ArrayReduceCtx {
         ))
     }
 }
+
+#[derive(Debug)]
+struct ScalarFnUnaryFilterPushDownRule;
+
+impl ArrayParentReduceRule<ScalarFnVTable> for ScalarFnUnaryFilterPushDownRule {
+    type Parent = Exact<FilterVTable>;
+
+    fn parent(&self) -> Self::Parent {
+        Exact::from(&FilterVTable)
+    }
+
+    fn reduce_parent(
+        &self,
+        child: &ScalarFnArray,
+        parent: &FilterArray,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        // If we only have one non-constant child, then it is _always_ cheaper to push down the
+        // filter over the children of the scalar function array.
+        if child
+            .children
+            .iter()
+            .filter(|c| !c.is::<ConstantVTable>())
+            .count()
+            == 1
+        {
+            let new_children: Vec<_> = child
+                .children
+                .iter()
+                .map(|c| match c.as_opt::<ConstantVTable>() {
+                    Some(array) => {
+                        ConstantArray::new(array.scalar().clone(), parent.len()).into_array()
+                    }
+                    None => FilterArray::new(c.clone(), parent.filter_mask().clone()).into_array(),
+                })
+                .collect();
+
+            let new_array =
+                ScalarFnArray::try_new(child.scalar_fn.clone(), new_children, parent.len())?
+                    .into_array();
+
+            return Ok(Some(new_array));
+        }
+
+        Ok(None)
+    }
+}

vortex-array/src/arrow/executor/struct_.rs

@@ -6,15 +6,20 @@ use std::sync::Arc;
 use arrow_array::ArrayRef as ArrowArrayRef;
 use arrow_array::StructArray;
 use arrow_buffer::NullBuffer;
+use arrow_schema::DataType;
 use arrow_schema::Fields;
-use itertools::Itertools;
+use vortex_compute::arrow::IntoArrow;
+use vortex_dtype::DType;
+use vortex_dtype::StructFields;
+use vortex_dtype::arrow::FromArrowType;
+use vortex_error::VortexError;
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_session::VortexSession;
 
 use crate::Array;
 use crate::ArrayRef;
+use crate::VectorExecutor;
 use crate::arrays::ScalarFnVTable;
 use crate::arrays::StructVTable;
 use crate::arrow::ArrowArrayExecutor;
@@ -52,28 +57,19 @@ pub(super) fn to_arrow_struct(
         );
     }
 
-    // Otherwise, we have some options:
-    //  1. Use get_item expression to extract each field? This is a bit sad because get_item
-    //     will perform the validity masking again.
-    //  2. Execute a full struct vector. But this may do unnecessary work on fields that may
-    //    have a more direct conversion to the desired Arrow field type.
-    //  3. Something else?
-    //
-    // For now, we go with option 1. Although we really ought to figure out CSE for this.
-    let field_arrays = fields
-        .iter()
-        .map(|f| array.get_item(f.name().as_str()))
-        .try_collect()?;
+    // Otherwise, we fall back to executing the full struct vector.
+    // First we apply a cast to ensure we push down casting where possible into the struct fields.
+    let vx_fields = StructFields::from_arrow(fields);
+    let array = array.cast(DType::Struct(
+        vx_fields,
+        vortex_dtype::Nullability::Nullable,
+    ))?;
 
-    if !array.all_valid() {
-        // TODO(ngates): we should grab the nullability using the is_not_null expression.
-        vortex_bail!(
-            "Cannot convert nullable Struct array with nulls to Arrow\n{}",
-            array.display_tree()
-        );
-    }
+    let struct_array = array.execute_vector(session)?.into_struct().into_arrow()?;
 
-    create_from_fields(fields, field_arrays, None, len, session)
+    // Finally, we cast to Arrow to ensure any types not representable by Vortex (e.g. Dictionary)
+    // are properly converted.
+    arrow_cast::cast(&struct_array, &DataType::Struct(fields.clone())).map_err(VortexError::from)
 }
 
 fn create_from_fields(

vortex-layout/src/layouts/flat/reader.rs

@@ -219,14 +219,17 @@ impl LayoutReader for FlatReader {
             }
 
             Ok(if *USE_VORTEX_OPERATORS {
-                // Evaluate the projection expression.
-                array = array.apply(&expr)?;
-
-                // Filter the array based on the row mask.
+                // First apply the filter to the array.
+                // NOTE(ngates): we *must* filter first before applying the expression, as the
+                // expression may depend on the filtered rows being removed e.g.
+                //  `CAST(a, u8) WHERE a < 256`
                 if !mask.all_true() {
                     array = array.filter(mask)?;
                 }
 
+                // Evaluate the projection expression.
+                array = array.apply(&expr)?;
+
                 array
             } else {
                 // Filter the array based on the row mask.