diff --git a/vortex-layout/src/plan/mod.rs b/vortex-layout/src/plan/mod.rs index 06f7acd0031..802e210b07e 100644 --- a/vortex-layout/src/plan/mod.rs +++ b/vortex-layout/src/plan/mod.rs @@ -11,6 +11,7 @@ mod children; mod display; mod lower; mod optimize; +pub mod optimizer; mod plans; mod typed; mod vtable; @@ -38,12 +39,13 @@ pub use plans::PackPlan; pub use plans::RowIdx; pub use plans::RowIdxData; pub use plans::RowIdxPlan; -pub use plans::RowIdxPlanMetadata; pub use plans::SegmentScan; pub use plans::SegmentScanData; pub use plans::SegmentScanPlan; pub use plans::Take; pub use plans::TakePlan; +pub use plans::plan_row_idx_expression; +pub use plans::row_idx_dtype; pub use typed::DynPlan; pub use typed::Plan; pub use typed::PlanParts; diff --git a/vortex-layout/src/plan/optimize.rs b/vortex-layout/src/plan/optimize.rs index d78798cdb65..f471b7744f9 100644 --- a/vortex-layout/src/plan/optimize.rs +++ b/vortex-layout/src/plan/optimize.rs @@ -1,15 +1,35 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Generic bottom-up optimization over physical plans. +//! Plan optimization. +//! +//! The optimizer applies static rewrites top-down, optimizes children, then retries rewrites +//! exposed by the optimized children. use vortex_error::VortexResult; -use crate::plan::Eval; use crate::plan::PlanRef; +use crate::plan::optimizer::reduce_parent; +use crate::plan::optimizer::reduce_plan; + +fn reduce(plan: &PlanRef) -> VortexResult> { + if let Some(rewritten) = reduce_plan(plan)? { + return Ok(Some(rewritten)); + } + for child_idx in 0..plan.child_count() { + if let Some(rewritten) = reduce_parent(plan, child_idx)? { + return Ok(Some(rewritten)); + } + } + Ok(None) +} /// Optimizes `plan`, preserving its dtype and row domain. pub fn optimize(plan: PlanRef) -> VortexResult { + if let Some(rewritten) = reduce(&plan)? { + return optimize(rewritten); + } + let mut children = Vec::with_capacity(plan.child_count()); let mut changed = false; for child in plan.children().iter() { @@ -19,17 +39,13 @@ pub fn optimize(plan: PlanRef) -> VortexResult { children.push(optimized); } - let plan = if changed { - plan.with_children(children)? - } else { - plan - }; - - let Some(eval) = plan.as_opt::() else { + if !changed { return Ok(plan); - }; - if eval.expression().is_root() { - return eval.child_plan(); + } + + let plan = plan.with_children(children)?; + if let Some(rewritten) = reduce(&plan)? { + return optimize(rewritten); } Ok(plan) } diff --git a/vortex-layout/src/plan/optimizer/mod.rs b/vortex-layout/src/plan/optimizer/mod.rs new file mode 100644 index 00000000000..a63006be62e --- /dev/null +++ b/vortex-layout/src/plan/optimizer/mod.rs @@ -0,0 +1,57 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Static rewrite rules for physical plans. + +mod rules; + +pub use rules::DynPlanParentReduceRule; +pub use rules::DynPlanReduceRule; +pub use rules::PlanParentReduceRule; +pub use rules::PlanParentReduceRuleAdapter; +pub use rules::PlanParentRuleSet; +pub use rules::PlanReduceRule; +pub use rules::PlanReduceRuleAdapter; +pub use rules::PlanRuleSet; +use vortex_error::VortexResult; + +use super::Concat; +use super::Eval; +use super::Pack; +use super::PlanRef; +use super::Take; +use super::plans::EvalIdentityRule; +use super::plans::ExpressionConcatRule; +use super::plans::ExpressionPackRule; +use super::plans::ExpressionTakeRule; + +static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter = + PlanReduceRuleAdapter::new(EvalIdentityRule); + +static PLAN_RULES: PlanRuleSet = PlanRuleSet::new(&[&EVAL_IDENTITY_RULE]); + +static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionConcatRule); +static EXPRESSION_TAKE_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionTakeRule); +static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter = + PlanParentReduceRuleAdapter::new(ExpressionPackRule); + +static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[ + &EXPRESSION_CONCAT_RULE, + &EXPRESSION_TAKE_RULE, + &EXPRESSION_PACK_RULE, +]); + +/// Attempts a static rewrite for `plan`. +pub(crate) fn reduce_plan(plan: &PlanRef) -> VortexResult> { + PLAN_RULES.evaluate(plan) +} + +/// 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/optimizer/rules.rs b/vortex-layout/src/plan/optimizer/rules.rs new file mode 100644 index 00000000000..3276aab22f6 --- /dev/null +++ b/vortex-layout/src/plan/optimizer/rules.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Typed and type-erased interfaces for plan rewrites. + +use std::any::type_name; +use std::fmt::Debug; +use std::marker::PhantomData; + +use vortex_error::VortexResult; + +use crate::plan::Plan; +use crate::plan::PlanRef; +use crate::plan::PlanVTable; + +/// A rewrite over one concrete plan operator. +/// +/// Rules return one rewrite without recursively optimizing the replacement. The plan optimizer +/// owns traversal and drives further rewrites. +pub trait PlanReduceRule: Debug + Send + Sync + 'static { + /// Attempts to replace `plan`. + fn reduce(&self, plan: &Plan

) -> VortexResult>; +} + +/// Type-erased interface used by [`PlanRuleSet`]. +pub trait DynPlanReduceRule: Debug + Send + Sync + 'static { + /// Returns whether this rule supports the concrete plan operator. + fn matches(&self, plan: &PlanRef) -> bool; + + /// Attempts to replace `plan`. + fn reduce(&self, plan: &PlanRef) -> VortexResult>; +} + +/// Bridges a typed [`PlanReduceRule`] to a type-erased static registry. +pub struct PlanReduceRuleAdapter { + rule: R, + _plan: PhantomData P>, +} + +impl PlanReduceRuleAdapter { + /// Creates an adapter for a typed plan rule. + pub const fn new(rule: R) -> Self { + Self { + rule, + _plan: PhantomData, + } + } +} + +impl Debug for PlanReduceRuleAdapter +where + P: PlanVTable, + R: PlanReduceRule

