diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index 3e625324a18..64be1e281d4 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -210,6 +210,7 @@ mod tests { use vortex_error::vortex_err; use crate::dtype::DType; + use crate::dtype::FieldName; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; @@ -218,8 +219,13 @@ mod tests { use crate::expr::get_item; use crate::expr::lit; use crate::expr::lt_eq; + use crate::expr::merge; use crate::expr::or; + use crate::expr::pack; use crate::expr::root; + use crate::expr::select; + use crate::expr::transform::replace; + use crate::expr::transform::replace_root_fields; use crate::scalar::Scalar; use crate::scalar_fn::fns::literal::Literal; @@ -272,4 +278,78 @@ mod tests { assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable)); Ok(()) } + + /// `Select` over a `pack` rewrites to a `pack` of the selected child fields. It must read + /// those fields out of the child `pack` itself, because the optimizer visits children before + /// their parent and so never revisits the children a rule introduces. + #[test] + fn optimize_select_of_pack_reads_pack_fields() -> VortexResult<()> { + let field_names = (0..4) + .map(|idx| FieldName::from(format!("field_{idx}"))) + .collect::>(); + let fields = StructFields::new( + field_names.clone().into(), + vec![DType::Primitive(PType::U64, Nullability::NonNullable); field_names.len()], + ); + let scope = DType::Struct(fields.clone(), Nullability::NonNullable); + + let expr = select( + field_names + .iter() + .cloned() + .chain(["input_row_idx".into(), "input_file_idx".into()]) + .collect::>(), + merge([ + root(), + pack( + [ + ("input_row_idx", lit(0_u64)), + ("input_file_idx", lit(1_u64)), + ], + Nullability::NonNullable, + ), + ]), + ); + let expr = replace(expr, &root(), replace_root_fields(root(), &fields)); + + let expected = pack( + field_names + .into_iter() + .map(|name| (name.clone(), get_item(name, root()))) + .chain([ + ("input_row_idx".into(), lit(0_u64)), + ("input_file_idx".into(), lit(1_u64)), + ]) + .collect::>(), + Nullability::NonNullable, + ); + + let optimized = expr.optimize_recursive(&scope)?; + assert_eq!(optimized, expected, "got {optimized}"); + assert_eq!(optimized.optimize_recursive(&scope)?, expected); + + Ok(()) + } + + /// `Merge` reduces to a `pack` that reads each field out of its children, so it must also + /// read straight out of a `pack` child. + #[test] + fn optimize_merge_of_packs_reads_pack_fields() -> VortexResult<()> { + let scope = DType::Struct(StructFields::empty(), Nullability::NonNullable); + let expr = merge([ + pack([("a", lit(0_u64))], Nullability::NonNullable), + pack([("b", lit(1_u64))], Nullability::NonNullable), + ]); + + let expected = pack( + [("a", lit(0_u64)), ("b", lit(1_u64))], + Nullability::NonNullable, + ); + + let optimized = expr.optimize_recursive(&scope)?; + assert_eq!(optimized, expected, "got {optimized}"); + assert_eq!(optimized.optimize_recursive(&scope)?, expected); + + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/fns/get_item.rs index 986dd620460..17927e7cbbd 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/fns/get_item.rs @@ -3,6 +3,7 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::slice; use prost::Message; use vortex_error::VortexResult; @@ -51,6 +52,57 @@ impl GetItem { ) -> VortexResult { ScalarFnArray::try_new(GetItem.bind(field_name.into()), vec![input]) } + + /// Read `field_name` out of `struct_node`, without a `get_item` node where possible. + /// + /// Rules that read a field out of their own child use this instead of building the + /// `get_item` node themselves. The optimizer visits children before their parent, so a + /// `get_item` that a rule introduces is never simplified again. Removing it here keeps the + /// output of the rule a fixed point. + pub(crate) fn project_node( + struct_node: &T, + field_name: &FieldName, + ) -> VortexResult { + if let Some(field) = Self::project_pack_node(struct_node, field_name)? { + return Ok(field); + } + + struct_node.new_node( + GetItem.bind(field_name.clone()), + slice::from_ref(struct_node), + ) + } + + /// Read `field_name` straight out of a `pack` node, which holds each field as a child. + /// + /// Returns `None` if `struct_node` is not a `pack`, or does not hold `field_name`. + fn project_pack_node( + struct_node: &T, + field_name: &FieldName, + ) -> VortexResult> { + let Some(pack) = struct_node + .scalar_fn() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + else { + return Ok(None); + }; + let Some(idx) = pack.names.find(field_name) else { + return Ok(None); + }; + + let mut field = struct_node.child(idx); + + // A nullable pack intersects its validity into the field, so mask the field with an + // all-true literal to make its dtype nullable. + if pack.nullability.is_nullable() { + field = struct_node.new_node( + Mask.bind(EmptyOptions), + &[field, struct_node.new_node(Literal.bind(true.into()), &[])?], + )?; + } + + Ok(Some(field)) + } } impl ScalarFnVTable for GetItem { @@ -136,25 +188,7 @@ impl ScalarFnVTable for GetItem { } fn reduce(&self, field_name: &FieldName, node: &T) -> VortexResult> { - let child = node.child(0); - if let Some(child_fn) = child.scalar_fn() - && let Some(pack) = child_fn.as_opt::() - && let Some(idx) = pack.names.find(field_name) - { - let mut field = child.child(idx); - - // Possibly mask the field if the pack is nullable - if pack.nullability.is_nullable() { - field = node.new_node( - Mask.bind(EmptyOptions), - &[field, node.new_node(Literal.bind(true.into()), &[])?], - )?; - } - - return Ok(Some(field)); - } - - Ok(None) + Self::project_pack_node(&node.child(0), field_name) } fn simplify_untyped( diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 178075e936c..8e20c9c2197 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -175,12 +175,12 @@ impl ScalarFnVTable for Merge { fn reduce(&self, options: &Self::Options, node: &T) -> VortexResult> { let mut names = Vec::with_capacity(node.child_count() * 2); - let mut children = Vec::with_capacity(node.child_count() * 2); + let mut fields: Vec = Vec::with_capacity(node.child_count() * 2); let mut duplicate_names = HashSet::<_>::new(); for child in (0..node.child_count()).map(|i| node.child(i)) { let child_dtype = child.node_dtype()?; - if !child_dtype.is_struct() { + if !child_dtype.is_struct() || child_dtype.is_nullable() { vortex_bail!( "Merge child must return a non-nullable struct dtype, got {}", child_dtype @@ -192,12 +192,18 @@ impl ScalarFnVTable for Merge { .vortex_expect("expected struct"); for name in child_dtype.names().iter() { + // `project_node` reads straight out of a `pack` child, so this reduction leaves + // behind no `get_item` for a later pass to simplify. The child is non-nullable, + // so the read cannot intersect a struct validity into the field, which is the + // only way it could change the field dtype. + let field = GetItem::project_node(&child, name)?; + if let Some(idx) = names.iter().position(|n| n == name) { duplicate_names.insert(name.clone()); - children[idx] = child.clone(); + fields[idx] = field; } else { names.push(name.clone()); - children.push(child.clone()); + fields.push(field); } } @@ -209,18 +215,12 @@ impl ScalarFnVTable for Merge { } } - let pack_children: Vec<_> = names - .iter() - .zip(children) - .map(|(name, child)| node.new_node(GetItem.bind(name.clone()), &[child])) - .try_collect()?; - let pack_expr = node.new_node( Pack.bind(PackOptions { names: FieldNames::from(names), nullability: node.node_dtype()?.nullability(), }), - &pack_children, + &fields, )?; Ok(Some(pack_expr)) diff --git a/vortex-array/src/scalar_fn/fns/select.rs b/vortex-array/src/scalar_fn/fns/select.rs index b0ac654a2f1..939fd0d0e99 100644 --- a/vortex-array/src/scalar_fn/fns/select.rs +++ b/vortex-array/src/scalar_fn/fns/select.rs @@ -6,6 +6,7 @@ use std::fmt::Formatter; use itertools::Itertools; use prost::Message; +use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -27,7 +28,6 @@ use crate::dtype::FieldNames; use crate::expr::display::ExprDisplay; use crate::expr::expression::Expression; use crate::expr::field::DisplayFieldNames; -use crate::expr::get_item; use crate::expr::pack; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -205,29 +205,41 @@ impl ScalarFnVTable for Select { return Ok(Some(pack(empty, struct_nullability))); } - // We cannot always convert a `select` into a `pack(get_item(f1), get_item(f2), ...)`. - // This is because `get_item` does a validity intersection of the struct validity with its - // fields, which is not the same as just "masking" out the unwanted fields (a selection). + // We cannot always convert a `select` into a projection of its child fields. This is + // because reading a field does a validity intersection of the struct validity with the + // field, which is not the same as just "masking" out the unwanted fields (a selection). // // We can, however, make this simplification when the child of the `select` is already a - // `pack` and we know that `get_item` will do no validity intersections. - let child_is_pack = child_struct.is::(); - - // `get_item` only performs validity intersection when the struct is nullable but the field - // is not. This would change the semantics of a `select`, so we can only simplify when this - // won't happen. + // `pack` and we know that reading a field will do no validity intersections. + // + // The intersection only happens when the struct is nullable but the field is not. That + // would change the semantics of a `select`, so we can only simplify when it won't happen. let would_intersect_validity = struct_nullability.is_nullable() && !all_included_fields_are_nullable; - if child_is_pack && !would_intersect_validity { - let pack_expr = pack( - included_fields - .into_iter() - .map(|name| (name.clone(), get_item(name, child_struct.clone()))), - struct_nullability, - ); - - return Ok(Some(pack_expr)); + if let Some(child_pack) = child_struct.as_opt::() + && !would_intersect_validity + { + // Take the field expressions straight out of the child `pack` rather than wrap the + // `pack` in `get_item`. The guard above proves the two are equivalent here, and the + // optimizer visits children before their parent, so a `get_item` introduced now + // would never be simplified away. + let fields: Vec<_> = included_fields + .iter() + .map(|name| { + let idx = child_pack.names.find(name).ok_or_else(|| { + vortex_err!( + "Cannot find field {} in pack fields {:?}", + name, + child_pack.names + ) + })?; + + Ok((name.clone(), child_struct.child(idx).clone())) + }) + .try_collect::<_, _, VortexError>()?; + + return Ok(Some(pack(fields, struct_nullability))); } Ok(None)