From 65a1ca61fe8091c7f1db581061915cc27cb7557e Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 4 Aug 2026 11:38:25 +0100 Subject: [PATCH 1/3] Push expressions through layout plans Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/mod.rs | 11 +++ vortex-layout/src/plan/plans/dict.rs | 33 ++++++++ vortex-layout/src/plan/plans/expression.rs | 8 +- vortex-layout/src/plan/plans/struct_.rs | 42 ++++++++++ vortex-layout/src/plan/tests.rs | 89 ++++++++++++++++++++++ 5 files changed, 179 insertions(+), 4 deletions(-) diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 4ccd5a9f325..7ef3a49209c 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -27,6 +27,7 @@ pub use plans::ListPlan; pub use plans::RowIdxPlan; pub use plans::StructPlan; use vortex_array::dtype::DType; +use vortex_array::expr::Expression; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -63,6 +64,16 @@ pub trait Plan: 'static + Send + Sync { /// Implementations may defer child optimization until the child is accessed. fn optimize(&self) -> VortexResult; + /// Attempts to rewrite `expression` through this plan. + /// + /// Returns `None` when this plan has no applicable expression rewrite. Implementations may + /// request only the children needed by the rewrite and should preserve all other lazy child + /// slots. + fn optimize_expression(&self, expression: &Expression) -> VortexResult> { + let _ = expression; + Ok(None) + } + /// Returns the dtype produced by this plan. fn dtype(&self) -> &DType; diff --git a/vortex-layout/src/plan/plans/dict.rs b/vortex-layout/src/plan/plans/dict.rs index 47827c6950c..a89d936bd4c 100644 --- a/vortex-layout/src/plan/plans/dict.rs +++ b/vortex-layout/src/plan/plans/dict.rs @@ -4,10 +4,16 @@ use std::borrow::Cow; use std::sync::Arc; +use vortex_array::expr::Expression; +use vortex_array::expr::is_root; +use vortex_array::expr::label_is_fallible; +use vortex_array::expr::label_strict; +use vortex_array::expr::label_tree; use vortex_error::VortexResult; use vortex_error::vortex_bail; use crate::layouts::dict::DictLayout; +use crate::plan::ExpressionPlan; use crate::plan::Plan; use crate::plan::PlanRef; use crate::plan::new_plan; @@ -67,6 +73,33 @@ impl Plan for DictPlan { Ok(Arc::new(self.with_children(codes, values))) } + fn optimize_expression(&self, expression: &Expression) -> VortexResult> { + if !expression.return_dtype(&self.dtype)?.is_boolean() { + return Ok(None); + } + let references_root = label_tree(expression, is_root, |acc, &child| acc | child) + .get(expression) + .copied() + .unwrap_or(false); + let is_strict = label_strict(expression) + .get(expression) + .copied() + .unwrap_or(false); + let is_fallible = label_is_fallible(expression) + .get(expression) + .copied() + .unwrap_or(true); + if !references_root || !is_strict || is_fallible { + return Ok(None); + } + + let values = + ExpressionPlan::try_new(expression.clone(), Arc::clone(&self.values))?.optimize()?; + Ok(Some(Arc::new( + self.with_children(Arc::clone(&self.codes), values), + ))) + } + fn dtype(&self) -> &vortex_array::dtype::DType { &self.dtype } diff --git a/vortex-layout/src/plan/plans/expression.rs b/vortex-layout/src/plan/plans/expression.rs index 686ad81b007..82cd6f939dc 100644 --- a/vortex-layout/src/plan/plans/expression.rs +++ b/vortex-layout/src/plan/plans/expression.rs @@ -62,10 +62,10 @@ impl Plan for ExpressionPlan { } if let Some(inner) = child.as_any().downcast_ref::() { let expression = replace(expression, &root(), inner.expression.clone()); - return Ok(Arc::new(Self::try_new( - expression, - Arc::clone(&inner.child), - )?)); + return Self::try_new(expression, Arc::clone(&inner.child))?.optimize(); + } + if let Some(rewritten) = child.optimize_expression(&expression)? { + return Ok(rewritten); } Ok(Arc::new(Self::try_new(expression, child)?)) } diff --git a/vortex-layout/src/plan/plans/struct_.rs b/vortex-layout/src/plan/plans/struct_.rs index e216a3f56c3..8b6be1b0cf1 100644 --- a/vortex-layout/src/plan/plans/struct_.rs +++ b/vortex-layout/src/plan/plans/struct_.rs @@ -5,9 +5,18 @@ use std::borrow::Cow; use std::sync::Arc; use vortex_array::dtype::DType; +use vortex_array::expr::Expression; +use vortex_array::expr::col; +use vortex_array::expr::make_free_field_annotator; +use vortex_array::expr::root; +use vortex_array::expr::transform::partition; +use vortex_array::expr::transform::replace; +use vortex_array::expr::transform::replace_root_fields; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::layouts::struct_::StructLayout; +use crate::plan::ExpressionPlan; use crate::plan::LazyPlanChildren; use crate::plan::Plan; use crate::plan::PlanRef; @@ -68,6 +77,39 @@ impl Plan for StructPlan { Ok(Arc::new(self.with_children(children))) } + fn optimize_expression(&self, expression: &Expression) -> VortexResult> { + if self.dtype.is_nullable() { + return Ok(None); + } + + let fields = self.layout.struct_fields(); + let expanded = + replace_root_fields(expression.clone(), fields).optimize_recursive(&self.dtype)?; + let partitioned = partition( + expanded.clone(), + &self.dtype, + make_free_field_annotator(fields), + )?; + if partitioned.partition_names.len() != 1 { + return Ok(None); + } + + let field_name = partitioned + .partition_names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition has no field"))?; + let field_index = fields.find(field_name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{field_name}'") + })?; + let field = self + .children + .get(field_index)? + .ok_or_else(|| vortex_err!("Struct field '{field_name}' has no plan"))?; + let lowered = replace(expanded, &col(field_name.clone()), root()); + + Ok(Some(ExpressionPlan::try_new(lowered, field)?.optimize()?)) + } + fn dtype(&self) -> &DType { &self.dtype } diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 6d5bc469ac9..a88622af8d1 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -8,7 +8,11 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; +use vortex_array::expr::checked_add; use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::is_null; +use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -454,3 +458,88 @@ fn row_idx_plan_preserves_row_index_expressions() -> VortexResult<()> { assert_eq!(expression.row_count(), 3); Ok(()) } + +#[test] +fn expression_pushes_through_struct_field_and_dictionary_values() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let codes_dtype = primitive(PType::U8, Nullability::NonNullable); + let dictionary = + DictLayout::new(flat(2, value_dtype.clone(), 0), flat(3, codes_dtype, 1)).into_layout(); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), + Nullability::NonNullable, + ), + vec![dictionary, unsupported(3, value_dtype)], + ) + .into_layout(); + let plan: PlanRef = Arc::new(ExpressionPlan::try_new( + gt(get_item("a", root()), lit(5_i32)), + make_plan(layout)?, + )?); + + let optimized = plan.optimize()?; + + insta::assert_snapshot!(optimized.tree_display(), @r" + root: DictPlan(bool, rows=3) + codes: FlatPlan(u8, rows=3) + values: ExpressionPlan(bool, rows=2) expr=($ > 5i32) + child: FlatPlan(i32, rows=2) + "); + Ok(()) +} + +#[test] +fn dictionary_pushdown_rejects_unsafe_expressions() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let dictionary = DictLayout::new( + flat(2, value_dtype, 0), + flat(3, primitive(PType::U8, Nullability::NonNullable), 1), + ) + .into_layout(); + + for expression in [ + lit(false), + is_null(root()), + gt(checked_add(root(), lit(1_i32)), lit(5_i32)), + ] { + let plan = + ExpressionPlan::try_new(expression.clone(), make_plan(Arc::clone(&dictionary))?)?; + let optimized = plan.optimize()?; + let expression_plan = optimized + .as_any() + .downcast_ref::() + .ok_or_else(|| { + vortex_err!("Expression unexpectedly pushed into dictionary: {expression}") + })?; + assert!(expression_plan.child_plan().as_any().is::()); + } + Ok(()) +} + +#[test] +fn nullable_struct_keeps_expression_above_parent_validity() -> VortexResult<()> { + let field_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", field_dtype.clone())]), + Nullability::Nullable, + ), + vec![ + flat(3, DType::Bool(Nullability::NonNullable), 0), + flat(3, field_dtype, 1), + ], + ) + .into_layout(); + let plan = ExpressionPlan::try_new(gt(get_item("a", root()), lit(5_i32)), make_plan(layout)?)?; + + let optimized = plan.optimize()?; + let expression_plan = optimized + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("Nullable struct expression unexpectedly pushed down"))?; + assert!(expression_plan.child_plan().as_any().is::()); + Ok(()) +} From 87243b886817e9e0f345db635daa3c7e99522694 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 4 Aug 2026 17:47:45 +0100 Subject: [PATCH 2/3] Push expressions through chunked and row-index plans Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/children.rs | 29 +- vortex-layout/src/plan/mod.rs | 12 +- vortex-layout/src/plan/plans/chunked.rs | 32 +- vortex-layout/src/plan/plans/mod.rs | 2 + vortex-layout/src/plan/plans/row_idx.rs | 273 +++++++++++++ vortex-layout/src/plan/plans/struct_.rs | 2 +- vortex-layout/src/plan/tests.rs | 490 ++++++++---------------- 7 files changed, 484 insertions(+), 356 deletions(-) diff --git a/vortex-layout/src/plan/children.rs b/vortex-layout/src/plan/children.rs index e74f7f67dbb..a9342b2d609 100644 --- a/vortex-layout/src/plan/children.rs +++ b/vortex-layout/src/plan/children.rs @@ -46,17 +46,24 @@ impl LazyPlanChildren { Ok(cell.get_or_try_init(|| (self.initializer)(index))?.clone()) } - /// Lazily transforms each present child into a new child collection. - pub(crate) fn map( + /// Eagerly transforms each present child into a new child collection. + pub(crate) fn try_map( &self, - transform: impl Fn(usize, PlanRef) -> VortexResult + 'static + Send + Sync, - ) -> Self { - let source = self.clone(); - Self::new(self.len(), move |index| { - source - .get(index)? - .map(|child| transform(index, child)) - .transpose() - }) + transform: impl Fn(usize, PlanRef) -> VortexResult, + ) -> VortexResult { + // TODO: Make recursive child optimization lazy again once the optimizer API can + // explicitly distinguish fully optimized plans from plans with deferred optimizer work. + let children = (0..self.len()) + .map(|index| { + self.get(index)? + .map(|child| transform(index, child)) + .transpose() + }) + .collect::>>()?; + let children: Arc<[Option]> = children.into(); + let len = children.len(); + Ok(Self::new(len, move |index| { + Ok(children.get(index).cloned().flatten()) + })) } } diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 7ef3a49209c..bb06b3e26b2 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -24,7 +24,9 @@ pub use plans::DictPlan; pub use plans::ExpressionPlan; pub use plans::FlatPlan; pub use plans::ListPlan; +pub use plans::RowIdxPartitionPlan; pub use plans::RowIdxPlan; +pub use plans::RowIdxValuesPlan; pub use plans::StructPlan; use vortex_array::dtype::DType; use vortex_array::expr::Expression; @@ -59,16 +61,14 @@ pub trait Plan: 'static + Send + Sync { std::any::type_name::() } - /// Optimizes this plan while preserving its dtype and row domain. - /// - /// Implementations may defer child optimization until the child is accessed. + /// Recursively optimizes this plan and all of its children while preserving its dtype and row + /// domain. fn optimize(&self) -> VortexResult; /// Attempts to rewrite `expression` through this plan. /// - /// Returns `None` when this plan has no applicable expression rewrite. Implementations may - /// request only the children needed by the rewrite and should preserve all other lazy child - /// slots. + /// Returns `None` when this plan has no applicable expression rewrite. Implementations should + /// preserve child slots that the rewrite does not change. fn optimize_expression(&self, expression: &Expression) -> VortexResult> { let _ = expression; Ok(None) diff --git a/vortex-layout/src/plan/plans/chunked.rs b/vortex-layout/src/plan/plans/chunked.rs index fd145ec5221..5d5d7c0403b 100644 --- a/vortex-layout/src/plan/plans/chunked.rs +++ b/vortex-layout/src/plan/plans/chunked.rs @@ -5,9 +5,13 @@ use std::borrow::Cow; use std::sync::Arc; use vortex_array::dtype::DType; +use vortex_array::expr::Expression; +use vortex_array::expr::label_tree; use vortex_error::VortexResult; use crate::layouts::chunked::ChunkedLayout; +use crate::layouts::row_idx::RowIdx; +use crate::plan::ExpressionPlan; use crate::plan::LazyPlanChildren; use crate::plan::Plan; use crate::plan::PlanRef; @@ -36,10 +40,10 @@ impl ChunkedPlan { } } - fn with_chunks(&self, chunks: LazyPlanChildren) -> Self { + fn with_chunks(&self, dtype: DType, chunks: LazyPlanChildren) -> Self { Self { layout: self.layout.clone(), - dtype: self.dtype.clone(), + dtype, chunks, } } @@ -55,8 +59,28 @@ impl Plan for ChunkedPlan { } fn optimize(&self) -> VortexResult { - let chunks = self.chunks.map(|_, chunk| chunk.optimize()); - Ok(Arc::new(self.with_chunks(chunks))) + let chunks = self.chunks.try_map(|_, chunk| chunk.optimize())?; + Ok(Arc::new(self.with_chunks(self.dtype.clone(), chunks))) + } + + fn optimize_expression(&self, expression: &Expression) -> VortexResult> { + let references_row_idx = label_tree( + expression, + |node| node.is::(), + |acc, &child| acc | child, + ) + .get(expression) + .copied() + .unwrap_or(false); + if references_row_idx { + return Ok(None); + } + + let dtype = expression.return_dtype(&self.dtype)?; + let chunks = self + .chunks + .try_map(|_, chunk| ExpressionPlan::try_new(expression.clone(), chunk)?.optimize())?; + Ok(Some(Arc::new(self.with_chunks(dtype, chunks)))) } fn dtype(&self) -> &DType { diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index a6ffe427d3c..1bc5e1ddce9 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -14,5 +14,7 @@ pub use dict::DictPlan; pub use expression::ExpressionPlan; pub use flat::FlatPlan; pub use list::ListPlan; +pub use row_idx::RowIdxPartitionPlan; pub use row_idx::RowIdxPlan; +pub use row_idx::RowIdxValuesPlan; pub use struct_::StructPlan; diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/plans/row_idx.rs index 58d9370122b..73315645657 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/plans/row_idx.rs @@ -3,12 +3,27 @@ use std::any::Any; use std::borrow::Cow; +use std::fmt::Display; +use std::fmt::Formatter; use std::sync::Arc; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::Expression; +use vortex_array::expr::get_item; +use vortex_array::expr::root; +use vortex_array::expr::transform::partition; +use vortex_array::expr::transform::replace; +use vortex_array::scalar_fn::fns::pack::Pack; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use crate::layouts::row_idx::RowIdx; +use crate::layouts::row_idx::row_idx; +use crate::plan::ExpressionPlan; use crate::plan::Plan; use crate::plan::PlanRef; @@ -38,6 +53,102 @@ impl Plan for RowIdxPlan { Ok(Self::new_ref(self.row_offset, self.child.optimize()?)) } + fn optimize_expression(&self, expression: &Expression) -> VortexResult> { + let partitioned = partition(expression.clone(), self.dtype(), |node| { + if node.is::() { + vec![RowIdxExpressionPartition::RowIdx] + } else if vortex_array::expr::is_root(node) { + vec![RowIdxExpressionPartition::Child] + } else { + vec![] + } + })?; + + if partitioned.partition_annotations.len() == 1 { + return match partitioned.partition_annotations[0] { + RowIdxExpressionPartition::RowIdx => { + let expression = replace(expression.clone(), &row_idx(), root()); + let values = RowIdxValuesPlan::new_ref(self.row_offset, self.row_count()); + Ok(Some( + ExpressionPlan::try_new(expression, values)?.optimize()?, + )) + } + RowIdxExpressionPartition::Child => Ok(Some( + ExpressionPlan::try_new(expression.clone(), Arc::clone(&self.child))? + .optimize()?, + )), + }; + } + + if partitioned.partition_annotations.len() != 2 { + return Ok(None); + } + let Some(row_idx_index) = partitioned + .partition_annotations + .iter() + .position(|partition| *partition == RowIdxExpressionPartition::RowIdx) + else { + return Ok(None); + }; + let Some(child_index) = partitioned + .partition_annotations + .iter() + .position(|partition| *partition == RowIdxExpressionPartition::Child) + else { + return Ok(None); + }; + + let row_idx_partition = &partitioned.partitions[row_idx_index]; + let child_partition = &partitioned.partitions[child_index]; + let (Some(row_idx_pack), Some(child_pack)) = ( + row_idx_partition.as_opt::(), + child_partition.as_opt::(), + ) else { + return Ok(None); + }; + // A general pack plan is needed to expose more than one result from either branch. + if row_idx_partition.children().len() != 1 || child_partition.children().len() != 1 { + return Ok(None); + } + + let (Some(row_idx_value_name), Some(child_value_name)) = + (row_idx_pack.names.get(0), child_pack.names.get(0)) + else { + return Ok(None); + }; + let row_idx_value_name = row_idx_value_name.clone(); + let child_value_name = child_value_name.clone(); + let row_idx_partition_name = partitioned.partition_names[row_idx_index].clone(); + let child_partition_name = partitioned.partition_names[child_index].clone(); + let row_idx_expression = row_idx_partition.child(0).clone(); + let child_expression = child_partition.child(0).clone(); + + let residual = replace( + partitioned.root, + &get_item(row_idx_value_name, get_item(row_idx_partition_name, root())), + get_item(RowIdxExpressionPartition::RowIdx.name(), root()), + ); + let residual = replace( + residual, + &get_item(child_value_name, get_item(child_partition_name, root())), + get_item(RowIdxExpressionPartition::Child.name(), root()), + ); + + let row_idx_expression = replace(row_idx_expression, &row_idx(), root()); + let row_idx_plan = ExpressionPlan::try_new( + row_idx_expression, + RowIdxValuesPlan::new_ref(self.row_offset, self.row_count()), + )? + .optimize()?; + let child_plan = + ExpressionPlan::try_new(child_expression, Arc::clone(&self.child))?.optimize()?; + let partitions = RowIdxPartitionPlan::try_new(row_idx_plan, child_plan)?; + + Ok(Some( + ExpressionPlan::try_new(residual, partitions)?.optimize()?, + )) + } + fn dtype(&self) -> &DType { self.child.dtype() } @@ -65,3 +176,165 @@ impl Plan for RowIdxPlan { } } } + +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +enum RowIdxExpressionPartition { + RowIdx, + Child, +} + +impl RowIdxExpressionPartition { + fn name(self) -> &'static str { + match self { + Self::RowIdx => "row_idx", + Self::Child => "child", + } + } +} + +impl Display for RowIdxExpressionPartition { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.name()) + } +} + +impl From for FieldName { + fn from(partition: RowIdxExpressionPartition) -> Self { + partition.name().into() + } +} + +/// A plan that generates the global row-index values for a row domain. +pub struct RowIdxValuesPlan { + row_offset: u64, + row_count: u64, + dtype: DType, +} + +impl RowIdxValuesPlan { + /// Creates a shared row-index values plan starting at `row_offset`. + pub fn new_ref(row_offset: u64, row_count: u64) -> PlanRef { + Arc::new(Self { + row_offset, + row_count, + dtype: DType::Primitive(PType::U64, Nullability::NonNullable), + }) + } + + /// Returns the global row index assigned to the first row. + pub fn row_offset(&self) -> u64 { + self.row_offset + } +} + +impl Plan for RowIdxValuesPlan { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &'static str { + "RowIdxValuesPlan" + } + + fn optimize(&self) -> VortexResult { + Ok(Self::new_ref(self.row_offset, self.row_count)) + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.row_count + } +} + +/// A plan that combines independently evaluated row-index and data expression partitions. +pub struct RowIdxPartitionPlan { + row_idx: PlanRef, + child: PlanRef, + dtype: DType, +} + +impl RowIdxPartitionPlan { + /// Creates a shared partition plan whose branches have the same row domain. + pub fn try_new(row_idx: PlanRef, child: PlanRef) -> VortexResult { + if row_idx.row_count() != child.row_count() { + vortex_bail!( + "Row-index partition row count {} does not match child row count {}", + row_idx.row_count(), + child.row_count() + ) + } + let dtype = DType::Struct( + StructFields::from_iter([ + ( + RowIdxExpressionPartition::RowIdx.name(), + row_idx.dtype().clone(), + ), + ( + RowIdxExpressionPartition::Child.name(), + child.dtype().clone(), + ), + ]), + Nullability::NonNullable, + ); + Ok(Arc::new(Self { + row_idx, + child, + dtype, + })) + } + + /// Returns the plan that evaluates the row-index expression partition. + pub fn row_idx_plan(&self) -> &PlanRef { + &self.row_idx + } + + /// Returns the plan that evaluates the data-child expression partition. + pub fn child_plan(&self) -> &PlanRef { + &self.child + } +} + +impl Plan for RowIdxPartitionPlan { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &'static str { + "RowIdxPartitionPlan" + } + + fn optimize(&self) -> VortexResult { + Self::try_new(self.row_idx.optimize()?, self.child.optimize()?) + } + + fn dtype(&self) -> &DType { + &self.dtype + } + + fn row_count(&self) -> u64 { + self.child.row_count() + } + + fn child_count(&self) -> usize { + 2 + } + + fn child(&self, index: usize) -> VortexResult> { + match index { + 0 => Ok(Some(Arc::clone(&self.row_idx))), + 1 => Ok(Some(Arc::clone(&self.child))), + _ => vortex_bail!("Row-index partition plan has no child {index}"), + } + } + + fn child_name(&self, index: usize) -> Cow<'_, str> { + match index { + 0 => Cow::Borrowed(RowIdxExpressionPartition::RowIdx.name()), + 1 => Cow::Borrowed(RowIdxExpressionPartition::Child.name()), + _ => Cow::Owned(format!("child[{index}]")), + } + } +} diff --git a/vortex-layout/src/plan/plans/struct_.rs b/vortex-layout/src/plan/plans/struct_.rs index 8b6be1b0cf1..ff2e573cf09 100644 --- a/vortex-layout/src/plan/plans/struct_.rs +++ b/vortex-layout/src/plan/plans/struct_.rs @@ -73,7 +73,7 @@ impl Plan for StructPlan { } fn optimize(&self) -> VortexResult { - let children = self.children.map(|_, child| child.optimize()); + let children = self.children.try_map(|_, child| child.optimize())?; Ok(Arc::new(self.with_children(children))) } diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index a88622af8d1..6cdd04c3b05 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fmt; use std::sync::Arc; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; +use vortex_array::expr::and; use vortex_array::expr::checked_add; use vortex_array::expr::get_item; use vortex_array::expr::gt; @@ -26,7 +26,6 @@ use crate::layouts::chunked::ChunkedLayout; use crate::layouts::dict::DictLayout; use crate::layouts::flat::FlatLayout; use crate::layouts::foreign::new_foreign_layout; -use crate::layouts::list::ListLayout; use crate::layouts::row_idx::row_idx; use crate::layouts::struct_::StructLayout; use crate::segments::SegmentId; @@ -55,12 +54,21 @@ fn make_plan(layout: LayoutRef) -> VortexResult { } #[test] -fn unsupported_layout_has_no_plan() -> VortexResult<()> { - let layout = unsupported(3, DType::Null); - - let error = new_plan(&layout) +fn struct_plan_optimization_visits_all_fields() -> VortexResult<()> { + let field_dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = StructLayout::new( + 1, + DType::Struct( + StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]), + Nullability::NonNullable, + ), + vec![flat(1, field_dtype.clone(), 0), unsupported(1, field_dtype)], + ) + .into_layout(); + let error = make_plan(layout)? + .optimize() .err() - .ok_or_else(|| vortex_err!("unsupported layout unexpectedly produced a plan"))?; + .ok_or_else(|| vortex_err!("unsupported field was not visited during optimization"))?; assert!( error .to_string() @@ -70,45 +78,7 @@ fn unsupported_layout_has_no_plan() -> VortexResult<()> { } #[test] -fn flat_plan_has_no_children() -> VortexResult<()> { - let plan = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?; - - assert!(plan.as_any().is::()); - assert_eq!(plan.child_count(), 0); - assert!(plan.child(0).is_err()); - Ok(()) -} - -#[test] -fn chunked_plan_exposes_chunks() -> VortexResult<()> { - let dtype = primitive(PType::I32, Nullability::NonNullable); - let layout = ChunkedLayout::new( - 3, - dtype.clone(), - OwnedLayoutChildren::layout_children(vec![flat(2, dtype.clone(), 0), flat(1, dtype, 1)]), - ) - .into_layout(); - let plan = make_plan(layout)?; - - assert!(plan.as_any().is::()); - assert_eq!(plan.child_count(), 2); - assert_eq!( - plan.child(0)? - .ok_or_else(|| vortex_err!("missing first chunk"))? - .row_count(), - 2 - ); - assert_eq!( - plan.child(1)? - .ok_or_else(|| vortex_err!("missing second chunk"))? - .row_count(), - 1 - ); - Ok(()) -} - -#[test] -fn chunked_plan_defers_unrequested_chunks_through_optimization() -> VortexResult<()> { +fn chunked_plan_optimization_visits_all_chunks() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); let layout = ChunkedLayout::new( 2, @@ -119,21 +89,11 @@ fn chunked_plan_defers_unrequested_chunks_through_optimization() -> VortexResult ]), ) .into_layout(); - let plan = make_plan(layout)?.optimize()?; - - let first = plan - .child(0)? - .ok_or_else(|| vortex_err!("missing first chunk"))?; - assert!(first.as_any().is::()); - let cached = plan - .child(0)? - .ok_or_else(|| vortex_err!("missing cached first chunk"))?; - assert!(Arc::ptr_eq(&first, &cached)); - let error = plan - .child(1) + let error = make_plan(layout)? + .optimize() .err() - .ok_or_else(|| vortex_err!("unsupported chunk unexpectedly produced a plan"))?; + .ok_or_else(|| vortex_err!("unsupported chunk was not visited during optimization"))?; assert!( error .to_string() @@ -143,335 +103,183 @@ fn chunked_plan_defers_unrequested_chunks_through_optimization() -> VortexResult } #[test] -fn dict_plan_orders_codes_before_values() -> VortexResult<()> { - let values_dtype = primitive(PType::I32, Nullability::NonNullable); - let codes_dtype = primitive(PType::U8, Nullability::NonNullable); - let layout = DictLayout::new( - flat(2, values_dtype.clone(), 0), - flat(3, codes_dtype.clone(), 1), - ) - .into_layout(); - let plan = make_plan(layout)?; - - assert!(plan.as_any().is::()); - assert_eq!(plan.child_count(), 2); - assert_eq!( - plan.child(0)? - .ok_or_else(|| vortex_err!("missing codes"))? - .dtype(), - &codes_dtype - ); - assert_eq!( - plan.child(1)? - .ok_or_else(|| vortex_err!("missing values"))? - .dtype(), - &values_dtype - ); - Ok(()) -} - -#[test] -fn list_plan_has_stable_optional_validity_slot() -> VortexResult<()> { - let element_dtype = primitive(PType::I32, Nullability::NonNullable); - let offsets_dtype = primitive(PType::U32, Nullability::NonNullable); - let non_nullable = ListLayout::new( - DType::List(Arc::new(element_dtype.clone()), Nullability::NonNullable), - flat(3, element_dtype.clone(), 0), - flat(3, offsets_dtype.clone(), 1), - None, - ) - .into_layout(); - let plan = make_plan(non_nullable)?; +fn row_idx_only_expression_uses_generated_values_plan() -> VortexResult<()> { + let layout = flat(3, primitive(PType::I32, Nullability::NonNullable), 0); + let plan: PlanRef = Arc::new(ExpressionPlan::try_new( + row_idx(), + RowIdxPlan::new_ref(10, make_plan(layout)?), + )?); - assert!(plan.as_any().is::()); - assert_eq!(plan.child_count(), 3); - assert_eq!( - plan.child(0)? - .ok_or_else(|| vortex_err!("missing elements"))? - .dtype(), - &element_dtype - ); - assert_eq!( - plan.child(1)? - .ok_or_else(|| vortex_err!("missing offsets"))? - .dtype(), - &offsets_dtype - ); - assert!(plan.child(2)?.is_none()); + insta::assert_snapshot!(plan.tree_display(), @r" + root: ExpressionPlan(u64, rows=3) expr=#row_idx + child: RowIdxPlan(i32, rows=3) + child: FlatPlan(i32, rows=3) + "); - let nullable = ListLayout::new( - DType::List(Arc::new(element_dtype.clone()), Nullability::Nullable), - flat(3, element_dtype, 2), - flat(3, offsets_dtype, 3), - Some(flat(2, DType::Bool(Nullability::NonNullable), 4)), - ) - .into_layout(); - let nullable_plan = make_plan(nullable)?; - assert_eq!( - nullable_plan - .child(2)? - .ok_or_else(|| vortex_err!("missing validity"))? - .dtype(), - &DType::Bool(Nullability::NonNullable) + let optimized = plan.optimize()?; + insta::assert_snapshot!( + optimized.tree_display(), + @"root: RowIdxValuesPlan(u64, rows=3)" ); + let values = optimized + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("optimized plan does not generate row-index values"))?; + assert_eq!(values.row_offset(), 10); Ok(()) } #[test] -fn struct_plan_orders_fields_before_optional_validity() -> VortexResult<()> { - let field_dtype = primitive(PType::I32, Nullability::NonNullable); - let fields = StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]); - let non_nullable = StructLayout::new( - 3, - DType::Struct(fields.clone(), Nullability::NonNullable), - vec![ - flat(3, field_dtype.clone(), 0), - flat(3, field_dtype.clone(), 1), - ], - ) - .into_layout(); - let plan = make_plan(non_nullable)?; - - assert!(plan.as_any().is::()); - assert_eq!(plan.child_count(), 3); - assert_eq!( - plan.child(0)? - .ok_or_else(|| vortex_err!("missing field a"))? - .dtype(), - &field_dtype - ); - assert_eq!( - plan.child(1)? - .ok_or_else(|| vortex_err!("missing field b"))? - .dtype(), - &field_dtype - ); - assert!(plan.child(2)?.is_none()); - - let nullable = StructLayout::new( - 3, - DType::Struct(fields, Nullability::Nullable), - vec![ - flat(3, DType::Bool(Nullability::NonNullable), 2), - flat(3, field_dtype.clone(), 3), - flat(3, field_dtype, 4), - ], +fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let dictionary = DictLayout::new( + flat(2, value_dtype.clone(), 0), + flat(3, primitive(PType::U8, Nullability::NonNullable), 1), ) .into_layout(); - let nullable_plan = make_plan(nullable)?; - assert_eq!( - nullable_plan - .child(2)? - .ok_or_else(|| vortex_err!("missing validity"))? - .dtype(), - &DType::Bool(Nullability::NonNullable) - ); - Ok(()) -} - -#[test] -fn struct_plan_defers_unrequested_fields_through_optimization() -> VortexResult<()> { - let field_dtype = primitive(PType::I32, Nullability::NonNullable); - let layout = StructLayout::new( - 1, - DType::Struct( - StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]), - Nullability::NonNullable, - ), - vec![flat(1, field_dtype.clone(), 0), unsupported(1, field_dtype)], - ) - .into_layout(); - let plan = make_plan(layout)?.optimize()?; - - assert!( - plan.child(0)? - .ok_or_else(|| vortex_err!("missing field a"))? - .as_any() - .is::() - ); - let error = plan - .child(1) - .err() - .ok_or_else(|| vortex_err!("unsupported field unexpectedly produced a plan"))?; - assert!( - error - .to_string() - .contains("No physical plan implementation for layout 'vortex.test.unsupported'") - ); - Ok(()) -} - -#[test] -fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { - let field_dtype = primitive(PType::I32, Nullability::NonNullable); let layout = StructLayout::new( 3, DType::Struct( - StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]), + StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), Nullability::NonNullable, ), - vec![flat(3, field_dtype.clone(), 0), flat(3, field_dtype, 1)], + vec![dictionary, flat(3, value_dtype, 2)], ) .into_layout(); + let expression = and( + gt(row_idx(), lit(11_u64)), + gt(get_item("a", root()), lit(5_i32)), + ); let plan: PlanRef = Arc::new(ExpressionPlan::try_new( - get_item("a", root()), - make_plan(layout)?, + expression, + RowIdxPlan::new_ref(10, make_plan(layout)?), )?); - assert_eq!(plan.to_string(), "ExpressionPlan(i32, rows=3)"); insta::assert_snapshot!(plan.tree_display(), @r" - root: ExpressionPlan(i32, rows=3) expr=$.a - child: StructPlan({a=i32, b=i32}, rows=3) - a: FlatPlan(i32, rows=3) - b: FlatPlan(i32, rows=3) - "); - - struct DepthExtractor; - - impl PlanTreeExtractor for DepthExtractor { - fn write_header( - &self, - _plan: &dyn Plan, - context: &PlanTreeContext, - formatter: &mut fmt::Formatter<'_>, - ) -> fmt::Result { - write!(formatter, " depth={}", context.depth()) - } - } - - insta::assert_snapshot!(plan.tree_display_builder().with(DepthExtractor), @r" - root: depth=0 - child: depth=1 - a: depth=2 - b: depth=2 + root: ExpressionPlan(bool, rows=3) expr=((#row_idx > 11u64) and ($.a > 5i32)) + child: RowIdxPlan({a=i32, b=i32}, rows=3) + child: StructPlan({a=i32, b=i32}, rows=3) + a: DictPlan(i32, rows=3) + codes: FlatPlan(u8, rows=3) + values: FlatPlan(i32, rows=2) + b: FlatPlan(i32, rows=3) "); - let nullable_fields = StructFields::from_iter([ - ("a", primitive(PType::I32, Nullability::NonNullable)), - ("b", primitive(PType::I32, Nullability::NonNullable)), - ]); - let nullable_layout = StructLayout::new( - 3, - DType::Struct(nullable_fields, Nullability::Nullable), - vec![ - flat(3, DType::Bool(Nullability::NonNullable), 2), - flat(3, primitive(PType::I32, Nullability::NonNullable), 3), - flat(3, primitive(PType::I32, Nullability::NonNullable), 4), - ], - ) - .into_layout(); - let nullable = make_plan(nullable_layout)?; - insta::assert_snapshot!(nullable.tree_display_builder(), @r" - root: - a: - b: - validity: + let optimized = plan.optimize()?; + insta::assert_snapshot!(optimized.tree_display(), @r" + root: ExpressionPlan(bool, rows=3) expr=($.row_idx and $.child) + child: RowIdxPartitionPlan({row_idx=bool, child=bool}, rows=3) + row_idx: ExpressionPlan(bool, rows=3) expr=($ > 11u64) + child: RowIdxValuesPlan(u64, rows=3) + child: DictPlan(bool, rows=3) + codes: FlatPlan(u8, rows=3) + values: ExpressionPlan(bool, rows=2) expr=($ > 5i32) + child: FlatPlan(i32, rows=2) "); + let residual = optimized + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("optimized plan has no residual expression"))?; + let partitions = residual + .child_plan() + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("optimized plan has no row-index partitions"))?; + let row_idx_expression = partitions + .row_idx_plan() + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("row-index partition has no expression"))?; + let values = row_idx_expression + .child_plan() + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("row-index partition has no generated values"))?; + assert_eq!(values.row_offset(), 10); Ok(()) } #[test] -fn chunked_plan_display_names_chunks() -> VortexResult<()> { +fn chunked_plan_preserves_global_row_index_expressions() -> VortexResult<()> { let dtype = primitive(PType::I32, Nullability::NonNullable); let layout = ChunkedLayout::new( 3, dtype.clone(), - OwnedLayoutChildren::layout_children(vec![flat(2, dtype.clone(), 0), flat(1, dtype, 1)]), + OwnedLayoutChildren::layout_children(vec![flat(1, dtype.clone(), 0), flat(2, dtype, 1)]), ) .into_layout(); - let plan = make_plan(layout)?; + let plan = ExpressionPlan::try_new(row_idx(), make_plan(layout)?)?.optimize()?; + let expression = plan + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("Row-index expression unexpectedly pushed into chunks"))?; - insta::assert_snapshot!(plan.display_tree(), @r" - root: ChunkedPlan(i32, rows=3) - chunks[0]: FlatPlan(i32, rows=2) - chunks[1]: FlatPlan(i32, rows=1) - "); + assert_eq!(expression.expression(), &row_idx()); + assert!(expression.child_plan().as_any().is::()); Ok(()) } #[test] -fn dict_plan_display_names_logical_children() -> VortexResult<()> { - let layout = DictLayout::new( - flat(2, primitive(PType::I32, Nullability::NonNullable), 0), - flat(3, primitive(PType::U8, Nullability::NonNullable), 1), +fn expression_pushes_through_struct_field_and_dictionary_values() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let codes_dtype = primitive(PType::U8, Nullability::NonNullable); + let dictionary = + DictLayout::new(flat(2, value_dtype.clone(), 0), flat(3, codes_dtype, 1)).into_layout(); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), + Nullability::NonNullable, + ), + vec![dictionary, flat(3, value_dtype, 2)], ) .into_layout(); - let plan = make_plan(layout)?; + let plan: PlanRef = Arc::new(ExpressionPlan::try_new( + gt(get_item("a", root()), lit(5_i32)), + make_plan(layout)?, + )?); insta::assert_snapshot!(plan.tree_display(), @r" - root: DictPlan(i32, rows=3) + root: ExpressionPlan(bool, rows=3) expr=($.a > 5i32) + child: StructPlan({a=i32, b=i32}, rows=3) + a: DictPlan(i32, rows=3) + codes: FlatPlan(u8, rows=3) + values: FlatPlan(i32, rows=2) + b: FlatPlan(i32, rows=3) + "); + + let optimized = plan.optimize()?; + + insta::assert_snapshot!(optimized.tree_display(), @r" + root: DictPlan(bool, rows=3) codes: FlatPlan(u8, rows=3) - values: FlatPlan(i32, rows=2) + values: ExpressionPlan(bool, rows=2) expr=($ > 5i32) + child: FlatPlan(i32, rows=2) "); Ok(()) } #[test] -fn list_plan_display_handles_optional_validity() -> VortexResult<()> { - let element_dtype = primitive(PType::I32, Nullability::NonNullable); - let offsets_dtype = primitive(PType::U32, Nullability::NonNullable); - let non_nullable_layout = ListLayout::new( - DType::List(Arc::new(element_dtype.clone()), Nullability::NonNullable), - flat(4, element_dtype.clone(), 0), - flat(3, offsets_dtype.clone(), 1), - None, +fn expression_pushes_through_struct_field_with_heterogeneous_chunks() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let dictionary = DictLayout::new( + flat(2, value_dtype.clone(), 0), + flat(3, primitive(PType::U8, Nullability::NonNullable), 1), ) .into_layout(); - let non_nullable = make_plan(non_nullable_layout)?; - - insta::assert_snapshot!(non_nullable.tree_display(), @r" - root: ListPlan(list(i32), rows=2) - elements: FlatPlan(i32, rows=4) - offsets: FlatPlan(u32, rows=3) - "); - - let nullable_layout = ListLayout::new( - DType::List(Arc::new(element_dtype.clone()), Nullability::Nullable), - flat(4, element_dtype, 2), - flat(3, offsets_dtype, 3), - Some(flat(2, DType::Bool(Nullability::NonNullable), 4)), + let chunks = ChunkedLayout::new( + 5, + value_dtype.clone(), + OwnedLayoutChildren::layout_children(vec![dictionary, flat(2, value_dtype.clone(), 2)]), ) .into_layout(); - let nullable = make_plan(nullable_layout)?; - - insta::assert_snapshot!(nullable.tree_display(), @r" - root: ListPlan(list(i32)?, rows=2) - elements: FlatPlan(i32, rows=4) - offsets: FlatPlan(u32, rows=3) - validity: FlatPlan(bool, rows=2) - "); - Ok(()) -} - -#[test] -fn row_idx_plan_preserves_row_index_expressions() -> VortexResult<()> { - let layout = flat(3, primitive(PType::I32, Nullability::NonNullable), 0); - let plan = RowIdxPlan::new_ref(10, make_plan(layout)?); - let plan = ExpressionPlan::try_new(row_idx(), plan)?.optimize()?; - let expression = plan - .as_any() - .downcast_ref::() - .ok_or_else(|| vortex_err!("optimized plan is not an expression plan"))?; - - assert_eq!(expression.expression(), &row_idx()); - assert!(expression.child_plan().as_any().is::()); - assert_eq!(expression.row_count(), 3); - Ok(()) -} - -#[test] -fn expression_pushes_through_struct_field_and_dictionary_values() -> VortexResult<()> { - let value_dtype = primitive(PType::I32, Nullability::NonNullable); - let codes_dtype = primitive(PType::U8, Nullability::NonNullable); - let dictionary = - DictLayout::new(flat(2, value_dtype.clone(), 0), flat(3, codes_dtype, 1)).into_layout(); let layout = StructLayout::new( - 3, + 5, DType::Struct( StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), Nullability::NonNullable, ), - vec![dictionary, unsupported(3, value_dtype)], + vec![chunks, flat(5, value_dtype, 3)], ) .into_layout(); let plan: PlanRef = Arc::new(ExpressionPlan::try_new( @@ -479,12 +287,26 @@ fn expression_pushes_through_struct_field_and_dictionary_values() -> VortexResul make_plan(layout)?, )?); + insta::assert_snapshot!(plan.tree_display(), @r" + root: ExpressionPlan(bool, rows=5) expr=($.a > 5i32) + child: StructPlan({a=i32, b=i32}, rows=5) + a: ChunkedPlan(i32, rows=5) + chunks[0]: DictPlan(i32, rows=3) + codes: FlatPlan(u8, rows=3) + values: FlatPlan(i32, rows=2) + chunks[1]: FlatPlan(i32, rows=2) + b: FlatPlan(i32, rows=5) + "); + let optimized = plan.optimize()?; insta::assert_snapshot!(optimized.tree_display(), @r" - root: DictPlan(bool, rows=3) - codes: FlatPlan(u8, rows=3) - values: ExpressionPlan(bool, rows=2) expr=($ > 5i32) + root: ChunkedPlan(bool, rows=5) + chunks[0]: DictPlan(bool, rows=3) + codes: FlatPlan(u8, rows=3) + values: ExpressionPlan(bool, rows=2) expr=($ > 5i32) + child: FlatPlan(i32, rows=2) + chunks[1]: ExpressionPlan(bool, rows=2) expr=($ > 5i32) child: FlatPlan(i32, rows=2) "); Ok(()) From a9043d58b8709c12c6242603caa91fdfa65980dc Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 5 Aug 2026 15:30:39 +0100 Subject: [PATCH 3/3] Use static parent-reduction rules for plan expressions Signed-off-by: Joe Isaacs --- vortex-layout/src/plan/mod.rs | 10 -- vortex-layout/src/plan/optimizer/mod.rs | 35 +++++ vortex-layout/src/plan/plans/chunked.rs | 56 +++++--- vortex-layout/src/plan/plans/dict.rs | 70 ++++++---- vortex-layout/src/plan/plans/expression.rs | 6 +- vortex-layout/src/plan/plans/mod.rs | 4 + vortex-layout/src/plan/plans/row_idx.rs | 133 ++++++++++-------- vortex-layout/src/plan/plans/struct_.rs | 149 +++++++++++++++------ vortex-layout/src/plan/tests.rs | 120 +++++++++++++++-- 9 files changed, 417 insertions(+), 166 deletions(-) diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index bb06b3e26b2..64473c3d95f 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -29,7 +29,6 @@ pub use plans::RowIdxPlan; pub use plans::RowIdxValuesPlan; pub use plans::StructPlan; use vortex_array::dtype::DType; -use vortex_array::expr::Expression; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -65,15 +64,6 @@ pub trait Plan: 'static + Send + Sync { /// domain. fn optimize(&self) -> VortexResult; - /// Attempts to rewrite `expression` through this plan. - /// - /// Returns `None` when this plan has no applicable expression rewrite. Implementations should - /// preserve child slots that the rewrite does not change. - fn optimize_expression(&self, expression: &Expression) -> VortexResult> { - let _ = expression; - Ok(None) - } - /// Returns the dtype produced by this plan. fn dtype(&self) -> &DType; diff --git a/vortex-layout/src/plan/optimizer/mod.rs b/vortex-layout/src/plan/optimizer/mod.rs index c9df9fb547e..8c2ee5ba051 100644 --- a/vortex-layout/src/plan/optimizer/mod.rs +++ b/vortex-layout/src/plan/optimizer/mod.rs @@ -9,3 +9,38 @@ pub use rules::DynPlanParentReduceRule; pub use rules::PlanParentReduceRule; pub use rules::PlanParentReduceRuleAdapter; pub use rules::PlanParentRuleSet; +use vortex_error::VortexResult; + +use super::ChunkedPlan; +use super::DictPlan; +use super::PlanRef; +use super::RowIdxPlan; +use super::StructPlan; +use super::plans::ExpressionChunkedRule; +use super::plans::ExpressionDictRule; +use super::plans::ExpressionRowIdxRule; +use super::plans::ExpressionStructRule; + +static EXPRESSION_CHUNKED_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionChunkedRule); +static EXPRESSION_DICT_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionDictRule); +static EXPRESSION_ROW_IDX_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionRowIdxRule); +static EXPRESSION_STRUCT_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionStructRule); + +static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[ + &EXPRESSION_CHUNKED_RULE, + &EXPRESSION_DICT_RULE, + &EXPRESSION_ROW_IDX_RULE, + &EXPRESSION_STRUCT_RULE, +]); + +/// Attempts a static rewrite for `parent` and its child at `child_idx`. +pub(crate) fn reduce_parent(parent: &PlanRef, child_idx: usize) -> VortexResult> { + let Some(child) = parent.child(child_idx)? else { + return Ok(None); + }; + PARENT_RULES.evaluate(&child, parent, child_idx) +} diff --git a/vortex-layout/src/plan/plans/chunked.rs b/vortex-layout/src/plan/plans/chunked.rs index 5d5d7c0403b..853e23987e5 100644 --- a/vortex-layout/src/plan/plans/chunked.rs +++ b/vortex-layout/src/plan/plans/chunked.rs @@ -5,7 +5,6 @@ use std::borrow::Cow; use std::sync::Arc; use vortex_array::dtype::DType; -use vortex_array::expr::Expression; use vortex_array::expr::label_tree; use vortex_error::VortexResult; @@ -16,6 +15,7 @@ use crate::plan::LazyPlanChildren; use crate::plan::Plan; use crate::plan::PlanRef; use crate::plan::new_plan; +use crate::plan::optimizer::PlanParentReduceRule; /// A physical plan with one child per row chunk. pub struct ChunkedPlan { @@ -63,26 +63,6 @@ impl Plan for ChunkedPlan { Ok(Arc::new(self.with_chunks(self.dtype.clone(), chunks))) } - fn optimize_expression(&self, expression: &Expression) -> VortexResult> { - let references_row_idx = label_tree( - expression, - |node| node.is::(), - |acc, &child| acc | child, - ) - .get(expression) - .copied() - .unwrap_or(false); - if references_row_idx { - return Ok(None); - } - - let dtype = expression.return_dtype(&self.dtype)?; - let chunks = self - .chunks - .try_map(|_, chunk| ExpressionPlan::try_new(expression.clone(), chunk)?.optimize())?; - Ok(Some(Arc::new(self.with_chunks(dtype, chunks)))) - } - fn dtype(&self) -> &DType { &self.dtype } @@ -106,3 +86,37 @@ impl Plan for ChunkedPlan { Cow::Owned(format!("chunks[{index}]")) } } + +/// Pushes an expression through every chunk of a chunked plan. +#[derive(Debug)] +pub(crate) struct ExpressionChunkedRule; + +impl PlanParentReduceRule for ExpressionChunkedRule { + type Parent = ExpressionPlan; + + fn reduce_parent( + &self, + child: &ChunkedPlan, + parent: &ExpressionPlan, + _child_idx: usize, + ) -> VortexResult> { + let expression = parent.expression(); + let references_row_idx = label_tree( + expression, + |node| node.is::(), + |acc, &child| acc | child, + ) + .get(expression) + .copied() + .unwrap_or(false); + if references_row_idx { + return Ok(None); + } + + let dtype = expression.return_dtype(&child.dtype)?; + let chunks = child + .chunks + .try_map(|_, chunk| ExpressionPlan::try_new(expression.clone(), chunk)?.optimize())?; + Ok(Some(Arc::new(child.with_chunks(dtype, chunks)))) + } +} diff --git a/vortex-layout/src/plan/plans/dict.rs b/vortex-layout/src/plan/plans/dict.rs index a89d936bd4c..365b210be54 100644 --- a/vortex-layout/src/plan/plans/dict.rs +++ b/vortex-layout/src/plan/plans/dict.rs @@ -4,7 +4,6 @@ use std::borrow::Cow; use std::sync::Arc; -use vortex_array::expr::Expression; use vortex_array::expr::is_root; use vortex_array::expr::label_is_fallible; use vortex_array::expr::label_strict; @@ -17,6 +16,7 @@ use crate::plan::ExpressionPlan; use crate::plan::Plan; use crate::plan::PlanRef; use crate::plan::new_plan; +use crate::plan::optimizer::PlanParentReduceRule; /// A physical dictionary plan with children ordered as `[codes, values]`. pub struct DictPlan { @@ -73,33 +73,6 @@ impl Plan for DictPlan { Ok(Arc::new(self.with_children(codes, values))) } - fn optimize_expression(&self, expression: &Expression) -> VortexResult> { - if !expression.return_dtype(&self.dtype)?.is_boolean() { - return Ok(None); - } - let references_root = label_tree(expression, is_root, |acc, &child| acc | child) - .get(expression) - .copied() - .unwrap_or(false); - let is_strict = label_strict(expression) - .get(expression) - .copied() - .unwrap_or(false); - let is_fallible = label_is_fallible(expression) - .get(expression) - .copied() - .unwrap_or(true); - if !references_root || !is_strict || is_fallible { - return Ok(None); - } - - let values = - ExpressionPlan::try_new(expression.clone(), Arc::clone(&self.values))?.optimize()?; - Ok(Some(Arc::new( - self.with_children(Arc::clone(&self.codes), values), - ))) - } - fn dtype(&self) -> &vortex_array::dtype::DType { &self.dtype } @@ -128,3 +101,44 @@ impl Plan for DictPlan { } } } + +/// Pushes a safe boolean expression into dictionary values. +#[derive(Debug)] +pub(crate) struct ExpressionDictRule; + +impl PlanParentReduceRule for ExpressionDictRule { + type Parent = ExpressionPlan; + + fn reduce_parent( + &self, + child: &DictPlan, + parent: &ExpressionPlan, + _child_idx: usize, + ) -> VortexResult> { + let expression = parent.expression(); + if !expression.return_dtype(&child.dtype)?.is_boolean() { + return Ok(None); + } + let references_root = label_tree(expression, is_root, |acc, &child| acc | child) + .get(expression) + .copied() + .unwrap_or(false); + let is_strict = label_strict(expression) + .get(expression) + .copied() + .unwrap_or(false); + let is_fallible = label_is_fallible(expression) + .get(expression) + .copied() + .unwrap_or(true); + if !references_root || !is_strict || is_fallible { + return Ok(None); + } + + let values = + ExpressionPlan::try_new(expression.clone(), Arc::clone(&child.values))?.optimize()?; + Ok(Some(Arc::new( + child.with_children(Arc::clone(&child.codes), values), + ))) + } +} diff --git a/vortex-layout/src/plan/plans/expression.rs b/vortex-layout/src/plan/plans/expression.rs index 82cd6f939dc..a98f1147872 100644 --- a/vortex-layout/src/plan/plans/expression.rs +++ b/vortex-layout/src/plan/plans/expression.rs @@ -15,6 +15,7 @@ use vortex_error::vortex_bail; use crate::plan::Plan; use crate::plan::PlanRef; +use crate::plan::optimizer::reduce_parent; /// A physical plan that applies an expression to the output of `child`. pub struct ExpressionPlan { @@ -64,10 +65,11 @@ impl Plan for ExpressionPlan { let expression = replace(expression, &root(), inner.expression.clone()); return Self::try_new(expression, Arc::clone(&inner.child))?.optimize(); } - if let Some(rewritten) = child.optimize_expression(&expression)? { + let parent: PlanRef = Arc::new(Self::try_new(expression, child)?); + if let Some(rewritten) = reduce_parent(&parent, 0)? { return Ok(rewritten); } - Ok(Arc::new(Self::try_new(expression, child)?)) + Ok(parent) } fn dtype(&self) -> &DType { diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index 1bc5e1ddce9..ae1e7a96124 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -10,11 +10,15 @@ mod row_idx; mod struct_; pub use chunked::ChunkedPlan; +pub(crate) use chunked::ExpressionChunkedRule; pub use dict::DictPlan; +pub(crate) use dict::ExpressionDictRule; pub use expression::ExpressionPlan; pub use flat::FlatPlan; pub use list::ListPlan; +pub(crate) use row_idx::ExpressionRowIdxRule; pub use row_idx::RowIdxPartitionPlan; pub use row_idx::RowIdxPlan; pub use row_idx::RowIdxValuesPlan; +pub(crate) use struct_::ExpressionStructRule; pub use struct_::StructPlan; diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/plans/row_idx.rs index 73315645657..baa6aadad6d 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/plans/row_idx.rs @@ -12,7 +12,6 @@ use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; -use vortex_array::expr::Expression; use vortex_array::expr::get_item; use vortex_array::expr::root; use vortex_array::expr::transform::partition; @@ -26,6 +25,7 @@ use crate::layouts::row_idx::row_idx; use crate::plan::ExpressionPlan; use crate::plan::Plan; use crate::plan::PlanRef; +use crate::plan::optimizer::PlanParentReduceRule; /// A physical plan that adds row-index expression support to its child. pub struct RowIdxPlan { @@ -53,8 +53,49 @@ impl Plan for RowIdxPlan { Ok(Self::new_ref(self.row_offset, self.child.optimize()?)) } - fn optimize_expression(&self, expression: &Expression) -> VortexResult> { - let partitioned = partition(expression.clone(), self.dtype(), |node| { + fn dtype(&self) -> &DType { + self.child.dtype() + } + + fn row_count(&self) -> u64 { + self.child.row_count() + } + + fn child_count(&self) -> usize { + 1 + } + + fn child(&self, index: usize) -> VortexResult> { + if index != 0 { + vortex_bail!("Row-index plan has no child {index}") + } + Ok(Some(Arc::clone(&self.child))) + } + + fn child_name(&self, index: usize) -> Cow<'_, str> { + if index == 0 { + Cow::Borrowed("child") + } else { + Cow::Owned(format!("child[{index}]")) + } + } +} + +/// Partitions an expression between generated row indices and the data child. +#[derive(Debug)] +pub(crate) struct ExpressionRowIdxRule; + +impl PlanParentReduceRule for ExpressionRowIdxRule { + type Parent = ExpressionPlan; + + fn reduce_parent( + &self, + child: &RowIdxPlan, + parent: &ExpressionPlan, + _child_idx: usize, + ) -> VortexResult> { + let expression = parent.expression(); + let partitioned = partition(expression.clone(), child.dtype(), |node| { if node.is::() { vec![RowIdxExpressionPartition::RowIdx] } else if vortex_array::expr::is_root(node) { @@ -68,13 +109,13 @@ impl Plan for RowIdxPlan { return match partitioned.partition_annotations[0] { RowIdxExpressionPartition::RowIdx => { let expression = replace(expression.clone(), &row_idx(), root()); - let values = RowIdxValuesPlan::new_ref(self.row_offset, self.row_count()); + let values = RowIdxValuesPlan::new_ref(child.row_offset, child.row_count()); Ok(Some( ExpressionPlan::try_new(expression, values)?.optimize()?, )) } RowIdxExpressionPartition::Child => Ok(Some( - ExpressionPlan::try_new(expression.clone(), Arc::clone(&self.child))? + ExpressionPlan::try_new(expression.clone(), Arc::clone(&child.child))? .optimize()?, )), }; @@ -106,75 +147,51 @@ impl Plan for RowIdxPlan { ) else { return Ok(None); }; - // A general pack plan is needed to expose more than one result from either branch. - if row_idx_partition.children().len() != 1 || child_partition.children().len() != 1 { - return Ok(None); - } - - let (Some(row_idx_value_name), Some(child_value_name)) = - (row_idx_pack.names.get(0), child_pack.names.get(0)) - else { - return Ok(None); - }; - let row_idx_value_name = row_idx_value_name.clone(); - let child_value_name = child_value_name.clone(); let row_idx_partition_name = partitioned.partition_names[row_idx_index].clone(); let child_partition_name = partitioned.partition_names[child_index].clone(); - let row_idx_expression = row_idx_partition.child(0).clone(); - let child_expression = child_partition.child(0).clone(); + let mut residual = partitioned.root; - let residual = replace( - partitioned.root, - &get_item(row_idx_value_name, get_item(row_idx_partition_name, root())), - get_item(RowIdxExpressionPartition::RowIdx.name(), root()), - ); - let residual = replace( - residual, - &get_item(child_value_name, get_item(child_partition_name, root())), - get_item(RowIdxExpressionPartition::Child.name(), root()), - ); + let row_idx_expression = if row_idx_partition.children().len() == 1 { + let Some(value_name) = row_idx_pack.names.get(0) else { + return Ok(None); + }; + residual = replace( + residual, + &get_item(value_name.clone(), get_item(row_idx_partition_name, root())), + get_item(RowIdxExpressionPartition::RowIdx.name(), root()), + ); + row_idx_partition.child(0).clone() + } else { + row_idx_partition.clone() + }; + let child_expression = if child_partition.children().len() == 1 { + let Some(value_name) = child_pack.names.get(0) else { + return Ok(None); + }; + residual = replace( + residual, + &get_item(value_name.clone(), get_item(child_partition_name, root())), + get_item(RowIdxExpressionPartition::Child.name(), root()), + ); + child_partition.child(0).clone() + } else { + child_partition.clone() + }; let row_idx_expression = replace(row_idx_expression, &row_idx(), root()); let row_idx_plan = ExpressionPlan::try_new( row_idx_expression, - RowIdxValuesPlan::new_ref(self.row_offset, self.row_count()), + RowIdxValuesPlan::new_ref(child.row_offset, child.row_count()), )? .optimize()?; let child_plan = - ExpressionPlan::try_new(child_expression, Arc::clone(&self.child))?.optimize()?; + ExpressionPlan::try_new(child_expression, Arc::clone(&child.child))?.optimize()?; let partitions = RowIdxPartitionPlan::try_new(row_idx_plan, child_plan)?; Ok(Some( ExpressionPlan::try_new(residual, partitions)?.optimize()?, )) } - - fn dtype(&self) -> &DType { - self.child.dtype() - } - - fn row_count(&self) -> u64 { - self.child.row_count() - } - - fn child_count(&self) -> usize { - 1 - } - - fn child(&self, index: usize) -> VortexResult> { - if index != 0 { - vortex_bail!("Row-index plan has no child {index}") - } - Ok(Some(Arc::clone(&self.child))) - } - - fn child_name(&self, index: usize) -> Cow<'_, str> { - if index == 0 { - Cow::Borrowed("child") - } else { - Cow::Owned(format!("child[{index}]")) - } - } } #[derive(Clone, Copy, PartialEq, Eq, Hash)] diff --git a/vortex-layout/src/plan/plans/struct_.rs b/vortex-layout/src/plan/plans/struct_.rs index ff2e573cf09..e5ba5eb5d4e 100644 --- a/vortex-layout/src/plan/plans/struct_.rs +++ b/vortex-layout/src/plan/plans/struct_.rs @@ -5,13 +5,15 @@ use std::borrow::Cow; use std::sync::Arc; use vortex_array::dtype::DType; -use vortex_array::expr::Expression; +use vortex_array::dtype::StructFields; use vortex_array::expr::col; +use vortex_array::expr::get_item; use vortex_array::expr::make_free_field_annotator; use vortex_array::expr::root; use vortex_array::expr::transform::partition; use vortex_array::expr::transform::replace; use vortex_array::expr::transform::replace_root_fields; +use vortex_array::scalar_fn::fns::pack::Pack; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -21,6 +23,7 @@ use crate::plan::LazyPlanChildren; use crate::plan::Plan; use crate::plan::PlanRef; use crate::plan::new_plan; +use crate::plan::optimizer::PlanParentReduceRule; /// A physical struct plan with children ordered as `[field(0), ..., field(n - 1), validity?]`. pub struct StructPlan { @@ -54,12 +57,24 @@ impl StructPlan { } } - fn with_children(&self, children: LazyPlanChildren) -> Self { - Self { + fn with_children(&self, children: LazyPlanChildren) -> VortexResult { + let fields = self.layout.struct_fields(); + let dtypes = (0..fields.nfields()) + .map(|index| { + children + .get(index)? + .map(|child| child.dtype().clone()) + .ok_or_else(|| vortex_err!("Struct field {index} has no plan")) + }) + .collect::>>()?; + Ok(Self { layout: self.layout.clone(), - dtype: self.dtype.clone(), + dtype: DType::Struct( + StructFields::new(fields.names().clone(), dtypes), + self.layout.dtype().nullability(), + ), children, - } + }) } } @@ -74,40 +89,7 @@ impl Plan for StructPlan { fn optimize(&self) -> VortexResult { let children = self.children.try_map(|_, child| child.optimize())?; - Ok(Arc::new(self.with_children(children))) - } - - fn optimize_expression(&self, expression: &Expression) -> VortexResult> { - if self.dtype.is_nullable() { - return Ok(None); - } - - let fields = self.layout.struct_fields(); - let expanded = - replace_root_fields(expression.clone(), fields).optimize_recursive(&self.dtype)?; - let partitioned = partition( - expanded.clone(), - &self.dtype, - make_free_field_annotator(fields), - )?; - if partitioned.partition_names.len() != 1 { - return Ok(None); - } - - let field_name = partitioned - .partition_names - .get(0) - .ok_or_else(|| vortex_err!("Struct expression partition has no field"))?; - let field_index = fields.find(field_name).ok_or_else(|| { - vortex_err!("Struct expression references unknown field '{field_name}'") - })?; - let field = self - .children - .get(field_index)? - .ok_or_else(|| vortex_err!("Struct field '{field_name}' has no plan"))?; - let lowered = replace(expanded, &col(field_name.clone()), root()); - - Ok(Some(ExpressionPlan::try_new(lowered, field)?.optimize()?)) + Ok(Arc::new(self.with_children(children)?)) } fn dtype(&self) -> &DType { @@ -136,3 +118,92 @@ impl Plan for StructPlan { Cow::Owned(format!("child[{index}]")) } } + +/// Partitions an expression across the fields of a struct plan. +#[derive(Debug)] +pub(crate) struct ExpressionStructRule; + +impl PlanParentReduceRule for ExpressionStructRule { + type Parent = ExpressionPlan; + + fn reduce_parent( + &self, + child: &StructPlan, + parent: &ExpressionPlan, + _child_idx: usize, + ) -> VortexResult> { + if child.dtype.is_nullable() { + return Ok(None); + } + + let expression = parent.expression(); + let fields = child.layout.struct_fields(); + let expanded = + replace_root_fields(expression.clone(), fields).optimize_recursive(&child.dtype)?; + let partitioned = partition( + expanded.clone(), + &child.dtype, + make_free_field_annotator(fields), + )?; + if partitioned.partition_names.is_empty() { + return Ok(None); + } + + if partitioned.partition_names.len() == 1 { + let field_name = partitioned + .partition_names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition has no field"))?; + let field_index = fields.find(field_name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{field_name}'") + })?; + let field = child + .children + .get(field_index)? + .ok_or_else(|| vortex_err!("Struct field '{field_name}' has no plan"))?; + let lowered = replace(expanded, &col(field_name.clone()), root()); + + return Ok(Some(ExpressionPlan::try_new(lowered, field)?.optimize()?)); + } + + let mut residual = partitioned.root; + let mut field_expressions = vec![None; fields.nfields()]; + for index in 0..partitioned.partitions.len() { + let field_name = &partitioned.partition_names[index]; + let partition = &partitioned.partitions[index]; + let field_index = fields.find(field_name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{field_name}'") + })?; + + let lowered = if let Some(pack) = partition.as_opt::() + && partition.children().len() == 1 + { + let value_name = pack + .names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition pack is empty"))?; + residual = replace( + residual, + &get_item(value_name.clone(), get_item(field_name.clone(), root())), + get_item(field_name.clone(), root()), + ); + replace(partition.child(0).clone(), &col(field_name.clone()), root()) + } else { + replace(partition.clone(), &col(field_name.clone()), root()) + }; + field_expressions[field_index] = Some(lowered); + } + + let children = child.children.try_map(|index, field| { + let Some(expression) = field_expressions.get(index).and_then(Option::as_ref) else { + return Ok(field); + }; + ExpressionPlan::try_new(expression.clone(), field)?.optimize() + })?; + let rewritten: PlanRef = Arc::new(child.with_children(children)?); + + Ok(Some(Arc::new(ExpressionPlan::try_new( + residual, rewritten, + )?))) + } +} diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 6cdd04c3b05..01fa7fcecfb 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -148,7 +148,10 @@ fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { .into_layout(); let expression = and( gt(row_idx(), lit(11_u64)), - gt(get_item("a", root()), lit(5_i32)), + and( + gt(get_item("a", root()), lit(5_i32)), + gt(get_item("b", root()), lit(7_i32)), + ), ); let plan: PlanRef = Arc::new(ExpressionPlan::try_new( expression, @@ -156,7 +159,7 @@ fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { )?); insta::assert_snapshot!(plan.tree_display(), @r" - root: ExpressionPlan(bool, rows=3) expr=((#row_idx > 11u64) and ($.a > 5i32)) + root: ExpressionPlan(bool, rows=3) expr=((#row_idx > 11u64) and (($.a > 5i32) and ($.b > 7i32))) child: RowIdxPlan({a=i32, b=i32}, rows=3) child: StructPlan({a=i32, b=i32}, rows=3) a: DictPlan(i32, rows=3) @@ -167,14 +170,18 @@ fn expression_partitions_across_row_idx_and_struct() -> VortexResult<()> { let optimized = plan.optimize()?; insta::assert_snapshot!(optimized.tree_display(), @r" - root: ExpressionPlan(bool, rows=3) expr=($.row_idx and $.child) - child: RowIdxPartitionPlan({row_idx=bool, child=bool}, rows=3) + root: ExpressionPlan(bool, rows=3) expr=(($.row_idx and $.child.child_0) and $.child.child_1) + child: RowIdxPartitionPlan({row_idx=bool, child={child_0=bool, child_1=bool}}, rows=3) row_idx: ExpressionPlan(bool, rows=3) expr=($ > 11u64) child: RowIdxValuesPlan(u64, rows=3) - child: DictPlan(bool, rows=3) - codes: FlatPlan(u8, rows=3) - values: ExpressionPlan(bool, rows=2) expr=($ > 5i32) - child: FlatPlan(i32, rows=2) + child: ExpressionPlan({child_0=bool, child_1=bool}, rows=3) expr=pack(child_0: $.a, child_1: $.b) + child: StructPlan({a=bool, b=bool}, rows=3) + a: DictPlan(bool, rows=3) + codes: FlatPlan(u8, rows=3) + values: ExpressionPlan(bool, rows=2) expr=($ > 5i32) + child: FlatPlan(i32, rows=2) + b: ExpressionPlan(bool, rows=3) expr=($ > 7i32) + child: FlatPlan(i32, rows=3) "); let residual = optimized .as_any() @@ -312,6 +319,103 @@ fn expression_pushes_through_struct_field_with_heterogeneous_chunks() -> VortexR Ok(()) } +#[test] +fn multi_field_struct_expression_pushes_into_each_field() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let dictionary = DictLayout::new( + flat(2, value_dtype.clone(), 0), + flat(3, primitive(PType::U8, Nullability::NonNullable), 1), + ) + .into_layout(); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([ + ("a", value_dtype.clone()), + ("b", value_dtype.clone()), + ("c", value_dtype.clone()), + ]), + Nullability::NonNullable, + ), + vec![ + dictionary, + flat(3, value_dtype.clone(), 2), + flat(3, value_dtype, 3), + ], + ) + .into_layout(); + let expression = and( + gt(get_item("a", root()), lit(5_i32)), + gt(get_item("b", root()), lit(7_i32)), + ); + let plan: PlanRef = Arc::new(ExpressionPlan::try_new(expression, make_plan(layout)?)?); + + insta::assert_snapshot!(plan.tree_display(), @r" + root: ExpressionPlan(bool, rows=3) expr=(($.a > 5i32) and ($.b > 7i32)) + child: StructPlan({a=i32, b=i32, c=i32}, rows=3) + a: DictPlan(i32, rows=3) + codes: FlatPlan(u8, rows=3) + values: FlatPlan(i32, rows=2) + b: FlatPlan(i32, rows=3) + c: FlatPlan(i32, rows=3) + "); + + let optimized = plan.optimize()?; + + insta::assert_snapshot!(optimized.tree_display(), @r" + root: ExpressionPlan(bool, rows=3) expr=($.a and $.b) + child: StructPlan({a=bool, b=bool, c=i32}, rows=3) + a: DictPlan(bool, rows=3) + codes: FlatPlan(u8, rows=3) + values: ExpressionPlan(bool, rows=2) expr=($ > 5i32) + child: FlatPlan(i32, rows=2) + b: ExpressionPlan(bool, rows=3) expr=($ > 7i32) + child: FlatPlan(i32, rows=3) + c: FlatPlan(i32, rows=3) + "); + assert_eq!( + optimized.tree_display().to_string(), + optimized.optimize()?.tree_display().to_string() + ); + Ok(()) +} + +#[test] +fn multi_field_struct_expression_keeps_cross_field_refinement() -> VortexResult<()> { + let value_dtype = primitive(PType::I32, Nullability::NonNullable); + let dictionary = DictLayout::new( + flat(2, value_dtype.clone(), 0), + flat(3, primitive(PType::U8, Nullability::NonNullable), 1), + ) + .into_layout(); + let layout = StructLayout::new( + 3, + DType::Struct( + StructFields::from_iter([("a", value_dtype.clone()), ("b", value_dtype.clone())]), + Nullability::NonNullable, + ), + vec![dictionary, flat(3, value_dtype, 2)], + ) + .into_layout(); + let expression = gt( + checked_add(get_item("a", root()), get_item("b", root())), + lit(10_i32), + ); + let plan = ExpressionPlan::try_new(expression, make_plan(layout)?)?; + + let optimized = plan.optimize()?; + + insta::assert_snapshot!(optimized.tree_display(), @r" + root: ExpressionPlan(bool, rows=3) expr=(($.a + $.b) > 10i32) + child: StructPlan({a=i32, b=i32}, rows=3) + a: DictPlan(i32, rows=3) + codes: FlatPlan(u8, rows=3) + values: FlatPlan(i32, rows=2) + b: FlatPlan(i32, rows=3) + "); + Ok(()) +} + #[test] fn dictionary_pushdown_rejects_unsafe_expressions() -> VortexResult<()> { let value_dtype = primitive(PType::I32, Nullability::NonNullable);