diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index 65ef2a3ceb216..07252655434f8 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -104,3 +104,7 @@ name = "string_concat" [package.metadata.cargo-machete] ignored = ["half"] + +[[bench]] +harness = false +name = "equivalence_properties" diff --git a/datafusion/physical-expr/benches/equivalence_properties.rs b/datafusion/physical-expr/benches/equivalence_properties.rs new file mode 100644 index 0000000000000..f014ed4d7429b --- /dev/null +++ b/datafusion/physical-expr/benches/equivalence_properties.rs @@ -0,0 +1,150 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for the ordering satisfaction checks on [`EquivalenceProperties`]. +//! +//! These are called repeatedly during physical optimization (sort removal, +//! `EnforceSorting`, `EnforceDistribution`, and the requirement checks for +//! windows, joins and aggregates), so their cost shows up directly in planning +//! time. +//! +//! The benchmarks are parameterized by the number of equivalence classes, since +//! that -- not the schema width, which is behind an `Arc` -- is what these +//! checks carry around. + +use std::sync::Arc; + +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, + PhysicalSortRequirement, +}; + +fn schema(n_cols: usize) -> SchemaRef { + Arc::new(Schema::new( + (0..n_cols) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, true)) + .collect::>(), + )) +} + +fn col(i: usize) -> Arc { + Arc::new(Column::new(&format!("c{i}"), i)) +} + +fn asc(i: usize) -> PhysicalSortExpr { + PhysicalSortExpr::new(col(i), SortOptions::default()) +} + +/// Properties with three equivalent orderings and `n_classes` equivalence +/// classes, i.e. roughly what a scan feeding a join and a window function looks +/// like. Columns `c0..c7` carry the orderings; the equivalence classes are built +/// from the columns above them. +fn properties(n_classes: usize) -> EquivalenceProperties { + let schema = schema(8 + 2 * n_classes); + let mut props = EquivalenceProperties::new(schema); + props.add_orderings([ + vec![asc(0), asc(1), asc(2), asc(3)], + vec![asc(4), asc(5)], + vec![asc(6)], + ]); + for i in 0..n_classes { + props + .add_equal_conditions(col(8 + 2 * i), col(9 + 2 * i)) + .unwrap(); + } + props +} + +fn bench_ordering_satisfaction(c: &mut Criterion) { + let mut group = c.benchmark_group("equivalence_properties"); + + for n_classes in [2, 8, 32] { + let props = properties(n_classes); + + // A single sort key: the most common shape by far. + group.bench_with_input( + BenchmarkId::new("ordering_satisfy/1_key", n_classes), + &n_classes, + |b, _| b.iter(|| props.ordering_satisfy([asc(0)]).unwrap()), + ); + // A single sort key that is not satisfied: exits on the first key. + group.bench_with_input( + BenchmarkId::new("ordering_satisfy/1_key_unsatisfied", n_classes), + &n_classes, + |b, _| b.iter(|| props.ordering_satisfy([asc(7)]).unwrap()), + ); + // Four sort keys: exercises the per-key constant registration. + group.bench_with_input( + BenchmarkId::new("ordering_satisfy/4_keys", n_classes), + &n_classes, + |b, _| { + b.iter(|| { + props + .ordering_satisfy([asc(0), asc(1), asc(2), asc(3)]) + .unwrap() + }) + }, + ); + group.bench_with_input( + BenchmarkId::new("ordering_satisfy_requirement/1_key", n_classes), + &n_classes, + |b, _| { + b.iter(|| { + props + .ordering_satisfy_requirement([PhysicalSortRequirement::new( + col(0), + None, + )]) + .unwrap() + }) + }, + ); + group.bench_with_input( + BenchmarkId::new("ordering_satisfy_requirement/4_keys", n_classes), + &n_classes, + |b, _| { + b.iter(|| { + props + .ordering_satisfy_requirement( + (0..4).map(|i| PhysicalSortRequirement::new(col(i), None)), + ) + .unwrap() + }) + }, + ); + group.bench_with_input( + BenchmarkId::new("extract_common_sort_prefix/4_keys", n_classes), + &n_classes, + |b, _| { + b.iter(|| { + let ordering = + LexOrdering::new((0..4).map(asc).collect::>()).unwrap(); + props.extract_common_sort_prefix(ordering).unwrap() + }) + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_ordering_satisfaction); +criterion_main!(benches); diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 08c05efe0ccc0..c68157fecbd8c 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -640,8 +640,13 @@ impl EquivalenceProperties { return Ok(true); } let schema = self.schema(); - let mut eq_properties = self.clone(); - for element in normal_reqs { + // Registering satisfied keys as constants mutates the state, so it + // needs an owned copy -- but only from the second requirement onwards. + // Single-element requirements (the common case) never pay for the clone. + let last_idx = normal_reqs.len() - 1; + let mut owned = None::; + for (idx, element) in normal_reqs.into_iter().enumerate() { + let eq_properties = owned.as_ref().unwrap_or(self); // Check whether given requirement is satisfied: let ExprProperties { sort_properties, .. @@ -658,10 +663,16 @@ impl EquivalenceProperties { if !satisfy { return Ok(false); } + if idx == last_idx { + // Nothing left to check, so no need to update the state: + break; + } // Treat satisfied keys (and the sub-expressions they pin down) as // constants in subsequent iterations. See // [`Self::add_satisfied_key_constants`] for the rationale. - eq_properties.add_satisfied_key_constants(element.expr)?; + owned + .get_or_insert_with(|| self.clone()) + .add_satisfied_key_constants(element.expr)?; } Ok(true) } @@ -718,8 +729,12 @@ impl EquivalenceProperties { return Ok(full_length); } let schema = self.schema(); - let mut eq_properties = self.clone(); + // Registering satisfied keys as constants mutates the state, so it + // needs an owned copy -- but only from the second sort expression + // onwards. Single-element orderings never pay for the clone. + let mut owned = None::; for (idx, element) in normal_ordering.into_iter().enumerate() { + let eq_properties = owned.as_ref().unwrap_or(self); // Check whether given ordering is satisfied: let ExprProperties { sort_properties, .. @@ -739,10 +754,16 @@ impl EquivalenceProperties { // many we've satisfied so far: return Ok(idx); } + if idx + 1 == full_length { + // Nothing left to check, so no need to update the state: + break; + } // Treat satisfied keys (and the sub-expressions they pin down) as // constants in subsequent iterations. See // [`Self::add_satisfied_key_constants`] for the rationale. - eq_properties.add_satisfied_key_constants(Arc::clone(&element.expr))?; + owned + .get_or_insert_with(|| self.clone()) + .add_satisfied_key_constants(Arc::clone(&element.expr))?; } // All sort expressions are satisfied, return full length: Ok(full_length)