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 4ccd5a9f325..64473c3d95f 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_error::VortexResult; @@ -58,9 +60,8 @@ 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; /// Returns the dtype produced by this plan. 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 fd145ec5221..853e23987e5 100644 --- a/vortex-layout/src/plan/plans/chunked.rs +++ b/vortex-layout/src/plan/plans/chunked.rs @@ -5,13 +5,17 @@ use std::borrow::Cow; use std::sync::Arc; use vortex_array::dtype::DType; +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; use crate::plan::new_plan; +use crate::plan::optimizer::PlanParentReduceRule; /// A physical plan with one child per row chunk. pub struct ChunkedPlan { @@ -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,8 @@ 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 dtype(&self) -> &DType { @@ -82,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 47827c6950c..365b210be54 100644 --- a/vortex-layout/src/plan/plans/dict.rs +++ b/vortex-layout/src/plan/plans/dict.rs @@ -4,13 +4,19 @@ use std::borrow::Cow; use std::sync::Arc; +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; +use crate::plan::optimizer::PlanParentReduceRule; /// A physical dictionary plan with children ordered as `[codes, values]`. pub struct DictPlan { @@ -95,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 686ad81b007..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 { @@ -62,12 +63,13 @@ 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(); } - Ok(Arc::new(Self::try_new(expression, child)?)) + let parent: PlanRef = Arc::new(Self::try_new(expression, child)?); + if let Some(rewritten) = reduce_parent(&parent, 0)? { + return Ok(rewritten); + } + 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 a6ffe427d3c..ae1e7a96124 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -10,9 +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 58d9370122b..baa6aadad6d 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/vortex-layout/src/plan/plans/row_idx.rs @@ -3,14 +3,29 @@ 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::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; +use crate::plan::optimizer::PlanParentReduceRule; /// A physical plan that adds row-index expression support to its child. pub struct RowIdxPlan { @@ -65,3 +80,278 @@ impl Plan for RowIdxPlan { } } } + +/// 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) { + 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(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(&child.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); + }; + let row_idx_partition_name = partitioned.partition_names[row_idx_index].clone(); + let child_partition_name = partitioned.partition_names[child_index].clone(); + let mut residual = partitioned.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(child.row_offset, child.row_count()), + )? + .optimize()?; + let child_plan = + 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()?, + )) + } +} + +#[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 e216a3f56c3..e5ba5eb5d4e 100644 --- a/vortex-layout/src/plan/plans/struct_.rs +++ b/vortex-layout/src/plan/plans/struct_.rs @@ -5,13 +5,25 @@ use std::borrow::Cow; use std::sync::Arc; use vortex_array::dtype::DType; +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; use crate::layouts::struct_::StructLayout; +use crate::plan::ExpressionPlan; 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 { @@ -45,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, - } + }) } } @@ -64,8 +88,8 @@ impl Plan for StructPlan { } fn optimize(&self) -> VortexResult { - let children = self.children.map(|_, child| child.optimize()); - Ok(Arc::new(self.with_children(children))) + let children = self.children.try_map(|_, child| child.optimize())?; + Ok(Arc::new(self.with_children(children)?)) } fn dtype(&self) -> &DType { @@ -94,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 6d5bc469ac9..01fa7fcecfb 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -1,14 +1,18 @@ // 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; +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; @@ -22,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; @@ -51,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() @@ -66,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, @@ -115,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() @@ -139,318 +103,369 @@ 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)?; +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(), 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 + 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 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 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, +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 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 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()); - - 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)), + 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 nullable_plan = make_plan(nullable)?; - assert_eq!( - nullable_plan - .child(2)? - .ok_or_else(|| vortex_err!("missing validity"))? - .dtype(), - &DType::Bool(Nullability::NonNullable) + let expression = and( + gt(row_idx(), lit(11_u64)), + 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, + RowIdxPlan::new_ref(10, make_plan(layout)?), + )?); + + insta::assert_snapshot!(plan.tree_display(), @r" + 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) + 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: 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: 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() + .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 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( +fn chunked_plan_preserves_global_row_index_expressions() -> VortexResult<()> { + let dtype = primitive(PType::I32, Nullability::NonNullable); + let layout = ChunkedLayout::new( 3, - DType::Struct(fields.clone(), Nullability::NonNullable), - vec![ - flat(3, field_dtype.clone(), 0), - flat(3, field_dtype.clone(), 1), - ], + dtype.clone(), + OwnedLayoutChildren::layout_children(vec![flat(1, dtype.clone(), 0), flat(2, dtype, 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 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"))?; - 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), - ], - ) - .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) - ); + assert_eq!(expression.expression(), &row_idx()); + assert!(expression.child_plan().as_any().is::()); Ok(()) } #[test] -fn struct_plan_defers_unrequested_fields_through_optimization() -> VortexResult<()> { - let field_dtype = primitive(PType::I32, Nullability::NonNullable); +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( - 1, + 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(1, field_dtype.clone(), 0), unsupported(1, field_dtype)], + vec![dictionary, flat(3, value_dtype, 2)], ) .into_layout(); - let plan = make_plan(layout)?.optimize()?; + let plan: PlanRef = Arc::new(ExpressionPlan::try_new( + gt(get_item("a", root()), lit(5_i32)), + make_plan(layout)?, + )?); - 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'") - ); + insta::assert_snapshot!(plan.tree_display(), @r" + 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: ExpressionPlan(bool, rows=2) expr=($ > 5i32) + child: FlatPlan(i32, rows=2) + "); Ok(()) } #[test] -fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> { - let field_dtype = primitive(PType::I32, Nullability::NonNullable); +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 chunks = ChunkedLayout::new( + 5, + value_dtype.clone(), + OwnedLayoutChildren::layout_children(vec![dictionary, flat(2, value_dtype.clone(), 2)]), + ) + .into_layout(); let layout = StructLayout::new( - 3, + 5, 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![chunks, flat(5, value_dtype, 3)], ) .into_layout(); let plan: PlanRef = Arc::new(ExpressionPlan::try_new( - get_item("a", root()), + gt(get_item("a", root()), lit(5_i32)), 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) + 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) "); - 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()) - } - } + let optimized = plan.optimize()?; - insta::assert_snapshot!(plan.tree_display_builder().with(DepthExtractor), @r" - root: depth=0 - child: depth=1 - a: depth=2 - b: depth=2 + insta::assert_snapshot!(optimized.tree_display(), @r" + 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(()) +} - let nullable_fields = StructFields::from_iter([ - ("a", primitive(PType::I32, Nullability::NonNullable)), - ("b", primitive(PType::I32, Nullability::NonNullable)), - ]); - let nullable_layout = StructLayout::new( +#[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(nullable_fields, Nullability::Nullable), + DType::Struct( + StructFields::from_iter([ + ("a", value_dtype.clone()), + ("b", value_dtype.clone()), + ("c", value_dtype.clone()), + ]), + Nullability::NonNullable, + ), vec![ - flat(3, DType::Bool(Nullability::NonNullable), 2), - flat(3, primitive(PType::I32, Nullability::NonNullable), 3), - flat(3, primitive(PType::I32, Nullability::NonNullable), 4), + dictionary, + flat(3, value_dtype.clone(), 2), + flat(3, value_dtype, 3), ], ) .into_layout(); - let nullable = make_plan(nullable_layout)?; - insta::assert_snapshot!(nullable.tree_display_builder(), @r" - root: - a: - b: - validity: + 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 chunked_plan_display_names_chunks() -> VortexResult<()> { - let dtype = primitive(PType::I32, Nullability::NonNullable); - let layout = ChunkedLayout::new( +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.clone(), - OwnedLayoutChildren::layout_children(vec![flat(2, dtype.clone(), 0), flat(1, dtype, 1)]), + 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 expression = gt( + checked_add(get_item("a", root()), get_item("b", root())), + lit(10_i32), + ); + let plan = ExpressionPlan::try_new(expression, make_plan(layout)?)?; - insta::assert_snapshot!(plan.display_tree(), @r" - root: ChunkedPlan(i32, rows=3) - chunks[0]: FlatPlan(i32, rows=2) - chunks[1]: FlatPlan(i32, rows=1) + 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 dict_plan_display_names_logical_children() -> VortexResult<()> { - let layout = DictLayout::new( - flat(2, primitive(PType::I32, Nullability::NonNullable), 0), +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(); - let plan = make_plan(layout)?; - insta::assert_snapshot!(plan.tree_display(), @r" - root: DictPlan(i32, rows=3) - codes: FlatPlan(u8, rows=3) - values: FlatPlan(i32, rows=2) - "); + 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 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, - ) - .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)), +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 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(()) -} + let plan = ExpressionPlan::try_new(gt(get_item("a", root()), lit(5_i32)), make_plan(layout)?)?; -#[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 + let optimized = plan.optimize()?; + let expression_plan = optimized .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_or_else(|| vortex_err!("Nullable struct expression unexpectedly pushed down"))?; + assert!(expression_plan.child_plan().as_any().is::()); Ok(()) }