From 7124a9d40a74a8cae99a98c0b063075fe87a1c1f Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 3 Aug 2026 18:04:50 +0100 Subject: [PATCH] refactor: partition bound expressions Signed-off-by: Joe Isaacs --- vortex-array/src/expr/analysis/annotation.rs | 160 +++++- .../src/expr/analysis/immediate_access.rs | 30 +- vortex-array/src/expr/analysis/labeling.rs | 101 +++- .../src/expr/transform/bound_partition.rs | 471 ++++++++++++++++++ vortex-array/src/expr/transform/mod.rs | 4 +- vortex-array/src/expr/traversal/mod.rs | 70 +++ 6 files changed, 798 insertions(+), 38 deletions(-) create mode 100644 vortex-array/src/expr/transform/bound_partition.rs diff --git a/vortex-array/src/expr/analysis/annotation.rs b/vortex-array/src/expr/analysis/annotation.rs index 32608f38d1e..f2c63d7c4d7 100644 --- a/vortex-array/src/expr/analysis/annotation.rs +++ b/vortex-array/src/expr/analysis/annotation.rs @@ -8,7 +8,10 @@ use vortex_error::VortexResult; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; +use crate::expr::BoundExpression; +use crate::expr::ExactBoundExpr; use crate::expr::Expression; +use crate::expr::traversal::Node; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeVisitor; use crate::expr::traversal::TraversalOrder; @@ -17,28 +20,38 @@ pub trait Annotation: Clone + Hash + Eq {} impl Annotation for A where A: Clone + Hash + Eq {} -pub trait AnnotationFn: Fn(&Expression) -> Vec { +pub trait AnnotationFn: Fn(&N) -> Vec { type Annotation: Annotation; } -impl AnnotationFn for F +impl AnnotationFn for F where A: Annotation, - F: Fn(&Expression) -> Vec, + F: Fn(&N) -> Vec, { type Annotation = A; } -pub type Annotations<'a, A> = HashMap<&'a Expression, HashSet>; +pub type Annotations<'a, A, N = Expression> = HashMap<&'a N, HashSet>; + +/// Annotations keyed by bound-tree identity. +/// +/// Identity keys avoid structurally hashing every node's dtype. That matters when a bound root +/// carries a lazy schema whose structural hash would deserialize every field. +pub type BoundAnnotations = HashMap>; /// Walk the expression tree and annotate each expression with zero or more annotations. /// /// Returns a map of each expression to all annotations that any of its descendent (child) /// expressions are annotated with. -pub fn descendent_annotations( - expr: &Expression, +pub fn descendent_annotations<'a, N, A>( + expr: &'a N, annotate: A, -) -> Annotations<'_, A::Annotation> { +) -> Annotations<'a, A::Annotation, N> +where + N: Node + Eq + Hash, + A: AnnotationFn, +{ let mut visitor = AnnotationVisitor { annotations: Default::default(), annotate, @@ -53,10 +66,11 @@ pub fn descendent_annotations( /// /// Returns a map of each expression to all annotations. Annotations of /// children are not propagated to parents. -pub fn direct_annotations( - expr: &Expression, - annotate: A, -) -> Annotations<'_, A::Annotation> { +pub fn direct_annotations<'a, N, A>(expr: &'a N, annotate: A) -> Annotations<'a, A::Annotation, N> +where + N: Node + Eq + Hash, + A: AnnotationFn, +{ let mut visitor = AnnotationVisitor { annotations: Default::default(), annotate, @@ -66,14 +80,66 @@ pub fn direct_annotations( visitor.annotations } -struct AnnotationVisitor<'a, A: AnnotationFn> { - annotations: Annotations<'a, A::Annotation>, +/// Annotate a bound expression and propagate each annotation to its ancestors. +/// +/// Unlike [`descendent_annotations`], this uses [`ExactBoundExpr`] keys to preserve the cheap +/// identity semantics of an already-bound tree. +pub fn descendent_bound_annotations( + expr: &BoundExpression, + annotate: A, +) -> BoundAnnotations +where + A: AnnotationFn, +{ + bound_annotations(expr, annotate, true) +} + +/// Annotate each bound-expression node without propagating annotations to its ancestors. +/// +/// The returned map uses [`ExactBoundExpr`] keys so lookups do not structurally hash node dtypes. +pub fn direct_bound_annotations( + expr: &BoundExpression, + annotate: A, +) -> BoundAnnotations +where + A: AnnotationFn, +{ + bound_annotations(expr, annotate, false) +} + +fn bound_annotations( + expr: &BoundExpression, + annotate: A, + propagate_up: bool, +) -> BoundAnnotations +where + A: AnnotationFn, +{ + let mut visitor = BoundAnnotationVisitor { + annotations: Default::default(), + annotate, + propagate_up, + }; + expr.accept(&mut visitor).vortex_expect("Infallible"); + visitor.annotations +} + +struct AnnotationVisitor<'a, N, A> +where + N: Node + Eq + Hash, + A: AnnotationFn, +{ + annotations: Annotations<'a, A::Annotation, N>, annotate: A, propagate_up: bool, } -impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> { - type NodeTy = Expression; +impl<'a, N, A> NodeVisitor<'a> for AnnotationVisitor<'a, N, A> +where + N: Node + Eq + Hash, + A: AnnotationFn, +{ + type NodeTy = N; fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult { let annotations = (self.annotate)(node); @@ -89,20 +155,74 @@ impl<'a, A: AnnotationFn> NodeVisitor<'a> for AnnotationVisitor<'a, A> { } } - fn visit_up(&mut self, node: &'a Expression) -> VortexResult { + fn visit_up(&mut self, node: &'a N) -> VortexResult { if !self.propagate_up { return Ok(TraversalOrder::Continue); } + let child_annotations = node.iter_children(|children| { + children + .filter_map(|child| self.annotations.get(child).cloned()) + .collect::>() + }); + + let annotations = self.annotations.entry(node).or_default(); + child_annotations + .into_iter() + .for_each(|ps| annotations.extend(ps.iter().cloned())); + + Ok(TraversalOrder::Continue) + } +} + +struct BoundAnnotationVisitor +where + A: AnnotationFn, +{ + annotations: BoundAnnotations, + annotate: A, + propagate_up: bool, +} + +impl<'a, A> NodeVisitor<'a> for BoundAnnotationVisitor +where + A: AnnotationFn, +{ + type NodeTy = BoundExpression; + + fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult { + let annotations = (self.annotate)(node); + if annotations.is_empty() { + return Ok(TraversalOrder::Continue); + } + + self.annotations + .entry(ExactBoundExpr(node.clone())) + .or_default() + .extend(annotations); + Ok(TraversalOrder::Skip) + } + + fn visit_up(&mut self, node: &'a Self::NodeTy) -> VortexResult { + if !self.propagate_up { + return Ok(TraversalOrder::Continue); + } + let child_annotations = node .children() .iter() - .filter_map(|c| self.annotations.get(c).cloned()) + .filter_map(|child| { + self.annotations + .get(&ExactBoundExpr(child.clone())) + .cloned() + }) .collect::>(); - - let annotations = self.annotations.entry(node).or_default(); + let annotations = self + .annotations + .entry(ExactBoundExpr(node.clone())) + .or_default(); child_annotations .into_iter() - .for_each(|ps| annotations.extend(ps.iter().cloned())); + .for_each(|child| annotations.extend(child)); Ok(TraversalOrder::Continue) } diff --git a/vortex-array/src/expr/analysis/immediate_access.rs b/vortex-array/src/expr/analysis/immediate_access.rs index 805400be793..8df9e056f7e 100644 --- a/vortex-array/src/expr/analysis/immediate_access.rs +++ b/vortex-array/src/expr/analysis/immediate_access.rs @@ -6,6 +6,7 @@ use vortex_utils::aliases::hash_set::HashSet; use crate::dtype::FieldName; use crate::dtype::StructFields; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::analysis::AnnotationFn; use crate::expr::analysis::Annotations; @@ -43,7 +44,7 @@ pub type FieldAccesses<'a> = Annotations<'a, FieldName>; /// - The full expression has free fields `{a, d}` (not `b`, only top-level fields are tracked). pub fn make_free_field_annotator( scope: &StructFields, -) -> impl AnnotationFn { +) -> impl AnnotationFn { move |expr: &Expression| { if let Some(selection) = expr.as_opt::() { + if expr.children()[0].is_root() { + return selection + .normalize_to_included_fields(scope.names()) + .vortex_expect("Select fields must be valid for scope") + .into_iter() + .collect(); + } + } else if let Some(field_name) = scalar_fn.as_opt::() + && expr.children()[0].is_root() + { + return vec![field_name.clone()]; + } + + vec![] + } +} + /// For all subexpressions in an expression, find the fields that are accessed directly from the /// scope, but not any fields in those fields /// e.g. scope = {a: {b: .., c: ..}, d: ..}, expr = root().a.b + root().d accesses {a,d} (not b). diff --git a/vortex-array/src/expr/analysis/labeling.rs b/vortex-array/src/expr/analysis/labeling.rs index 85bb7e5e5f9..51957f26190 100644 --- a/vortex-array/src/expr/analysis/labeling.rs +++ b/vortex-array/src/expr/analysis/labeling.rs @@ -1,17 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::hash::Hash; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_utils::aliases::hash_map::HashMap; +use crate::expr::BoundExpression; +use crate::expr::ExactBoundExpr; use crate::expr::Expression; +use crate::expr::traversal::Node; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeVisitor; use crate::expr::traversal::TraversalOrder; /// Boolean labels keyed by each expression node in a tree. -pub type BooleanLabels<'a> = HashMap<&'a Expression, bool>; +pub type BooleanLabels<'a, N = Expression> = HashMap<&'a N, bool>; + +/// Labels keyed by bound-tree identity. +pub type BoundLabels = HashMap; /// Label each node in an expression tree using a bottom-up traversal. /// @@ -32,11 +40,14 @@ pub type BooleanLabels<'a> = HashMap<&'a Expression, bool>; /// - `merge_child`: Mutable function that folds child labels into an accumulator. /// Takes `(self_label, child_label)` and returns the updated accumulator. /// Called once per child, with the initial accumulator being the node's self-label. -pub fn label_tree( - expr: &Expression, - self_label: impl Fn(&Expression) -> L, +pub fn label_tree( + expr: &N, + self_label: impl Fn(&N) -> L, mut merge_child: impl FnMut(L, &L) -> L, -) -> HashMap<&Expression, L> { +) -> HashMap<&N, L> +where + N: Node + Eq + Hash, +{ let mut visitor = LabelingVisitor { labels: Default::default(), self_label, @@ -47,40 +58,98 @@ pub fn label_tree( visitor.labels } -struct LabelingVisitor<'a, 'b, L, F, G> +/// Label each node in a bound expression using identity-keyed lookups. +/// +/// This avoids structurally hashing bound dtypes, which may deserialize a lazy schema. +pub fn label_bound_tree( + expr: &BoundExpression, + self_label: impl Fn(&BoundExpression) -> L, + mut merge_child: impl FnMut(L, &L) -> L, +) -> BoundLabels { + let mut visitor = BoundLabelingVisitor { + labels: Default::default(), + self_label, + merge_child: &mut merge_child, + }; + expr.accept(&mut visitor) + .vortex_expect("BoundLabelingVisitor is infallible"); + visitor.labels +} + +struct LabelingVisitor<'a, 'b, N, L, F, G> where - F: Fn(&Expression) -> L, + N: Node + Eq + Hash, + F: Fn(&N) -> L, G: FnMut(L, &L) -> L, { - labels: HashMap<&'a Expression, L>, + labels: HashMap<&'a N, L>, self_label: F, merge_child: &'b mut G, } -impl<'a, 'b, L: Clone, F, G> NodeVisitor<'a> for LabelingVisitor<'a, 'b, L, F, G> +impl<'a, 'b, N, L: Clone, F, G> NodeVisitor<'a> for LabelingVisitor<'a, 'b, N, L, F, G> where - F: Fn(&Expression) -> L, + N: Node + Eq + Hash, + F: Fn(&N) -> L, G: FnMut(L, &L) -> L, { - type NodeTy = Expression; + type NodeTy = N; fn visit_down(&mut self, _node: &'a Self::NodeTy) -> VortexResult { Ok(TraversalOrder::Continue) } - fn visit_up(&mut self, node: &'a Expression) -> VortexResult { + fn visit_up(&mut self, node: &'a N) -> VortexResult { let self_label = (self.self_label)(node); + let final_label = node.iter_children(|children| { + children.fold(self_label, |acc, child| { + let child_label = self + .labels + .get(child) + .vortex_expect("child must have label"); + (self.merge_child)(acc, child_label) + }) + }); + + self.labels.insert(node, final_label); + + Ok(TraversalOrder::Continue) + } +} + +struct BoundLabelingVisitor<'a, L, F, G> +where + F: Fn(&BoundExpression) -> L, + G: FnMut(L, &L) -> L, +{ + labels: BoundLabels, + self_label: F, + merge_child: &'a mut G, +} + +impl<'node, 'visitor, L: Clone, F, G> NodeVisitor<'node> for BoundLabelingVisitor<'visitor, L, F, G> +where + F: Fn(&BoundExpression) -> L, + G: FnMut(L, &L) -> L, +{ + type NodeTy = BoundExpression; + + fn visit_down(&mut self, _node: &'node Self::NodeTy) -> VortexResult { + Ok(TraversalOrder::Continue) + } + + fn visit_up(&mut self, node: &'node Self::NodeTy) -> VortexResult { + let self_label = (self.self_label)(node); let final_label = node.children().iter().fold(self_label, |acc, child| { let child_label = self .labels - .get(child) + .get(&ExactBoundExpr(child.clone())) .vortex_expect("child must have label"); (self.merge_child)(acc, child_label) }); - - self.labels.insert(node, final_label); - + self.labels + .insert(ExactBoundExpr(node.clone()), final_label); Ok(TraversalOrder::Continue) } } diff --git a/vortex-array/src/expr/transform/bound_partition.rs b/vortex-array/src/expr/transform/bound_partition.rs new file mode 100644 index 00000000000..9f1c59bcd6f --- /dev/null +++ b/vortex-array/src/expr/transform/bound_partition.rs @@ -0,0 +1,471 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::dtype::DType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::StructFields; +use crate::expr::BoundExpression; +use crate::expr::ExactBoundExpr; +use crate::expr::analysis::Annotation; +use crate::expr::analysis::AnnotationFn; +use crate::expr::analysis::BoundAnnotations; +use crate::expr::analysis::descendent_bound_annotations; +use crate::expr::traversal::NodeExt; +use crate::expr::traversal::NodeRewriter; +use crate::expr::traversal::Transformed; +use crate::expr::traversal::TraversalOrder; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::fns::pack::Pack; +use crate::scalar_fn::fns::pack::PackOptions; + +/// Partition an expression into sub-expressions that are uniquely associated with an annotation. +/// A root expression is also returned that can be used to recombine the results of the partitions +/// into the result of the original expression. +/// +/// ## 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 `BoundPartitionedExpr` to include a +/// partition expression for computing the validity, or to include that expression as part of the +/// root. +/// +/// See . +pub fn partition_bound>( + expr: BoundExpression, + annotate_fn: A, +) -> VortexResult> +where + A::Annotation: Display, + FieldName: From, +{ + // Annotate each expression with the annotations that any of its descendent expressions have. + let annotations = descendent_bound_annotations(&expr, annotate_fn); + partition_bound_annotations(expr, annotations) +} + +/// Partition an already-annotated bound expression tree. +/// +/// Prefer [`partition_bound`] when annotations can be derived by an [`AnnotationFn`]. +pub fn partition_bound_annotations( + expr: BoundExpression, + annotations: BoundAnnotations, +) -> VortexResult> +where + A: Display + Clone + Eq + Hash, + FieldName: From, +{ + let mut collector = PartitionCollector::::new(&annotations); + expr.clone().rewrite(&mut collector)?; + + let mut partitions = Vec::with_capacity(collector.sub_expressions.len()); + let mut partition_annotations = Vec::with_capacity(collector.sub_expressions.len()); + + for (annotation, exprs) in collector.sub_expressions { + // We pack all sub-expressions for the same annotation into a single expression. + let names = exprs + .iter() + .enumerate() + .map(|(idx, _)| PartitionCollector::field_name(&annotation, idx)) + .collect(); + let expr = bound_pack(names, exprs)?; + + partitions.push(expr); + partition_annotations.push(annotation); + } + + let partition_names = partition_annotations + .iter() + .map(|id| FieldName::from(id.clone())) + .collect::(); + let root_scope = partition_root_dtype(&partition_names, &partitions); + let mut rewriter = PartitionRootRewriter::new(&annotations, root_scope); + let root = expr.rewrite(&mut rewriter)?.value; + + Ok(BoundPartitionedExpr { + root, + partitions: partitions.into_boxed_slice(), + partition_names, + partition_annotations: partition_annotations.into_boxed_slice(), + }) +} + +/// The result of partitioning an expression. +#[derive(Debug)] +pub struct BoundPartitionedExpr { + /// The root expression used to re-assemble the results. + pub root: BoundExpression, + /// The partition expressions themselves. + pub partitions: Box<[BoundExpression]>, + /// The field name of each partition as referenced in the root expression. + pub partition_names: FieldNames, + /// The annotation associated with each partition. + pub partition_annotations: Box<[A]>, +} + +impl Display for BoundPartitionedExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "root: {} {{{}}}", + self.root, + self.partition_names + .iter() + .zip(self.partitions.iter()) + .map(|(name, partition)| format!("{name}: {partition}")) + .join(", ") + ) + } +} + +impl BoundPartitionedExpr +where + FieldName: From, +{ + /// Return the partition for a given field, if it exists. + // FIXME(ngates): this should return an iterator since an annotation may have multiple partitions. + pub fn find_partition(&self, id: &A) -> Option<&BoundExpression> { + let id = FieldName::from(id.clone()); + self.partition_names + .iter() + .position(|field| field == id) + .map(|idx| &self.partitions[idx]) + } + + /// Replace the partition expressions and update every root dtype in the recombination tree. + pub fn replace_partitions(&mut self, partitions: Box<[BoundExpression]>) -> VortexResult<()> { + vortex_ensure!( + partitions.len() == self.partition_names.len(), + "Expected {} partitions, got {}", + self.partition_names.len(), + partitions.len() + ); + + let root_dtype = partition_root_dtype(&self.partition_names, &partitions); + let root = replace_root_dtype(self.root.clone(), root_dtype)?; + self.partitions = partitions; + self.root = root; + Ok(()) + } +} + +#[derive(Debug)] +struct PartitionCollector<'a, A: Annotation> { + annotations: &'a BoundAnnotations, + sub_expressions: HashMap>, +} + +impl<'a, A: Annotation + Display> PartitionCollector<'a, A> { + fn new(annotations: &'a BoundAnnotations) -> Self { + Self { + sub_expressions: HashMap::new(), + annotations, + } + } + + /// Each annotation may be associated with multiple sub-expressions, so we need to + /// a unique name for each sub-expression. + fn field_name(annotation: &A, idx: usize) -> FieldName { + format!("{annotation}_{idx}").into() + } +} + +impl NodeRewriter for PartitionCollector<'_, A> +where + FieldName: From, +{ + type NodeTy = BoundExpression; + + fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult> { + match self.annotations.get(&ExactBoundExpr(node.clone())) { + // If this expression only accesses a single field, then we can skip the children + Some(annotations) if annotations.len() == 1 => { + let annotation = annotations + .iter() + .next() + .vortex_expect("expected one field"); + let sub_exprs = self.sub_expressions.entry(annotation.clone()).or_default(); + sub_exprs.push(node.clone()); + Ok(Transformed { + value: node, + changed: false, + order: TraversalOrder::Skip, + }) + } + + // Otherwise, continue traversing. + _ => Ok(Transformed::no(node)), + } + } + + fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult> { + Ok(Transformed::no(node)) + } +} + +struct PartitionRootRewriter<'a, A: Annotation> { + annotations: &'a BoundAnnotations, + partition_offsets: HashMap, + root_dtype: DType, +} + +impl<'a, A: Annotation> PartitionRootRewriter<'a, A> { + fn new(annotations: &'a BoundAnnotations, root_dtype: DType) -> Self { + Self { + annotations, + partition_offsets: HashMap::new(), + root_dtype, + } + } +} + +impl NodeRewriter for PartitionRootRewriter<'_, A> +where + FieldName: From, +{ + type NodeTy = BoundExpression; + + fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult> { + let Some(annotations) = self.annotations.get(&ExactBoundExpr(node.clone())) else { + return Ok(Transformed::no(node)); + }; + if annotations.len() != 1 { + return Ok(Transformed::no(node)); + } + + let annotation = annotations + .iter() + .next() + .vortex_expect("expected one annotation"); + let offset = self + .partition_offsets + .entry(annotation.clone()) + .or_default(); + let field_name = PartitionCollector::field_name(annotation, *offset); + *offset += 1; + + let partition = bound_get_item( + FieldName::from(annotation.clone()), + BoundExpression::new_root(self.root_dtype.clone()), + )?; + let value = bound_get_item(field_name, partition)?; + + Ok(Transformed { + value, + changed: true, + order: TraversalOrder::Skip, + }) + } +} + +fn bound_get_item(field_name: FieldName, child: BoundExpression) -> VortexResult { + BoundExpression::try_new(GetItem.bind(field_name), [child]) +} + +fn bound_pack(names: FieldNames, children: Vec) -> VortexResult { + BoundExpression::try_new( + Pack.bind(PackOptions { + names, + nullability: Nullability::NonNullable, + }), + children, + ) +} + +fn partition_root_dtype(names: &FieldNames, partitions: &[BoundExpression]) -> DType { + DType::Struct( + StructFields::new( + names.clone(), + partitions + .iter() + .map(|partition| partition.dtype().clone()) + .collect(), + ), + Nullability::NonNullable, + ) +} + +fn replace_root_dtype(expr: BoundExpression, root_dtype: DType) -> VortexResult { + Ok(expr + .transform_down(|node| { + if node.is_root() { + Ok(Transformed { + value: BoundExpression::new_root(root_dtype.clone()), + changed: true, + order: TraversalOrder::Skip, + }) + } else { + Ok(Transformed::no(node)) + } + })? + .into_inner()) +} + +#[cfg(test)] +mod tests { + use rstest::fixture; + use rstest::rstest; + + 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_bound_free_field_annotator; + use crate::expr::and; + use crate::expr::col; + use crate::expr::get_item; + use crate::expr::lit; + use crate::expr::merge; + use crate::expr::pack; + use crate::expr::root; + use crate::expr::transform::replace::replace_root_fields; + + #[fixture] + fn dtype() -> DType { + DType::Struct( + StructFields::from_iter([ + ( + "a", + DType::Struct( + StructFields::from_iter([("x", I32.into()), ("y", DType::from(I32))]), + NonNullable, + ), + ), + ("b", I32.into()), + ("c", I32.into()), + ]), + NonNullable, + ) + } + + fn partition_by_field( + expr: BoundExpression, + dtype: &DType, + ) -> VortexResult> { + let fields = dtype.as_struct_fields_opt().unwrap(); + partition_bound(expr, make_bound_free_field_annotator(fields)) + } + + #[rstest] + fn test_expr_top_level_ref(dtype: DType) { + let fields = dtype.as_struct_fields_opt().unwrap(); + + let expr = root(); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + // An un-expanded root expression is annotated by all fields, but since it is a single node + assert_eq!(partitioned.partitions.len(), 0); + assert_eq!(partitioned.root.unbind(), root()); + + // Instead, callers must expand the root expression themselves. + let expr = replace_root_fields(expr, fields); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + assert_eq!(partitioned.partitions.len(), fields.names().len()); + } + + #[rstest] + fn test_expr_top_level_ref_get_item_and_split(dtype: DType) { + let expr = get_item("y", get_item("a", root())); + + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + assert_eq!( + partitioned.root.unbind(), + get_item("a_0", get_item("a", root())) + ); + } + + #[rstest] + fn test_expr_top_level_ref_get_item_and_split_pack(dtype: DType) { + let expr = pack( + [ + ("x", get_item("x", get_item("a", root()))), + ("y", get_item("y", get_item("a", root()))), + ("c", get_item("c", root())), + ], + NonNullable, + ); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + let split_a = partitioned.find_partition(&"a".into()).unwrap(); + assert_eq!( + split_a.unbind(), + pack( + [ + ("a_0", get_item("x", get_item("a", root()))), + ("a_1", get_item("y", get_item("a", root()))) + ], + NonNullable + ) + ); + } + + #[rstest] + fn test_expr_top_level_ref_get_item_add(dtype: DType) { + let expr = and(get_item("y", get_item("a", root())), lit(1)); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + // Whole expr is a single split + assert_eq!(partitioned.partitions.len(), 1); + } + + #[rstest] + fn test_expr_top_level_ref_get_item_add_cannot_split(dtype: DType) { + let expr = and(get_item("y", get_item("a", root())), get_item("b", root())); + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + + // One for id.a and id.b + assert_eq!(partitioned.partitions.len(), 2); + } + + #[rstest] + fn test_expr_merge(dtype: DType) { + let expr = merge([col("a"), pack([("b", col("b"))], NonNullable)]); + + let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let expected = merge([get_item("a_0", col("a")), get_item("b_0", col("b"))]); + assert_eq!( + partitioned.root.unbind(), + expected, + "{} {}", + partitioned.root, + expected + ); + + assert_eq!(partitioned.partitions.len(), 2); + + let part_a = partitioned.find_partition(&"a".into()).unwrap(); + let expected_a = pack([("a_0", col("a"))], NonNullable); + assert_eq!(part_a.unbind(), expected_a, "{part_a} {expected_a}"); + + let part_b = partitioned.find_partition(&"b".into()).unwrap(); + let expected_b = pack([("b_0", pack([("b", col("b"))], NonNullable))], NonNullable); + assert_eq!(part_b.unbind(), expected_b, "{part_b} {expected_b}"); + } + + #[rstest] + fn replacing_partitions_refreshes_root_dtype(dtype: DType) -> VortexResult<()> { + let mut partitioned = partition_by_field(col("b").bind(&dtype)?, &dtype)?; + let field_dtype = DType::Primitive(I32, Nullable); + let replacement = pack([("b_0", root())], NonNullable).bind(&field_dtype)?; + + partitioned.replace_partitions(vec![replacement].into_boxed_slice())?; + + assert_eq!(partitioned.root.dtype(), &field_dtype); + Ok(()) + } +} diff --git a/vortex-array/src/expr/transform/mod.rs b/vortex-array/src/expr/transform/mod.rs index 6013241d971..0138cbe47c7 100644 --- a/vortex-array/src/expr/transform/mod.rs +++ b/vortex-array/src/expr/transform/mod.rs @@ -1,12 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A collection of transformations that can be applied to a [`crate::expr::Expression`]. +//! Transformations for [`crate::expr::Expression`] and [`crate::expr::BoundExpression`] trees. +mod bound_partition; mod coerce; pub(crate) mod match_between; mod partition; mod replace; +pub use bound_partition::*; pub use coerce::*; pub use partition::*; pub use replace::*; diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index b000439082c..c6000bfdcd4 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -23,7 +23,9 @@ pub use visitor::pre_order_visit_down; pub use visitor::pre_order_visit_up; use vortex_error::VortexResult; +use crate::expr::BoundExpression; use crate::expr::Expression; +use crate::expr::bound_expression::BoundKind; use crate::expr::traversal::fold::NodeFolderContextWrapper; /// Signal to control a traversal's flow @@ -533,6 +535,74 @@ impl Node for Expression { } } +impl Node for BoundExpression { + fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult>( + &'a self, + mut f: F, + ) -> VortexResult { + let BoundKind::Scalar { children, .. } = self.kind() else { + return Ok(TraversalOrder::Continue); + }; + + for child in children.iter() { + match f(child)? { + TraversalOrder::Continue | TraversalOrder::Skip => {} + TraversalOrder::Stop => return Ok(TraversalOrder::Stop), + } + } + + Ok(TraversalOrder::Continue) + } + + fn map_children VortexResult>>( + self, + mut f: F, + ) -> VortexResult> { + let BoundKind::Scalar { children, .. } = self.kind() else { + return Ok(Transformed::no(self)); + }; + + let mut order = TraversalOrder::Continue; + let mut changed = false; + let children = children + .iter() + .cloned() + .map(|child| match order { + TraversalOrder::Continue | TraversalOrder::Skip => f(child).map(|result| { + order = result.order; + changed |= result.changed; + result.value + }), + TraversalOrder::Stop => Ok(child), + }) + .collect::>>()?; + + if changed { + Ok(Transformed { + value: self.with_children(children)?, + order, + changed: true, + }) + } else { + Ok(Transformed::no(self)) + } + } + + fn iter_children(&self, f: impl FnOnce(&mut dyn Iterator) -> T) -> T { + match self.kind() { + BoundKind::Scalar { children, .. } => f(&mut children.iter()), + BoundKind::Root => f(&mut std::iter::empty()), + } + } + + fn children_count(&self) -> usize { + match self.kind() { + BoundKind::Scalar { children, .. } => children.len(), + BoundKind::Root => 0, + } + } +} + #[cfg(test)] mod tests { use vortex_error::VortexResult;