diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index c016944f76bec..06f384ac2db03 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -310,6 +310,33 @@ pub struct EquivalenceGroup { } impl EquivalenceGroup { + /// A cheap, deliberately conservative check that two groups hold the same + /// equivalence classes. + /// + /// This is not `PartialEq`, and the distinction is the point. `classes` is a + /// `Vec` whose order carries no meaning -- `remove_class_at_idx` uses + /// `swap_remove` -- so two groups describing exactly the same equalities can + /// hold their classes in different orders and this returns `false` for them. + /// Naming it `PartialEq` would invite callers to read it as semantic equality, + /// which it is not. + /// + /// The comparison is positional because it runs on a hot path: + /// `ProjectionExec` consults it every time a rule replaces its child. Set + /// semantics would mean scanning the other group once per class, and that + /// quadratic term costs more than the recomputation the caller is trying to + /// skip, by a margin that widens with the number of classes. + /// + /// Only the false direction is reachable: a group can be reported different + /// when it is not, never the same when it is not. Callers using this to skip + /// work must be built so that a `false` merely costs them that work, which is + /// exactly how the projection fast path uses it. + /// + /// `map` is an index into `classes` and carries no information the classes do + /// not already have, so it takes no part in the comparison. + pub fn has_same_classes(&self, other: &Self) -> bool { + self.classes == other.classes + } + /// Creates an equivalence group from the given equivalence classes. pub fn new(classes: impl IntoIterator) -> Self { classes.into_iter().collect::>().into() @@ -924,7 +951,7 @@ mod tests { use super::*; use crate::equivalence::tests::create_test_params; use crate::expressions::{BinaryExpr, Column, binary, col, lit}; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_expr::Operator; @@ -1245,4 +1272,92 @@ mod tests { Ok(()) } + + /// Builds a group from a list of equated column pairs. + fn group_of(schema: &SchemaRef, pairs: &[(&str, &str)]) -> Result { + let mut group = EquivalenceGroup::default(); + for (lhs, rhs) in pairs { + group.add_equal_conditions(col(lhs, schema)?, col(rhs, schema)?); + } + Ok(group) + } + + fn abc_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + Field::new("d", DataType::Int32, false), + ])) + } + + #[test] + fn test_has_same_classes_is_conservative_about_class_order() -> Result<()> { + // `classes` is a `Vec` and `remove_class_at_idx` uses `swap_remove`, so + // the order two groups hold their classes in depends on how they were + // built. This check is positional and reports such a pair as different. + // + // That is deliberate, and it is why this is not `PartialEq`: set + // semantics would mean scanning the other group per class, and the + // quadratic cost dwarfs the recomputation the caller is trying to skip. + // The error can only go this way -- different when they match, never the + // reverse -- so a caller only forfeits an optimization. + let schema = abc_schema(); + let ab_then_cd = group_of(&schema, &[("a", "b"), ("c", "d")])?; + let cd_then_ab = group_of(&schema, &[("c", "d"), ("a", "b")])?; + + assert_eq!(ab_then_cd.len(), 2, "expected two disjoint classes"); + assert!(!ab_then_cd.has_same_classes(&cd_then_ab)); + + Ok(()) + } + + #[test] + fn test_has_same_classes_compares_classes() -> Result<()> { + let schema = abc_schema(); + + // Two empty groups agree. + assert!(group_of(&schema, &[])?.has_same_classes(&group_of(&schema, &[])?)); + // The same class, built the same way. + assert!( + group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "b")])?) + ); + // A class is a set, so the order within a pair is immaterial. + assert!( + group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("b", "a")])?) + ); + // A populated group is not an empty one. + assert!( + !group_of(&schema, &[("a", "b")])?.has_same_classes(&group_of(&schema, &[])?) + ); + // Equating a different pair yields a different group. + assert!( + !group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "c")])?) + ); + // Widening a class yields a different group. + assert!( + !group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "b"), ("b", "c")])?) + ); + + Ok(()) + } + + #[test] + fn test_has_same_classes_ignores_the_map() -> Result<()> { + // `map` indexes into `classes`, so equal classes must imply equal + // groups no matter how the classes were arrived at. Bridging `a = b` + // and `b = c` into one class must match stating `a = c` and `a = b`. + let schema = abc_schema(); + let bridged = group_of(&schema, &[("a", "b"), ("b", "c")])?; + let direct = group_of(&schema, &[("a", "c"), ("a", "b")])?; + + assert_eq!(bridged.len(), 1, "expected a single bridged class"); + assert!(bridged.has_same_classes(&direct)); + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index f52b320ed284f..08c05efe0ccc0 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -194,6 +194,31 @@ impl OrderingEquivalenceCache { } } +/// Counts how often [`EquivalenceProperties::project_reusing`] reused a cached +/// equivalence group instead of reprojecting it. +/// +/// Thread local rather than a global counter: the test binary runs tests in +/// parallel, and a shared count would let one test observe another's hits. +#[cfg(test)] +mod eq_group_reuse_probe { + use std::cell::Cell; + + thread_local! { + static HITS: Cell = const { Cell::new(0) }; + } + + pub(super) fn record_hit() { + HITS.with(|h| h.set(h.get() + 1)); + } + + /// Runs `f` and reports how many reuses happened while it did. + pub(super) fn count(f: impl FnOnce() -> T) -> (T, usize) { + let before = HITS.with(Cell::get); + let value = f(); + (value, HITS.with(Cell::get) - before) + } +} + impl EquivalenceProperties { /// Helper used by the ordering equivalence rule when considering whether /// an expression can replace an existing sort key without invalidating @@ -1188,6 +1213,48 @@ impl EquivalenceProperties { /// dropped. pub fn project(&self, mapping: &ProjectionMapping, output_schema: SchemaRef) -> Self { let eq_group = self.eq_group.project(mapping); + // Built here, so it satisfies the precondition by construction; going + // through the checked entry point would reproject it under + // `debug_assertions` for nothing. + self.project_with_eq_group_unchecked(mapping, output_schema, eq_group) + } + + /// Projects `self`, reusing `cached`'s already-projected equivalence group + /// when `self`'s group is unchanged from `previous`'s. + /// + /// [`EquivalenceGroup::project`] is a pure function of the group and the + /// mapping, so when the group is unchanged the previous result can be handed + /// back rather than recomputed. Orderings are still derived here: they are + /// precisely what changes when a sort is introduced below this node. + /// + /// Falls back to a full projection when the groups differ, so there is no + /// precondition to violate. The caller does still have to pass a `cached` + /// that came from projecting `previous` through this same `mapping`; that + /// part cannot be checked here, and in practice it holds because the caller + /// carries its projection over untouched. + pub fn project_reusing( + &self, + mapping: &ProjectionMapping, + output_schema: SchemaRef, + previous: &EquivalenceProperties, + cached: &EquivalenceProperties, + ) -> Self { + let eq_group = if self.eq_group.has_same_classes(&previous.eq_group) { + #[cfg(test)] + eq_group_reuse_probe::record_hit(); + cached.eq_group.clone() + } else { + self.eq_group.project(mapping) + }; + self.project_with_eq_group_unchecked(mapping, output_schema, eq_group) + } + + fn project_with_eq_group_unchecked( + &self, + mapping: &ProjectionMapping, + output_schema: SchemaRef, + eq_group: EquivalenceGroup, + ) -> Self { let orderings = self.projected_orderings(mapping, self.oeq_cache.normal_cls.clone()); let normal_orderings = orderings @@ -1559,3 +1626,127 @@ fn get_expr_properties( expr.get_properties(&child_states) } } + +#[cfg(test)] +mod tests { + use super::*; + + use crate::equivalence::tests::create_test_params; + use crate::expressions::col; + + use arrow::datatypes::{DataType, Field, Schema}; + + /// Renames `a`, `b`, `c` and `d` so the projection is non-trivial while + /// still carrying every ordering and the `a = c` class across. + fn renaming_mapping( + schema: &SchemaRef, + output_schema: &SchemaRef, + ) -> Result { + [ + ("a", "a1", 0), + ("b", "b1", 1), + ("c", "c1", 2), + ("d", "d1", 3), + ] + .into_iter() + .map(|(source, target, index)| { + Ok(( + col(source, schema)?, + vec![(col(target, output_schema)?, index)].into(), + )) + }) + .collect::>>() + .map(|entries| entries.into_iter().collect()) + } + + fn renamed_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, true), + Field::new("b1", DataType::Int32, true), + Field::new("c1", DataType::Int32, true), + Field::new("d1", DataType::Int32, true), + ])) + } + + #[test] + fn test_project_reusing_matches_project() -> Result<()> { + let (schema, eq_properties) = create_test_params()?; + let output_schema = renamed_schema(); + let mapping = renaming_mapping(&schema, &output_schema)?; + + let baseline = eq_properties.project(&mapping, Arc::clone(&output_schema)); + + // Reusing the projection of an identical group must reproduce it exactly. + let (reused, hits) = eq_group_reuse_probe::count(|| { + eq_properties.project_reusing( + &mapping, + Arc::clone(&output_schema), + &eq_properties, + &baseline, + ) + }); + assert_eq!( + hits, 1, + "the reuse path was not taken, so this compares nothing" + ); + + // Guard against a vacuous comparison. + assert!( + !baseline.eq_group().is_empty(), + "the projection dropped the equivalence class, nothing is being tested" + ); + assert!( + !baseline.oeq_class().is_empty(), + "the projection dropped every ordering, nothing is being tested" + ); + + assert!( + baseline.eq_group().has_same_classes(reused.eq_group()), + "equivalence group" + ); + assert_eq!(baseline.oeq_class(), reused.oeq_class(), "orderings"); + assert_eq!(baseline.constraints(), reused.constraints(), "constraints"); + assert_eq!(baseline.schema(), reused.schema(), "schema"); + + Ok(()) + } + + #[test] + fn test_project_reusing_falls_back_when_the_group_moved() -> Result<()> { + // A `previous` whose group differs must not have its projection carried + // over. There is no precondition to violate here: the fallback is what + // keeps a mismatched pair from producing properties that disagree with + // the projection. + let (schema, eq_properties) = create_test_params()?; + let output_schema = renamed_schema(); + let mapping = renaming_mapping(&schema, &output_schema)?; + + let baseline = eq_properties.project(&mapping, Arc::clone(&output_schema)); + let unrelated = EquivalenceProperties::new(Arc::clone(&schema)); + assert!( + !eq_properties + .eq_group() + .has_same_classes(unrelated.eq_group()), + "the two groups were meant to differ" + ); + + let (recomputed, hits) = eq_group_reuse_probe::count(|| { + eq_properties.project_reusing( + &mapping, + Arc::clone(&output_schema), + &unrelated, + &baseline, + ) + }); + assert_eq!(hits, 0, "a mismatched group was carried over anyway"); + + // Falling back must land on exactly what `project` would have produced. + assert!( + baseline.eq_group().has_same_classes(recomputed.eq_group()), + "equivalence group" + ); + assert_eq!(baseline.oeq_class(), recomputed.oeq_class(), "orderings"); + + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index cf362cdee55d3..1ec278eb9c377 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -52,6 +52,7 @@ use datafusion_common::tree_node::{ use datafusion_common::{DataFusionError, JoinSide, Result, internal_err, plan_err}; use datafusion_execution::TaskContext; use datafusion_expr::ExpressionPlacement; +use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_expr::equivalence::ProjectionMapping; use datafusion_physical_expr::projection::Projector; use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; @@ -177,6 +178,32 @@ impl ProjectionExec { fn try_from_projector( projector: Projector, input: Arc, + ) -> Result { + Self::try_from_projector_with_eq_group(projector, input, None) + } + + /// As [`Self::try_from_projector`], but `reuse_from` may carry the previous + /// child's equivalence properties together with the projection they produced, + /// letting [`EquivalenceProperties::project_reusing`] skip reprojecting a + /// group that has not changed. + /// + /// Projecting an equivalence group is a pure function of the group and the + /// mapping, so reuse is sound exactly when both are unchanged. + /// + /// The caller establishes the first by comparing the old and new child + /// groups. The second holds because the mapping comes from + /// `projector.projection()`, carried over untouched, and from the child's + /// schema, which `ProjectionMapping::try_new` consults only for field names + /// and indices -- never for types or nullability. So a child differing only + /// in nullability keeps the same mapping. A child that renamed or reordered + /// those fields would change the group too, since its members are `Column`s + /// carrying those names, and the comparison above would reject it; were one + /// to slip through anyway, `try_new`'s name assertion errors out rather than + /// letting a stale group into the plan. + fn try_from_projector_with_eq_group( + projector: Projector, + input: Arc, + reuse_from: Option<(&EquivalenceProperties, &EquivalenceProperties)>, ) -> Result { // Construct a map from the input expressions to the output expression of the Projection let projection_mapping = @@ -185,6 +212,7 @@ impl ProjectionExec { &input, &projection_mapping, Arc::clone(projector.output_schema()), + reuse_from, )?; Ok(Self { projector, @@ -214,10 +242,21 @@ impl ProjectionExec { input: &Arc, projection_mapping: &ProjectionMapping, schema: SchemaRef, + reuse_from: Option<(&EquivalenceProperties, &EquivalenceProperties)>, ) -> Result { - // Calculate equivalence properties: + // Calculate equivalence properties. Whether the group is reprojected or + // handed back is the only thing reuse changes; everything below is + // common, so the two paths cannot drift apart. let input_eq_properties = input.equivalence_properties(); - let eq_properties = input_eq_properties.project(projection_mapping, schema); + let eq_properties = match reuse_from { + Some((previous, cached)) => input_eq_properties.project_reusing( + projection_mapping, + schema, + previous, + cached, + ), + None => input_eq_properties.project(projection_mapping, schema), + }; // Calculate output partitioning, which needs to respect aliases: let output_partitioning = input .output_partitioning() @@ -352,11 +391,28 @@ impl ExecutionPlan for ProjectionExec { metrics: ExecutionPlanMetricsSet::new(), ..Self::clone(&*self) })), - ChildrenPropertiesMode::Recompute => ProjectionExec::try_from_projector( - self.projector.clone(), - children.swap_remove(0), - ) - .map(|p| Arc::new(p) as _), + ChildrenPropertiesMode::Recompute => { + // `Keep` above requires the child's properties to be unchanged + // outright. A rule that introduces a sort below this projection + // does not qualify, yet the child's *equivalence group* is still + // identical: sorting changes which orderings hold, not which + // expressions are equal to one another. Projecting that group + // again would reproduce the group already cached here, so reuse + // it and derive only the orderings. + // Hand over what this projection was built from and what that + // produced; `project_reusing` decides whether the group can be + // carried over and falls back to a full projection otherwise. + let reuse_from = Some(( + self.input.equivalence_properties(), + self.cache.equivalence_properties(), + )); + ProjectionExec::try_from_projector_with_eq_group( + self.projector.clone(), + children.swap_remove(0), + reuse_from, + ) + .map(|p| Arc::new(p) as _) + } } } @@ -1436,6 +1492,8 @@ mod tests { use crate::common::collect; use crate::empty::EmptyExec; + use crate::filter::FilterExec; + use crate::sorts::sort::SortExec; use crate::filter_pushdown::PushedDown; use crate::statistics::{StatisticsArgs, StatisticsContext}; @@ -2215,4 +2273,281 @@ mod tests { Ok(()) } + + /// `EmptyExec(a, b, c)` under a filter that equates `lhs` and `rhs`, so the + /// child carries a non-trivial equivalence group. + fn filtered_source(lhs: &str, rhs: &str) -> Result> { + filtered_source_with_nullability(lhs, rhs, false) + } + + /// As [`filtered_source`], but `nullable` varies the schema's nullability + /// while leaving field names and order alone. + fn filtered_source_with_nullability( + lhs: &str, + rhs: &str, + nullable: bool, + ) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, nullable), + Field::new("b", DataType::Int32, nullable), + Field::new("c", DataType::Int32, nullable), + ])); + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let predicate = binary( + col(lhs, &schema)?, + Operator::Eq, + col(rhs, &schema)?, + &schema, + )?; + Ok(Arc::new(FilterExec::try_new(predicate, input)?)) + } + + /// `[a AS x, b AS y, c AS z]` against `filtered_source`'s schema. + fn renaming_exprs(schema: &SchemaRef) -> Result> { + [("a", "x"), ("b", "y"), ("c", "z")] + .into_iter() + .map(|(source, alias)| { + Ok(ProjectionExpr { + expr: col(source, schema)?, + alias: alias.to_string(), + }) + }) + .collect() + } + + fn assert_same_properties(actual: &dyn ExecutionPlan, expected: &ProjectionExec) { + let actual_props = actual.properties().equivalence_properties(); + let expected_props = expected.properties().equivalence_properties(); + assert!( + actual_props + .eq_group() + .has_same_classes(expected_props.eq_group()), + "equivalence group: {:?} vs {:?}", + actual_props.eq_group(), + expected_props.eq_group() + ); + assert_eq!( + actual_props.oeq_class(), + expected_props.oeq_class(), + "orderings" + ); + assert_eq!( + actual_props.constraints(), + expected_props.constraints(), + "constraints" + ); + assert_eq!(actual_props.schema(), expected_props.schema(), "schema"); + // `Partitioning` has no `PartialEq`, so compare the partition count and + // the explicit `Display` form. Derived `Debug` would change with any + // field addition, making this brittle for no gain. + let actual_partitioning = actual.properties().output_partitioning(); + let expected_partitioning = expected.properties().output_partitioning(); + assert_eq!( + actual_partitioning.partition_count(), + expected_partitioning.partition_count(), + "partition count" + ); + assert_eq!( + actual_partitioning.to_string(), + expected_partitioning.to_string(), + "partitioning" + ); + } + + #[test] + fn test_sort_below_changes_orderings_but_not_the_equivalence_group() -> Result<()> { + // The premise the fast path rests on. If a sort ever starts altering + // the equivalence group, reusing the cached group becomes unsound and + // this test is the one that should fail first. + let child = filtered_source("a", "b")?; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "c", + &child.schema(), + )?)]) + .expect("non-empty ordering"); + let sorted = SortExec::new(ordering, Arc::clone(&child)); + + let child_props = child.properties().equivalence_properties(); + let sorted_props = sorted.properties().equivalence_properties(); + + assert!( + !child_props.eq_group().is_empty(), + "the filter did not produce an equivalence class" + ); + assert!( + child_props + .eq_group() + .has_same_classes(sorted_props.eq_group()), + "sorting altered the equivalence group" + ); + assert_ne!( + child_props.oeq_class(), + sorted_props.oeq_class(), + "sorting did not alter the orderings" + ); + + Ok(()) + } + + #[test] + fn test_replace_children_reuses_eq_group_when_only_orderings_change() -> Result<()> { + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = + Arc::new(ProjectionExec::try_new(exprs.clone(), Arc::clone(&child))?); + + // Sorting below the projection changes which orderings hold but leaves + // the equivalence group untouched -- the case the fast path targets. + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "c", + &child.schema(), + )?)]) + .expect("non-empty ordering"); + let sorted: Arc = + Arc::new(SortExec::new(ordering, Arc::clone(&child))); + + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&sorted)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + // Guard against vacuity: the group must be worth reusing, and the sort + // must genuinely have added an ordering the original did not have. + assert!( + !projection + .properties() + .equivalence_properties() + .eq_group() + .is_empty(), + "the projection carries no equivalence class, nothing is being reused" + ); + assert!( + projection + .properties() + .equivalence_properties() + .oeq_class() + .is_empty(), + "the unsorted projection was already ordered" + ); + assert!( + !replaced + .properties() + .equivalence_properties() + .oeq_class() + .is_empty(), + "the sort did not introduce an ordering" + ); + + // The fast path must agree with building the projection from scratch. + let expected = ProjectionExec::try_new(exprs, sorted)?; + assert_same_properties(replaced.as_ref(), &expected); + + Ok(()) + } + + #[test] + fn test_replace_children_reuses_eq_group_across_a_nullability_change() -> Result<()> { + // Reuse is sound only if the projection mapping is unchanged as well. + // `ProjectionMapping::try_new` reads the child schema for field names + // and indices alone, so a child differing only in nullability keeps the + // same mapping and must still take the fast path. + // + // The swap tightens nullability rather than loosening it, matching what + // `is_allowed_field_change` permits of a physical optimizer rule. + // + // The comparison is against `try_from_projector`, the path this one + // replaces, rather than a freshly built projection: `replace_children` + // carries the existing `Projector` over, so the output schema stays as + // it was, while `try_new` would derive a new one from the new child. + // That difference is inherent to `replace_children` and not something + // this fast path introduces, so the meaningful contract is that the two + // `replace_children` paths agree. + let child = filtered_source_with_nullability("a", "b", true)?; + let exprs = renaming_exprs(&child.schema())?; + let projection = Arc::new(ProjectionExec::try_new(exprs, Arc::clone(&child))?); + + let tightened_child = filtered_source_with_nullability("a", "b", false)?; + assert_ne!( + child.schema(), + tightened_child.schema(), + "the two children were meant to differ in nullability" + ); + assert!( + child + .properties() + .equivalence_properties() + .eq_group() + .has_same_classes( + tightened_child + .properties() + .equivalence_properties() + .eq_group() + ), + "nullability moved the equivalence group, so the fast path is no longer under test" + ); + + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&tightened_child)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let recomputed = ProjectionExec::try_from_projector( + projection.projector.clone(), + tightened_child, + )?; + + assert_same_properties(replaced.as_ref(), &recomputed); + + Ok(()) + } + + #[test] + fn test_replace_children_recomputes_when_eq_group_changes() -> Result<()> { + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = + Arc::new(ProjectionExec::try_new(exprs.clone(), Arc::clone(&child))?); + + // This child equates a different pair, so the cached group is stale and + // reusing it would be unsound: the guard has to fall through. + let other = filtered_source("a", "c")?; + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&other)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + assert!( + !replaced + .properties() + .equivalence_properties() + .eq_group() + .has_same_classes( + projection.properties().equivalence_properties().eq_group() + ), + "the projection kept the previous child's equivalence group" + ); + + let expected = ProjectionExec::try_new(exprs, other)?; + assert_same_properties(replaced.as_ref(), &expected); + + Ok(()) + } + + #[test] + fn test_replace_children_keep_mode_carries_properties_over() -> Result<()> { + // `Keep` is the caller's promise that the new child's properties match + // the old one's, so the cached properties must survive verbatim rather + // than being derived again. + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = Arc::new(ProjectionExec::try_new(exprs, Arc::clone(&child))?); + + let replaced = Arc::clone(&projection).replace_children( + vec![filtered_source("a", "b")?], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + )?; + + assert_same_properties(replaced.as_ref(), &projection); + + Ok(()) + } }