physical plan: add plan re-use benchmark

This patch adds a benchmark which is intended to measure overhead of actions,
required to perform  making an independent instance of the execution plan to
re-execute it, avoiding re-planning stage. There are several typical plans
that are tested, covering projection, aggregation, filtration, re-partition.

Also, the function `reset_plan_states(...)` is publically exported.

Author: askalt

datafusion/physical-plan/Cargo.toml

@@ -90,6 +90,10 @@ tokio = { workspace = true, features = [
 harness = false
 name = "partial_ordering"
 
+[[bench]]
+harness = false
+name = "plan_reuse"
+
 [[bench]]
 harness = false
 name = "spill_io"

datafusion/physical-plan/benches/plan_reuse.rs

@@ -0,0 +1,326 @@
+use std::sync::{Arc, LazyLock};
+
+use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef};
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion_common::{NullEquality, Result};
+use datafusion_execution::TaskContext;
+use datafusion_expr::{JoinType, Operator};
+use datafusion_functions_aggregate::average::avg_udaf;
+use datafusion_physical_expr::aggregate::AggregateExprBuilder;
+use datafusion_physical_expr::expressions::{self, binary, lit};
+use datafusion_physical_expr::projection::ProjectionExpr;
+use datafusion_physical_expr::{Partitioning, PhysicalExpr};
+use datafusion_physical_plan::aggregates::{
+    AggregateExec, AggregateMode, PhysicalGroupBy,
+};
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::execution_plan::reset_plan_states;
+use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode};
+use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::{
+    ExecutionPlan, execute_stream, filter::FilterExec, test::TestMemoryExec,
+};
+
+const NUM_FIELDS: usize = 1000;
+
+static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
+    Arc::new(Schema::new(
+        (0..NUM_FIELDS)
+            .map(|i| Arc::new(Field::new(format!("x_{i}"), DataType::Int64, false)))
+            .collect::<Fields>(),
+    ))
+});
+
+fn partitioning() -> Partitioning {
+    Partitioning::RoundRobinBatch(16)
+}
+
+fn col_name(i: usize) -> String {
+    format!("x_{i}")
+}
+
+fn aggr_name(i: usize) -> String {
+    format!("aggr({})", col_name(i))
+}
+
+fn col(i: usize) -> Arc<dyn PhysicalExpr> {
+    expressions::col(&col_name(i), &SCHEMA).unwrap()
+}
+
+/// Returns a typical plan for the query like:
+///
+/// ```sql
+/// SELECT aggr1(col1) as aggr1, aggr2(col2) as aggr2 FROM t
+/// WHERE p1
+/// HAVING p2
+/// ```
+///
+/// A plan looks like:
+///
+/// ```text
+/// FilterExec
+///   AggregateExec: mode=Final
+///     CoalescePartitionsExec
+///       AggregateExec: mode=Partial
+///         RepartitionExec
+///           FilterExec
+///             TestMemoryExec
+/// ```
+///
+fn query1_plan() -> Result<Arc<dyn ExecutionPlan>> {
+    let schema = Arc::clone(&SCHEMA);
+    let input = TestMemoryExec::try_new(&[vec![]], Arc::clone(&schema), None)?;
+
+    let plan = FilterExec::try_new(
+        // Some predicate.
+        binary(
+            binary(col(0), Operator::Eq, col(1), &schema)?,
+            Operator::And,
+            binary(col(2), Operator::Eq, lit(42_i64), &schema)?,
+            &schema,
+        )?,
+        Arc::new(input),
+    )?;
+
+    let plan = RepartitionExec::try_new(Arc::new(plan), partitioning())?;
+
+    let plan = {
+        // Partial aggregation.
+        let aggr_expr = (0..NUM_FIELDS)
+            .map(|i| {
+                AggregateExprBuilder::new(avg_udaf(), vec![col(i)])
+                    .schema(Arc::clone(&schema))
+                    .alias(aggr_name(i))
+                    .build()
+                    .map(Arc::new)
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let filter_expr = (0..aggr_expr.len()).map(|_| None).collect();
+
+        AggregateExec::try_new(
+            AggregateMode::Partial,
+            PhysicalGroupBy::new(vec![], vec![], vec![], false),
+            aggr_expr,
+            filter_expr,
+            Arc::new(plan),
+            Arc::clone(&schema),
+        )?
+    };
+
+    let plan = CoalescePartitionsExec::new(Arc::new(plan));
+
+    let schema = plan.schema();
+    let plan = {
+        // Final aggregation.
+        let aggr_expr = (0..NUM_FIELDS)
+            .map(|i| {
+                AggregateExprBuilder::new(
+                    avg_udaf(),
+                    vec![Arc::new(expressions::Column::new(&aggr_name(i), i))],
+                )
+                .schema(Arc::clone(&schema))
+                .alias(aggr_name(i))
+                .build()
+                .map(Arc::new)
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let filter_expr = (0..aggr_expr.len()).map(|_| None).collect();
+
+        AggregateExec::try_new(
+            AggregateMode::Partial,
+            PhysicalGroupBy::new(vec![], vec![], vec![], false),
+            aggr_expr,
+            filter_expr,
+            Arc::new(plan),
+            Arc::clone(&schema),
+        )?
+    };
+
+    let schema = plan.schema();
+    let plan = {
+        let predicate = (0..schema.fields.len()).fold(lit(true), |expr, i| {
+            binary(
+                expr,
+                Operator::And,
+                binary(
+                    Arc::new(expressions::Column::new(schema.field(i).name(), i)),
+                    Operator::Gt,
+                    lit(i as i64),
+                    &schema,
+                )
+                .unwrap(),
+                &schema,
+            )
+            .unwrap()
+        });
+
+        FilterExec::try_new(predicate, Arc::new(plan))?
+    };
+
+    Ok(Arc::new(plan))
+}
+
+/// Returns a typical plan for the query like:
+///
+/// ```sql
+/// SELECT projection FROM t1 JOIN t2 ON t1.a = t2.a
+/// WHERE p1
+/// ```
+///
+/// A plan looks like:
+///
+/// ```text
+/// HashJoinExec
+///   RepartitionExec
+///     FilterExec
+///       TestMemoryExec
+///   RepartitionExec
+///     FilterExec
+///       TestMemoryExec
+/// ```
+///
+fn query2_plan() -> Result<Arc<dyn ExecutionPlan>> {
+    let schema = Arc::clone(&SCHEMA);
+    let t: Arc<dyn ExecutionPlan> = Arc::new(TestMemoryExec::try_new(
+        &[vec![]],
+        Arc::clone(&schema),
+        None,
+    )?);
+
+    // Some predicate.
+    let predicate = (0..schema.fields.len().min(15)).fold(lit(true), |expr, i| {
+        binary(
+            expr,
+            Operator::And,
+            binary(
+                Arc::new(expressions::Column::new(schema.field(i).name(), i)),
+                Operator::Eq,
+                lit(i as i64),
+                &schema,
+            )
+            .unwrap(),
+            &schema,
+        )
+        .unwrap()
+    });
+
+    let left = RepartitionExec::try_new(
+        Arc::new(FilterExec::try_new(Arc::clone(&predicate), Arc::clone(&t))?),
+        partitioning(),
+    )?;
+
+    let right = RepartitionExec::try_new(
+        Arc::new(FilterExec::try_new(Arc::clone(&predicate), Arc::clone(&t))?),
+        partitioning(),
+    )?;
+
+    let projection = Some((0..schema.fields.len()).step_by(15).collect());
+    let join = HashJoinExec::try_new(
+        Arc::new(left),
+        Arc::new(right),
+        vec![(col(0), col(0))],
+        None,
+        &JoinType::Left,
+        projection,
+        PartitionMode::Partitioned,
+        NullEquality::NullEqualsNothing,
+        false,
+    )?;
+
+    Ok(Arc::new(join))
+}
+
+/// Returns a typical plan for the query like:
+///
+/// ```sql
+/// SELECT non_trivial_projection FROM t
+/// WHERE p
+/// ```
+///
+/// A plan looks like:
+///
+/// ```text
+/// ProjectionExec
+///   RepartitionExec
+///     FilterExec
+///       TestMemoryExec
+/// ```
+///
+fn query3_plan() -> Result<Arc<dyn ExecutionPlan>> {
+    let schema = Arc::clone(&SCHEMA);
+    let t: Arc<dyn ExecutionPlan> = Arc::new(TestMemoryExec::try_new(
+        &[vec![]],
+        Arc::clone(&schema),
+        None,
+    )?);
+
+    // Some predicate.
+    let predicate = (0..schema.fields.len().min(15)).fold(lit(true), |expr, i| {
+        binary(
+            expr,
+            Operator::And,
+            binary(
+                Arc::new(expressions::Column::new(schema.field(i).name(), i)),
+                Operator::Eq,
+                lit(i as i64),
+                &schema,
+            )
+            .unwrap(),
+            &schema,
+        )
+        .unwrap()
+    });
+
+    let plan = FilterExec::try_new(predicate, t)?;
+    let plan = RepartitionExec::try_new(Arc::new(plan), partitioning())?;
+    let plan = ProjectionExec::try_new(
+        (0..schema.fields.len() / 2).map(|i| {
+            ProjectionExpr::new(
+                binary(col(i * 2), Operator::Plus, col(i * 2 + 1), &schema).unwrap(),
+                col_name(i),
+            )
+        }),
+        Arc::new(plan),
+    )?;
+
+    Ok(Arc::new(plan))
+}
+
+fn reuse_plan(
+    b: &mut criterion::Bencher,
+    rt_handle: &tokio::runtime::Handle,
+    task_ctx: &Arc<TaskContext>,
+    plan: &Arc<dyn ExecutionPlan>,
+) {
+    b.iter(|| {
+        let plan = reset_plan_states(Arc::clone(plan)).unwrap();
+        std::hint::black_box(
+            rt_handle.block_on(async { execute_stream(plan, Arc::clone(task_ctx)) }),
+        )
+    });
+}
+
+/// Benchmark is intended to measure overhead of actions, required to perform
+/// making an independent instance of the execution plan to re-execute it, avoiding
+/// re-planning stage.
+fn bench_plan_reuse(c: &mut Criterion) {
+    let task_ctx = Arc::new(TaskContext::default());
+    let rt = tokio::runtime::Runtime::new().unwrap();
+
+    macro_rules! bench_query {
+        ($query_producer: expr) => {{
+            let plan = $query_producer().unwrap();
+            let task_ctx = Arc::clone(&task_ctx);
+            let h = rt.handle().clone();
+            move |b| reuse_plan(b, &h, &task_ctx, &plan)
+        }};
+    }
+
+    c.bench_function("query1", bench_query!(query1_plan));
+    c.bench_function("query2", bench_query!(query2_plan));
+    c.bench_function("query3", bench_query!(query3_plan));
+}
+
+criterion_group!(benches, bench_plan_reuse);
+criterion_main!(benches);

datafusion/physical-plan/src/execution_plan.rs

@@ -26,6 +26,7 @@ use crate::sort_pushdown::SortOrderPushdownResult;
 pub use crate::stream::EmptyRecordBatchStream;
 
 pub use datafusion_common::hash_utils;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 pub use datafusion_common::utils::project_schema;
 pub use datafusion_common::{ColumnStatistics, Statistics, internal_err};
 pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream};
@@ -1384,6 +1385,28 @@ pub fn check_not_null_constraints(
     Ok(batch)
 }
 
+/// Make plan ready to be re-executed returning its clone with state reset for all nodes.
+///
+/// Some plans will change their internal states after execution, making them unable to be executed again.
+/// This function uses [`ExecutionPlan::reset_state`] to reset any internal state within the plan.
+///
+/// An example is `CrossJoinExec`, which loads the left table into memory and stores it in the plan.
+/// However, if the data of the left table is derived from the work table, it will become outdated
+/// as the work table changes. When the next iteration executes this plan again, we must clear the left table.
+///
+/// # Note
+///
+/// While this function enables plan reuse, it does not allow the same plan to be executed concurrently if
+/// it uses dynamic filters or represents a recursive query.
+///
+pub fn reset_plan_states(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
+    plan.transform_up(|plan| {
+        let new_plan = Arc::clone(&plan).reset_state()?;
+        Ok(Transformed::yes(new_plan))
+    })
+    .data()
+}
+
 /// Utility function yielding a string representation of the given [`ExecutionPlan`].
 pub fn get_plan_string(plan: &Arc<dyn ExecutionPlan>) -> Vec<String> {
     let formatted = displayable(plan.as_ref()).indent(true).to_string();