, +{ + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PlanReduceRuleAdapter") + .field("plan", &type_name::

()) + .field("rule", &self.rule) + .finish() + } +} + +impl DynPlanReduceRule for PlanReduceRuleAdapter +where + P: PlanVTable, + R: PlanReduceRule

, +{ + fn matches(&self, plan: &PlanRef) -> bool { + plan.is::

() + } + + fn reduce(&self, plan: &PlanRef) -> VortexResult> { + let Some(plan) = plan.as_opt::

() else { + return Ok(None); + }; + self.rule.reduce(plan) + } +} + +/// An ordered static collection of single-plan rewrite rules. +pub struct PlanRuleSet { + rules: &'static [&'static dyn DynPlanReduceRule], +} + +impl PlanRuleSet { + /// Creates a rule set whose first successful rewrite wins. + pub const fn new(rules: &'static [&'static dyn DynPlanReduceRule]) -> Self { + Self { rules } + } + + /// Evaluates rules registered for the concrete plan operator. + pub fn evaluate(&self, plan: &PlanRef) -> VortexResult> { + for rule in self.rules { + if !rule.matches(plan) { + continue; + } + let Some(reduced) = rule.reduce(plan)? else { + continue; + }; + + #[cfg(debug_assertions)] + { + vortex_error::vortex_ensure!( + reduced.row_count() == plan.row_count(), + "Plan rewrite from {rule:?} changed row count from {} to {}", + plan.row_count(), + reduced.row_count() + ); + vortex_error::vortex_ensure!( + reduced.dtype() == plan.dtype(), + "Plan rewrite from {rule:?} changed dtype from {} to {}", + plan.dtype(), + reduced.dtype() + ); + } + + return Ok(Some(reduced)); + } + Ok(None) + } +} + +/// A metadata-only rewrite where a child plan rewrites its parent plan. +/// +/// Rules return one rewrite without recursively optimizing the replacement. The plan optimizer +/// owns traversal and drives further rewrites. +pub trait PlanParentReduceRule: Debug + Send + Sync + 'static { + /// The concrete parent operator matched by this rule. + type Parent: PlanVTable; + + /// Attempts to replace `parent` based on its child at `child_idx`. + fn reduce_parent( + &self, + child: &Plan, + parent: &Plan, + child_idx: usize, + ) -> VortexResult>; +} + +/// Type-erased interface used by [`PlanParentRuleSet`]. +pub trait DynPlanParentReduceRule: Debug + Send + Sync + 'static { + /// Returns whether this rule supports the concrete child and parent operators. + fn matches(&self, child: &PlanRef, parent: &PlanRef) -> bool; + + /// Attempts to replace `parent` based on `child` at `child_idx`. + fn reduce_parent( + &self, + child: &PlanRef, + parent: &PlanRef, + child_idx: usize, + ) -> VortexResult>; +} + +/// Bridges a typed [`PlanParentReduceRule`] to a type-erased static registry. +pub struct PlanParentReduceRuleAdapter { + rule: R, + _child: PhantomData C>, +} + +impl PlanParentReduceRuleAdapter { + /// Creates an adapter for a typed parent-child rule. + pub const fn new(rule: R) -> Self { + Self { + rule, + _child: PhantomData, + } + } +} + +impl Debug for PlanParentReduceRuleAdapter +where + C: PlanVTable, + R: PlanParentReduceRule, +{ + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PlanParentReduceRuleAdapter") + .field("parent", &type_name::()) + .field("child", &type_name::()) + .field("rule", &self.rule) + .finish() + } +} + +impl DynPlanParentReduceRule for PlanParentReduceRuleAdapter +where + C: PlanVTable, + R: PlanParentReduceRule, +{ + fn matches(&self, child: &PlanRef, parent: &PlanRef) -> bool { + child.is::() && parent.is::() + } + + fn reduce_parent( + &self, + child: &PlanRef, + parent: &PlanRef, + child_idx: usize, + ) -> VortexResult> { + let Some(child) = child.as_opt::() else { + return Ok(None); + }; + let Some(parent) = parent.as_opt::() else { + return Ok(None); + }; + self.rule.reduce_parent(child, parent, child_idx) + } +} + +/// An ordered static collection of parent-child plan rewrite rules. +pub struct PlanParentRuleSet { + rules: &'static [&'static dyn DynPlanParentReduceRule], +} + +impl PlanParentRuleSet { + /// Creates a rule set whose first successful rewrite wins. + pub const fn new(rules: &'static [&'static dyn DynPlanParentReduceRule]) -> Self { + Self { rules } + } + + /// Evaluates rules registered for the concrete `(parent, child)` pair. + pub fn evaluate( + &self, + child: &PlanRef, + parent: &PlanRef, + child_idx: usize, + ) -> VortexResult> { + for rule in self.rules { + if !rule.matches(child, parent) { + continue; + } + let Some(reduced) = rule.reduce_parent(child, parent, child_idx)? else { + continue; + }; + + #[cfg(debug_assertions)] + { + vortex_error::vortex_ensure!( + reduced.row_count() == parent.row_count(), + "Plan rewrite from {rule:?} changed row count from {} to {}", + parent.row_count(), + reduced.row_count() + ); + vortex_error::vortex_ensure!( + reduced.dtype() == parent.dtype(), + "Plan rewrite from {rule:?} changed dtype from {} to {}", + parent.dtype(), + reduced.dtype() + ); + } + + return Ok(Some(reduced)); + } + Ok(None) + } +} diff --git a/vortex-layout/src/plan/plans/concat.rs b/vortex-layout/src/plan/plans/concat.rs index 2e2893fde5f..2e814156309 100644 --- a/vortex-layout/src/plan/plans/concat.rs +++ b/vortex-layout/src/plan/plans/concat.rs @@ -10,12 +10,15 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_session::registry::CachedId; +use crate::plan::Eval; +use crate::plan::EvalPlan; use crate::plan::Plan; use crate::plan::PlanChildren; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; +use crate::plan::optimizer::PlanParentReduceRule; /// Concatenates its children row-wise. #[derive(Clone, Debug)] @@ -140,3 +143,28 @@ impl PlanVTable for Concat { Cow::Owned(format!("chunks[{index}]")) } } + +/// Pushes an expression into every chunk of a [`Concat`]. +#[derive(Debug)] +pub(crate) struct ExpressionConcatRule; + +impl PlanParentReduceRule for ExpressionConcatRule { + type Parent = Eval; + + fn reduce_parent( + &self, + child: &Plan, + parent: &Plan, + _child_idx: usize, + ) -> VortexResult> { + let expression = parent.expression(); + let chunks = child + .children() + .iter() + .map(|chunk| Ok(EvalPlan::try_new(expression.clone(), chunk?)?.into_plan())) + .collect::>>()?; + Ok(Some( + ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(), + )) + } +} diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/plans/eval.rs index 70a645da687..f1edb967e41 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/vortex-layout/src/plan/plans/eval.rs @@ -17,6 +17,7 @@ use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; use crate::plan::check_child_count; +use crate::plan::optimizer::PlanReduceRule; /// Applies an expression to the output of its child. #[derive(Clone, Debug)] @@ -123,3 +124,17 @@ fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> V } Ok(()) } + +/// Removes an [`Eval`] whose expression is the identity expression. +#[derive(Debug)] +pub(crate) struct EvalIdentityRule; + +impl PlanReduceRule for EvalIdentityRule { + fn reduce(&self, plan: &Plan) -> VortexResult> { + if plan.expression().is_root() { + Ok(Some(plan.child_plan()?)) + } else { + Ok(None) + } + } +} diff --git a/vortex-layout/src/plan/plans/mod.rs b/vortex-layout/src/plan/plans/mod.rs index 2e4a6dbad5e..7b6a1893d0b 100644 --- a/vortex-layout/src/plan/plans/mod.rs +++ b/vortex-layout/src/plan/plans/mod.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod concat; -mod eval; +pub(crate) mod eval; mod list_pack; mod pack; mod row_idx; @@ -12,21 +12,26 @@ mod take; pub use concat::Concat; pub use concat::ConcatData; pub use concat::ConcatPlan; +pub(crate) use concat::ExpressionConcatRule; pub use eval::Eval; pub use eval::EvalData; +pub(crate) use eval::EvalIdentityRule; pub use eval::EvalPlan; pub use list_pack::ListPack; pub use list_pack::ListPackData; pub use list_pack::ListPackPlan; +pub(crate) use pack::ExpressionPackRule; pub use pack::Pack; pub use pack::PackData; pub use pack::PackPlan; pub use row_idx::RowIdx; pub use row_idx::RowIdxData; pub use row_idx::RowIdxPlan; -pub use row_idx::RowIdxPlanMetadata; +pub use row_idx::plan_row_idx_expression; +pub use row_idx::row_idx_dtype; pub use segment_scan::SegmentScan; pub use segment_scan::SegmentScanData; pub use segment_scan::SegmentScanPlan; +pub(crate) use take::ExpressionTakeRule; pub use take::Take; pub use take::TakePlan; diff --git a/vortex-layout/src/plan/plans/pack.rs b/vortex-layout/src/plan/plans/pack.rs index 9cf47ebe4b7..6497cc63d45 100644 --- a/vortex-layout/src/plan/plans/pack.rs +++ b/vortex-layout/src/plan/plans/pack.rs @@ -5,20 +5,39 @@ use std::borrow::Cow; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::descendent_bound_annotations; +use vortex_array::expr::make_bound_free_field_annotator; +use vortex_array::expr::transform::partition_bound; +use vortex_array::expr::traversal::NodeExt; +use vortex_array::expr::traversal::Transformed; +use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::get_item::GetItem; +use vortex_array::scalar_fn::fns::pack::Pack as PackFn; +use vortex_array::scalar_fn::fns::pack::PackOptions; +use vortex_array::scalar_fn::fns::select::Select; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::registry::CachedId; +use crate::plan::Eval; +use crate::plan::EvalPlan; use crate::plan::Plan; use crate::plan::PlanChildren; use crate::plan::PlanId; use crate::plan::PlanParts; use crate::plan::PlanRef; use crate::plan::PlanVTable; +use crate::plan::optimizer::PlanParentReduceRule; /// Assembles a struct from one child per field, plus an optional trailing validity child. #[derive(Clone, Debug)] @@ -224,3 +243,339 @@ fn validate_validity_child(expected_row_count: u64, child: &PlanRef) -> VortexRe } Ok(()) } + +impl PackPlan { + /// Rebuilds this plan with only `fields`, which must be a subset of the current fields. + /// + /// Pruning is only sound for a non-nullable struct: dropping a field of a nullable struct + /// would drop the validity child that the remaining fields depend on. + pub(crate) fn with_pruned_fields( + &self, + fields: Vec<(FieldName, PlanRef)>, + ) -> VortexResult { + vortex_ensure!( + !self.dtype().is_nullable(), + "Cannot prune fields from a nullable Pack" + ); + let struct_fields = StructFields::from_iter( + fields + .iter() + .map(|(name, plan)| (name.clone(), plan.dtype().clone())), + ); + let field_plans = fields.into_iter().map(|(_, plan)| plan).collect::>(); + PackPlan::try_new( + struct_fields, + Nullability::NonNullable, + self.row_count(), + field_plans, + None, + ) + } +} + +/// Pushes an expression into the referenced fields of a [`Pack`], pruning the rest. +#[derive(Debug)] +pub(crate) struct ExpressionPackRule; + +impl PlanParentReduceRule for ExpressionPackRule { + type Parent = Eval; + + fn reduce_parent( + &self, + child: &Plan, + parent: &Plan, + _child_idx: usize, + ) -> VortexResult> { + if child.dtype().is_nullable() { + return Ok(None); + } + + let expression = parent.expression(); + let fields = child.fields(); + let referenced_fields = + descendent_bound_annotations(expression, make_bound_free_field_annotator(fields)) + .get(&ExactBoundExpr(expression.clone())) + .vortex_expect("Bound expression missing free-field annotations") + .clone(); + let expanded_root = expanded_struct_root(child.dtype(), fields)?; + let expanded = expand_struct_root(expression.clone(), &expanded_root, fields)?; + let partitioned = + partition_bound(expanded.clone(), make_bound_free_field_annotator(fields))?; + + if partitioned.partition_names.is_empty() { + let selected_indices = fields + .names() + .iter() + .enumerate() + .filter_map(|(index, name)| referenced_fields.contains(name).then_some(index)) + .collect::>(); + if selected_indices.len() == fields.nfields() { + return Ok(None); + } + + let pruned_fields = selected_indices + .into_iter() + .map(|field_index| { + Ok(( + field_name(fields, field_index)?, + field_plan(child, field_index)?, + )) + }) + .collect::>>()?; + let rewritten = child.with_pruned_fields(pruned_fields)?.into_plan(); + return Ok(Some( + EvalPlan::try_new(expression.clone(), rewritten)?.into_plan(), + )); + } + + if partitioned.partition_names.len() == 1 { + let name = partitioned + .partition_names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition has no field"))?; + let index = fields.find(name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{name}'") + })?; + let field = field_plan(child, index)?; + let lowered = step_into_struct_field(expanded, name, field.dtype().clone())?; + return Ok(Some(EvalPlan::try_new(lowered, field)?.into_plan())); + } + + let residual = partitioned.root; + let mut collapsed = Vec::with_capacity(partitioned.partitions.len()); + let mut field_expressions = vec![None; fields.nfields()]; + for index in 0..partitioned.partitions.len() { + let name = &partitioned.partition_names[index]; + let partition = &partitioned.partitions[index]; + let field_index = fields.find(name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{name}'") + })?; + let field = field_plan(child, field_index)?; + let lowered = if let Some(pack) = partition + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + && partition.children().len() == 1 + { + let value_name = pack + .names + .get(0) + .ok_or_else(|| vortex_err!("Struct expression partition pack is empty"))?; + collapsed.push((name.clone(), value_name.clone())); + partition.children()[0].clone() + } else { + partition.clone() + }; + let lowered = step_into_struct_field(lowered, name, field.dtype().clone())?; + field_expressions[field_index] = Some(lowered); + } + + let mut fields_changed = partitioned.partition_names.len() != fields.nfields(); + let mut pruned_fields = Vec::with_capacity(partitioned.partition_names.len()); + for (field_index, expression) in field_expressions.into_iter().enumerate() { + let Some(expression) = expression else { + continue; + }; + let field = field_plan(child, field_index)?; + let field = if is_identity_expression(&expression, field.dtype())? { + field + } else { + fields_changed = true; + EvalPlan::try_new(expression, field)?.into_plan() + }; + pruned_fields.push((field_name(fields, field_index)?, field)); + } + let rewritten = if fields_changed { + child.with_pruned_fields(pruned_fields)?.into_plan() + } else { + child.to_plan() + }; + let residual = rewrite_partition_root(residual, rewritten.dtype().clone(), &collapsed)?; + + if !fields_changed && residual == *expression { + return Ok(None); + } + + Ok(Some(EvalPlan::try_new(residual, rewritten)?.into_plan())) + } +} + +/// Rebinds a partitioned residual expression after collapsing single-value partitions. +/// +/// # Arguments +/// +/// * `expression` - The residual recombination expression returned by `partition_bound`. +/// * `root_dtype` - The dtype produced by the rewritten plan and used to rebind every root. +/// * `collapsed` - `(partition_name, value_name)` pairs whose one-field `Pack` was removed; +/// each `$.partition_name.value_name` access is rewritten to `$.partition_name`. +pub(super) fn rewrite_partition_root( + expression: BoundExpression, + root_dtype: DType, + collapsed: &[(FieldName, FieldName)], +) -> VortexResult { + Ok(expression + .transform_down(|node| { + if let Some(value_name) = node.as_opt::() { + let partition_access = &node.children()[0]; + if let Some(partition_name) = partition_access.as_opt::() + && partition_access.children()[0].is_root() + && collapsed.iter().any(|(partition, value)| { + partition == partition_name && value == value_name + }) + { + return Ok(Transformed { + value: BoundExpression::try_new( + GetItem.bind(partition_name.clone()), + [BoundExpression::new_root(root_dtype.clone())], + )?, + changed: true, + order: TraversalOrder::Skip, + }); + } + } + + if node.is_root() { + Ok(Transformed { + value: BoundExpression::new_root(root_dtype.clone()), + changed: true, + order: TraversalOrder::Skip, + }) + } else { + Ok(Transformed::no(node)) + } + })? + .into_inner()) +} + +fn field_name(fields: &StructFields, index: usize) -> VortexResult { + Ok(fields + .field_name(index) + .ok_or_else(|| vortex_err!("Struct field {index} has no name"))? + .clone()) +} + +fn field_plan(plan: &Plan, index: usize) -> VortexResult { + plan.child(index)? + .ok_or_else(|| vortex_err!("Struct field {index} has no plan")) +} + +fn expanded_struct_root( + root_dtype: &DType, + fields: &StructFields, +) -> VortexResult { + let root = BoundExpression::new_root(root_dtype.clone()); + let children = fields + .names() + .iter() + .map(|name| BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()])) + .collect::>>()?; + bound_pack(fields.names().clone(), children) +} + +fn is_identity_expression(expression: &BoundExpression, input_dtype: &DType) -> VortexResult { + if expression.is_root() { + return Ok(expression.dtype() == input_dtype); + } + if input_dtype.is_nullable() { + return Ok(false); + } + let Some(fields) = input_dtype.as_struct_fields_opt() else { + return Ok(false); + }; + Ok(expression == &expanded_struct_root(input_dtype, fields)?) +} + +fn expand_struct_root( + expression: BoundExpression, + expanded_root: &BoundExpression, + fields: &StructFields, +) -> VortexResult { + Ok(expression + .transform_down(|node| { + if node.is_root() { + return Ok(Transformed { + value: expanded_root.clone(), + changed: true, + order: TraversalOrder::Skip, + }); + } + + let Some(scalar_fn) = node.as_scalar() else { + return Ok(Transformed::no(node)); + }; + if !node + .children() + .first() + .is_some_and(BoundExpression::is_root) + { + return Ok(Transformed::no(node)); + } + + if let Some(field_name) = scalar_fn.as_opt::() { + let index = fields.find(field_name).ok_or_else(|| { + vortex_err!("Field {field_name} not found while expanding struct root") + })?; + return Ok(Transformed { + value: expanded_root.children()[index].clone(), + changed: true, + order: TraversalOrder::Skip, + }); + } + + if let Some(selection) = scalar_fn.as_opt::