feat(expr): partition struct validity as a first-class coordinate

vortex-array/src/expr/analysis/mod.rs

@@ -7,6 +7,7 @@ pub mod immediate_access;
 mod labeling;
 mod null_sensitive;
 mod referenced_field_paths;
+pub mod struct_part;
 
 pub use annotation::*;
 pub use fallible::label_is_fallible;
@@ -15,3 +16,4 @@ pub use labeling::*;
 pub use null_sensitive::BooleanLabels;
 pub use null_sensitive::label_null_sensitive;
 pub use referenced_field_paths::referenced_field_paths;
+pub use struct_part::*;

vortex-array/src/expr/analysis/struct_part.rs

@@ -0,0 +1,109 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt::Display;
+use std::fmt::Formatter;
+
+use vortex_error::VortexExpect;
+
+use crate::dtype::FieldName;
+use crate::dtype::Nullability;
+use crate::dtype::StructFields;
+use crate::expr::Expression;
+use crate::expr::analysis::AnnotationFn;
+use crate::scalar_fn::fns::get_item::GetItem;
+use crate::scalar_fn::fns::is_not_null::IsNotNull;
+use crate::scalar_fn::fns::root::Root;
+use crate::scalar_fn::fns::select::Select;
+
+/// A coordinate of a struct scope: either one of its fields, or — for a nullable struct —
+/// the struct's own validity.
+///
+/// A nullable struct with `n` fields decomposes into `n + 1` coordinates. Annotating
+/// expressions with `StructPart` instead of a bare [`FieldName`] lets the validity
+/// coordinate be routed like any field, and cannot collide with a real field that happens
+/// to be named "validity".
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum StructPart {
+    /// A top-level field of the struct scope.
+    Field(FieldName),
+    /// The validity of the struct scope itself, referenced as `is_not_null($)`.
+    Validity,
+}
+
+impl Display for StructPart {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            StructPart::Field(name) => write!(f, "{name}"),
+            StructPart::Validity => write!(f, "$validity"),
+        }
+    }
+}
+
+impl From<StructPart> for FieldName {
+    fn from(part: StructPart) -> Self {
+        match part {
+            StructPart::Field(name) => name,
+            // NOTE: this name is only used in the synthetic scope that recombines
+            // partitions. It is not distinguishable from a real field literally named
+            // "$validity"; routing must use the `StructPart` annotation, never this name.
+            StructPart::Validity => FieldName::from("$validity"),
+        }
+    }
+}
+
+/// Returns the free [`StructPart`] coordinates for an expression node.
+///
+/// This extends [`super::make_free_field_annotator`] with the validity coordinate of a
+/// nullable scope:
+///
+/// - **`is_not_null(root())`**: Returns `[Validity]` — the canonical reference to the
+///   scope's validity coordinate (it is also what [`crate::scalar_fn::fns::mask::Mask`]'s
+///   validity derivation produces for the expanded root, see
+///   [`crate::expr::transform::replace_root_scope`]).
+/// - **[`Select`] / [`GetItem`] on [`Root`]**: Returns the referenced fields, as
+///   `Field(..)`. These reference the field coordinates *without* the scope validity.
+/// - **[`Root`]**: Returns every field, plus `Validity` if the scope is nullable
+///   (conservative over-approximation: a bare `$` observes all coordinates).
+/// - **Everything else**: Returns empty (annotations aggregate from children).
+pub fn make_struct_part_annotator(
+    scope: &StructFields,
+    scope_nullability: Nullability,
+) -> impl AnnotationFn<Annotation = StructPart> {
+    move |expr: &Expression| {
+        if expr.is::<IsNotNull>() && expr.child(0).is::<Root>() {
+            if scope_nullability.is_nullable() {
+                return vec![StructPart::Validity];
+            }
+            // A non-nullable scope has no validity coordinate; `is_not_null($)` is a
+            // constant. Leave it unannotated so it does not force a partition.
+            return vec![];
+        } else if let Some(selection) = expr.as_opt::<Select>() {
+            if expr.child(0).is::<Root>() {
+                return selection
+                    .normalize_to_included_fields(scope.names())
+                    .vortex_expect("Select fields must be valid for scope")
+                    .into_iter()
+                    .map(StructPart::Field)
+                    .collect();
+            }
+        } else if let Some(field_name) = expr.as_opt::<GetItem>() {
+            if expr.child(0).is::<Root>() {
+                return vec![StructPart::Field(field_name.clone())];
+            }
+        } else if expr.is::<Root>() {
+            let mut parts: Vec<StructPart> = scope
+                .names()
+                .iter()
+                .cloned()
+                .map(StructPart::Field)
+                .collect();
+            if scope_nullability.is_nullable() {
+                parts.push(StructPart::Validity);
+            }
+            return parts;
+        }
+
+        vec![]
+    }
+}

vortex-array/src/expr/transform/partition.rs

@@ -34,11 +34,12 @@ use crate::expr::traversal::TraversalOrder;
 ///
 /// ## Note
 ///
