Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Cache output equivalence_properties in ProjectionExec #9097

Closed
wants to merge 6 commits into from
Closed

Cache output equivalence_properties in ProjectionExec #9097

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9d1ae0a
Keep track of output equivalence_properties in ProjectionExec
gruuya Feb 1, 2024
695d8b2
Amend CoalesceBatchesExec and RepartitionExec to the EXPLAIN for a fa…
gruuya Feb 1, 2024
438860d
Merge branch 'main' into add-eq-props-projection
gruuya Feb 1, 2024
40a2644
Call input's eqivalence_properties only in outer_partitioning computa…
gruuya Feb 1, 2024
71bf589
Merge branch 'main' into add-eq-props-projection
gruuya Feb 2, 2024
87958e0
Update datafusion/physical-plan/src/projection.rs
gruuya Feb 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions datafusion/physical-plan/src/aggregates/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -603,14 +603,12 @@ impl ExecutionPlan for AggregateExec {
// First stage aggregation will not change the output partitioning,
// but needs to respect aliases (e.g. mapping in the GROUP BY
// expression).
let input_eq_properties = self.input.equivalence_properties();
// First stage Aggregation will not change the output partitioning but need to respect the Alias
let input_partition = self.input.output_partitioning();
if let Partitioning::Hash(exprs, part) = input_partition {
let normalized_exprs = exprs
.into_iter()
.map(|expr| {
input_eq_properties
self.input
.equivalence_properties()
.project_expr(&expr, &self.projection_mapping)
.unwrap_or_else(|| {
Arc::new(UnKnownColumn::new(&expr.to_string()))
Expand Down
17 changes: 10 additions & 7 deletions datafusion/physical-plan/src/projection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,9 @@ pub struct ProjectionExec {
input: Arc<dyn ExecutionPlan>,
/// The output ordering
output_ordering: Option<Vec<PhysicalSortExpr>>,
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In theory this doesn't need to be stored now, since it can be obtained via self.equivalence_properties.oeq_class().output_ordering(), however the last method has some non-trivial computation which are probably best avoided anytime there's a call to ProjectionExec::output_ordering()

/// The output equivalence properties. Cached here as computing them
/// can be expensive, see <https://github.com/apache/arrow-datafusion/issues/9084>
equivalence_properties: EquivalenceProperties,
/// The mapping used to normalize expressions like Partitioning and
/// PhysicalSortExpr that maps input to output
projection_mapping: ProjectionMapping,
Expand Down Expand Up @@ -96,14 +99,16 @@ impl ProjectionExec {
let projection_mapping = ProjectionMapping::try_new(&expr, &input_schema)?;

let input_eqs = input.equivalence_properties();
let project_eqs = input_eqs.project(&projection_mapping, schema.clone());
let output_ordering = project_eqs.oeq_class().output_ordering();
let equivalence_properties =
input_eqs.project(&projection_mapping, schema.clone());
let output_ordering = equivalence_properties.oeq_class().output_ordering();

Ok(Self {
expr,
schema,
input,
output_ordering,
equivalence_properties,
projection_mapping,
metrics: ExecutionPlanMetricsSet::new(),
})
Expand Down Expand Up @@ -173,12 +178,12 @@ impl ExecutionPlan for ProjectionExec {
fn output_partitioning(&self) -> Partitioning {
// Output partition need to respect the alias
let input_partition = self.input.output_partitioning();
let input_eq_properties = self.input.equivalence_properties();
if let Partitioning::Hash(exprs, part) = input_partition {
let normalized_exprs = exprs
.into_iter()
.map(|expr| {
input_eq_properties
self.input
.equivalence_properties()
.project_expr(&expr, &self.projection_mapping)
.unwrap_or_else(|| {
Arc::new(UnKnownColumn::new(&expr.to_string()))
Expand All @@ -201,9 +206,7 @@ impl ExecutionPlan for ProjectionExec {
}

fn equivalence_properties(&self) -> EquivalenceProperties {
self.input
.equivalence_properties()
.project(&self.projection_mapping, self.schema())
self.equivalence_properties.clone()
}

fn with_new_children(
Expand Down
Loading