diff --git a/vortex-array/src/expr/analysis/mod.rs b/vortex-array/src/expr/analysis/mod.rs index f5208a31be8..c181b52d18c 100644 --- a/vortex-array/src/expr/analysis/mod.rs +++ b/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::*; diff --git a/vortex-array/src/expr/analysis/struct_part.rs b/vortex-array/src/expr/analysis/struct_part.rs new file mode 100644 index 00000000000..7195cb2745e --- /dev/null +++ b/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 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 { + move |expr: &Expression| { + if expr.is::() && expr.child(0).is::() { + 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::