-/// This function currently respects the validity of each field in the scope, but the not validity
-/// of the scope itself. The fix would be for the returned `PartitionedExpr` to include a partition
-/// expression for computing the validity, or to include that expression as part of the root.
-///
-/// See <https://github.com/vortex-data/vortex/issues/1907>.
+/// This function is agnostic to what the annotations denote. To partition over a *nullable*
+/// struct scope, expand the root with
+/// [`replace_root_scope`][crate::expr::transform::replace_root_scope] and annotate with
+/// [`make_struct_part_annotator`][crate::expr::analysis::make_struct_part_annotator], which
+/// treats the scope's own validity as one more coordinate alongside its fields
+/// (see <https://github.com/vortex-data/vortex/issues/1907>).
 pub fn partition<A: AnnotationFn>(
     expr: Expression,
     scope: &DType,
@@ -217,21 +218,28 @@ where
 mod tests {
     use rstest::fixture;
     use rstest::rstest;
+    use vortex_utils::aliases::hash_set::HashSet;
 
     use super::*;
     use crate::dtype::DType;
     use crate::dtype::Nullability::NonNullable;
+    use crate::dtype::Nullability::Nullable;
     use crate::dtype::PType::I32;
     use crate::dtype::StructFields;
     use crate::expr::analysis::make_free_field_annotator;
+    use crate::expr::analysis::make_struct_part_annotator;
+    use crate::expr::analysis::struct_part::StructPart;
     use crate::expr::and;
     use crate::expr::col;
     use crate::expr::get_item;
+    use crate::expr::gt;
+    use crate::expr::is_not_null;
     use crate::expr::lit;
     use crate::expr::merge;
     use crate::expr::pack;
     use crate::expr::root;
     use crate::expr::transform::replace::replace_root_fields;
+    use crate::expr::transform::replace::replace_root_scope;
 
     #[fixture]
     fn dtype() -> DType {
@@ -329,6 +337,33 @@ mod tests {
         assert_eq!(partitioned.partitions.len(), 2);
     }
 
+    #[rstest]
+    fn test_partition_validity_coordinate() -> VortexResult<()> {
+        // A nullable struct scope has a validity coordinate alongside its fields. After
+        // expanding the root, a validity observer partitions to the validity coordinate and
+        // a field predicate partitions to the field — the validity is never guessed at a
+        // fixed position.
+        let dtype = DType::Struct(
+            StructFields::from_iter([("a", I32.into()), ("b", DType::from(I32))]),
+            Nullable,
+        );
+        let fields = dtype.as_struct_fields_opt().unwrap();
+
+        let expr = and(is_not_null(root()), gt(get_item("a", root()), lit(5)));
+        let expr = replace_root_scope(expr, &dtype).optimize_recursive(&dtype)?;
+
+        let partitioned = partition(expr, &dtype, make_struct_part_annotator(fields, Nullable))?;
+
+        let annotations: HashSet<StructPart> =
+            partitioned.partition_annotations.iter().cloned().collect();
+        assert_eq!(
+            annotations,
+            HashSet::from_iter([StructPart::Validity, StructPart::Field("a".into())]),
+            "{partitioned}"
+        );
+        Ok(())
+    }
+
     #[rstest]
     fn test_expr_merge(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();

vortex-array/src/expr/transform/replace.rs

@@ -3,10 +3,13 @@
 
 use vortex_error::VortexExpect;
 
+use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::StructFields;
 use crate::expr::Expression;
 use crate::expr::col;
+use crate::expr::is_not_null;
+use crate::expr::mask;
 use crate::expr::pack;
 use crate::expr::root;
 use crate::expr::traversal::NodeExt;
@@ -33,26 +36,61 @@ pub fn replace(expr: Expression, needle: &Expression, replacement: Expression) -
 
 /// Expand the `root` expression with a pack of the given struct fields.
 pub fn replace_root_fields(expr: Expression, fields: &StructFields) -> Expression {
-    replace(
-        expr,
-        &root(),
-        pack(
-            fields
-                .names()
-                .iter()
-                .map(|name| (name.clone(), col(name.clone()))),
-            Nullability::NonNullable,
-        ),
+    replace(expr, &root(), root_fields_expansion(fields))
+}
+
+/// Expand the `root` expression of a struct scope into an expression over its coordinates.
+///
+/// For a non-nullable struct scope this is a `pack` of every field, as in
+/// [`replace_root_fields`]. For a nullable struct scope, the scope has one more coordinate
+/// than it has fields — its own validity — so the expansion is the identity
+///
+/// ```text
+/// $ == mask(pack(f1: $.f1, ..., fn: $.fn), is_not_null($))
+/// ```
+///
+/// where the surviving `$.f` references denote the fields *without* the struct's own
+/// validity applied (`pack` reassembles the values, `mask` re-applies the struct validity).
+/// This keeps the validity coordinate in the expression term, so downstream analyses
+/// (e.g. partitioning) can route it like any other coordinate instead of re-applying it
+/// out-of-band.
+pub fn replace_root_scope(expr: Expression, scope: &DType) -> Expression {
+    let fields = scope
+        .as_struct_fields_opt()
+        .vortex_expect("replace_root_scope requires a struct scope");
+    let expansion = match scope.nullability() {
+        Nullability::NonNullable => root_fields_expansion(fields),
+        Nullability::Nullable => mask(root_fields_expansion(fields), is_not_null(root())),
+    };
+    replace(expr, &root(), expansion)
+}
+
+fn root_fields_expansion(fields: &StructFields) -> Expression {
+    pack(
+        fields
+            .names()
+            .iter()
+            .map(|name| (name.clone(), col(name.clone()))),
+        Nullability::NonNullable,
     )
 }
 
 #[cfg(test)]
 mod test {
     use super::replace;
+    use super::replace_root_scope;
+    use crate::dtype::DType;
     use crate::dtype::Nullability::NonNullable;
+    use crate::dtype::Nullability::Nullable;
+    use crate::dtype::PType::I32;
+    use crate::dtype::StructFields;
+    use crate::expr::col;
     use crate::expr::get_item;
+    use crate::expr::is_not_null;
     use crate::expr::lit;
+    use crate::expr::mask;
     use crate::expr::pack;
+    use crate::expr::root;
 
     #[test]
     fn test_replace_full_tree() {
@@ -71,4 +109,30 @@ mod test {
         let replaced_expr = replace(e, &needle, replacement);
         assert_eq!(replaced_expr.to_string(), "pack(a: 1i32, b: 42i32)");
     }
+
+    #[test]
+    fn test_replace_root_scope_non_nullable() {
+        let dtype = DType::Struct(
+            StructFields::from_iter([("a", I32.into()), ("b", DType::from(I32))]),
+            NonNullable,
+        );
+        let expanded = replace_root_scope(root(), &dtype);
+        let expected = pack([("a", col("a")), ("b", col("b"))], NonNullable);
+        assert_eq!(&expanded, &expected);
+    }
+
+    #[test]
+    fn test_replace_root_scope_nullable() {
+        let dtype = DType::Struct(
+            StructFields::from_iter([("a", I32.into()), ("b", DType::from(I32))]),
+            Nullable,
+        );
+        let expanded = replace_root_scope(root(), &dtype);
+        // A nullable struct scope has n + 1 coordinates: its fields, and its own validity.
+        let expected = mask(
+            pack([("a", col("a")), ("b", col("b"))], NonNullable),
+            is_not_null(root()),
+        );
+        assert_eq!(&expanded, &expected);
+    }
 }

vortex-array/src/scalar_fn/erased.rs

@@ -138,6 +138,15 @@ impl ScalarFnRef {
         }))
     }
 
+    /// Like [`Self::validity`], but returns `None` when the vtable does not define a validity
+    /// derivation, instead of falling back to `is_not_null(expr)`.
+    ///
+    /// Rewrites that turn `is_not_null(expr)` into the derived validity must use this method:
+    /// the fallback would make such a rewrite a self-referential no-op.
+    pub fn validity_opt(&self, expr: &Expression) -> VortexResult<Option<Expression>> {
+        self.0.validity(expr)
+    }
+
     /// Execute the expression given the input arguments.
     pub fn execute(
         &self,

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

@@ -154,6 +154,15 @@ impl ScalarFnVTable for GetItem {
     ) -> VortexResult<Option<Expression>> {
         let child = expr.child(0);
 
+        // `get_item` distributes over `mask`: `get_item(f, mask(s, m)) == mask(get_item(f, s), m)`
+        // (a field of a masked struct is the field masked by the same mask). This exposes the
+        // field access to a `pack` beneath the mask, letting it collapse to a single coordinate.
+        if child.is::<Mask>() {
+            let input = child.child(0).clone();
+            let m = child.child(1).clone();
+            return Ok(Some(GetItem.new_expr(field_name.clone(), [input]).mask(m)?));
+        }
+
         // If the child is a Pack expression, we can directly return the corresponding child.
         if let Some(pack) = child.as_opt::<Pack>() {
             let idx = pack
@@ -258,6 +267,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_get_item_distributes_over_mask() -> vortex_error::VortexResult<()> {
+        use crate::expr::mask;
+        use crate::expr::test_harness;
+
+        // get_item(a, mask(pack(..), m)) == mask(get_item(a, pack(..)), m) == mask(a-child, m)
+        let m = get_item("bool1", root());
+        let expr = get_item(
+            "a",
+            mask(pack([("a", lit(1)), ("b", lit(2))], NonNullable), m.clone()),
+        );
+        let simplified = expr.optimize_recursive(&test_harness::struct_dtype())?;
+        assert_eq!(&simplified, &mask(lit(1), m));
+        Ok(())
+    }
+
     #[test]
     fn test_pack_get_item_rule() {
         // Create: pack(a: lit(1), b: lit(2)).get_item("b")