diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index b276ae32cf247..7abb39e1a7130 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/datafusion-examples/examples/custom_data_source/custom_datasource.rs @@ -192,7 +192,7 @@ impl TableProvider for CustomDataSource { struct CustomExec { db: CustomDataSource, projected_schema: SchemaRef, - cache: PlanProperties, + cache: Arc, } impl CustomExec { @@ -207,7 +207,7 @@ impl CustomExec { Self { db, projected_schema, - cache, + cache: Arc::new(cache), } } @@ -238,7 +238,7 @@ impl ExecutionPlan for CustomExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs index 48475acbb1542..6f377ea1ce3b9 100644 --- a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs +++ b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs @@ -199,7 +199,7 @@ impl ExternalBatchBufferer { struct BufferingExecutionPlan { schema: SchemaRef, input: Arc, - properties: PlanProperties, + properties: Arc, } impl BufferingExecutionPlan { @@ -233,7 +233,7 @@ impl ExecutionPlan for BufferingExecutionPlan { self.schema.clone() } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index f3910d461b6a8..b4f3d4f098996 100644 --- a/datafusion-examples/examples/proto/composed_extension_codec.rs +++ b/datafusion-examples/examples/proto/composed_extension_codec.rs @@ -106,7 +106,7 @@ impl ExecutionPlan for ParentExec { self } - fn properties(&self) -> &datafusion::physical_plan::PlanProperties { + fn properties(&self) -> &Arc { unreachable!() } @@ -182,7 +182,7 @@ impl ExecutionPlan for ChildExec { self } - fn properties(&self) -> &datafusion::physical_plan::PlanProperties { + fn properties(&self) -> &Arc { unreachable!() } diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index 657432ef31362..895f2fdd4ff3a 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -618,7 +618,7 @@ pub struct SampleExec { upper_bound: f64, seed: u64, metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, + cache: Arc, } impl SampleExec { @@ -656,7 +656,7 @@ impl SampleExec { upper_bound, seed, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -686,7 +686,7 @@ impl ExecutionPlan for SampleExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 7865eb016bee1..484b5f805e547 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -549,7 +549,7 @@ fn evaluate_filters_to_mask( struct DmlResultExec { rows_affected: u64, schema: SchemaRef, - properties: PlanProperties, + properties: Arc, } impl DmlResultExec { @@ -570,7 +570,7 @@ impl DmlResultExec { Self { rows_affected, schema, - properties, + properties: Arc::new(properties), } } } @@ -604,7 +604,7 @@ impl ExecutionPlan for DmlResultExec { Arc::clone(&self.schema) } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } diff --git a/datafusion/common/src/stats.rs b/datafusion/common/src/stats.rs index ba13ef392d912..cecf1d03418d7 100644 --- a/datafusion/common/src/stats.rs +++ b/datafusion/common/src/stats.rs @@ -391,8 +391,13 @@ impl Statistics { /// For example, if we had statistics for columns `{"a", "b", "c"}`, /// projecting to `vec![2, 1]` would return statistics for columns `{"c", /// "b"}`. - pub fn project(mut self, projection: Option<&Vec>) -> Self { - let Some(projection) = projection else { + pub fn project(self, projection: Option<&impl AsRef<[usize]>>) -> Self { + let projection = projection.map(AsRef::as_ref); + self.project_impl(projection) + } + + fn project_impl(mut self, projection: Option<&[usize]>) -> Self { + let Some(projection) = projection.map(AsRef::as_ref) else { return self; }; @@ -410,7 +415,7 @@ impl Statistics { .map(Slot::Present) .collect(); - for idx in projection { + for idx in projection.iter() { let next_idx = self.column_statistics.len(); let slot = std::mem::replace( columns.get_mut(*idx).expect("projection out of bounds"), @@ -1066,7 +1071,7 @@ mod tests { #[test] fn test_project_none() { - let projection = None; + let projection: Option> = None; let stats = make_stats(vec![10, 20, 30]).project(projection.as_ref()); assert_eq!(stats, make_stats(vec![10, 20, 30])); } diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 03310a7bde193..016b188e3d6b8 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -70,10 +70,10 @@ use std::thread::available_parallelism; /// ``` pub fn project_schema( schema: &SchemaRef, - projection: Option<&Vec>, + projection: Option<&impl AsRef<[usize]>>, ) -> Result { let schema = match projection { - Some(columns) => Arc::new(schema.project(columns)?), + Some(columns) => Arc::new(schema.project(columns.as_ref())?), None => Arc::clone(schema), }; Ok(schema) diff --git a/datafusion/core/benches/reset_plan_states.rs b/datafusion/core/benches/reset_plan_states.rs index f2f81f755b96e..56ad0a359d0eb 100644 --- a/datafusion/core/benches/reset_plan_states.rs +++ b/datafusion/core/benches/reset_plan_states.rs @@ -15,15 +15,20 @@ // specific language governing permissions and limitations // under the License. +use std::cell::OnceCell; use std::sync::{Arc, LazyLock}; use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; +use criterion::measurement::WallTime; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion::prelude::SessionContext; use datafusion_catalog::MemTable; +use datafusion_common::metadata::ScalarAndMetadata; +use datafusion_common::{ParamValues, ScalarValue}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::displayable; use datafusion_physical_plan::execution_plan::reset_plan_states; +use datafusion_physical_plan::reuse::ReusableExecutionPlan; use tokio::runtime::Runtime; const NUM_FIELDS: usize = 1000; @@ -37,6 +42,44 @@ static SCHEMA: LazyLock = LazyLock::new(|| { )) }); +/// Decides when to generate placeholders, helping to form a query +/// with a certain placeholders percent. +struct PlaceholderCounter { + placeholders_percent: usize, + c: usize, + placeholder_idx: usize, + num_placeholders: usize, +} + +impl PlaceholderCounter { + fn new(placeholders_percent: usize) -> Self { + Self { + placeholders_percent, + c: 0, + placeholder_idx: 0, + num_placeholders: 0, + } + } + + fn placeholder(&mut self) -> Option { + let is_placeholder = self.c < self.placeholders_percent; + self.c += 1; + if self.c >= 100 { + self.c = 0; + } + if is_placeholder { + self.num_placeholders += 1; + Some("$1".to_owned()) + } else { + None + } + } + + fn placeholder_or(&mut self, f: impl FnOnce() -> String) -> String { + self.placeholder().unwrap_or_else(f) + } +} + fn col_name(i: usize) -> String { format!("x_{i}") } @@ -47,7 +90,7 @@ fn aggr_name(i: usize) -> String { fn physical_plan( ctx: &SessionContext, - rt: &Runtime, + rt: &tokio::runtime::Handle, sql: &str, ) -> Arc { rt.block_on(async { @@ -60,15 +103,16 @@ fn physical_plan( }) } -fn predicate(col_name: impl Fn(usize) -> String, len: usize) -> String { +fn predicate(mut comparee: impl FnMut(usize) -> (String, String), len: usize) -> String { let mut predicate = String::new(); for i in 0..len { if i > 0 { predicate.push_str(" AND "); } - predicate.push_str(&col_name(i)); + let (lhs, rhs) = comparee(i); + predicate.push_str(&lhs); predicate.push_str(" = "); - predicate.push_str(&i.to_string()); + predicate.push_str(&rhs); } predicate } @@ -83,7 +127,8 @@ fn predicate(col_name: impl Fn(usize) -> String, len: usize) -> String { /// /// Where `p1` and `p2` some long predicates. /// -fn query1() -> String { +fn query0(placeholders_percent: usize) -> (String, usize) { + let mut plc = PlaceholderCounter::new(placeholders_percent); let mut query = String::new(); query.push_str("SELECT "); for i in 0..NUM_FIELDS { @@ -91,15 +136,37 @@ fn query1() -> String { query.push_str(", "); } query.push_str("AVG("); - query.push_str(&col_name(i)); + + if let Some(placeholder) = plc.placeholder() { + query.push_str(&format!("{}+{}", placeholder, col_name(i))); + } else { + query.push_str(&col_name(i)); + } + query.push_str(") AS "); query.push_str(&aggr_name(i)); } query.push_str(" FROM t WHERE "); - query.push_str(&predicate(col_name, PREDICATE_LEN)); + query.push_str(&predicate( + |i| { + ( + plc.placeholder_or(|| col_name(i)), + plc.placeholder_or(|| col_name(i + 1)), + ) + }, + PREDICATE_LEN, + )); query.push_str(" HAVING "); - query.push_str(&predicate(aggr_name, PREDICATE_LEN)); - query + query.push_str(&predicate( + |i| { + ( + plc.placeholder_or(|| aggr_name(i)), + plc.placeholder_or(|| aggr_name(i + 1)), + ) + }, + PREDICATE_LEN, + )); + (query, plc.num_placeholders) } /// Returns a typical plan for the query like: @@ -109,27 +176,35 @@ fn query1() -> String { /// WHERE p1 /// ``` /// -fn query2() -> String { +fn query1(placeholders_percent: usize) -> (String, usize) { + let mut plc = PlaceholderCounter::new(placeholders_percent); let mut query = String::new(); query.push_str("SELECT "); for i in (0..NUM_FIELDS).step_by(2) { if i > 0 { query.push_str(", "); } - if (i / 2) % 2 == 0 { - query.push_str(&format!("t.{}", col_name(i))); + let col = if (i / 2) % 2 == 0 { + format!("t.{}", col_name(i)) } else { - query.push_str(&format!("v.{}", col_name(i))); - } + format!("v.{}", col_name(i)) + }; + let add = plc.placeholder_or(|| "1".to_owned()); + let proj = format!("{col} + {add}"); + query.push_str(&proj); } query.push_str(" FROM t JOIN v ON t.x_0 = v.x_0 WHERE "); - fn qualified_name(i: usize) -> String { - format!("t.{}", col_name(i)) - } - - query.push_str(&predicate(qualified_name, PREDICATE_LEN)); - query + query.push_str(&predicate( + |i| { + ( + plc.placeholder_or(|| format!("t.{}", col_name(i))), + plc.placeholder_or(|| i.to_string()), + ) + }, + PREDICATE_LEN, + )); + (query, plc.num_placeholders) } /// Returns a typical plan for the query like: @@ -139,7 +214,8 @@ fn query2() -> String { /// WHERE p /// ``` /// -fn query3() -> String { +fn query2(placeholders_percent: usize) -> (String, usize) { + let mut plc = PlaceholderCounter::new(placeholders_percent); let mut query = String::new(); query.push_str("SELECT "); @@ -150,22 +226,23 @@ fn query3() -> String { } query.push_str(&col_name(i * 2)); query.push_str(" + "); - query.push_str(&col_name(i * 2 + 1)); + query.push_str(&plc.placeholder_or(|| col_name(i * 2 + 1))); } query.push_str(" FROM t WHERE "); - query.push_str(&predicate(col_name, PREDICATE_LEN)); - query -} - -fn run_reset_states(b: &mut criterion::Bencher, plan: &Arc) { - b.iter(|| std::hint::black_box(reset_plan_states(Arc::clone(plan)).unwrap())); + query.push_str(&predicate( + |i| { + ( + plc.placeholder_or(|| col_name(i)), + plc.placeholder_or(|| i.to_string()), + ) + }, + PREDICATE_LEN, + )); + (query, plc.num_placeholders) } -/// 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_reset_plan_states(c: &mut Criterion) { +fn init() -> (SessionContext, Runtime) { let rt = Runtime::new().unwrap(); let ctx = SessionContext::new(); ctx.register_table( @@ -179,20 +256,73 @@ fn bench_reset_plan_states(c: &mut Criterion) { Arc::new(MemTable::try_new(Arc::clone(&SCHEMA), vec![vec![], vec![]]).unwrap()), ) .unwrap(); + (ctx, rt) +} - macro_rules! bench_query { - ($query_producer: expr) => {{ - let sql = $query_producer(); - let plan = physical_plan(&ctx, &rt, &sql); - log::debug!("plan:\n{}", displayable(plan.as_ref()).indent(true)); - move |b| run_reset_states(b, &plan) - }}; - } +/// 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_reset( + g: &mut criterion::BenchmarkGroup<'_, WallTime>, + query_fn: impl FnOnce() -> String, +) { + let (ctx, rt) = init(); + let query = query_fn(); + let rt = rt.handle(); + let plan: OnceCell> = OnceCell::new(); + g.bench_function("reset", |b| { + let plan = plan.get_or_init(|| { + log::info!("sql:\n{}\n\n", query); + let plan = physical_plan(&ctx, rt, &query); + log::info!("plan:\n{}", displayable(plan.as_ref()).indent(true)); + plan + }); + b.iter(|| std::hint::black_box(reset_plan_states(Arc::clone(&plan)).unwrap())) + }); +} + +/// The same as [`bench_reset`] for placeholdered plans. +/// `placeholders_percent` is a percent of placeholders that must be used in generated queries. +fn bench_bind( + g: &mut criterion::BenchmarkGroup<'_, WallTime>, + placeholders_percent: usize, + query_fn: impl FnOnce(usize) -> (String, usize), +) { + let (ctx, rt) = init(); + let params = ParamValues::List(vec![ScalarAndMetadata::new( + ScalarValue::Int64(Some(42)), + None, + )]); + let (query, num_placeholders) = query_fn(placeholders_percent); + let rt = rt.handle(); + let plan: OnceCell = OnceCell::new(); + g.bench_function(format!("{num_placeholders}_placeholders"), move |b| { + let plan = plan.get_or_init(|| { + log::info!("sql:\n{}\n\n", query); + let plan = physical_plan(&ctx, rt, &query); + log::info!("plan:\n{}", displayable(plan.as_ref()).indent(true)); + plan.into() + }); + b.iter(|| std::hint::black_box(plan.bind(Some(¶ms)))) + }); +} - c.bench_function("query1", bench_query!(query1)); - c.bench_function("query2", bench_query!(query2)); - c.bench_function("query3", bench_query!(query3)); +fn criterion_benchmark(c: &mut Criterion) { + env_logger::init(); + + for (query_idx, query_fn) in [query0, query1, query2].iter().enumerate() { + { + let mut g = c.benchmark_group(format!("reset_query{query_idx}")); + bench_reset(&mut g, || query_fn(0).0); + } + { + let mut g = c.benchmark_group(format!("bind_query{query_idx}")); + for placeholders_percent in [0, 1, 10, 50, 100] { + bench_bind(&mut g, placeholders_percent, query_fn); + } + } + } } -criterion_group!(benches, bench_reset_plan_states); +criterion_group!(benches, criterion_benchmark); criterion_main!(benches); diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 54c5867c1804a..ce070b4b0203e 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1007,7 +1007,7 @@ impl DefaultPhysicalPlanner { // project the output columns excluding the async functions // The async functions are always appended to the end of the schema. .apply_projection(Some( - (0..input.schema().fields().len()).collect(), + (0..input.schema().fields().len()).collect::>(), ))? .with_batch_size(session_state.config().batch_size()) .build()? @@ -3722,13 +3722,15 @@ mod tests { #[derive(Debug)] struct NoOpExecutionPlan { - cache: PlanProperties, + cache: Arc, } impl NoOpExecutionPlan { fn new(schema: SchemaRef) -> Self { let cache = Self::compute_properties(schema); - Self { cache } + Self { + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -3766,7 +3768,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -3920,7 +3922,7 @@ digraph { fn children(&self) -> Vec<&Arc> { self.0.iter().collect::>() } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } fn execute( @@ -3969,7 +3971,7 @@ digraph { fn children(&self) -> Vec<&Arc> { unimplemented!() } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } fn execute( @@ -4090,7 +4092,7 @@ digraph { fn children(&self) -> Vec<&Arc> { vec![] } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } fn execute( diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index 8453615c2886b..6dd09ebd8832f 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -79,7 +79,7 @@ struct CustomTableProvider; #[derive(Debug, Clone)] struct CustomExecutionPlan { projection: Option>, - cache: PlanProperties, + cache: Arc, } impl CustomExecutionPlan { @@ -88,7 +88,10 @@ impl CustomExecutionPlan { let schema = project_schema(&schema, projection.as_ref()).expect("projected schema"); let cache = Self::compute_properties(schema); - Self { projection, cache } + Self { + projection, + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -157,7 +160,7 @@ impl ExecutionPlan for CustomExecutionPlan { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs index ca1eaa1f958ea..7a624d0636530 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -62,13 +62,16 @@ fn create_batch(value: i32, num_rows: usize) -> Result { #[derive(Debug)] struct CustomPlan { batches: Vec, - cache: PlanProperties, + cache: Arc, } impl CustomPlan { fn new(schema: SchemaRef, batches: Vec) -> Self { let cache = Self::compute_properties(schema); - Self { batches, cache } + Self { + batches, + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -109,7 +112,7 @@ impl ExecutionPlan for CustomPlan { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index 820c2a470b376..6964d72eacffe 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -45,7 +45,7 @@ use async_trait::async_trait; struct StatisticsValidation { stats: Statistics, schema: Arc, - cache: PlanProperties, + cache: Arc, } impl StatisticsValidation { @@ -59,7 +59,7 @@ impl StatisticsValidation { Self { stats, schema, - cache, + cache: Arc::new(cache), } } @@ -158,7 +158,7 @@ impl ExecutionPlan for StatisticsValidation { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/fuzz_cases/once_exec.rs b/datafusion/core/tests/fuzz_cases/once_exec.rs index 49e2caaa7417c..69edf9be1d825 100644 --- a/datafusion/core/tests/fuzz_cases/once_exec.rs +++ b/datafusion/core/tests/fuzz_cases/once_exec.rs @@ -32,7 +32,7 @@ use std::sync::{Arc, Mutex}; pub struct OnceExec { /// the results to send back stream: Mutex>, - cache: PlanProperties, + cache: Arc, } impl Debug for OnceExec { @@ -46,7 +46,7 @@ impl OnceExec { let cache = Self::compute_properties(stream.schema()); Self { stream: Mutex::new(Some(stream)), - cache, + cache: Arc::new(cache), } } @@ -83,7 +83,7 @@ impl ExecutionPlan for OnceExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 94ae82a9ad755..b3a3d29d070b1 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -120,7 +120,7 @@ macro_rules! assert_plan { struct SortRequiredExec { input: Arc, expr: LexOrdering, - cache: PlanProperties, + cache: Arc, } impl SortRequiredExec { @@ -132,7 +132,7 @@ impl SortRequiredExec { Self { input, expr: requirement, - cache, + cache: Arc::new(cache), } } @@ -174,7 +174,7 @@ impl ExecutionPlan for SortRequiredExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 9234a95591baa..ab35f399e2b11 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -762,7 +762,7 @@ async fn test_hash_join_swap_on_joins_with_projections( "ProjectionExec won't be added above if HashJoinExec contains embedded projection", ); - assert_eq!(swapped_join.projection, Some(vec![0_usize])); + assert_eq!(swapped_join.projection.as_deref().unwrap(), &[0_usize]); assert_eq!(swapped.schema().fields.len(), 1); assert_eq!(swapped.schema().fields[0].name(), "small_col"); Ok(()) @@ -979,7 +979,7 @@ impl RecordBatchStream for UnboundedStream { pub struct UnboundedExec { batch_produce: Option, batch: RecordBatch, - cache: PlanProperties, + cache: Arc, } impl UnboundedExec { @@ -995,7 +995,7 @@ impl UnboundedExec { Self { batch_produce, batch, - cache, + cache: Arc::new(cache), } } @@ -1052,7 +1052,7 @@ impl ExecutionPlan for UnboundedExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -1091,7 +1091,7 @@ pub enum SourceType { pub struct StatisticsExec { stats: Statistics, schema: Arc, - cache: PlanProperties, + cache: Arc, } impl StatisticsExec { @@ -1105,7 +1105,7 @@ impl StatisticsExec { Self { stats, schema: Arc::new(schema), - cache, + cache: Arc::new(cache), } } @@ -1153,7 +1153,7 @@ impl ExecutionPlan for StatisticsExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/core/tests/physical_optimizer/mod.rs b/datafusion/core/tests/physical_optimizer/mod.rs index 7d22a8c25f209..cf179cb727cf1 100644 --- a/datafusion/core/tests/physical_optimizer/mod.rs +++ b/datafusion/core/tests/physical_optimizer/mod.rs @@ -30,8 +30,6 @@ mod join_selection; mod limit_pushdown; mod limited_distinct_aggregation; mod partition_statistics; -#[expect(clippy::needless_pass_by_value)] -mod physical_expr_resolver; mod projection_pushdown; mod pushdown_sort; mod replace_with_order_preserving_variants; diff --git a/datafusion/core/tests/physical_optimizer/physical_expr_resolver.rs b/datafusion/core/tests/physical_optimizer/physical_expr_resolver.rs deleted file mode 100644 index 43013cd62c65f..0000000000000 --- a/datafusion/core/tests/physical_optimizer/physical_expr_resolver.rs +++ /dev/null @@ -1,276 +0,0 @@ -// 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. - -//! Integration tests for [`PhysicalExprResolver`] optimizer rule. - -use std::{collections::HashMap, sync::Arc}; - -use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use datafusion::config::ConfigOptions; -use datafusion_common::{ParamValues, Result, ScalarValue}; -use datafusion_execution::TaskContext; -use datafusion_expr::Operator; -use datafusion_physical_expr::{ - Partitioning, - expressions::{BinaryExpr, col, lit, placeholder}, -}; -use datafusion_physical_optimizer::{ - PhysicalOptimizerRule, physical_expr_resolver::PhysicalExprResolver, -}; -use datafusion_physical_plan::{ - ExecutionPlan, filter::FilterExec, get_plan_string, - plan_transformer::TransformPlanExec, repartition::RepartitionExec, -}; - -use crate::physical_optimizer::test_utils::{ - coalesce_partitions_exec, global_limit_exec, resolve_placeholders_exec, stream_exec, -}; - -fn create_schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new("c1", DataType::Int32, true), - Field::new("c2", DataType::Int32, true), - Field::new("c3", DataType::Int32, true), - ])) -} - -fn filter_exec( - schema: SchemaRef, - input: Arc, -) -> Result> { - Ok(Arc::new(FilterExec::try_new( - Arc::new(BinaryExpr::new( - col("c3", schema.as_ref()).unwrap(), - Operator::Gt, - lit(0), - )), - input, - )?)) -} - -fn filter_exec_with_placeholders( - schema: SchemaRef, - input: Arc, -) -> Result> { - Ok(Arc::new(FilterExec::try_new( - Arc::new(BinaryExpr::new( - col("c3", schema.as_ref()).unwrap(), - Operator::Gt, - placeholder("$foo", DataType::Int32), - )), - input, - )?)) -} - -fn repartition_exec( - streaming_table: Arc, -) -> Result> { - Ok(Arc::new(RepartitionExec::try_new( - streaming_table, - Partitioning::RoundRobinBatch(8), - )?)) -} - -#[test] -fn test_noop_if_no_placeholders_found() -> Result<()> { - let schema = create_schema(); - let streaming_table = stream_exec(&schema); - let repartition = repartition_exec(streaming_table)?; - let filter = filter_exec(schema, repartition)?; - let coalesce_partitions = coalesce_partitions_exec(filter); - let plan = global_limit_exec(coalesce_partitions, 0, Some(5)); - - let initial = get_plan_string(&plan); - let expected_initial = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > 0", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(initial, expected_initial); - - let after_optimize = PhysicalExprResolver::new_post_optimization() - .optimize(plan, &ConfigOptions::new())?; - - let optimized_plan_string = get_plan_string(&after_optimize); - assert_eq!(initial, optimized_plan_string); - - Ok(()) -} - -#[test] -fn test_wrap_plan_with_transformer() -> Result<()> { - let schema = create_schema(); - let streaming_table = stream_exec(&schema); - let repartition = repartition_exec(streaming_table)?; - let filter = filter_exec_with_placeholders(schema, repartition)?; - let coalesce_partitions = coalesce_partitions_exec(filter); - let plan = global_limit_exec(coalesce_partitions, 0, Some(5)); - - let initial = get_plan_string(&plan); - let expected_initial = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > $foo", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(initial, expected_initial); - - let after_optimize = PhysicalExprResolver::new_post_optimization() - .optimize(plan, &ConfigOptions::new())?; - - let expected_optimized = [ - "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]", - " GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > $foo", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - let optimized_plan_string = get_plan_string(&after_optimize); - assert_eq!(optimized_plan_string, expected_optimized); - - let transformer = after_optimize - .as_ref() - .as_any() - .downcast_ref::() - .expect("should downcast"); - - let param_values = ParamValues::Map(HashMap::from_iter([( - "foo".to_string(), - ScalarValue::Int32(Some(100)).into(), - )])); - - let ctx = Arc::new(TaskContext::default().with_param_values(param_values)); - let resolved_plan = transformer.transform(&ctx)?; - let resolved_plan_string = get_plan_string(&resolved_plan); - let expected_resolved = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > 100", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(resolved_plan_string, expected_resolved); - - Ok(()) -} - -#[test] -fn test_remove_useless_transformers() -> Result<()> { - let schema = create_schema(); - let streaming_table = stream_exec(&schema); - let repartition = repartition_exec(streaming_table)?; - let filter = filter_exec(schema, repartition)?; - let transformer = resolve_placeholders_exec(filter); - let coalesce_partitions = coalesce_partitions_exec(transformer); - let global_limit = global_limit_exec(coalesce_partitions, 0, Some(5)); - let plan = resolve_placeholders_exec(global_limit); - - let initial = get_plan_string(&plan); - let expected_initial = [ - "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=0]", - " GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=0]", - " FilterExec: c3@2 > 0", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(initial, expected_initial); - - let after_optimize = - PhysicalExprResolver::new().optimize(plan, &ConfigOptions::new())?; - - let expected_optimized = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > 0", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - let optimized_plan_string = get_plan_string(&after_optimize); - assert_eq!(optimized_plan_string, expected_optimized); - - Ok(()) -} - -#[test] -fn test_combine_transformers() -> Result<()> { - let schema = create_schema(); - let streaming_table = stream_exec(&schema); - let repartition = repartition_exec(streaming_table)?; - let transformer = resolve_placeholders_exec(repartition); - let filter = filter_exec_with_placeholders(schema, transformer)?; - let transformer = resolve_placeholders_exec(filter); - let coalesce_partitions = coalesce_partitions_exec(transformer); - let global_limit = global_limit_exec(coalesce_partitions, 0, Some(5)); - let plan = resolve_placeholders_exec(global_limit); - - let initial = get_plan_string(&plan); - let expected_initial = [ - "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=0]", - " GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]", - " FilterExec: c3@2 > $foo", - " TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=0]", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - assert_eq!(initial, expected_initial); - - let after_pre_optimization = - PhysicalExprResolver::new().optimize(plan, &ConfigOptions::new())?; - - let expected_optimized = [ - "GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > $foo", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - let optimized_plan_string = get_plan_string(&after_pre_optimization); - assert_eq!(optimized_plan_string, expected_optimized); - - let after_post_optimization = PhysicalExprResolver::new_post_optimization() - .optimize(after_pre_optimization, &ConfigOptions::new())?; - - let expected_optimized = [ - "TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1]", - " GlobalLimitExec: skip=0, fetch=5", - " CoalescePartitionsExec", - " FilterExec: c3@2 > $foo", - " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", - " StreamingTableExec: partition_sizes=1, projection=[c1, c2, c3], infinite_source=true", - ]; - - let optimized_plan_string = get_plan_string(&after_post_optimization); - assert_eq!(optimized_plan_string, expected_optimized); - - Ok(()) -} diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 6db9d816c40aa..02a3f7d0f67dd 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -484,7 +484,7 @@ impl ExecutionPlan for TestNode { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { self.input.properties() } diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 582506457560c..f8c91ba272a9f 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -59,9 +59,6 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{JoinFilter, JoinOn}; use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; -use datafusion_physical_plan::plan_transformer::{ - ResolvePlaceholdersRule, TransformPlanExec, -}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; @@ -396,15 +393,6 @@ pub fn projection_exec( Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) } -pub fn resolve_placeholders_exec( - input: Arc, -) -> Arc { - Arc::new( - TransformPlanExec::try_new(input, vec![Box::new(ResolvePlaceholdersRule::new())]) - .unwrap(), - ) -} - /// A test [`ExecutionPlan`] whose requirements can be configured. #[derive(Debug)] pub struct RequirementsTestExec { @@ -466,7 +454,7 @@ impl ExecutionPlan for RequirementsTestExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { self.input.properties() } @@ -837,7 +825,7 @@ pub fn sort_expr_named(name: &str, index: usize) -> PhysicalSortExpr { pub struct TestScan { schema: SchemaRef, output_ordering: Vec, - plan_properties: PlanProperties, + plan_properties: Arc, // Store the requested ordering for display requested_ordering: Option, } @@ -871,7 +859,7 @@ impl TestScan { Self { schema, output_ordering, - plan_properties, + plan_properties: Arc::new(plan_properties), requested_ordering: None, } } @@ -927,7 +915,7 @@ impl ExecutionPlan for TestScan { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.plan_properties } diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index 5ba8545766e32..3c2e256a09bc4 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -19,7 +19,7 @@ use std::collections::HashMap; use super::*; use datafusion_common::{ParamValues, ScalarValue, metadata::ScalarAndMetadata}; -use datafusion_execution::TaskContext; +use datafusion_physical_plan::reuse::ReusableExecutionPlan; use insta::assert_snapshot; #[tokio::test] @@ -439,15 +439,15 @@ async fn test_resolve_window_function() -> Result<()> { let batch = record_batch!(("id", Int32, [1, 2]), ("name", Utf8, ["Alex", "Bob"]))?; ctx.register_batch("t1", batch)?; - let plan = ctx - .sql("SELECT id, SUM(id + $1) OVER (PARTITION BY name ORDER BY id) FROM t1") - .await? - .create_physical_plan() - .await?; + let plan = ReusableExecutionPlan::new( + ctx.sql("SELECT id, SUM(id + $1) OVER (PARTITION BY name ORDER BY id) FROM t1") + .await? + .create_physical_plan() + .await?, + ); let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(100)).into()]); - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(plan, task_ctx).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +----+--------------------------------------------------------------------------------------------------------------------------+ @@ -479,15 +479,15 @@ async fn test_resolve_join() -> Result<()> { ctx.register_batch("t1", batch_1)?; ctx.register_batch("t2", batch_2)?; - let plan = ctx - .sql("SELECT t1.name, t2.age FROM t1 JOIN t2 ON t1.id + $1 = t2.id;") - .await? - .create_physical_plan() - .await?; + let plan = ReusableExecutionPlan::new( + ctx.sql("SELECT t1.name, t2.age FROM t1 JOIN t2 ON t1.id + $1 = t2.id;") + .await? + .create_physical_plan() + .await?, + ); let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(8)).into()]); - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(plan, task_ctx).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +------+-----+ @@ -503,24 +503,21 @@ async fn test_resolve_join() -> Result<()> { #[tokio::test] async fn test_resolve_cast() -> Result<()> { let ctx = SessionContext::new(); - let plan = ctx - .sql("SELECT CAST($1 as INT)") - .await? - .create_physical_plan() - .await?; + let plan = ReusableExecutionPlan::new( + ctx.sql("SELECT CAST($1 as INT)") + .await? + .create_physical_plan() + .await?, + ); let param_values = ParamValues::List(vec![ ScalarValue::Utf8(Some("not a number".to_string())).into(), ]); - - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let result = collect(Arc::clone(&plan), task_ctx).await; - assert!(result.is_err()); + assert!(plan.bind(Some(¶m_values)).is_err()); let param_values = ParamValues::List(vec![ScalarValue::Utf8(Some("200".to_string())).into()]); - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(plan, task_ctx).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +-----+ @@ -536,18 +533,17 @@ async fn test_resolve_cast() -> Result<()> { #[tokio::test] async fn test_resolve_try_cast() -> Result<()> { let ctx = SessionContext::new(); - let plan = ctx - .sql("SELECT TRY_CAST($1 as INT)") - .await? - .create_physical_plan() - .await?; + let plan = ReusableExecutionPlan::new( + ctx.sql("SELECT TRY_CAST($1 as INT)") + .await? + .create_physical_plan() + .await?, + ); let param_values = ParamValues::List(vec![ ScalarValue::Utf8(Some("not a number".to_string())).into(), ]); - - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(Arc::clone(&plan), task_ctx).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +----+ @@ -559,8 +555,7 @@ async fn test_resolve_try_cast() -> Result<()> { let param_values = ParamValues::List(vec![ScalarValue::Utf8(Some("200".to_string())).into()]); - let task_ctx = Arc::new(TaskContext::from(&ctx).with_param_values(param_values)); - let batches = collect(plan, task_ctx).await?; + let batches = collect(plan.bind(Some(¶m_values))?.plan(), ctx.task_ctx()).await?; assert_snapshot!(batches_to_sort_string(&batches), @r" +-----+ diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index 7ad00dece1b24..4d2a31ca1f960 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -122,20 +122,22 @@ impl TableProvider for TestInsertTableProvider { #[derive(Debug)] struct TestInsertExec { op: InsertOp, - plan_properties: PlanProperties, + plan_properties: Arc, } impl TestInsertExec { fn new(op: InsertOp) -> Self { Self { op, - plan_properties: PlanProperties::new( - EquivalenceProperties::new(make_count_schema()), - Partitioning::UnknownPartitioning(1), - EmissionType::Incremental, - Boundedness::Bounded, - ) - .with_scheduling_type(SchedulingType::Cooperative), + plan_properties: Arc::new( + PlanProperties::new( + EquivalenceProperties::new(make_count_schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ) + .with_scheduling_type(SchedulingType::Cooperative), + ), } } } @@ -159,7 +161,7 @@ impl ExecutionPlan for TestInsertExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.plan_properties } diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index d53e076739608..c2533e73d2be9 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -653,13 +653,17 @@ struct TopKExec { input: Arc, /// The maximum number of values k: usize, - cache: PlanProperties, + cache: Arc, } impl TopKExec { fn new(input: Arc, k: usize) -> Self { let cache = Self::compute_properties(input.schema()); - Self { input, k, cache } + Self { + input, + k, + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -704,7 +708,7 @@ impl ExecutionPlan for TopKExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 5acc89722b200..f149109dff5cc 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -89,7 +89,7 @@ pub struct DataSinkExec { count_schema: SchemaRef, /// Optional required sort order for output data. sort_order: Option, - cache: PlanProperties, + cache: Arc, } impl Debug for DataSinkExec { @@ -117,7 +117,7 @@ impl DataSinkExec { sink, count_schema: make_count_schema(), sort_order, - cache, + cache: Arc::new(cache), } } @@ -174,7 +174,7 @@ impl ExecutionPlan for DataSinkExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 75413c743b8e1..cffb4f41e61c9 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -267,7 +267,7 @@ pub struct DataSourceExec { /// The source of the data -- for example, `FileScanConfig` or `MemorySourceConfig` data_source: Arc, /// Cached plan properties such as sort order - cache: PlanProperties, + cache: Arc, } impl DisplayAs for DataSourceExec { @@ -291,7 +291,7 @@ impl ExecutionPlan for DataSourceExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -361,7 +361,7 @@ impl ExecutionPlan for DataSourceExec { fn with_fetch(&self, limit: Option) -> Option> { let data_source = self.data_source.with_fetch(limit)?; - let cache = self.cache.clone(); + let cache = Arc::clone(&self.cache); Some(Arc::new(Self { data_source, cache })) } @@ -405,7 +405,8 @@ impl ExecutionPlan for DataSourceExec { let mut new_node = self.clone(); new_node.data_source = data_source; // Re-compute properties since we have new filters which will impact equivalence info - new_node.cache = Self::compute_properties(&new_node.data_source); + new_node.cache = + Arc::new(Self::compute_properties(&new_node.data_source)); Ok(FilterPushdownPropagation { filters: res.filters, @@ -470,7 +471,10 @@ impl DataSourceExec { // Default constructor for `DataSourceExec`, setting the `cooperative` flag to `true`. pub fn new(data_source: Arc) -> Self { let cache = Self::compute_properties(&data_source); - Self { data_source, cache } + Self { + data_source, + cache: Arc::new(cache), + } } /// Return the source object @@ -479,20 +483,20 @@ impl DataSourceExec { } pub fn with_data_source(mut self, data_source: Arc) -> Self { - self.cache = Self::compute_properties(&data_source); + self.cache = Arc::new(Self::compute_properties(&data_source)); self.data_source = data_source; self } /// Assign constraints pub fn with_constraints(mut self, constraints: Constraints) -> Self { - self.cache = self.cache.with_constraints(constraints); + Arc::make_mut(&mut self.cache).set_constraints(constraints); self } /// Assign output partitioning pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self { - self.cache = self.cache.with_partitioning(partitioning); + Arc::make_mut(&mut self.cache).partitioning = partitioning; self } diff --git a/datafusion/datasource/src/values.rs b/datafusion/datasource/src/values.rs index bea1f871d76a2..348e26fef177a 100644 --- a/datafusion/datasource/src/values.rs +++ b/datafusion/datasource/src/values.rs @@ -341,8 +341,7 @@ mod tests { use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, lit, placeholder}; use datafusion_physical_plan::{ - ExecutionPlan, collect, - plan_transformer::{ResolvePlaceholdersRule, TransformPlanExec}, + ExecutionPlan, collect, reuse::ReusableExecutionPlan, }; #[test] @@ -437,13 +436,17 @@ mod tests { // Should be ValuesSource because of placeholder. assert!(values_exec.data_source().as_any().is::()); - let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>]; - let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?); - let task_ctx = Arc::new(TaskContext::default().with_param_values( - ParamValues::List(vec![ScalarValue::Int32(Some(10)).into()]), - )); + let exec = ReusableExecutionPlan::new(values_exec as _); + let task_ctx = Arc::new(TaskContext::default()); - let batch = collect(exec, task_ctx).await?; + let batch = collect( + exec.bind(Some(&ParamValues::List(vec![ + ScalarValue::Int32(Some(10)).into(), + ])))? + .plan(), + task_ctx, + ) + .await?; let expected = [ "+-----+----+", "| a | b |", @@ -469,18 +472,20 @@ mod tests { placeholder("$2", DataType::Int32), ]]; - let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)?; - let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>]; - let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?) as Arc<_>; + let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)? as _; + let exec = ReusableExecutionPlan::new(values_exec); - let task_ctx = Arc::new(TaskContext::default().with_param_values( - ParamValues::List(vec![ + let task_ctx = Arc::new(TaskContext::default()); + + let batch = collect( + exec.bind(Some(&ParamValues::List(vec![ ScalarValue::Int32(Some(10)).into(), ScalarValue::Int32(Some(20)).into(), - ]), - )); - - let batch = collect(Arc::clone(&exec), task_ctx).await?; + ])))? + .plan(), + Arc::clone(&task_ctx), + ) + .await?; let expected = [ "+----+----+", "| a | b |", @@ -490,14 +495,15 @@ mod tests { ]; assert_batches_eq!(expected, &batch); - let task_ctx = Arc::new(TaskContext::default().with_param_values( - ParamValues::List(vec![ + let batch = collect( + exec.bind(Some(&ParamValues::List(vec![ ScalarValue::Int32(Some(30)).into(), ScalarValue::Int32(Some(40)).into(), - ]), - )); - - let batch = collect(exec, task_ctx).await?; + ])))? + .plan(), + task_ctx, + ) + .await?; let expected = [ "+----+----+", "| a | b |", @@ -527,22 +533,24 @@ mod tests { let data: Vec>> = vec![vec![lit(10), placeholder("$foo", DataType::Int32)]]; - let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)?; - let rules = vec![Box::new(ResolvePlaceholdersRule::new()) as Box<_>]; - let exec = Arc::new(TransformPlanExec::try_new(values_exec, rules)?) as Arc<_>; + let values_exec = ValuesSource::try_new_exec(Arc::clone(&schema), data)? as _; let task_ctx = Arc::new(TaskContext::default()); - let result = collect(Arc::clone(&exec), task_ctx).await; + let result = collect(Arc::clone(&values_exec), task_ctx).await; assert!(result.is_err()); - let task_ctx = Arc::new(TaskContext::default().with_param_values( - ParamValues::Map(HashMap::from_iter([( + let exec = ReusableExecutionPlan::new(values_exec); + + let task_ctx = Arc::new(TaskContext::default()); + let batch = collect( + exec.bind(Some(&ParamValues::Map(HashMap::from_iter([( "foo".to_string(), ScalarValue::Int32(Some(20)).into(), - )])), - )); - - let batch = collect(Arc::clone(&exec), task_ctx).await?; + )]))))? + .plan(), + task_ctx, + ) + .await?; let expected = [ "+----+----+", "| a | b |", diff --git a/datafusion/execution/src/task.rs b/datafusion/execution/src/task.rs index 2c191fbc03dd4..38f31cf4629eb 100644 --- a/datafusion/execution/src/task.rs +++ b/datafusion/execution/src/task.rs @@ -19,9 +19,7 @@ use crate::{ config::SessionConfig, memory_pool::MemoryPool, registry::FunctionRegistry, runtime_env::RuntimeEnv, }; -use datafusion_common::{ - ParamValues, Result, internal_datafusion_err, plan_datafusion_err, -}; +use datafusion_common::{Result, internal_datafusion_err, plan_datafusion_err}; use datafusion_expr::planner::ExprPlanner; use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use std::collections::HashSet; @@ -50,8 +48,6 @@ pub struct TaskContext { window_functions: HashMap>, /// Runtime environment associated with this task context runtime: Arc, - /// External query parameters - param_values: Option, } impl Default for TaskContext { @@ -67,7 +63,6 @@ impl Default for TaskContext { aggregate_functions: HashMap::new(), window_functions: HashMap::new(), runtime, - param_values: None, } } } @@ -95,7 +90,6 @@ impl TaskContext { aggregate_functions, window_functions, runtime, - param_values: None, } } @@ -124,11 +118,6 @@ impl TaskContext { Arc::clone(&self.runtime) } - /// Return param values associated with this [`TaskContext`] - pub fn param_values(&self) -> &Option { - &self.param_values - } - pub fn scalar_functions(&self) -> &HashMap> { &self.scalar_functions } @@ -152,12 +141,6 @@ impl TaskContext { self.runtime = runtime; self } - - /// Update the [`ParamValues`] - pub fn with_param_values(mut self, param_values: ParamValues) -> Self { - self.param_values = Some(param_values); - self - } } impl FunctionRegistry for TaskContext { diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index c879b022067c3..ec8f7538fee1a 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -90,7 +90,7 @@ impl FFI_ExecutionPlan { unsafe extern "C" fn properties_fn_wrapper( plan: &FFI_ExecutionPlan, ) -> FFI_PlanProperties { - plan.inner().properties().into() + plan.inner().properties().as_ref().into() } unsafe extern "C" fn children_fn_wrapper( @@ -192,7 +192,7 @@ impl Drop for FFI_ExecutionPlan { pub struct ForeignExecutionPlan { name: String, plan: FFI_ExecutionPlan, - properties: PlanProperties, + properties: Arc, children: Vec>, } @@ -244,7 +244,7 @@ impl TryFrom<&FFI_ExecutionPlan> for Arc { let plan = ForeignExecutionPlan { name, plan: plan.clone(), - properties, + properties: Arc::new(properties), children, }; @@ -262,7 +262,7 @@ impl ExecutionPlan for ForeignExecutionPlan { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } @@ -278,7 +278,7 @@ impl ExecutionPlan for ForeignExecutionPlan { plan: self.plan.clone(), name: self.name.clone(), children, - properties: self.properties.clone(), + properties: Arc::clone(&self.properties), })) } @@ -305,19 +305,19 @@ pub(crate) mod tests { #[derive(Debug)] pub struct EmptyExec { - props: PlanProperties, + props: Arc, children: Vec>, } impl EmptyExec { pub fn new(schema: arrow::datatypes::SchemaRef) -> Self { Self { - props: PlanProperties::new( + props: Arc::new(PlanProperties::new( datafusion::physical_expr::EquivalenceProperties::new(schema), Partitioning::UnknownPartitioning(3), EmissionType::Incremental, Boundedness::Bounded, - ), + )), children: Vec::default(), } } @@ -342,7 +342,7 @@ pub(crate) mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.props } @@ -355,7 +355,7 @@ pub(crate) mod tests { children: Vec>, ) -> Result> { Ok(Arc::new(EmptyExec { - props: self.props.clone(), + props: Arc::clone(&self.props), children, })) } diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index 6149736c58555..8370cf19e6589 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -162,7 +162,7 @@ impl Drop for AsyncTableProvider { #[derive(Debug)] struct AsyncTestExecutionPlan { - properties: datafusion_physical_plan::PlanProperties, + properties: Arc, batch_request: mpsc::Sender, batch_receiver: broadcast::Receiver>, } @@ -173,12 +173,12 @@ impl AsyncTestExecutionPlan { batch_receiver: broadcast::Receiver>, ) -> Self { Self { - properties: datafusion_physical_plan::PlanProperties::new( + properties: Arc::new(datafusion_physical_plan::PlanProperties::new( EquivalenceProperties::new(super::create_test_schema()), Partitioning::UnknownPartitioning(3), datafusion_physical_plan::execution_plan::EmissionType::Incremental, datafusion_physical_plan::execution_plan::Boundedness::Bounded, - ), + )), batch_request, batch_receiver, } @@ -194,7 +194,7 @@ impl ExecutionPlan for AsyncTestExecutionPlan { self } - fn properties(&self) -> &datafusion_physical_plan::PlanProperties { + fn properties(&self) -> &Arc { &self.properties } diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 996bc4b08fcd2..a98341b10765a 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -207,8 +207,13 @@ impl EquivalenceProperties { } /// Adds constraints to the properties. - pub fn with_constraints(mut self, constraints: Constraints) -> Self { + pub fn set_constraints(&mut self, constraints: Constraints) { self.constraints = constraints; + } + + /// Adds constraints to the properties. + pub fn with_constraints(mut self, constraints: Constraints) -> Self { + self.set_constraints(constraints); self } diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 34314f1c80d5f..8312d6df1f48c 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -55,6 +55,8 @@ pub use literal::{Literal, lit}; pub use negative::{NegativeExpr, negative}; pub use no_op::NoOp; pub use not::{NotExpr, not}; -pub use placeholder::{PlaceholderExpr, has_placeholders, placeholder}; +pub use placeholder::{ + PlaceholderExpr, has_placeholders, placeholder, resolve_expr_placeholders, +}; pub use try_cast::{TryCastExpr, try_cast}; pub use unknown_column::UnKnownColumn; diff --git a/datafusion/physical-expr/src/expressions/placeholder.rs b/datafusion/physical-expr/src/expressions/placeholder.rs index 9149517a0a57a..74cb81550ffc9 100644 --- a/datafusion/physical-expr/src/expressions/placeholder.rs +++ b/datafusion/physical-expr/src/expressions/placeholder.rs @@ -28,12 +28,43 @@ use arrow::{ datatypes::{DataType, Field, FieldRef, Schema}, }; use datafusion_common::{ - DataFusionError, Result, exec_datafusion_err, tree_node::TreeNode, + DataFusionError, ParamValues, Result, exec_datafusion_err, exec_err, + tree_node::{Transformed, TreeNode}, }; use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use std::hash::Hash; +use crate::{expressions::Literal, simplifier::const_evaluator::simplify_const_expr}; + +/// Resolves [`PlaceholderExpr`] in the physical expression using the provided [`ParamValues`]. +pub fn resolve_expr_placeholders( + expr: Arc, + param_values: Option<&ParamValues>, +) -> Result> { + let expr = expr.transform_up(|node| { + let Some(placeholder) = node.as_any().downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(ref field) = placeholder.field else { + return Ok(Transformed::no(node)); + }; + let Some(param_values) = param_values else { + return exec_err!("value for placeholder {} is not found", placeholder.id); + }; + let scalar = param_values.get_placeholders_with_values(&placeholder.id)?; + let value = scalar.value.cast_to(field.data_type())?; + let literal = Literal::new_with_metadata(value, scalar.metadata); + Ok(Transformed::yes(Arc::new(literal))) + })?; + + if expr.transformed { + simplify_const_expr(expr.data).map(|t| t.data) + } else { + Ok(expr.data) + } +} + /// Physical expression representing a placeholder parameter (e.g., $1, $2, or named parameters) in /// the physical plan. /// diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index a84cb618333d3..c6adccc9b49de 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -29,7 +29,8 @@ use arrow::datatypes::{Field, Schema, SchemaRef}; use datafusion_common::stats::{ColumnStatistics, Precision}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{ - Result, ScalarValue, assert_or_internal_err, internal_datafusion_err, plan_err, + Result, ScalarValue, Statistics, assert_or_internal_err, internal_datafusion_err, + plan_err, }; use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; @@ -125,7 +126,8 @@ impl From for (Arc, String) { /// indices. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ProjectionExprs { - exprs: Vec, + /// [`Arc`] used for a cheap clone, which improves physical plan optimization performance. + exprs: Arc<[ProjectionExpr]>, } impl std::fmt::Display for ProjectionExprs { @@ -137,14 +139,16 @@ impl std::fmt::Display for ProjectionExprs { impl From> for ProjectionExprs { fn from(value: Vec) -> Self { - Self { exprs: value } + Self { + exprs: value.into(), + } } } impl From<&[ProjectionExpr]> for ProjectionExprs { fn from(value: &[ProjectionExpr]) -> Self { Self { - exprs: value.to_vec(), + exprs: value.iter().cloned().collect(), } } } @@ -152,7 +156,7 @@ impl From<&[ProjectionExpr]> for ProjectionExprs { impl FromIterator for ProjectionExprs { fn from_iter>(exprs: T) -> Self { Self { - exprs: exprs.into_iter().collect::>(), + exprs: exprs.into_iter().collect(), } } } @@ -164,12 +168,17 @@ impl AsRef<[ProjectionExpr]> for ProjectionExprs { } impl ProjectionExprs { - pub fn new(exprs: I) -> Self - where - I: IntoIterator, - { + /// Make a new [`ProjectionExprs`] from expressions iterator. + pub fn new(exprs: impl IntoIterator) -> Self { + Self { + exprs: exprs.into_iter().collect(), + } + } + + /// Make a new [`ProjectionExprs`] from expressions. + pub fn from_expressions(exprs: impl Into>) -> Self { Self { - exprs: exprs.into_iter().collect::>(), + exprs: exprs.into(), } } @@ -285,13 +294,14 @@ impl ProjectionExprs { { let exprs = self .exprs - .into_iter() + .iter() + .cloned() .map(|mut proj| { proj.expr = f(proj.expr)?; Ok(proj) }) - .collect::>>()?; - Ok(Self::new(exprs)) + .collect::>>()?; + Ok(Self::from_expressions(exprs)) } /// Apply another projection on top of this projection, returning the combined projection. @@ -361,7 +371,7 @@ impl ProjectionExprs { /// applied on top of this projection. pub fn try_merge(&self, other: &ProjectionExprs) -> Result { let mut new_exprs = Vec::with_capacity(other.exprs.len()); - for proj_expr in &other.exprs { + for proj_expr in other.exprs.iter() { let new_expr = update_expr(&proj_expr.expr, &self.exprs, true)? .ok_or_else(|| { internal_datafusion_err!( @@ -614,12 +624,12 @@ impl ProjectionExprs { /// ``` pub fn project_statistics( &self, - mut stats: datafusion_common::Statistics, + mut stats: Statistics, output_schema: &Schema, - ) -> Result { + ) -> Result { let mut column_statistics = vec![]; - for proj_expr in &self.exprs { + for proj_expr in self.exprs.iter() { let expr = &proj_expr.expr; let col_stats = if let Some(col) = expr.as_any().downcast_ref::() { std::mem::take(&mut stats.column_statistics[col.index()]) @@ -766,13 +776,52 @@ impl Projector { } } -impl IntoIterator for ProjectionExprs { - type Item = ProjectionExpr; - type IntoIter = std::vec::IntoIter; +/// Describes an immutable reference counted projection. +/// +/// This structure represents projecting a set of columns by index. +/// [`Arc`] is used to make it cheap to clone. +pub type ProjectionRef = Arc<[usize]>; - fn into_iter(self) -> Self::IntoIter { - self.exprs.into_iter() - } +/// Combine two projections. +/// +/// If `p1` is [`None`] then there are no changes. +/// Otherwise, if passed `p2` is not [`None`] then it is remapped +/// according to the `p1`. Otherwise, there are no changes. +/// +/// # Example +/// +/// If stored projection is [0, 2] and we call `apply_projection([0, 2, 3])`, +/// then the resulting projection will be [0, 3]. +/// +/// # Error +/// +/// Returns an internal error if `p1` contains index that is greater than `p2` len. +/// +pub fn combine_projections( + p1: Option<&ProjectionRef>, + p2: Option<&ProjectionRef>, +) -> Result> { + let Some(p1) = p1 else { + return Ok(None); + }; + let Some(p2) = p2 else { + return Ok(Some(Arc::clone(p1))); + }; + + Ok(Some( + p1.iter() + .map(|i| { + let idx = *i; + assert_or_internal_err!( + idx < p2.len(), + "unable to apply projection: index {} is greater than new projection len {}", + idx, + p2.len(), + ); + Ok(p2[*i]) + }) + .collect::>>()?, + )) } /// The function operates in two modes: diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index acb1c588097ee..790669b5c9dbf 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -49,7 +49,7 @@ use datafusion_physical_plan::aggregates::{ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::execution_plan::EmissionType; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, PartitionMode, SortMergeJoinExec, + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, PartitionMode, SortMergeJoinExec, }; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -305,18 +305,19 @@ pub fn adjust_input_keys_ordering( Vec<(PhysicalExprRef, PhysicalExprRef)>, Vec, )| { - HashJoinExec::try_new( + HashJoinExecBuilder::new( Arc::clone(left), Arc::clone(right), new_conditions.0, - filter.clone(), - join_type, - // TODO: although projection is not used in the join here, because projection pushdown is after enforce_distribution. Maybe we need to handle it later. Same as filter. - projection.clone(), - PartitionMode::Partitioned, - *null_equality, - *null_aware, + *join_type, ) + .with_filter(filter.clone()) + // TODO: although projection is not used in the join here, because projection pushdown is after enforce_distribution. Maybe we need to handle it later. Same as filter. + .with_projection_ref(projection.clone()) + .with_partition_mode(PartitionMode::Partitioned) + .with_null_equality(*null_equality) + .with_null_aware(*null_aware) + .build() .map(|e| Arc::new(e) as _) }; return reorder_partitioned_join_keys( @@ -638,17 +639,20 @@ pub fn reorder_join_keys_to_inputs( right_keys, } = join_keys; let new_join_on = new_join_conditions(&left_keys, &right_keys); - return Ok(Arc::new(HashJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - new_join_on, - filter.clone(), - join_type, - projection.clone(), - PartitionMode::Partitioned, - *null_equality, - *null_aware, - )?)); + return Ok(Arc::new( + HashJoinExecBuilder::new( + Arc::clone(left), + Arc::clone(right), + new_join_on, + *join_type, + ) + .with_filter(filter.clone()) + .with_projection_ref(projection.clone()) + .with_partition_mode(PartitionMode::Partitioned) + .with_null_equality(*null_equality) + .with_null_aware(*null_aware) + .build()?, + )); } } } else if let Some(SortMergeJoinExec { diff --git a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs index 698fdea8e766e..2dc61ba2453fb 100644 --- a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs @@ -723,7 +723,7 @@ fn handle_hash_join( .collect(); let column_indices = build_join_column_index(plan); - let projected_indices: Vec<_> = if let Some(projection) = &plan.projection { + let projected_indices: Vec<_> = if let Some(projection) = plan.projection.as_ref() { projection.iter().map(|&i| &column_indices[i]).collect() } else { column_indices.iter().collect() diff --git a/datafusion/physical-optimizer/src/ensure_coop.rs b/datafusion/physical-optimizer/src/ensure_coop.rs index 5d00d00bce21d..ef8946f9a49d1 100644 --- a/datafusion/physical-optimizer/src/ensure_coop.rs +++ b/datafusion/physical-optimizer/src/ensure_coop.rs @@ -281,7 +281,7 @@ mod tests { input: Arc, scheduling_type: SchedulingType, evaluation_type: EvaluationType, - properties: PlanProperties, + properties: Arc, } impl DummyExec { @@ -305,7 +305,7 @@ mod tests { input, scheduling_type, evaluation_type, - properties, + properties: Arc::new(properties), } } } @@ -327,7 +327,7 @@ mod tests { fn as_any(&self) -> &dyn Any { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } fn children(&self) -> Vec<&Arc> { diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 7412d0ba97812..02ef378d704a0 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -34,7 +34,7 @@ use datafusion_physical_expr::expressions::Column; use datafusion_physical_plan::execution_plan::EmissionType; use datafusion_physical_plan::joins::utils::ColumnIndex; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, + CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; @@ -191,30 +191,18 @@ pub(crate) fn try_collect_left( { Ok(Some(hash_join.swap_inputs(PartitionMode::CollectLeft)?)) } else { - Ok(Some(Arc::new(HashJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - hash_join.on().to_vec(), - hash_join.filter().cloned(), - hash_join.join_type(), - hash_join.projection.clone(), - PartitionMode::CollectLeft, - hash_join.null_equality(), - hash_join.null_aware, - )?))) + Ok(Some(Arc::new( + HashJoinExecBuilder::from(hash_join) + .with_partition_mode(PartitionMode::CollectLeft) + .build()?, + ))) } } - (true, false) => Ok(Some(Arc::new(HashJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - hash_join.on().to_vec(), - hash_join.filter().cloned(), - hash_join.join_type(), - hash_join.projection.clone(), - PartitionMode::CollectLeft, - hash_join.null_equality(), - hash_join.null_aware, - )?))), + (true, false) => Ok(Some(Arc::new( + HashJoinExecBuilder::from(hash_join) + .with_partition_mode(PartitionMode::CollectLeft) + .build()?, + ))), (false, true) => { // Don't swap null-aware anti joins as they have specific side requirements if hash_join.join_type().supports_swap() && !hash_join.null_aware { @@ -254,17 +242,11 @@ pub(crate) fn partitioned_hash_join( PartitionMode::Partitioned }; - Ok(Arc::new(HashJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - hash_join.on().to_vec(), - hash_join.filter().cloned(), - hash_join.join_type(), - hash_join.projection.clone(), - partition_mode, - hash_join.null_equality(), - hash_join.null_aware, - )?)) + Ok(Arc::new( + HashJoinExecBuilder::from(hash_join) + .with_partition_mode(partition_mode) + .build()?, + )) } } diff --git a/datafusion/physical-optimizer/src/lib.rs b/datafusion/physical-optimizer/src/lib.rs index 195b77c9b4b70..3a0d79ae2d234 100644 --- a/datafusion/physical-optimizer/src/lib.rs +++ b/datafusion/physical-optimizer/src/lib.rs @@ -37,7 +37,6 @@ pub mod limit_pushdown_past_window; pub mod limited_distinct_aggregation; pub mod optimizer; pub mod output_requirements; -pub mod physical_expr_resolver; pub mod projection_pushdown; pub use datafusion_pruning as pruning; pub mod pushdown_sort; diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index f8a10e3a2727e..ff71c9ec64385 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -30,7 +30,6 @@ use crate::join_selection::JoinSelection; use crate::limit_pushdown::LimitPushdown; use crate::limited_distinct_aggregation::LimitedDistinctAggregation; use crate::output_requirements::OutputRequirements; -use crate::physical_expr_resolver::PhysicalExprResolver; use crate::projection_pushdown::ProjectionPushdown; use crate::sanity_checker::SanityCheckPlan; use crate::topk_aggregation::TopKAggregation; @@ -87,8 +86,6 @@ impl PhysicalOptimizer { // If there is a output requirement of the query, make sure that // this information is not lost across different rules during optimization. Arc::new(OutputRequirements::new_add_mode()), - // This rule removes all existing `TransformPlanExec` nodes from the plan tree. - Arc::new(PhysicalExprResolver::new()), Arc::new(AggregateStatistics::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the @@ -148,10 +145,6 @@ impl PhysicalOptimizer { // PushdownSort: Detect sorts that can be pushed down to data sources. Arc::new(PushdownSort::new()), Arc::new(EnsureCooperative::new()), - // This rule prepares the physical plan for placeholder resolution by wrapping it in a - // `TransformPlanExec` with a `ResolvePlaceholdersRule` if it contains any unresolved - // placeholders. - Arc::new(PhysicalExprResolver::new_post_optimization()), // This FilterPushdown handles dynamic filters that may have references to the source ExecutionPlan. // Therefore it should be run at the end of the optimization process since any changes to the plan may break the dynamic filter's references. // See `FilterPushdownPhase` for more details. diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 0dc6a25fbc0b7..9c4169ec654f8 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -98,7 +98,7 @@ pub struct OutputRequirementExec { input: Arc, order_requirement: Option, dist_requirement: Distribution, - cache: PlanProperties, + cache: Arc, fetch: Option, } @@ -114,7 +114,7 @@ impl OutputRequirementExec { input, order_requirement: requirements, dist_requirement, - cache, + cache: Arc::new(cache), fetch, } } @@ -200,7 +200,7 @@ impl ExecutionPlan for OutputRequirementExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-optimizer/src/physical_expr_resolver.rs b/datafusion/physical-optimizer/src/physical_expr_resolver.rs deleted file mode 100644 index a18a6d4d6bbc9..0000000000000 --- a/datafusion/physical-optimizer/src/physical_expr_resolver.rs +++ /dev/null @@ -1,138 +0,0 @@ -// 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. - -//! [`PhysicalExprResolver`] ensures that the physical plan is prepared for placeholder resolution -//! by wrapping it in a [`TransformPlanExec`] with a [`ResolvePlaceholdersRule`] if the plan -//! contains any unresolved placeholders. The actual resolution happens during execution. - -use std::sync::Arc; - -use datafusion_common::{ - Result, - config::ConfigOptions, - tree_node::{Transformed, TreeNode}, -}; -use datafusion_physical_plan::{ - ExecutionPlan, - plan_transformer::{ResolvePlaceholdersRule, TransformPlanExec}, -}; - -use crate::PhysicalOptimizerRule; - -/// The phase in which the [`PhysicalExprResolver`] rule is applied. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum PhysicalExprResolverPhase { - /// Optimization that happens before most other optimizations. - /// This optimization removes all [`TransformPlanExec`] execution plans from the plan - /// tree. - Pre, - /// Optimization that happens after most other optimizations. - /// This optimization checks for the presence of placeholders in the optimized plan, and if - /// they are present, wraps the plan in a [`TransformPlanExec`] with a [`ResolvePlaceholdersRule`]. - Post, -} - -/// Physical optimizer rule that prepares the plan for placeholder resolution during execution. -#[derive(Debug)] -pub struct PhysicalExprResolver { - phase: PhysicalExprResolverPhase, -} - -impl PhysicalExprResolver { - /// Creates a new [`PhysicalExprResolver`] optimizer rule that runs in the pre-optimization - /// phase. In this phase, the rule removes any existing [`TransformPlanExec`] from the - /// plan tree. - pub fn new() -> Self { - Self { - phase: PhysicalExprResolverPhase::Pre, - } - } - - /// Creates a new [`PhysicalExprResolver`] optimizer rule that runs in the post-optimization - /// phase. In this phase, the rule wraps the physical plan in a [`TransformPlanExec`] with a - /// [`ResolvePlaceholdersRule`] if the plan contains any unresolved placeholders. - pub fn new_post_optimization() -> Self { - Self { - phase: PhysicalExprResolverPhase::Post, - } - } -} - -impl Default for PhysicalExprResolver { - fn default() -> Self { - Self::new() - } -} - -impl PhysicalOptimizerRule for PhysicalExprResolver { - fn name(&self) -> &str { - match self.phase { - PhysicalExprResolverPhase::Pre => "PhysicalExprResolver", - PhysicalExprResolverPhase::Post => "PhysicalExprResolver(Post)", - } - } - - fn optimize( - &self, - plan: Arc, - _config: &ConfigOptions, - ) -> Result> { - match self.phase { - PhysicalExprResolverPhase::Pre => plan - .transform_up(|plan| { - if let Some(plan) = plan.as_any().downcast_ref::() - { - Ok(Transformed::yes(Arc::clone(plan.input()))) - } else { - Ok(Transformed::no(plan)) - } - }) - .map(|t| t.data), - PhysicalExprResolverPhase::Post => { - if let Some(transformer) = - plan.as_any().downcast_ref::() - { - let resolves_placeholders = - transformer.has_rule::(); - - if resolves_placeholders { - Ok(plan) - } else { - transformer - .add_rule(Box::new(ResolvePlaceholdersRule::new())) - .map(|r| Arc::new(r) as Arc<_>) - } - } else { - let transformer = TransformPlanExec::try_new( - plan, - vec![Box::new(ResolvePlaceholdersRule::new())], - )?; - - if transformer.plans_to_transform() > 0 { - Ok(Arc::new(transformer)) - } else { - Ok(Arc::clone(transformer.input())) - } - } - } - } - } - - fn schema_check(&self) -> bool { - true - } -} diff --git a/datafusion/physical-optimizer/src/projection_pushdown.rs b/datafusion/physical-optimizer/src/projection_pushdown.rs index 99922ba075cc0..44d0926a8b250 100644 --- a/datafusion/physical-optimizer/src/projection_pushdown.rs +++ b/datafusion/physical-optimizer/src/projection_pushdown.rs @@ -135,7 +135,7 @@ fn try_push_down_join_filter( ); let new_lhs_length = lhs_rewrite.data.0.schema().fields.len(); - let projections = match projections { + let projections = match projections.as_ref() { None => match join.join_type() { JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { // Build projections that ignore the newly projected columns. diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 27bd8610b46e3..13f91fd7d4ea2 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -102,7 +102,3 @@ name = "sort_preserving_merge" harness = false name = "aggregate_vectorized" required-features = ["test_utils"] - -[[bench]] -harness = false -name = "plan_transformer" diff --git a/datafusion/physical-plan/benches/plan_transformer.rs b/datafusion/physical-plan/benches/plan_transformer.rs deleted file mode 100644 index 42816bd92d5cc..0000000000000 --- a/datafusion/physical-plan/benches/plan_transformer.rs +++ /dev/null @@ -1,313 +0,0 @@ -// 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. - -use std::any::Any; -use std::sync::Arc; - -use arrow_schema::Schema; -use criterion::measurement::WallTime; -use datafusion_common::Result; -use datafusion_common::tree_node::{ - Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, -}; -use datafusion_execution::TaskContext; -use datafusion_physical_plan::ExecutionPlan; -use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion_physical_plan::empty::EmptyExec; -use datafusion_physical_plan::plan_transformer::{ - ExecutionTransformationRule, TransformPlanExec, -}; - -use criterion::{BatchSize, BenchmarkGroup, Criterion, criterion_group, criterion_main}; - -#[derive(Debug, Clone)] -struct ResetAllRule {} - -impl ExecutionTransformationRule for ResetAllRule { - fn name(&self) -> &str { - "ResetAllRule" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } - - fn matches(&mut self, _node: &Arc) -> Result { - Ok(true) - } - - fn rewrite( - &self, - node: Arc, - _ctx: &TaskContext, - ) -> Result>> { - node.reset_state().map(Transformed::yes) - } -} - -#[derive(Debug, Clone)] -struct ResetByNameRule { - node_name: String, -} - -impl ExecutionTransformationRule for ResetByNameRule { - fn name(&self) -> &str { - "ResetByNameRule" - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } - - fn matches(&mut self, node: &Arc) -> Result { - Ok(node.name() == self.node_name) - } - - fn rewrite( - &self, - node: Arc, - _ctx: &TaskContext, - ) -> Result>> { - node.reset_state().map(Transformed::yes) - } -} - -struct ResetAllRewriter; - -impl TreeNodeRewriter for ResetAllRewriter { - type Node = Arc; - - fn f_up(&mut self, node: Self::Node) -> Result> { - // This part is needed for handling nested rewriters. - if node.as_any().is::() { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Jump)); - } - node.reset_state().map(Transformed::yes) - } -} - -struct ResetByNameRewriter { - node_name: String, -} - -impl TreeNodeRewriter for ResetByNameRewriter { - type Node = Arc; - - fn f_up(&mut self, node: Self::Node) -> Result> { - // This part is needed for handling nested rewriters. - if node.as_any().is::() { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Jump)); - } - - if node.name() == self.node_name { - node.reset_state().map(Transformed::yes) - } else { - Ok(Transformed::no(node)) - } - } -} - -struct MockRewriterExec { - input: Arc, -} - -impl MockRewriterExec { - fn execute( - &self, - rewriter: &mut impl TreeNodeRewriter>, - ) -> Result> { - let input = Arc::clone(&self.input); - input.rewrite(rewriter).map(|t| t.data) - } -} - -fn create_deep_tree(depth: usize) -> Arc { - let schema = Arc::new(Schema::empty()); - let mut node: Arc = Arc::new(EmptyExec::new(schema)); - for _ in 0..depth { - node = Arc::new(CoalescePartitionsExec::new(node)); - } - node -} - -fn nodes_amount(plan: &Arc) -> usize { - let mut amount = 0; - plan.apply(|_| { - amount += 1; - Ok(TreeNodeRecursion::Continue) - }) - .unwrap(); - amount -} - -fn benchmark_with_transformer_exec( - group: &mut BenchmarkGroup<'_, WallTime>, - batch_label: &str, - plan: &Arc, - rules: &[Box], - batch_size: BatchSize, -) { - let ctx = Arc::new(TaskContext::default()); - let nodes_amount = nodes_amount(plan); - - group.bench_function( - format!( - "transform_plan_exec_two_phases_{batch_label}_{nodes_amount}_nodes_{}_rule(s)", - rules.len() - ), - |b| { - b.iter_batched( - || { - ( - Arc::clone(plan), - rules.iter().map(|r| r.clone_box()).collect(), - ) - }, - |(plan, rules)| { - let transformer = TransformPlanExec::try_new(plan, rules).unwrap(); - transformer.transform(&ctx).unwrap(); - }, - batch_size, - ) - }, - ); - - group.bench_function( - format!( - "transform_plan_exec_second_phase_{batch_label}_{nodes_amount}_nodes_{}_rule(s)", - rules.len() - ), - |b| { - let plan = Arc::clone(plan); - let rules = rules.iter().map(|r| r.clone_box()).collect(); - let transformer = Arc::new(TransformPlanExec::try_new(plan, rules).unwrap()); - - b.iter_batched( - || Arc::clone(&transformer), - |transformer| { - transformer.transform(&ctx).unwrap(); - }, - batch_size, - ) - }, - ); -} - -fn benchmark_with_tree_node_rewriter( - group: &mut BenchmarkGroup<'_, WallTime>, - batch_label: &str, - plan: &Arc, - mut rewriter: impl TreeNodeRewriter>, - rewrites_amount: usize, - batch_size: BatchSize, -) { - let nodes_amount = nodes_amount(plan); - let mock_rewriter_exec = Arc::new(MockRewriterExec { - input: Arc::clone(plan), - }); - - group.bench_function( - format!( - "tree_node_rewriter_{batch_label}_{nodes_amount}_nodes_{rewrites_amount}_iteration(s)" - ), - |b| { - b.iter_batched( - || Arc::clone(&mock_rewriter_exec), - |plan| { - for _ in 0..rewrites_amount { - plan.execute(&mut rewriter).unwrap(); - } - }, - batch_size, - ) - }, - ); -} - -fn criterion_benchmark(c: &mut Criterion) { - let depths = [5, 30]; - let rules_count = [1, 2]; - let batch_size = BatchSize::SmallInput; - let mut group = c.benchmark_group("plan_transformation"); - - for depth in depths { - let plan = create_deep_tree(depth); - - for count in rules_count { - let reset_all_rules = (0..count) - .map(|_| Box::new(ResetAllRule {}) as Box<_>) - .collect::>(); - - let reset_one_rules = (0..count) - .map(|_| { - Box::new(ResetByNameRule { - node_name: "EmptyExec".to_string(), - }) as Box<_> - }) - .collect::>(); - - benchmark_with_transformer_exec( - &mut group, - "reset_all", - &plan, - &reset_all_rules, - batch_size, - ); - - benchmark_with_tree_node_rewriter( - &mut group, - "reset_all", - &plan, - ResetAllRewriter, - count, - batch_size, - ); - - benchmark_with_transformer_exec( - &mut group, - "reset_one", - &plan, - &reset_one_rules, - batch_size, - ); - - benchmark_with_tree_node_rewriter( - &mut group, - "reset_one", - &plan, - ResetByNameRewriter { - node_name: "EmptyExec".to_string(), - }, - count, - batch_size, - ); - } - } - - group.finish(); -} - -criterion_group!(benches, criterion_benchmark); -criterion_main!(benches); diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 464e5ce6c7c1a..1d85703e9166f 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -25,7 +25,9 @@ use crate::aggregates::{ no_grouping::AggregateStream, row_hash::GroupedHashAggregateStream, topk_stream::GroupedTopKAggregateStream, }; -use crate::execution_plan::{CardinalityEffect, EmissionType, ReplacePhysicalExpr}; +use crate::execution_plan::{ + CardinalityEffect, EmissionType, ReplacePhysicalExpr, has_same_children_properties, +}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDownPredicate, @@ -33,7 +35,7 @@ use crate::filter_pushdown::{ use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, InputOrderMode, - SendableRecordBatchStream, Statistics, + SendableRecordBatchStream, Statistics, check_if_same_properties, }; use datafusion_common::config::ConfigOptions; use datafusion_physical_expr::utils::collect_columns; @@ -626,11 +628,14 @@ pub struct AggregateExec { /// Aggregation mode (full, partial) mode: AggregateMode, /// Group by expressions - group_by: PhysicalGroupBy, + /// [`Arc`] used for a cheap clone, which improves physical plan optimization performance. + group_by: Arc, /// Aggregate expressions - aggr_expr: Vec>, + /// The same reason to [`Arc`] it as for [`Self::group_by`]. + aggr_expr: Arc<[Arc]>, /// FILTER (WHERE clause) expression for each aggregate expression - filter_expr: Vec>>, + /// The same reason to [`Arc`] it as for [`Self::group_by`]. + filter_expr: Arc<[Option>]>, /// Configuration for limit-based optimizations limit_options: Option, /// Input plan, could be a partial aggregate or the input to the aggregate @@ -648,7 +653,7 @@ pub struct AggregateExec { required_input_ordering: Option, /// Describes how the input is ordered relative to the group by columns input_order_mode: InputOrderMode, - cache: PlanProperties, + cache: Arc, /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]), /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node. /// @@ -664,18 +669,18 @@ impl AggregateExec { /// Rewrites aggregate exec with new aggregate expressions. pub fn with_new_aggr_exprs( &self, - aggr_expr: Vec>, + aggr_expr: impl Into]>>, ) -> Self { Self { - aggr_expr, + aggr_expr: aggr_expr.into(), // clone the rest of the fields required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), mode: self.mode, - group_by: self.group_by.clone(), - filter_expr: self.filter_expr.clone(), + group_by: Arc::clone(&self.group_by), + filter_expr: Arc::clone(&self.filter_expr), limit_options: self.limit_options, input: Arc::clone(&self.input), schema: Arc::clone(&self.schema), @@ -692,11 +697,11 @@ impl AggregateExec { required_input_ordering: self.required_input_ordering.clone(), metrics: ExecutionPlanMetricsSet::new(), input_order_mode: self.input_order_mode.clone(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), mode: self.mode, - group_by: self.group_by.clone(), - aggr_expr: self.aggr_expr.clone(), - filter_expr: self.filter_expr.clone(), + group_by: Arc::clone(&self.group_by), + aggr_expr: Arc::clone(&self.aggr_expr), + filter_expr: Arc::clone(&self.filter_expr), input: Arc::clone(&self.input), schema: Arc::clone(&self.schema), input_schema: Arc::clone(&self.input_schema), @@ -711,12 +716,13 @@ impl AggregateExec { /// Create a new hash aggregate execution plan pub fn try_new( mode: AggregateMode, - group_by: PhysicalGroupBy, + group_by: impl Into>, aggr_expr: Vec>, filter_expr: Vec>>, input: Arc, input_schema: SchemaRef, ) -> Result { + let group_by = group_by.into(); let schema = create_schema(&input.schema(), &group_by, &aggr_expr, mode)?; let schema = Arc::new(schema); @@ -741,13 +747,16 @@ impl AggregateExec { /// the schema in such cases. fn try_new_with_schema( mode: AggregateMode, - group_by: PhysicalGroupBy, + group_by: impl Into>, mut aggr_expr: Vec>, - filter_expr: Vec>>, + filter_expr: impl Into>]>>, input: Arc, input_schema: SchemaRef, schema: SchemaRef, ) -> Result { + let group_by = group_by.into(); + let filter_expr = filter_expr.into(); + // Make sure arguments are consistent in size assert_eq_or_internal_err!( aggr_expr.len(), @@ -814,13 +823,13 @@ impl AggregateExec { &group_expr_mapping, &mode, &input_order_mode, - aggr_expr.as_slice(), + aggr_expr.as_ref(), )?; let mut exec = AggregateExec { mode, group_by, - aggr_expr, + aggr_expr: aggr_expr.into(), filter_expr, input, schema, @@ -829,7 +838,7 @@ impl AggregateExec { required_input_ordering, limit_options: None, input_order_mode, - cache, + cache: Arc::new(cache), dynamic_filter: None, }; @@ -1187,6 +1196,17 @@ impl AggregateExec { _ => Precision::Absent, } } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for AggregateExec { @@ -1325,7 +1345,7 @@ impl ExecutionPlan for AggregateExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -1368,11 +1388,13 @@ impl ExecutionPlan for AggregateExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); + let mut me = AggregateExec::try_new_with_schema( self.mode, - self.group_by.clone(), - self.aggr_expr.clone(), - self.filter_expr.clone(), + Arc::clone(&self.group_by), + self.aggr_expr.to_vec(), + Arc::clone(&self.filter_expr), Arc::clone(&children[0]), Arc::clone(&self.input_schema), Arc::clone(&self.schema), @@ -1645,9 +1667,9 @@ impl ExecutionPlan for AggregateExec { } Ok(Some(Arc::new(Self { - group_by, - aggr_expr, - filter_expr, + group_by: Arc::new(group_by), + aggr_expr: aggr_expr.into(), + filter_expr: filter_expr.into(), dynamic_filter: None, metrics: ExecutionPlanMetricsSet::new(), ..self.clone() @@ -2507,14 +2529,17 @@ mod tests { struct TestYieldingExec { /// True if this exec should yield back to runtime the first time it is polled pub yield_first: bool, - cache: PlanProperties, + cache: Arc, } impl TestYieldingExec { fn new(yield_first: bool) -> Self { let schema = some_data().0; let cache = Self::compute_properties(schema); - Self { yield_first, cache } + Self { + yield_first, + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -2555,7 +2580,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/aggregates/no_grouping.rs b/datafusion/physical-plan/src/aggregates/no_grouping.rs index eb9b6766ab8ee..fe8942097ac93 100644 --- a/datafusion/physical-plan/src/aggregates/no_grouping.rs +++ b/datafusion/physical-plan/src/aggregates/no_grouping.rs @@ -62,7 +62,7 @@ struct AggregateStreamInner { mode: AggregateMode, input: SendableRecordBatchStream, aggregate_expressions: Vec>>, - filter_expressions: Vec>>, + filter_expressions: Arc<[Option>]>, // ==== Runtime States/Buffers ==== accumulators: Vec, @@ -277,7 +277,7 @@ impl AggregateStream { partition: usize, ) -> Result { let agg_schema = Arc::clone(&agg.schema); - let agg_filter_expr = agg.filter_expr.clone(); + let agg_filter_expr = Arc::clone(&agg.filter_expr); let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let input = agg.input.execute(partition, Arc::clone(context))?; @@ -285,7 +285,7 @@ impl AggregateStream { let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?; let filter_expressions = match agg.mode.input_mode() { AggregateInputMode::Raw => agg_filter_expr, - AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()], + AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(), }; let accumulators = create_accumulators(&agg.aggr_expr)?; diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/row_hash.rs index b2cf396b1500d..de857370ce285 100644 --- a/datafusion/physical-plan/src/aggregates/row_hash.rs +++ b/datafusion/physical-plan/src/aggregates/row_hash.rs @@ -377,10 +377,10 @@ pub(crate) struct GroupedHashAggregateStream { /// /// For example, for an aggregate like `SUM(x) FILTER (WHERE x >= 100)`, /// the filter expression is `x > 100`. - filter_expressions: Vec>>, + filter_expressions: Arc<[Option>]>, /// GROUP BY expressions - group_by: PhysicalGroupBy, + group_by: Arc, /// max rows in output RecordBatches batch_size: usize, @@ -465,8 +465,8 @@ impl GroupedHashAggregateStream { ) -> Result { debug!("Creating GroupedHashAggregateStream"); let agg_schema = Arc::clone(&agg.schema); - let agg_group_by = agg.group_by.clone(); - let agg_filter_expr = agg.filter_expr.clone(); + let agg_group_by = Arc::clone(&agg.group_by); + let agg_filter_expr = Arc::clone(&agg.filter_expr); let batch_size = context.session_config().batch_size(); let input = agg.input.execute(partition, Arc::clone(context))?; @@ -475,7 +475,7 @@ impl GroupedHashAggregateStream { let timer = baseline_metrics.elapsed_compute().timer(); - let aggregate_exprs = agg.aggr_expr.clone(); + let aggregate_exprs = Arc::clone(&agg.aggr_expr); // arguments for each aggregate, one vec of expressions per // aggregate @@ -493,7 +493,7 @@ impl GroupedHashAggregateStream { let filter_expressions = match agg.mode.input_mode() { AggregateInputMode::Raw => agg_filter_expr, - AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()], + AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(), }; // Instantiate the accumulators diff --git a/datafusion/physical-plan/src/aggregates/topk_stream.rs b/datafusion/physical-plan/src/aggregates/topk_stream.rs index 72c5d0c86745d..4aa566ccfcd0a 100644 --- a/datafusion/physical-plan/src/aggregates/topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/topk_stream.rs @@ -50,7 +50,7 @@ pub struct GroupedTopKAggregateStream { baseline_metrics: BaselineMetrics, group_by_metrics: GroupByMetrics, aggregate_arguments: Vec>>, - group_by: PhysicalGroupBy, + group_by: Arc, priority_map: PriorityMap, } @@ -62,7 +62,7 @@ impl GroupedTopKAggregateStream { limit: usize, ) -> Result { let agg_schema = Arc::clone(&aggr.schema); - let group_by = aggr.group_by.clone(); + let group_by = Arc::clone(&aggr.group_by); let input = aggr.input.execute(partition, Arc::clone(context))?; let baseline_metrics = BaselineMetrics::new(&aggr.metrics, partition); let group_by_metrics = GroupByMetrics::new(&aggr.metrics, partition); diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 1fb8f93a38782..eca31ea0e194f 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -51,7 +51,7 @@ pub struct AnalyzeExec { pub(crate) input: Arc, /// The output schema for RecordBatches of this exec node schema: SchemaRef, - cache: PlanProperties, + cache: Arc, } impl AnalyzeExec { @@ -70,7 +70,7 @@ impl AnalyzeExec { metric_types, input, schema, - cache, + cache: Arc::new(cache), } } @@ -131,7 +131,7 @@ impl ExecutionPlan for AnalyzeExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index a61fd95949d1a..93701a25b3bf4 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -16,10 +16,12 @@ // under the License. use crate::coalesce::LimitedBatchCoalescer; +use crate::execution_plan::has_same_children_properties; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::stream::RecordBatchStreamAdapter; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + check_if_same_properties, }; use arrow::array::RecordBatch; use arrow_schema::{Fields, Schema, SchemaRef}; @@ -45,12 +47,12 @@ use std::task::{Context, Poll, ready}; /// /// The schema of the output of the AsyncFuncExec is: /// Input columns followed by one column for each async expression -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct AsyncFuncExec { /// The async expressions to evaluate async_exprs: Vec>, input: Arc, - cache: PlanProperties, + cache: Arc, metrics: ExecutionPlanMetricsSet, } @@ -84,7 +86,7 @@ impl AsyncFuncExec { Ok(Self { input, async_exprs, - cache, + cache: Arc::new(cache), metrics: ExecutionPlanMetricsSet::new(), }) } @@ -113,6 +115,17 @@ impl AsyncFuncExec { pub fn input(&self) -> &Arc { &self.input } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for AsyncFuncExec { @@ -149,7 +162,7 @@ impl ExecutionPlan for AsyncFuncExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -159,16 +172,17 @@ impl ExecutionPlan for AsyncFuncExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { assert_eq_or_internal_err!( children.len(), 1, "AsyncFuncExec wrong number of children" ); + check_if_same_properties!(self, children); Ok(Arc::new(AsyncFuncExec::try_new( self.async_exprs.clone(), - Arc::clone(&children[0]), + children.swap_remove(0), )?)) } diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 3b80f9924e311..ae8d015988228 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -18,7 +18,9 @@ //! [`BufferExec`] decouples production and consumption on messages by buffering the input in the //! background up to a certain capacity. -use crate::execution_plan::{CardinalityEffect, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, SchedulingType, has_same_children_properties, +}; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -27,6 +29,7 @@ use crate::projection::ProjectionExec; use crate::stream::RecordBatchStreamAdapter; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SortOrderPushdownResult, + check_if_same_properties, }; use arrow::array::RecordBatch; use datafusion_common::config::ConfigOptions; @@ -92,7 +95,7 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; #[derive(Debug, Clone)] pub struct BufferExec { input: Arc, - properties: PlanProperties, + properties: Arc, capacity: usize, metrics: ExecutionPlanMetricsSet, } @@ -100,14 +103,12 @@ pub struct BufferExec { impl BufferExec { /// Builds a new [BufferExec] with the provided capacity in bytes. pub fn new(input: Arc, capacity: usize) -> Self { - let properties = input - .properties() - .clone() + let properties = PlanProperties::clone(input.properties()) .with_scheduling_type(SchedulingType::Cooperative); Self { input, - properties, + properties: Arc::new(properties), capacity, metrics: ExecutionPlanMetricsSet::new(), } @@ -122,6 +123,17 @@ impl BufferExec { pub fn capacity(&self) -> usize { self.capacity } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for BufferExec { @@ -146,7 +158,7 @@ impl ExecutionPlan for BufferExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } @@ -166,6 +178,7 @@ impl ExecutionPlan for BufferExec { self: Arc, mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); if children.len() != 1 { return plan_err!("BufferExec can only have one child"); } diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index dfcd3cb0bcae7..27a521ec6bc5d 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -27,6 +27,7 @@ use super::{DisplayAs, ExecutionPlanProperties, PlanProperties, Statistics}; use crate::projection::ProjectionExec; use crate::{ DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream, + check_if_same_properties, }; use arrow::datatypes::SchemaRef; @@ -36,7 +37,7 @@ use datafusion_execution::TaskContext; use datafusion_physical_expr::PhysicalExpr; use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; -use crate::execution_plan::CardinalityEffect; +use crate::execution_plan::{CardinalityEffect, has_same_children_properties}; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -71,7 +72,7 @@ pub struct CoalesceBatchesExec { fetch: Option, /// Execution metrics metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, + cache: Arc, } #[expect(deprecated)] @@ -84,7 +85,7 @@ impl CoalesceBatchesExec { target_batch_size, fetch: None, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), } } @@ -115,6 +116,17 @@ impl CoalesceBatchesExec { input.boundedness(), ) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } #[expect(deprecated)] @@ -159,7 +171,7 @@ impl ExecutionPlan for CoalesceBatchesExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -177,10 +189,11 @@ impl ExecutionPlan for CoalesceBatchesExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new( - CoalesceBatchesExec::new(Arc::clone(&children[0]), self.target_batch_size) + CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) .with_fetch(self.fetch), )) } @@ -222,7 +235,7 @@ impl ExecutionPlan for CoalesceBatchesExec { target_batch_size: self.target_batch_size, fetch: limit, metrics: self.metrics.clone(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), })) } diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 22dcc85d6ea3a..a7a2ce8a18b5c 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -27,11 +27,13 @@ use super::{ DisplayAs, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, Statistics, }; -use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, has_same_children_properties, +}; use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase}; use crate::projection::{ProjectionExec, make_with_child}; use crate::sort_pushdown::SortOrderPushdownResult; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; +use crate::{DisplayFormatType, ExecutionPlan, Partitioning, check_if_same_properties}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_common::config::ConfigOptions; @@ -47,7 +49,7 @@ pub struct CoalescePartitionsExec { input: Arc, /// Execution metrics metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, + cache: Arc, /// Optional number of rows to fetch. Stops producing rows after this fetch pub(crate) fetch: Option, } @@ -59,7 +61,7 @@ impl CoalescePartitionsExec { CoalescePartitionsExec { input, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), fetch: None, } } @@ -100,6 +102,17 @@ impl CoalescePartitionsExec { .with_evaluation_type(drive) .with_scheduling_type(scheduling) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for CoalescePartitionsExec { @@ -135,7 +148,7 @@ impl ExecutionPlan for CoalescePartitionsExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -149,9 +162,10 @@ impl ExecutionPlan for CoalescePartitionsExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { - let mut plan = CoalescePartitionsExec::new(Arc::clone(&children[0])); + check_if_same_properties!(self, children); + let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); plan.fetch = self.fetch; Ok(Arc::new(plan)) } @@ -274,7 +288,7 @@ impl ExecutionPlan for CoalescePartitionsExec { input: Arc::clone(&self.input), fetch: limit, metrics: self.metrics.clone(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), })) } diff --git a/datafusion/physical-plan/src/common.rs b/datafusion/physical-plan/src/common.rs index 32dc60b56ad48..590f6f09e8b9e 100644 --- a/datafusion/physical-plan/src/common.rs +++ b/datafusion/physical-plan/src/common.rs @@ -181,7 +181,7 @@ pub fn compute_record_batch_statistics( /// Checks if the given projection is valid for the given schema. pub fn can_project( schema: &arrow::datatypes::SchemaRef, - projection: Option<&Vec>, + projection: Option<&[usize]>, ) -> Result<()> { match projection { Some(columns) => { diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index ce54a451ac4d1..acc79ee009690 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -87,14 +87,14 @@ use crate::filter_pushdown::{ use crate::projection::ProjectionExec; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, SortOrderPushdownResult, + SendableRecordBatchStream, SortOrderPushdownResult, check_if_same_properties, }; use arrow::record_batch::RecordBatch; use arrow_schema::Schema; use datafusion_common::{Result, Statistics, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; -use crate::execution_plan::SchedulingType; +use crate::execution_plan::{SchedulingType, has_same_children_properties}; use crate::stream::RecordBatchStreamAdapter; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use futures::{Stream, StreamExt}; @@ -217,16 +217,15 @@ where #[derive(Debug, Clone)] pub struct CooperativeExec { input: Arc, - properties: PlanProperties, + properties: Arc, } impl CooperativeExec { /// Creates a new `CooperativeExec` operator that wraps the given input execution plan. pub fn new(input: Arc) -> Self { - let properties = input - .properties() - .clone() - .with_scheduling_type(SchedulingType::Cooperative); + let properties = PlanProperties::clone(input.properties()) + .with_scheduling_type(SchedulingType::Cooperative) + .into(); Self { input, properties } } @@ -235,6 +234,16 @@ impl CooperativeExec { pub fn input(&self) -> &Arc { &self.input } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + ..Self::clone(self) + } + } } impl DisplayAs for CooperativeExec { @@ -260,7 +269,7 @@ impl ExecutionPlan for CooperativeExec { self.input.schema() } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.properties } @@ -281,6 +290,7 @@ impl ExecutionPlan for CooperativeExec { 1, "CooperativeExec requires exactly one child" ); + check_if_same_properties!(self, children); Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) } diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 52c37a106b39e..d525f44541d5f 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -1153,7 +1153,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index fcfbcfa3e8277..d3e62f91c0570 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -43,7 +43,7 @@ pub struct EmptyExec { schema: SchemaRef, /// Number of partitions partitions: usize, - cache: PlanProperties, + cache: Arc, } impl EmptyExec { @@ -53,7 +53,7 @@ impl EmptyExec { EmptyExec { schema, partitions: 1, - cache, + cache: Arc::new(cache), } } @@ -62,7 +62,7 @@ impl EmptyExec { self.partitions = partitions; // Changing partitions may invalidate output partitioning, so update it: let output_partitioning = Self::output_partitioning_helper(self.partitions); - self.cache = self.cache.with_partitioning(output_partitioning); + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; self } @@ -114,7 +114,7 @@ impl ExecutionPlan for EmptyExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index abfa1df828224..6ad1afa051e8d 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -128,7 +128,7 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync { /// /// This information is available via methods on [`ExecutionPlanProperties`] /// trait, which is implemented for all `ExecutionPlan`s. - fn properties(&self) -> &PlanProperties; + fn properties(&self) -> &Arc; /// Returns an error if this individual node does not conform to its invariants. /// These invariants are typically only checked in debug mode. @@ -1105,12 +1105,17 @@ impl PlanProperties { self } - /// Overwrite equivalence properties with its new value. - pub fn with_eq_properties(mut self, eq_properties: EquivalenceProperties) -> Self { + /// Set equivalence properties having mut reference. + pub fn set_eq_properties(&mut self, eq_properties: EquivalenceProperties) { // Changing equivalence properties also changes output ordering, so // make sure to overwrite it: self.output_ordering = eq_properties.output_ordering(); self.eq_properties = eq_properties; + } + + /// Overwrite equivalence properties with its new value. + pub fn with_eq_properties(mut self, eq_properties: EquivalenceProperties) -> Self { + self.set_eq_properties(eq_properties); self } @@ -1142,9 +1147,14 @@ impl PlanProperties { self } + /// Set constraints having mut reference. + pub fn set_constraints(&mut self, constraints: Constraints) { + self.eq_properties.set_constraints(constraints); + } + /// Overwrite constraints with its new value. pub fn with_constraints(mut self, constraints: Constraints) -> Self { - self.eq_properties = self.eq_properties.with_constraints(constraints); + self.set_constraints(constraints); self } @@ -1556,6 +1566,41 @@ pub fn reset_plan_states(plan: Arc) -> Result, + children: &[Arc], +) -> Result { + let old_children = plan.children(); + assert_eq_or_internal_err!( + children.len(), + old_children.len(), + "Wrong number of children" + ); + for (lhs, rhs) in plan.children().iter().zip(children.iter()) { + if !Arc::ptr_eq(lhs.properties(), rhs.properties()) { + return Ok(false); + } + } + Ok(true) +} + +/// Helper macro to avoid properties re-computation if passed children properties +/// the same as plan already has. Could be used to implement fast-path for method +/// [`ExecutionPlan::with_new_children`]. +#[macro_export] +macro_rules! check_if_same_properties { + ($plan: expr, $children: expr) => { + if has_same_children_properties(&$plan, &$children)? { + let plan = $plan.with_new_children_and_same_properties($children); + return Ok(Arc::new(plan)); + } + }; +} + /// Utility function yielding a string representation of the given [`ExecutionPlan`]. pub fn get_plan_string(plan: &Arc) -> Vec { let formatted = displayable(plan.as_ref()).indent(true).to_string(); @@ -1618,7 +1663,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } @@ -1685,7 +1730,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unimplemented!() } diff --git a/datafusion/physical-plan/src/explain.rs b/datafusion/physical-plan/src/explain.rs index aa3c0afefe8b5..bf21b0484689a 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -44,7 +44,7 @@ pub struct ExplainExec { stringified_plans: Vec, /// control which plans to print verbose: bool, - cache: PlanProperties, + cache: Arc, } impl ExplainExec { @@ -59,7 +59,7 @@ impl ExplainExec { schema, stringified_plans, verbose, - cache, + cache: Arc::new(cache), } } @@ -112,7 +112,7 @@ impl ExecutionPlan for ExplainExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 79e3c2960c2ad..5393753aa509e 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -20,15 +20,19 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll, ready}; +use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; use itertools::Itertools; use super::{ ColumnStatistics, DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; +use crate::check_if_same_properties; use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use crate::common::can_project; -use crate::execution_plan::{CardinalityEffect, ReplacePhysicalExpr}; +use crate::execution_plan::{ + CardinalityEffect, ReplacePhysicalExpr, has_same_children_properties, +}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDown, PushedDownPredicate, @@ -84,9 +88,9 @@ pub struct FilterExec { /// Selectivity for statistics. 0 = no rows, 100 = all rows default_selectivity: u8, /// Properties equivalence properties, partitioning, etc. - cache: PlanProperties, + cache: Arc, /// The projection indices of the columns in the output schema of join - projection: Option>, + projection: Option, /// Target batch size for output batches batch_size: usize, /// Number of rows to fetch @@ -97,7 +101,7 @@ pub struct FilterExec { pub struct FilterExecBuilder { predicate: Arc, input: Arc, - projection: Option>, + projection: Option, default_selectivity: u8, batch_size: usize, fetch: Option, @@ -137,18 +141,19 @@ impl FilterExecBuilder { /// /// If no projection is currently set, the new projection is used directly. /// If `None` is passed, the projection is cleared. - pub fn apply_projection(mut self, projection: Option>) -> Result { + pub fn apply_projection(self, projection: Option>) -> Result { + let projection = projection.map(Into::into); + self.apply_projection_by_ref(projection.as_ref()) + } + + /// The same as [`Self::apply_projection`] but takes projection shared reference. + pub fn apply_projection_by_ref( + mut self, + projection: Option<&ProjectionRef>, + ) -> Result { // Check if the projection is valid against current output schema - can_project(&self.input.schema(), projection.as_ref())?; - self.projection = match projection { - Some(new_proj) => match &self.projection { - Some(existing_proj) => { - Some(new_proj.iter().map(|i| existing_proj[*i]).collect()) - } - None => Some(new_proj), - }, - None => None, - }; + can_project(&self.input.schema(), projection.map(AsRef::as_ref))?; + self.projection = combine_projections(projection, self.projection.as_ref())?; Ok(self) } @@ -190,16 +195,14 @@ impl FilterExecBuilder { } // Validate projection if provided - if let Some(ref proj) = self.projection { - can_project(&self.input.schema(), Some(proj))?; - } + can_project(&self.input.schema(), self.projection.as_deref())?; // Compute properties once with all parameters let cache = FilterExec::compute_properties( &self.input, &self.predicate, self.default_selectivity, - self.projection.as_ref(), + self.projection.as_deref(), )?; Ok(FilterExec { @@ -207,7 +210,7 @@ impl FilterExecBuilder { input: self.input, metrics: ExecutionPlanMetricsSet::new(), default_selectivity: self.default_selectivity, - cache, + cache: Arc::new(cache), projection: self.projection, batch_size: self.batch_size, fetch: self.fetch, @@ -280,7 +283,7 @@ impl FilterExec { input: Arc::clone(&self.input), metrics: self.metrics.clone(), default_selectivity: self.default_selectivity, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), projection: self.projection.clone(), batch_size, fetch: self.fetch, @@ -303,8 +306,8 @@ impl FilterExec { } /// Projection - pub fn projection(&self) -> Option<&Vec> { - self.projection.as_ref() + pub fn projection(&self) -> &Option { + &self.projection } /// Calculates `Statistics` for `FilterExec`, by applying selectivity (either default, or estimated) to input statistics. @@ -381,7 +384,7 @@ impl FilterExec { input: &Arc, predicate: &Arc, default_selectivity: u8, - projection: Option<&Vec>, + projection: Option<&[usize]>, ) -> Result { // Combine the equal predicates with the input equivalence properties // to construct the equivalence properties: @@ -420,7 +423,7 @@ impl FilterExec { if let Some(projection) = projection { let schema = eq_properties.schema(); let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; - let out_schema = project_schema(schema, Some(projection))?; + let out_schema = project_schema(schema, Some(&projection))?; output_partitioning = output_partitioning.project(&projection_mapping, &eq_properties); eq_properties = eq_properties.project(&projection_mapping, out_schema); @@ -433,6 +436,17 @@ impl FilterExec { input.boundedness(), )) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for FilterExec { @@ -487,7 +501,7 @@ impl ExecutionPlan for FilterExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -504,6 +518,7 @@ impl ExecutionPlan for FilterExec { self: Arc, mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); let new_input = children.swap_remove(0); FilterExecBuilder::from(&*self) .with_input(new_input) @@ -665,7 +680,7 @@ impl ExecutionPlan for FilterExec { let new_predicate = conjunction(unhandled_filters); let updated_node = if new_predicate.eq(&lit(true)) { // FilterExec is no longer needed, but we may need to leave a projection in place - match self.projection() { + match self.projection().as_ref() { Some(projection_indices) => { let filter_child_schema = filter_input.schema(); let proj_exprs = projection_indices @@ -697,12 +712,12 @@ impl ExecutionPlan for FilterExec { input: Arc::clone(&filter_input), metrics: self.metrics.clone(), default_selectivity: self.default_selectivity, - cache: Self::compute_properties( + cache: Arc::new(Self::compute_properties( &filter_input, &new_predicate, self.default_selectivity, - self.projection.as_ref(), - )?, + self.projection.as_deref(), + )?), projection: self.projection.clone(), batch_size: self.batch_size, fetch: self.fetch, @@ -722,7 +737,7 @@ impl ExecutionPlan for FilterExec { input: Arc::clone(&self.input), metrics: self.metrics.clone(), default_selectivity: self.default_selectivity, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), projection: self.projection.clone(), batch_size: self.batch_size, fetch, @@ -839,7 +854,7 @@ struct FilterExecStream { /// Runtime metrics recording metrics: FilterExecMetrics, /// The projection indices of the columns in the input schema - projection: Option>, + projection: Option, /// Batch coalescer to combine small batches batch_coalescer: LimitedBatchCoalescer, } @@ -930,8 +945,8 @@ impl Stream for FilterExecStream { .evaluate(&batch) .and_then(|v| v.into_array(batch.num_rows())) .and_then(|array| { - Ok(match self.projection { - Some(ref projection) => { + Ok(match self.projection.as_ref() { + Some(projection) => { let projected_batch = batch.project(projection)?; (array, projected_batch) }, @@ -1752,7 +1767,7 @@ mod tests { .build()?; // Verify projection is set correctly - assert_eq!(filter.projection(), Some(&vec![0, 2])); + assert_eq!(filter.projection(), &Some([0, 2].into())); // Verify schema contains only projected columns let output_schema = filter.schema(); @@ -1782,7 +1797,7 @@ mod tests { let filter = FilterExecBuilder::new(predicate, input).build()?; // Verify no projection is set - assert_eq!(filter.projection(), None); + assert!(filter.projection().is_none()); // Verify schema contains all columns let output_schema = filter.schema(); @@ -1996,7 +2011,7 @@ mod tests { .build()?; // Verify composed projection is [0, 3] - assert_eq!(filter.projection(), Some(&vec![0, 3])); + assert_eq!(filter.projection(), &Some([0, 3].into())); // Verify schema contains only columns a and d let output_schema = filter.schema(); @@ -2030,7 +2045,7 @@ mod tests { .build()?; // Projection should be cleared - assert_eq!(filter.projection(), None); + assert_eq!(filter.projection(), &None); // Schema should have all columns let output_schema = filter.schema(); diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 7ada14be66543..c25408e49db64 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -25,7 +25,9 @@ use super::utils::{ OnceAsync, OnceFut, StatefulStreamResult, adjust_right_output_partitioning, reorder_output_after_swap, }; -use crate::execution_plan::{EmissionType, boundedness_from_children}; +use crate::execution_plan::{ + EmissionType, boundedness_from_children, has_same_children_properties, +}; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::projection::{ ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children, @@ -34,7 +36,7 @@ use crate::projection::{ use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, handle_state, + SendableRecordBatchStream, Statistics, check_if_same_properties, handle_state, }; use arrow::array::{RecordBatch, RecordBatchOptions}; @@ -94,7 +96,7 @@ pub struct CrossJoinExec { /// Execution plan metrics metrics: ExecutionPlanMetricsSet, /// Properties such as schema, equivalence properties, ordering, partitioning, etc. - cache: PlanProperties, + cache: Arc, } impl CrossJoinExec { @@ -125,7 +127,7 @@ impl CrossJoinExec { schema, left_fut: Default::default(), metrics: ExecutionPlanMetricsSet::default(), - cache, + cache: Arc::new(cache), } } @@ -192,6 +194,23 @@ impl CrossJoinExec { &self.right.schema(), ) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + + Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + left_fut: Default::default(), + cache: Arc::clone(&self.cache), + schema: Arc::clone(&self.schema), + } + } } /// Asynchronously collect the result of the left child @@ -256,7 +275,7 @@ impl ExecutionPlan for CrossJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -272,6 +291,7 @@ impl ExecutionPlan for CrossJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(CrossJoinExec::new( Arc::clone(&children[0]), Arc::clone(&children[1]), @@ -285,7 +305,7 @@ impl ExecutionPlan for CrossJoinExec { schema: Arc::clone(&self.schema), left_fut: Default::default(), // reset the build side! metrics: ExecutionPlanMetricsSet::default(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), }; Ok(Arc::new(new_exec)) } diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 98f7d04383402..c3811083b94e4 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -24,6 +24,7 @@ use std::{any::Any, vec}; use crate::ExecutionPlanProperties; use crate::execution_plan::{ EmissionType, ReplacePhysicalExpr, boundedness_from_children, + has_same_children_properties, }; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, @@ -83,6 +84,7 @@ use datafusion_physical_expr::equivalence::{ ProjectionMapping, join_equivalence_properties, }; use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; +use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; use ahash::RandomState; @@ -248,6 +250,172 @@ impl JoinLeftData { } } +/// Helps to build [`HashJoinExec`]. +pub struct HashJoinExecBuilder { + left: Arc, + right: Arc, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + join_type: JoinType, + filter: Option, + projection: Option, + partition_mode: PartitionMode, + null_equality: NullEquality, + null_aware: bool, +} + +impl HashJoinExecBuilder { + /// Make a new [`HashJoinExecBuilder`]. + pub fn new( + left: Arc, + right: Arc, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + join_type: JoinType, + ) -> Self { + Self { + left, + right, + on, + filter: None, + projection: None, + partition_mode: PartitionMode::Auto, + join_type, + null_equality: NullEquality::NullEqualsNothing, + null_aware: false, + } + } + + /// Set projection from the vector. + pub fn with_projection(self, projection: Option>) -> Self { + self.with_projection_ref(projection.map(Into::into)) + } + + /// Set projection from the shared reference. + pub fn with_projection_ref(mut self, projection: Option) -> Self { + self.projection = projection; + self + } + + /// Set optional filter. + pub fn with_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + + /// Set partition mode. + pub fn with_partition_mode(mut self, mode: PartitionMode) -> Self { + self.partition_mode = mode; + self + } + + /// Set null equality property. + pub fn with_null_equality(mut self, null_equality: NullEquality) -> Self { + self.null_equality = null_equality; + self + } + + /// Set null aware property. + pub fn with_null_aware(mut self, null_aware: bool) -> Self { + self.null_aware = null_aware; + self + } + + /// Build resulting execution plan. + pub fn build(self) -> Result { + let Self { + left, + right, + on, + join_type, + filter, + projection, + partition_mode, + null_equality, + null_aware, + } = self; + + let left_schema = left.schema(); + let right_schema = right.schema(); + if on.is_empty() { + return plan_err!("On constraints in HashJoinExec should be non-empty"); + } + + check_join_is_valid(&left_schema, &right_schema, &on)?; + + // Validate null_aware flag + if null_aware { + if !matches!(join_type, JoinType::LeftAnti) { + return plan_err!( + "null_aware can only be true for LeftAnti joins, got {join_type}" + ); + } + if on.len() != 1 { + return plan_err!( + "null_aware anti join only supports single column join key, got {} columns", + on.len() + ); + } + } + + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &join_type); + + let random_state = HASH_JOIN_SEED; + + let join_schema = Arc::new(join_schema); + + // check if the projection is valid + can_project(&join_schema, projection.as_deref())?; + + let cache = HashJoinExec::compute_properties( + &left, + &right, + &join_schema, + join_type, + &on, + partition_mode, + projection.as_deref(), + )?; + + // Initialize both dynamic filter and bounds accumulator to None + // They will be set later if dynamic filtering is enabled + + Ok(HashJoinExec { + left, + right, + on, + filter, + join_type, + join_schema, + left_fut: Default::default(), + random_state, + mode: partition_mode, + metrics: ExecutionPlanMetricsSet::new(), + projection, + column_indices, + null_equality, + null_aware, + cache: Arc::new(cache), + dynamic_filter: None, + }) + } +} + +impl From<&HashJoinExec> for HashJoinExecBuilder { + fn from(exec: &HashJoinExec) -> Self { + Self { + left: Arc::clone(exec.left()), + right: Arc::clone(exec.right()), + on: exec.on.clone(), + join_type: exec.join_type, + filter: exec.filter.clone(), + projection: exec.projection.clone(), + partition_mode: exec.mode, + null_equality: exec.null_equality, + null_aware: exec.null_aware, + } + } +} + #[expect(rustdoc::private_intra_doc_links)] /// Join execution plan: Evaluates equijoin predicates in parallel on multiple /// partitions using a hash table and an optional filter list to apply post @@ -468,7 +636,7 @@ pub struct HashJoinExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// The projection indices of the columns in the output schema of join - pub projection: Option>, + pub projection: Option, /// Information of index and left / right placement of columns column_indices: Vec, /// The equality null-handling behavior of the join algorithm. @@ -476,7 +644,7 @@ pub struct HashJoinExec { /// Flag to indicate if this is a null-aware anti join pub null_aware: bool, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// Dynamic filter for pushing down to the probe side /// Set when dynamic filter pushdown is detected in handle_child_pushdown_result. /// HashJoinExec also needs to keep a shared bounds accumulator for coordinating updates. @@ -537,70 +705,13 @@ impl HashJoinExec { null_equality: NullEquality, null_aware: bool, ) -> Result { - let left_schema = left.schema(); - let right_schema = right.schema(); - if on.is_empty() { - return plan_err!("On constraints in HashJoinExec should be non-empty"); - } - - check_join_is_valid(&left_schema, &right_schema, &on)?; - - // Validate null_aware flag - if null_aware { - if !matches!(join_type, JoinType::LeftAnti) { - return plan_err!( - "null_aware can only be true for LeftAnti joins, got {join_type}" - ); - } - if on.len() != 1 { - return plan_err!( - "null_aware anti join only supports single column join key, got {} columns", - on.len() - ); - } - } - - let (join_schema, column_indices) = - build_join_schema(&left_schema, &right_schema, join_type); - - let random_state = HASH_JOIN_SEED; - - let join_schema = Arc::new(join_schema); - - // check if the projection is valid - can_project(&join_schema, projection.as_ref())?; - - let cache = Self::compute_properties( - &left, - &right, - &join_schema, - *join_type, - &on, - partition_mode, - projection.as_ref(), - )?; - - // Initialize both dynamic filter and bounds accumulator to None - // They will be set later if dynamic filtering is enabled - - Ok(HashJoinExec { - left, - right, - on, - filter, - join_type: *join_type, - join_schema, - left_fut: Default::default(), - random_state, - mode: partition_mode, - metrics: ExecutionPlanMetricsSet::new(), - projection, - column_indices, - null_equality, - null_aware, - cache, - dynamic_filter: None, - }) + HashJoinExecBuilder::new(left, right, on, *join_type) + .with_filter(filter) + .with_projection(projection) + .with_partition_mode(partition_mode) + .with_null_equality(null_equality) + .with_null_aware(null_aware) + .build() } fn create_dynamic_filter(on: &JoinOn) -> Arc { @@ -689,26 +800,14 @@ impl HashJoinExec { /// Return new instance of [HashJoinExec] with the given projection. pub fn with_projection(&self, projection: Option>) -> Result { + let projection = projection.map(Into::into); // check if the projection is valid - can_project(&self.schema(), projection.as_ref())?; - let projection = match projection { - Some(projection) => match &self.projection { - Some(p) => Some(projection.iter().map(|i| p[*i]).collect()), - None => Some(projection), - }, - None => None, - }; - Self::try_new( - Arc::clone(&self.left), - Arc::clone(&self.right), - self.on.clone(), - self.filter.clone(), - &self.join_type, - projection, - self.mode, - self.null_equality, - self.null_aware, - ) + can_project(&self.schema(), projection.as_deref())?; + let projection = + combine_projections(projection.as_ref(), self.projection.as_ref())?; + HashJoinExecBuilder::from(self) + .with_projection_ref(projection) + .build() } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -719,7 +818,7 @@ impl HashJoinExec { join_type: JoinType, on: JoinOnRef, mode: PartitionMode, - projection: Option<&Vec>, + projection: Option<&[usize]>, ) -> Result { // Calculate equivalence properties: let mut eq_properties = join_equivalence_properties( @@ -771,7 +870,7 @@ impl HashJoinExec { if let Some(projection) = projection { // construct a map from the input expressions to the output expression of the Projection let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; - let out_schema = project_schema(schema, Some(projection))?; + let out_schema = project_schema(schema, Some(&projection))?; output_partitioning = output_partitioning.project(&projection_mapping, &eq_properties); eq_properties = eq_properties.project(&projection_mapping, out_schema); @@ -826,7 +925,7 @@ impl HashJoinExec { swap_join_projection( left.schema().fields().len(), right.schema().fields().len(), - self.projection.as_ref(), + self.projection.as_deref(), self.join_type(), ), partition_mode, @@ -939,7 +1038,7 @@ impl ExecutionPlan for HashJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -1000,6 +1099,20 @@ impl ExecutionPlan for HashJoinExec { self: Arc, children: Vec>, ) -> Result> { + let cache = if has_same_children_properties(&self, &children)? { + Arc::clone(&self.cache) + } else { + Arc::new(Self::compute_properties( + &children[0], + &children[1], + &self.join_schema, + self.join_type, + &self.on, + self.mode, + self.projection.as_deref(), + )?) + }; + Ok(Arc::new(HashJoinExec { left: Arc::clone(&children[0]), right: Arc::clone(&children[1]), @@ -1015,15 +1128,7 @@ impl ExecutionPlan for HashJoinExec { column_indices: self.column_indices.clone(), null_equality: self.null_equality, null_aware: self.null_aware, - cache: Self::compute_properties( - &children[0], - &children[1], - &self.join_schema, - self.join_type, - &self.on, - self.mode, - self.projection.as_ref(), - )?, + cache, // Keep the dynamic filter, bounds accumulator will be reset dynamic_filter: self.dynamic_filter.clone(), })) @@ -1046,7 +1151,7 @@ impl ExecutionPlan for HashJoinExec { column_indices: self.column_indices.clone(), null_equality: self.null_equality, null_aware: self.null_aware, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), // Reset dynamic filter and bounds accumulator to initial state dynamic_filter: None, })) @@ -1183,7 +1288,7 @@ impl ExecutionPlan for HashJoinExec { let right_stream = self.right.execute(partition, context)?; // update column indices to reflect the projection - let column_indices_after_projection = match &self.projection { + let column_indices_after_projection = match self.projection.as_ref() { Some(projection) => projection .iter() .map(|i| self.column_indices[*i].clone()) @@ -1375,7 +1480,7 @@ impl ExecutionPlan for HashJoinExec { column_indices: self.column_indices.clone(), null_equality: self.null_equality, null_aware: self.null_aware, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), dynamic_filter: Some(HashJoinExecDynamicFilter { filter: dynamic_filter, build_accumulator: OnceLock::new(), @@ -1432,7 +1537,7 @@ impl ExecutionPlan for HashJoinExec { column_indices: self.column_indices.clone(), null_equality: self.null_equality, null_aware: self.null_aware, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), // Reset dynamic filter and bounds accumulator to initial state dynamic_filter: None, }))) diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs index 8592e1d968535..b915802ea4015 100644 --- a/datafusion/physical-plan/src/joins/hash_join/mod.rs +++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs @@ -17,7 +17,7 @@ //! [`HashJoinExec`] Partitioned Hash Join Operator -pub use exec::HashJoinExec; +pub use exec::{HashJoinExec, HashJoinExecBuilder}; pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState}; mod exec; diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index 848d0472fe885..2cdfa1e6ac020 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -20,8 +20,10 @@ use arrow::array::BooleanBufferBuilder; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; -pub use hash_join::{HashExpr, HashJoinExec, HashTableLookupExpr, SeededRandomState}; -pub use nested_loop_join::NestedLoopJoinExec; +pub use hash_join::{ + HashExpr, HashJoinExec, HashJoinExecBuilder, HashTableLookupExpr, SeededRandomState, +}; +pub use nested_loop_join::{NestedLoopJoinExec, NestedLoopJoinExecBuilder}; use parking_lot::Mutex; // Note: SortMergeJoin is not used in plans yet pub use piecewise_merge_join::PiecewiseMergeJoinExec; diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index d4c0d1ac4a870..4ecb58d9ce077 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -31,6 +31,7 @@ use super::utils::{ use crate::common::can_project; use crate::execution_plan::{ EmissionType, ReplacePhysicalExpr, boundedness_from_children, + has_same_children_properties, }; use crate::joins::SharedBitmapBuilder; use crate::joins::utils::{ @@ -48,6 +49,7 @@ use crate::projection::{ use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, + check_if_same_properties, }; use arrow::array::{ @@ -73,7 +75,10 @@ use datafusion_physical_expr::equivalence::{ ProjectionMapping, join_equivalence_properties, }; -use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::{ + PhysicalExpr, + projection::{ProjectionRef, combine_projections}, +}; use futures::{Stream, StreamExt, TryStreamExt}; use log::debug; use parking_lot::Mutex; @@ -195,50 +200,120 @@ pub struct NestedLoopJoinExec { /// Information of index and left / right placement of columns column_indices: Vec, /// Projection to apply to the output of the join - projection: Option>, + projection: Option, /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } -impl NestedLoopJoinExec { - /// Try to create a new [`NestedLoopJoinExec`] - pub fn try_new( +/// Helps to build [`NestedLoopJoinExec`]. +pub struct NestedLoopJoinExecBuilder { + left: Arc, + right: Arc, + join_type: JoinType, + filter: Option, + projection: Option, +} + +impl NestedLoopJoinExecBuilder { + /// Make a new [`NestedLoopJoinExecBuilder`]. + pub fn new( left: Arc, right: Arc, - filter: Option, - join_type: &JoinType, - projection: Option>, - ) -> Result { + join_type: JoinType, + ) -> Self { + Self { + left, + right, + join_type, + filter: None, + projection: None, + } + } + + /// Set projection from the vector. + pub fn with_projection(self, projection: Option>) -> Self { + self.with_projection_ref(projection.map(Into::into)) + } + + /// Set projection from the shared reference. + pub fn with_projection_ref(mut self, projection: Option) -> Self { + self.projection = projection; + self + } + + /// Set optional filter. + pub fn with_filter(mut self, filter: Option) -> Self { + self.filter = filter; + self + } + + /// Build resulting execution plan. + pub fn build(self) -> Result { + let Self { + left, + right, + join_type, + filter, + projection, + } = self; + let left_schema = left.schema(); let right_schema = right.schema(); check_join_is_valid(&left_schema, &right_schema, &[])?; let (join_schema, column_indices) = - build_join_schema(&left_schema, &right_schema, join_type); + build_join_schema(&left_schema, &right_schema, &join_type); let join_schema = Arc::new(join_schema); - let cache = Self::compute_properties( + let cache = NestedLoopJoinExec::compute_properties( &left, &right, &join_schema, - *join_type, - projection.as_ref(), + join_type, + projection.as_deref(), )?; - Ok(NestedLoopJoinExec { left, right, filter, - join_type: *join_type, + join_type, join_schema, build_side_data: Default::default(), column_indices, projection, metrics: Default::default(), - cache, + cache: Arc::new(cache), }) } +} + +impl From<&NestedLoopJoinExec> for NestedLoopJoinExecBuilder { + fn from(exec: &NestedLoopJoinExec) -> Self { + Self { + left: Arc::clone(exec.left()), + right: Arc::clone(exec.right()), + join_type: exec.join_type, + filter: exec.filter.clone(), + projection: exec.projection.clone(), + } + } +} + +impl NestedLoopJoinExec { + /// Try to create a new [`NestedLoopJoinExec`] + pub fn try_new( + left: Arc, + right: Arc, + filter: Option, + join_type: &JoinType, + projection: Option>, + ) -> Result { + NestedLoopJoinExecBuilder::new(left, right, *join_type) + .with_projection(projection) + .with_filter(filter) + .build() + } /// left side pub fn left(&self) -> &Arc { @@ -260,8 +335,8 @@ impl NestedLoopJoinExec { &self.join_type } - pub fn projection(&self) -> Option<&Vec> { - self.projection.as_ref() + pub fn projection(&self) -> &Option { + &self.projection } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -270,7 +345,7 @@ impl NestedLoopJoinExec { right: &Arc, schema: &SchemaRef, join_type: JoinType, - projection: Option<&Vec>, + projection: Option<&[usize]>, ) -> Result { // Calculate equivalence properties: let mut eq_properties = join_equivalence_properties( @@ -313,7 +388,7 @@ impl NestedLoopJoinExec { if let Some(projection) = projection { // construct a map from the input expressions to the output expression of the Projection let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; - let out_schema = project_schema(schema, Some(projection))?; + let out_schema = project_schema(schema, Some(&projection))?; output_partitioning = output_partitioning.project(&projection_mapping, &eq_properties); eq_properties = eq_properties.project(&projection_mapping, out_schema); @@ -337,22 +412,14 @@ impl NestedLoopJoinExec { } pub fn with_projection(&self, projection: Option>) -> Result { + let projection = projection.map(Into::into); // check if the projection is valid - can_project(&self.schema(), projection.as_ref())?; - let projection = match projection { - Some(projection) => match &self.projection { - Some(p) => Some(projection.iter().map(|i| p[*i]).collect()), - None => Some(projection), - }, - None => None, - }; - Self::try_new( - Arc::clone(&self.left), - Arc::clone(&self.right), - self.filter.clone(), - &self.join_type, - projection, - ) + can_project(&self.schema(), projection.as_deref())?; + let projection = + combine_projections(projection.as_ref(), self.projection.as_ref())?; + NestedLoopJoinExecBuilder::from(self) + .with_projection_ref(projection) + .build() } /// Returns a new `ExecutionPlan` that runs NestedLoopsJoins with the left @@ -374,7 +441,7 @@ impl NestedLoopJoinExec { swap_join_projection( left.schema().fields().len(), right.schema().fields().len(), - self.projection.as_ref(), + self.projection.as_deref(), self.join_type(), ), )?; @@ -402,6 +469,27 @@ impl NestedLoopJoinExec { Ok(plan) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + + Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + build_side_data: Default::default(), + cache: Arc::clone(&self.cache), + filter: self.filter.clone(), + join_type: self.join_type, + join_schema: Arc::clone(&self.join_schema), + column_indices: self.column_indices.clone(), + projection: self.projection.clone(), + } + } } impl DisplayAs for NestedLoopJoinExec { @@ -456,7 +544,7 @@ impl ExecutionPlan for NestedLoopJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -479,13 +567,17 @@ impl ExecutionPlan for NestedLoopJoinExec { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(NestedLoopJoinExec::try_new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - self.filter.clone(), - &self.join_type, - self.projection.clone(), - )?)) + check_if_same_properties!(self, children); + Ok(Arc::new( + NestedLoopJoinExecBuilder::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.join_type, + ) + .with_filter(self.filter.clone()) + .with_projection_ref(self.projection.clone()) + .build()?, + )) } fn execute( @@ -524,7 +616,7 @@ impl ExecutionPlan for NestedLoopJoinExec { let probe_side_data = self.right.execute(partition, context)?; // update column indices to reflect the projection - let column_indices_after_projection = match &self.projection { + let column_indices_after_projection = match self.projection.as_ref() { Some(projection) => projection .iter() .map(|i| self.column_indices[*i].clone()) @@ -661,7 +753,7 @@ impl ExecutionPlan for NestedLoopJoinExec { column_indices: self.column_indices.clone(), projection: self.projection.clone(), metrics: Default::default(), - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), }))) } } diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index d7ece845e943c..014d4cb13e2b7 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -41,7 +41,9 @@ use std::fmt::Formatter; use std::sync::Arc; use std::sync::atomic::AtomicUsize; -use crate::execution_plan::{EmissionType, boundedness_from_children}; +use crate::execution_plan::{ + EmissionType, boundedness_from_children, has_same_children_properties, +}; use crate::joins::piecewise_merge_join::classic_join::{ ClassicPWMJStream, PiecewiseMergeJoinStreamState, @@ -51,7 +53,9 @@ use crate::joins::piecewise_merge_join::utils::{ }; use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; -use crate::{DisplayAs, DisplayFormatType, ExecutionPlanProperties}; +use crate::{ + DisplayAs, DisplayFormatType, ExecutionPlanProperties, check_if_same_properties, +}; use crate::{ ExecutionPlan, PlanProperties, joins::{ @@ -86,7 +90,7 @@ use crate::{ /// Both sides are sorted so that we can iterate from index 0 to the end on each side. This ordering ensures /// that when we find the first matching pair of rows, we can emit the current stream row joined with all remaining /// probe rows from the match position onward, without rescanning earlier probe rows. -/// +/// /// For `<` and `<=` operators, both inputs are sorted in **descending** order, while for `>` and `>=` operators /// they are sorted in **ascending** order. This choice ensures that the pointer on the buffered side can advance /// monotonically as we stream new batches from the stream side. @@ -129,34 +133,34 @@ use crate::{ /// /// Processing Row 1: /// -/// Sorted Buffered Side Sorted Streamed Side -/// ┌──────────────────┐ ┌──────────────────┐ -/// 1 │ 100 │ 1 │ 100 │ -/// ├──────────────────┤ ├──────────────────┤ -/// 2 │ 200 │ ─┐ 2 │ 200 │ -/// ├──────────────────┤ │ For row 1 on streamed side with ├──────────────────┤ -/// 3 │ 200 │ │ value 100, we emit rows 2 - 5. 3 │ 500 │ +/// Sorted Buffered Side Sorted Streamed Side +/// ┌──────────────────┐ ┌──────────────────┐ +/// 1 │ 100 │ 1 │ 100 │ +/// ├──────────────────┤ ├──────────────────┤ +/// 2 │ 200 │ ─┐ 2 │ 200 │ +/// ├──────────────────┤ │ For row 1 on streamed side with ├──────────────────┤ +/// 3 │ 200 │ │ value 100, we emit rows 2 - 5. 3 │ 500 │ /// ├──────────────────┤ │ as matches when the operator is └──────────────────┘ /// 4 │ 300 │ │ `Operator::Lt` (<) Emitting all /// ├──────────────────┤ │ rows after the first match (row /// 5 │ 400 │ ─┘ 2 buffered side; 100 < 200) -/// └──────────────────┘ +/// └──────────────────┘ /// /// Processing Row 2: /// By sorting the streamed side we know /// -/// Sorted Buffered Side Sorted Streamed Side -/// ┌──────────────────┐ ┌──────────────────┐ -/// 1 │ 100 │ 1 │ 100 │ -/// ├──────────────────┤ ├──────────────────┤ -/// 2 │ 200 │ <- Start here when probing for the 2 │ 200 │ -/// ├──────────────────┤ streamed side row 2. ├──────────────────┤ -/// 3 │ 200 │ 3 │ 500 │ +/// Sorted Buffered Side Sorted Streamed Side +/// ┌──────────────────┐ ┌──────────────────┐ +/// 1 │ 100 │ 1 │ 100 │ +/// ├──────────────────┤ ├──────────────────┤ +/// 2 │ 200 │ <- Start here when probing for the 2 │ 200 │ +/// ├──────────────────┤ streamed side row 2. ├──────────────────┤ +/// 3 │ 200 │ 3 │ 500 │ /// ├──────────────────┤ └──────────────────┘ -/// 4 │ 300 │ -/// ├──────────────────┤ +/// 4 │ 300 │ +/// ├──────────────────┤ /// 5 │ 400 │ -/// └──────────────────┘ +/// └──────────────────┘ /// ``` /// /// ## Existence Joins (Semi, Anti, Mark) @@ -202,10 +206,10 @@ use crate::{ /// 1 │ 100 │ 1 │ 500 │ /// ├──────────────────┤ ├──────────────────┤ /// 2 │ 200 │ 2 │ 200 │ -/// ├──────────────────┤ ├──────────────────┤ +/// ├──────────────────┤ ├──────────────────┤ /// 3 │ 200 │ 3 │ 300 │ /// ├──────────────────┤ └──────────────────┘ -/// 4 │ 300 │ ─┐ +/// 4 │ 300 │ ─┐ /// ├──────────────────┤ | We emit matches for row 4 - 5 /// 5 │ 400 │ ─┘ on the buffered side. /// └──────────────────┘ @@ -236,11 +240,11 @@ use crate::{ /// /// # Mark Join: /// Sorts the probe side, then computes the min/max range of the probe keys and scans the buffered side only -/// within that range. +/// within that range. /// Complexity: `O(|S| + scan(R[range]))`. /// /// ## Nested Loop Join -/// Compares every row from `S` with every row from `R`. +/// Compares every row from `S` with every row from `R`. /// Complexity: `O(|S| * |R|)`. /// /// ## Nested Loop Join @@ -273,13 +277,12 @@ pub struct PiecewiseMergeJoinExec { left_child_plan_required_order: LexOrdering, /// The right sort order, descending for `<`, `<=` operations + ascending for `>`, `>=` operations /// Unsorted for mark joins - #[expect(dead_code)] right_batch_required_orders: LexOrdering, /// This determines the sort order of all join columns used in sorting the stream and buffered execution plans. sort_options: SortOptions, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// Number of partitions to process num_partitions: usize, } @@ -373,7 +376,7 @@ impl PiecewiseMergeJoinExec { left_child_plan_required_order, right_batch_required_orders, sort_options, - cache, + cache: Arc::new(cache), num_partitions, }) } @@ -466,6 +469,31 @@ impl PiecewiseMergeJoinExec { pub fn swap_inputs(&self) -> Result> { todo!() } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let buffered = children.swap_remove(0); + let streamed = children.swap_remove(0); + Self { + buffered, + streamed, + on: self.on.clone(), + operator: self.operator, + join_type: self.join_type, + schema: Arc::clone(&self.schema), + left_child_plan_required_order: self.left_child_plan_required_order.clone(), + right_batch_required_orders: self.right_batch_required_orders.clone(), + sort_options: self.sort_options, + cache: Arc::clone(&self.cache), + num_partitions: self.num_partitions, + + // Re-set state. + metrics: ExecutionPlanMetricsSet::new(), + buffered_fut: Default::default(), + } + } } impl ExecutionPlan for PiecewiseMergeJoinExec { @@ -477,7 +505,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -511,6 +539,7 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); match &children[..] { [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( Arc::clone(left), @@ -527,6 +556,13 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } } + fn reset_state(self: Arc) -> Result> { + Ok(Arc::new(self.with_new_children_and_same_properties(vec![ + Arc::clone(&self.buffered), + Arc::clone(&self.streamed), + ]))) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index b3e87daff360a..83457806ca801 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -25,6 +25,7 @@ use std::sync::Arc; use crate::execution_plan::{ EmissionType, ReplacePhysicalExpr, boundedness_from_children, + has_same_children_properties, }; use crate::expressions::PhysicalSortExpr; use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; @@ -41,7 +42,7 @@ use crate::projection::{ }; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, SendableRecordBatchStream, Statistics, + PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties, }; use arrow::compute::SortOptions; @@ -130,7 +131,7 @@ pub struct SortMergeJoinExec { /// Defines the null equality for the join. pub null_equality: NullEquality, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl SortMergeJoinExec { @@ -201,7 +202,7 @@ impl SortMergeJoinExec { right_sort_exprs, sort_options, null_equality, - cache, + cache: Arc::new(cache), }) } @@ -343,6 +344,20 @@ impl SortMergeJoinExec { reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema()) } } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for SortMergeJoinExec { @@ -408,7 +423,7 @@ impl ExecutionPlan for SortMergeJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -443,6 +458,7 @@ impl ExecutionPlan for SortMergeJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); match &children[..] { [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( Arc::clone(left), diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 3adef3a2dd93f..41fed11f4dc0a 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -32,9 +32,11 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::vec; +use crate::check_if_same_properties; use crate::common::SharedMemoryReservation; use crate::execution_plan::{ ReplacePhysicalExpr, boundedness_from_children, emission_type_from_children, + has_same_children_properties, }; use crate::joins::stream_join_utils::{ PruningJoinHashMap, SortedFilterExpr, StreamJoinMetrics, @@ -200,7 +202,7 @@ pub struct SymmetricHashJoinExec { /// Partition Mode mode: StreamJoinPartitionMode, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl SymmetricHashJoinExec { @@ -256,7 +258,7 @@ impl SymmetricHashJoinExec { left_sort_exprs, right_sort_exprs, mode, - cache, + cache: Arc::new(cache), }) } @@ -363,6 +365,20 @@ impl SymmetricHashJoinExec { } Ok(false) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for SymmetricHashJoinExec { @@ -414,7 +430,7 @@ impl ExecutionPlan for SymmetricHashJoinExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -456,6 +472,7 @@ impl ExecutionPlan for SymmetricHashJoinExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(SymmetricHashJoinExec::try_new( Arc::clone(&children[0]), Arc::clone(&children[1]), diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index a9243fe04e28d..e709703e07d45 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1674,7 +1674,7 @@ fn swap_reverting_projection( pub fn swap_join_projection( left_schema_len: usize, right_schema_len: usize, - projection: Option<&Vec>, + projection: Option<&[usize]>, join_type: &JoinType, ) -> Option> { match join_type { @@ -1685,7 +1685,7 @@ pub fn swap_join_projection( | JoinType::RightAnti | JoinType::RightSemi | JoinType::LeftMark - | JoinType::RightMark => projection.cloned(), + | JoinType::RightMark => projection.map(|p| p.to_vec()), _ => projection.map(|p| { p.iter() .map(|i| { diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index da28951e6d897..e65ec96ad84eb 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -81,10 +81,10 @@ pub mod limit; pub mod memory; pub mod metrics; pub mod placeholder_row; -pub mod plan_transformer; pub mod projection; pub mod recursive_query; pub mod repartition; +pub mod reuse; pub mod sort_pushdown; pub mod sorts; pub mod spill; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index fea7acb221304..e1bd606fa25ba 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -27,8 +27,13 @@ use super::{ DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; -use crate::execution_plan::{Boundedness, CardinalityEffect}; -use crate::{DisplayFormatType, Distribution, ExecutionPlan, Partitioning}; +use crate::execution_plan::{ + Boundedness, CardinalityEffect, has_same_children_properties, +}; +use crate::{ + DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + check_if_same_properties, +}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -51,10 +56,10 @@ pub struct GlobalLimitExec { fetch: Option, /// Execution metrics metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, /// Does the limit have to preserve the order of its input, and if so what is it? /// Some optimizations may reorder the input if no particular sort is required required_ordering: Option, + cache: Arc, } impl GlobalLimitExec { @@ -66,8 +71,8 @@ impl GlobalLimitExec { skip, fetch, metrics: ExecutionPlanMetricsSet::new(), - cache, required_ordering: None, + cache: Arc::new(cache), } } @@ -106,6 +111,17 @@ impl GlobalLimitExec { pub fn set_required_ordering(&mut self, required_ordering: Option) { self.required_ordering = required_ordering; } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for GlobalLimitExec { @@ -144,7 +160,7 @@ impl ExecutionPlan for GlobalLimitExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -166,10 +182,11 @@ impl ExecutionPlan for GlobalLimitExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(GlobalLimitExec::new( - Arc::clone(&children[0]), + children.swap_remove(0), self.skip, self.fetch, ))) @@ -229,7 +246,7 @@ impl ExecutionPlan for GlobalLimitExec { } /// LocalLimitExec applies a limit to a single partition -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct LocalLimitExec { /// Input execution plan input: Arc, @@ -237,10 +254,10 @@ pub struct LocalLimitExec { fetch: usize, /// Execution metrics metrics: ExecutionPlanMetricsSet, - cache: PlanProperties, /// If the child plan is a sort node, after the sort node is removed during /// physical optimization, we should add the required ordering to the limit node required_ordering: Option, + cache: Arc, } impl LocalLimitExec { @@ -251,8 +268,8 @@ impl LocalLimitExec { input, fetch, metrics: ExecutionPlanMetricsSet::new(), - cache, required_ordering: None, + cache: Arc::new(cache), } } @@ -286,6 +303,17 @@ impl LocalLimitExec { pub fn set_required_ordering(&mut self, required_ordering: Option) { self.required_ordering = required_ordering; } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for LocalLimitExec { @@ -315,7 +343,7 @@ impl ExecutionPlan for LocalLimitExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -335,6 +363,7 @@ impl ExecutionPlan for LocalLimitExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); match children.len() { 1 => Ok(Arc::new(LocalLimitExec::new( Arc::clone(&children[0]), diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index 4a406ca648d57..5c684bec98f9b 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -161,7 +161,7 @@ pub struct LazyMemoryExec { /// Functions to generate batches for each partition batch_generators: Vec>>, /// Plan properties cache storing equivalence properties, partitioning, and execution mode - cache: PlanProperties, + cache: Arc, /// Execution metrics metrics: ExecutionPlanMetricsSet, } @@ -200,7 +200,8 @@ impl LazyMemoryExec { EmissionType::Incremental, boundedness, ) - .with_scheduling_type(SchedulingType::Cooperative); + .with_scheduling_type(SchedulingType::Cooperative) + .into(); Ok(Self { schema, @@ -215,9 +216,9 @@ impl LazyMemoryExec { match projection.as_ref() { Some(columns) => { let projected = Arc::new(self.schema.project(columns).unwrap()); - self.cache = self.cache.with_eq_properties(EquivalenceProperties::new( - Arc::clone(&projected), - )); + Arc::make_mut(&mut self.cache).set_eq_properties( + EquivalenceProperties::new(Arc::clone(&projected)), + ); self.schema = projected; self.projection = projection; self @@ -236,12 +237,12 @@ impl LazyMemoryExec { partition_count, generator_count ); - self.cache.partitioning = partitioning; + Arc::make_mut(&mut self.cache).partitioning = partitioning; Ok(()) } pub fn add_ordering(&mut self, ordering: impl IntoIterator) { - self.cache + Arc::make_mut(&mut self.cache) .eq_properties .add_orderings(std::iter::once(ordering)); } @@ -306,7 +307,7 @@ impl ExecutionPlan for LazyMemoryExec { Arc::clone(&self.schema) } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -365,7 +366,7 @@ impl ExecutionPlan for LazyMemoryExec { Ok(Arc::new(LazyMemoryExec { schema: Arc::clone(&self.schema), batch_generators: generators, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), metrics: ExecutionPlanMetricsSet::new(), projection: self.projection.clone(), })) diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 4d00b73cff39c..f95b10771c02f 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -43,7 +43,7 @@ pub struct PlaceholderRowExec { schema: SchemaRef, /// Number of partitions partitions: usize, - cache: PlanProperties, + cache: Arc, } impl PlaceholderRowExec { @@ -54,7 +54,7 @@ impl PlaceholderRowExec { PlaceholderRowExec { schema, partitions, - cache, + cache: Arc::new(cache), } } @@ -63,7 +63,7 @@ impl PlaceholderRowExec { self.partitions = partitions; // Update output partitioning when updating partitions: let output_partitioning = Self::output_partitioning_helper(self.partitions); - self.cache = self.cache.with_partitioning(output_partitioning); + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; self } @@ -132,7 +132,7 @@ impl ExecutionPlan for PlaceholderRowExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/plan_transformer/mod.rs b/datafusion/physical-plan/src/plan_transformer/mod.rs deleted file mode 100644 index 9834e3c5d0d84..0000000000000 --- a/datafusion/physical-plan/src/plan_transformer/mod.rs +++ /dev/null @@ -1,629 +0,0 @@ -// 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. - -//! This module provides a mechanism for two-pass [`ExecutionPlan`] tree transformations. -//! -//! In the first pass, the tree is visited (top-down) to identify nodes that require transformation -//! based on certain criteria. -//! -//! In the second pass, the stored plans are applied to the tree (top-down) using a -//! [`TreeNodeRewriter`]. -//! -//! This approach is beneficial because it allows making multiple transformations during a single -//! pass and caches node indices, ensuring that only nodes matching the criteria are transformed, -//! which can improve performance. - -mod resolve_placeholders; - -pub use resolve_placeholders::ResolvePlaceholdersRule; - -use std::any::Any; -use std::fmt::{self, Debug}; -use std::sync::Arc; - -use datafusion_common::{ - Result, - tree_node::{ - Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor, - }, -}; -use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use itertools::Itertools; - -use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; -use crate::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, Statistics}; - -/// A rule for transforming an [`ExecutionPlan`] node. -/// -/// This trait is used in the two-pass tree transformation process. Both passes use a top-down -/// traversal. In the first pass, the [`ExecutionTransformationRule::matches`] method is used to -/// identify nodes that require transformation. In the second pass, the -/// [`ExecutionTransformationRule::rewrite`] method is used to apply the transformation. -pub trait ExecutionTransformationRule: Send + Sync + Debug { - /// Returns the rule name. - fn name(&self) -> &str; - - /// Returns the rule as [`Any`] so that it can be downcast to a specific implementation. - fn as_any(&self) -> &dyn Any; - - /// Clones this rule. - fn clone_box(&self) -> Box; - - /// Checks if the given [`ExecutionPlan`] node matches the criteria for this rule. - fn matches(&mut self, _node: &Arc) -> Result { - Ok(false) - } - - /// Transforms the given [`ExecutionPlan`] node. - fn rewrite( - &self, - node: Arc, - _ctx: &TaskContext, - ) -> Result>> { - Ok(Transformed::no(node)) - } -} - -/// Stores the transformation operations to be applied to an [`ExecutionPlan`] node at a specific -/// index. -#[derive(Debug, Clone)] -struct TransformationPlan { - /// The index of the node in the tree traversal. - pub node_index: usize, - /// The list of transformation rule indices to apply. - pub rule_indices: Vec, -} - -/// Helper for building transformation plans for an [`ExecutionPlan`] tree during a single pass. -struct TransformationPlanner { - cursor: usize, - rules: Vec>, - plans: Vec, -} - -impl TransformationPlanner { - fn new(rules: Vec>) -> Self { - Self { - cursor: 0, - rules, - plans: Vec::new(), - } - } -} - -impl<'n> TreeNodeVisitor<'n> for TransformationPlanner { - type Node = Arc; - - fn f_down(&mut self, node: &'n Self::Node) -> Result { - if node.as_any().is::() { - return Ok(TreeNodeRecursion::Jump); - } - - let index = self.cursor; - self.cursor += 1; - - let mut rule_indices = Vec::new(); - for (rule_index, rule) in self.rules.iter_mut().enumerate() { - if rule.matches(node)? { - rule_indices.push(rule_index); - } - } - - if !rule_indices.is_empty() { - self.plans.push(TransformationPlan { - node_index: index, - rule_indices, - }); - } - - Ok(TreeNodeRecursion::Continue) - } -} - -/// Helper for applying transformation plans to an [`ExecutionPlan`] tree during a single pass. -/// -/// This applier uses a [`TaskContext`] to provide necessary information for the transformation -/// rules. -struct TransformationApplier<'rules, 'plans, 'ctx> { - cursor: usize, - rules: &'rules [Box], - plans: &'plans [TransformationPlan], - ctx: &'ctx TaskContext, -} - -impl<'rules, 'plans, 'ctx> TransformationApplier<'rules, 'plans, 'ctx> { - fn new( - rules: &'rules [Box], - plans: &'plans [TransformationPlan], - ctx: &'ctx TaskContext, - ) -> Self { - Self { - cursor: 0, - rules, - plans, - ctx, - } - } -} - -impl<'rules, 'plans, 'ctx> TreeNodeRewriter - for TransformationApplier<'rules, 'plans, 'ctx> -{ - type Node = Arc; - - fn f_down(&mut self, mut node: Self::Node) -> Result> { - let Some(plan) = self.plans.first() else { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Stop)); - }; - - if node.as_any().is::() { - return Ok(Transformed::new(node, false, TreeNodeRecursion::Jump)); - } - - let index = self.cursor; - self.cursor += 1; - - if index != plan.node_index { - return Ok(Transformed::no(node)); - } - - self.plans = &self.plans[1..]; - - let mut transformed = false; - for rule_index in plan.rule_indices.iter() { - let rule = &self.rules[*rule_index]; - let transform = rule.rewrite(node, self.ctx)?; - node = transform.data; - transformed |= transform.transformed; - } - - Ok(Transformed::new( - node, - transformed, - TreeNodeRecursion::Continue, - )) - } -} - -/// An [`ExecutionPlan`] that applies transformation rules during execution. -#[derive(Debug)] -pub struct TransformPlanExec { - /// The input execution plan. - input: Arc, - /// The transformation rules to apply. - rules: Vec>, - /// The pre-calculated transformation plans. - plans: Vec, - /// Execution metrics. - metrics: ExecutionPlanMetricsSet, -} - -impl TransformPlanExec { - /// Create a new [TransformPlanExec]. - /// - /// This method returns an error if any of the transformation rules return an error during the - /// initial traversal of the input plan. - pub fn try_new( - input: Arc, - rules: Vec>, - ) -> Result { - let mut planner = TransformationPlanner::new(rules); - input.visit(&mut planner)?; - - Ok(Self { - input, - rules: planner.rules, - plans: planner.plans, - metrics: ExecutionPlanMetricsSet::new(), - }) - } - - /// Returns the number of nodes that match at least one transformation rule. - pub fn plans_to_transform(&self) -> usize { - self.plans.len() - } - - pub fn transform( - &self, - context: &Arc, - ) -> Result> { - let mut applier = TransformationApplier::new(&self.rules, &self.plans, context); - let input = Arc::clone(&self.input); - input.rewrite(&mut applier).map(|t| t.data) - } - - /// Returns the input plan. - pub fn input(&self) -> &Arc { - &self.input - } - - /// Returns the transformation rules. - pub fn rules(&self) -> &[Box] { - &self.rules - } - - /// Checks if the transformation rules contains a rule of a specific type. - pub fn has_rule(&self) -> bool { - self.rules.iter().any(|r| r.as_any().is::()) - } - - /// Adds a new transformation rule and recalculates transformation plans. - pub fn add_rule( - &self, - new_rule: Box, - ) -> Result { - self.add_rules(vec![new_rule]) - } - - /// Adds new transformation rules and recalculates transformation plans. - pub fn add_rules( - &self, - new_rules: Vec>, - ) -> Result { - let mut planner = TransformationPlanner::new(new_rules); - self.input.visit(&mut planner)?; - let new_rules = planner.rules; - let new_plans = planner.plans; - - let mut current_rules = self - .rules - .iter() - .map(|rule| rule.clone_box()) - .collect::>(); - - let offset = current_rules.len(); - current_rules.extend(new_rules); - - let mut merged_plans = Vec::with_capacity(self.plans.len() + new_plans.len()); - let mut old_plans_iter = self.plans.iter().peekable(); - let mut new_plans_iter = new_plans.into_iter().peekable(); - - while old_plans_iter.peek().is_some() || new_plans_iter.peek().is_some() { - match (old_plans_iter.peek(), new_plans_iter.peek()) { - (Some(&old), Some(new)) if old.node_index == new.node_index => { - let mut merged_plan = old.clone(); - for &rule_index in &new.rule_indices { - merged_plan.rule_indices.push(offset + rule_index); - } - merged_plans.push(merged_plan); - old_plans_iter.next(); - new_plans_iter.next(); - } - (Some(&old), Some(new)) if old.node_index < new.node_index => { - merged_plans.push(old.clone()); - old_plans_iter.next(); - } - (Some(_), Some(_)) => { - // old_node_index > new.node_index - let mut new_plan = new_plans_iter.next().unwrap(); - for rule_index in new_plan.rule_indices.iter_mut() { - *rule_index += offset; - } - merged_plans.push(new_plan); - } - (Some(&old), None) => { - merged_plans.push(old.clone()); - old_plans_iter.next(); - } - (None, Some(_)) => { - let mut new_plan = new_plans_iter.next().unwrap(); - for rule_index in new_plan.rule_indices.iter_mut() { - *rule_index += offset; - } - merged_plans.push(new_plan); - } - (None, None) => unreachable!(), - } - } - - Ok(Self { - input: Arc::clone(&self.input), - rules: current_rules, - plans: merged_plans, - metrics: ExecutionPlanMetricsSet::new(), - }) - } -} - -impl DisplayAs for TransformPlanExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - let mut rule_to_nodes_count = vec![0; self.rules.len()]; - for plan in self.plans.iter() { - for rule_index in plan.rule_indices.iter() { - rule_to_nodes_count[*rule_index] += 1; - } - } - - let rules = rule_to_nodes_count - .into_iter() - .enumerate() - .map(|(rule_index, nodes_count)| { - let rule_name = self.rules[rule_index].name(); - format!("{rule_name}: plans_to_modify={nodes_count}") - }) - .join(", "); - - write!(f, "TransformPlanExec: rules=[{rules}]") - } - DisplayFormatType::TreeRender => Ok(()), - } - } -} - -impl ExecutionPlan for TransformPlanExec { - fn static_name() -> &'static str - where - Self: Sized, - { - "TransformPlanExec" - } - - fn name(&self) -> &str { - Self::static_name() - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn properties(&self) -> &PlanProperties { - self.input.properties() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - let rules = self.rules.iter().map(|r| r.clone_box()).collect(); - Ok(Arc::new(TransformPlanExec::try_new( - Arc::clone(&children[0]), - rules, - )?)) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let metric = MetricBuilder::new(&self.metrics).elapsed_compute(partition); - let _timer = metric.timer(); - - let transformed = self.transform(&context)?; - transformed.execute(partition, context) - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } - - fn statistics(&self) -> Result { - self.partition_statistics(None) - } - - fn partition_statistics(&self, partition: Option) -> Result { - self.input.partition_statistics(partition) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::coalesce_partitions::CoalescePartitionsExec; - use crate::empty::EmptyExec; - use crate::limit::GlobalLimitExec; - use crate::placeholder_row::PlaceholderRowExec; - use crate::plan_transformer::resolve_placeholders::ResolvePlaceholdersRule; - use crate::projection::ProjectionExec; - use crate::{get_plan_string, test}; - use arrow_schema::{DataType, Schema}; - use datafusion_common::{ParamValues, ScalarValue, tree_node::Transformed}; - use datafusion_expr::Operator; - use datafusion_physical_expr::expressions::{binary, lit, placeholder}; - use datafusion_physical_expr::projection::ProjectionExpr; - use insta::assert_snapshot; - use std::sync::Arc; - - #[derive(Debug, Clone)] - struct MockRule { - name: String, - match_name: String, - } - - impl ExecutionTransformationRule for MockRule { - fn name(&self) -> &str { - &self.name - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn clone_box(&self) -> Box { - Box::new(self.clone()) - } - - fn matches(&mut self, node: &Arc) -> Result { - Ok(node.name() == self.match_name || self.match_name == "all") - } - - fn rewrite( - &self, - node: Arc, - _ctx: &TaskContext, - ) -> Result>> { - Ok(Transformed::no(node)) - } - } - - #[test] - fn test_has_rule() -> Result<()> { - let schema = Arc::new(Schema::empty()); - let input = Arc::new(EmptyExec::new(schema)); - - let resolve_rule = Box::new(ResolvePlaceholdersRule::new()); - let exec = TransformPlanExec::try_new(input, vec![resolve_rule])?; - assert!(exec.has_rule::()); - assert!(!exec.has_rule::()); - - let exec = exec.add_rule(Box::new(MockRule { - name: "mock".to_string(), - match_name: "any".to_string(), - }))?; - - assert!(exec.has_rule::()); - assert!(exec.has_rule::()); - - Ok(()) - } - - #[test] - fn test_add_rules_merge() -> Result<()> { - let schema = Arc::new(Schema::empty()); - let empty = Arc::new(EmptyExec::new(schema)); - let coalesce = Arc::new(CoalescePartitionsExec::new(empty)); - let input = Arc::new(GlobalLimitExec::new(coalesce, 0, None)) as Arc<_>; - - // Node 0: GlobalLimitExec - // Node 1: CoalescePartitionsExec - // Node 2: EmptyExec - - let rule_a = Box::new(MockRule { - name: "ruleA".to_string(), - match_name: "CoalescePartitionsExec".to_string(), - }); - let exec = TransformPlanExec::try_new(Arc::clone(&input), vec![rule_a.clone()])?; - - assert_eq!(exec.rules.len(), 1); - assert_eq!(exec.plans.len(), 1); - assert_eq!(exec.plans[0].node_index, 1); - assert_eq!(exec.plans[0].rule_indices, vec![0]); - - let rule_b = Box::new(MockRule { - name: "ruleB".to_string(), - match_name: "GlobalLimitExec".to_string(), - }); - let exec = exec.add_rules(vec![rule_b.clone()])?; - - assert_eq!(exec.rules.len(), 2); - assert_eq!(exec.plans.len(), 2); - assert_eq!(exec.plans[0].node_index, 0); - assert_eq!(exec.plans[0].rule_indices, vec![1]); - assert_eq!(exec.plans[1].node_index, 1); - assert_eq!(exec.plans[1].rule_indices, vec![0]); - - let rule_c = Box::new(MockRule { - name: "ruleC".to_string(), - match_name: "EmptyExec".to_string(), - }); - let exec = exec.add_rules(vec![rule_c.clone()])?; - - assert_eq!(exec.rules.len(), 3); - assert_eq!(exec.plans.len(), 3); - assert_eq!(exec.plans[0].node_index, 0); - assert_eq!(exec.plans[0].rule_indices, vec![1]); - assert_eq!(exec.plans[1].node_index, 1); - assert_eq!(exec.plans[1].rule_indices, vec![0]); - assert_eq!(exec.plans[2].node_index, 2); - assert_eq!(exec.plans[2].rule_indices, vec![2]); - - let rule_d = Box::new(MockRule { - name: "ruleD".to_string(), - match_name: "all".to_string(), - }); - - let exec = exec.add_rules(vec![rule_d.clone()])?; - let check_full_plan = |exec: TransformPlanExec| { - assert_eq!(exec.rules.len(), 4); - assert_eq!(exec.plans.len(), 3); - assert_eq!(exec.plans[0].node_index, 0); - assert_eq!(exec.plans[0].rule_indices, vec![1, 3]); - assert_eq!(exec.plans[1].node_index, 1); - assert_eq!(exec.plans[1].rule_indices, vec![0, 3]); - assert_eq!(exec.plans[2].node_index, 2); - assert_eq!(exec.plans[2].rule_indices, vec![2, 3]); - }; - - check_full_plan(exec); - - let exec = TransformPlanExec::try_new( - Arc::clone(&input), - vec![ - rule_a.clone(), - rule_b.clone(), - rule_c.clone(), - rule_d.clone(), - ], - )?; - - check_full_plan(exec); - - let exec = TransformPlanExec::try_new(Arc::clone(&input), vec![rule_a])? - .add_rules(vec![rule_b, rule_c, rule_d])?; - - check_full_plan(exec); - - Ok(()) - } - - #[tokio::test] - async fn test_consts_evaluation() -> Result<()> { - let schema = test::aggr_test_schema(); - let expr = binary( - lit(10), - Operator::Plus, - binary( - lit(5), - Operator::Multiply, - placeholder("$1", DataType::Int32), - &schema, - )?, - &schema, - )?; - - let projection_expr = ProjectionExpr { - expr, - alias: "sum".to_string(), - }; - - let row = Arc::new(PlaceholderRowExec::new(schema)); - let projection = ProjectionExec::try_new(vec![projection_expr], row)?; - let transformer = TransformPlanExec::try_new( - Arc::new(projection), - vec![Box::new(ResolvePlaceholdersRule::new())], - )?; - - let param_values = ParamValues::List(vec![ScalarValue::Int32(Some(20)).into()]); - let task_ctx = Arc::new(TaskContext::default().with_param_values(param_values)); - - let plan = transformer.transform(&task_ctx)?; - let plan_string = get_plan_string(&plan).join("\n"); - - assert_snapshot!(plan_string, @r" - ProjectionExec: expr=[110 as sum] - PlaceholderRowExec - "); - - Ok(()) - } -} diff --git a/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs b/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs deleted file mode 100644 index 115e9cc70ff4e..0000000000000 --- a/datafusion/physical-plan/src/plan_transformer/resolve_placeholders.rs +++ /dev/null @@ -1,126 +0,0 @@ -// 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. - -//! Rule to resolve placeholders in the physical plan with actual values. - -use std::sync::Arc; - -use datafusion_common::{ - ParamValues, Result, exec_err, - tree_node::{Transformed, TreeNode}, -}; -use datafusion_execution::TaskContext; -use datafusion_physical_expr::{ - PhysicalExpr, - expressions::{Literal, PlaceholderExpr, has_placeholders}, - simplifier::const_evaluator::simplify_const_expr, -}; - -use crate::{ExecutionPlan, plan_transformer::ExecutionTransformationRule}; - -/// A transformation rule that replaces [`PlaceholderExpr`] with actual values. -/// -/// This rule is applied to the physical plan when actual parameter values are provided in the -/// [`TaskContext`]. It traverses the plan and replaces any placeholders found in physical -/// expressions with their corresponding literal values. -#[derive(Debug, Clone, Default)] -pub struct ResolvePlaceholdersRule {} - -impl ResolvePlaceholdersRule { - /// Create a new [`ResolvePlaceholdersRule`]. - pub fn new() -> Self { - Self {} - } -} - -impl ExecutionTransformationRule for ResolvePlaceholdersRule { - fn name(&self) -> &str { - "ResolvePlaceholders" - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn clone_box(&self) -> Box { - Box::new(Self {}) - } - - fn matches(&mut self, node: &Arc) -> Result { - let Some(exprs) = node.physical_expressions() else { - return Ok(false); - }; - - for expr in exprs { - if has_placeholders(&expr) { - return Ok(true); - } - } - - Ok(false) - } - - fn rewrite( - &self, - node: Arc, - ctx: &TaskContext, - ) -> Result>> { - let Some(param_values) = ctx.param_values() else { - return Ok(Transformed::no(node)); - }; - - let Some(exprs) = node.physical_expressions() else { - return exec_err!("no physical expressions found"); - }; - - let new_exprs = exprs - .map(|expr| resolve_expr_placeholders(expr, param_values)) - .collect::>>()?; - - match node.with_physical_expressions(new_exprs.into())? { - Some(transformed_plan) => Ok(Transformed::yes(transformed_plan)), - None => exec_err!("failed to rewrite execution plan"), - } - } -} - -/// Resolves [`PlaceholderExpr`] in the physical expression using the provided [`ParamValues`]. -pub fn resolve_expr_placeholders( - expr: Arc, - param_values: &ParamValues, -) -> Result> { - let expr = expr.transform_up(|node| { - let Some(placeholder) = node.as_any().downcast_ref::() else { - return Ok(Transformed::no(node)); - }; - - if let Some(ref field) = placeholder.field { - let scalar = param_values.get_placeholders_with_values(&placeholder.id)?; - let value = scalar.value.cast_to(field.data_type())?; - let literal = Literal::new_with_metadata(value, scalar.metadata); - Ok(Transformed::yes(Arc::new(literal))) - } else { - Ok(Transformed::no(node)) - } - })?; - - if expr.transformed { - simplify_const_expr(expr.data).map(|t| t.data) - } else { - Ok(expr.data) - } -} diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index a4de3f0101ade..628a5722950bd 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -27,13 +27,15 @@ use super::{ SendableRecordBatchStream, SortOrderPushdownResult, Statistics, }; use crate::column_rewriter::PhysicalColumnRewriter; -use crate::execution_plan::{CardinalityEffect, ReplacePhysicalExpr}; +use crate::execution_plan::{ + CardinalityEffect, ReplacePhysicalExpr, has_same_children_properties, +}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterColumnChecker, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDownPredicate, }; use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef}; -use crate::{DisplayFormatType, ExecutionPlan, PhysicalExpr}; +use crate::{DisplayFormatType, ExecutionPlan, PhysicalExpr, check_if_same_properties}; use std::any::Any; use std::collections::HashMap; use std::pin::Pin; @@ -81,7 +83,7 @@ pub struct ProjectionExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl ProjectionExec { @@ -140,13 +142,19 @@ impl ProjectionExec { E: Into, { let input_schema = input.schema(); - // convert argument to Vec - let expr_vec = expr.into_iter().map(Into::into).collect::>(); - let projection = ProjectionExprs::new(expr_vec); + let expr_arc = expr.into_iter().map(Into::into).collect::>(); + let projection = ProjectionExprs::from_expressions(expr_arc); let projector = projection.make_projector(&input_schema)?; + Self::try_from_projector(projector, input) + } + fn try_from_projector( + projector: Projector, + input: Arc, + ) -> Result { // Construct a map from the input expressions to the output expression of the Projection - let projection_mapping = projection.projection_mapping(&input_schema)?; + let projection_mapping = + projector.projection().projection_mapping(&input.schema())?; let cache = Self::compute_properties( &input, &projection_mapping, @@ -156,7 +164,7 @@ impl ProjectionExec { projector, input, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -219,6 +227,17 @@ impl ProjectionExec { } Ok(alias_map) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for ProjectionExec { @@ -272,7 +291,7 @@ impl ExecutionPlan for ProjectionExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -307,8 +326,9 @@ impl ExecutionPlan for ProjectionExec { self: Arc, mut children: Vec>, ) -> Result> { - ProjectionExec::try_new( - self.projector.projection().clone(), + check_if_same_properties!(self, children); + ProjectionExec::try_from_projector( + self.projector.clone(), children.swap_remove(0), ) .map(|p| Arc::new(p) as _) diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 936a02581e89c..f8847cbacefb5 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -74,7 +74,7 @@ pub struct RecursiveQueryExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl RecursiveQueryExec { @@ -97,7 +97,7 @@ impl RecursiveQueryExec { is_distinct, work_table, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -143,7 +143,7 @@ impl ExecutionPlan for RecursiveQueryExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 612c7bb27ddf4..2e02cd3210786 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -31,7 +31,9 @@ use super::{ DisplayAs, ExecutionPlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use crate::coalesce::LimitedBatchCoalescer; -use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, has_same_children_properties, +}; use crate::hash_utils::create_hashes; use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; @@ -39,7 +41,10 @@ use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::spill::spill_manager::SpillManager; use crate::spill::spill_pool::{self, SpillPoolWriter}; use crate::stream::RecordBatchStreamAdapter; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics}; +use crate::{ + DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics, + check_if_same_properties, +}; use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions}; use arrow::compute::take_arrays; @@ -763,7 +768,7 @@ pub struct RepartitionExec { /// `SortPreservingRepartitionExec`, false means `RepartitionExec`. preserve_order: bool, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } #[derive(Debug, Clone)] @@ -832,6 +837,18 @@ impl RepartitionExec { pub fn name(&self) -> &str { "RepartitionExec" } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + state: Default::default(), + ..Self::clone(self) + } + } } impl DisplayAs for RepartitionExec { @@ -891,7 +908,7 @@ impl ExecutionPlan for RepartitionExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -903,6 +920,7 @@ impl ExecutionPlan for RepartitionExec { self: Arc, mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); let mut repartition = RepartitionExec::try_new( children.swap_remove(0), self.partitioning().clone(), @@ -1204,7 +1222,7 @@ impl ExecutionPlan for RepartitionExec { _config: &ConfigOptions, ) -> Result>> { use Partitioning::*; - let mut new_properties = self.cache.clone(); + let mut new_properties = PlanProperties::clone(&self.cache); new_properties.partitioning = match new_properties.partitioning { RoundRobinBatch(_) => RoundRobinBatch(target_partitions), Hash(hash, _) => Hash(hash, target_partitions), @@ -1215,7 +1233,7 @@ impl ExecutionPlan for RepartitionExec { state: Arc::clone(&self.state), metrics: self.metrics.clone(), preserve_order: self.preserve_order, - cache: new_properties, + cache: new_properties.into(), }))) } } @@ -1235,7 +1253,7 @@ impl RepartitionExec { state: Default::default(), metrics: ExecutionPlanMetricsSet::new(), preserve_order, - cache, + cache: Arc::new(cache), }) } @@ -1296,7 +1314,7 @@ impl RepartitionExec { // to maintain order self.input.output_partitioning().partition_count() > 1; let eq_properties = Self::eq_properties_helper(&self.input, self.preserve_order); - self.cache = self.cache.with_eq_properties(eq_properties); + Arc::make_mut(&mut self.cache).set_eq_properties(eq_properties); self } diff --git a/datafusion/physical-plan/src/reuse.rs b/datafusion/physical-plan/src/reuse.rs new file mode 100644 index 0000000000000..5f3536ccc11f5 --- /dev/null +++ b/datafusion/physical-plan/src/reuse.rs @@ -0,0 +1,198 @@ +// 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. + +use std::sync::Arc; + +use datafusion_common::{ + ParamValues, Result, exec_err, + tree_node::{Transformed, TreeNode, TreeNodeRecursion}, +}; +use datafusion_physical_expr::{ + PhysicalExpr, + expressions::{has_placeholders, resolve_expr_placeholders}, +}; + +use crate::{ExecutionPlan, execution_plan::ReplacePhysicalExpr}; + +/// Wraps an [`ExecutionPlan`] that may contain placeholders, making +/// it re-executable, avoiding re-planning stage. +/// +/// # Limitations +/// +/// Plan does not support re-execution if it (OR): +/// +/// * uses dynamic filters, +/// * represents a recursive query. +/// +/// This invariant is not enforced by [`ReusableExecutionPlan`] itself, so it +/// must be checked by user. +/// +pub struct ReusableExecutionPlan { + binder: Binder, + bound_plan: Option>, +} + +impl ReusableExecutionPlan { + /// Make a new [`ReusableExecutionPlan`] bound to the passed `plan`. + pub fn new(plan: Arc) -> Self { + let binder = Binder::new(plan); + Self { + binder, + bound_plan: None, + } + } + + /// Return a ready to execution instance of [`ReusableExecutionPlan`], + /// where placeholders are bound to the passed `params`. + pub fn bind(&self, params: Option<&ParamValues>) -> Result { + let bound_plan = self.binder.bind(params).map(Some)?; + Ok(Self { + binder: self.binder.clone(), + bound_plan, + }) + } + + /// Return an inner plan to execute. + /// + /// If this plan is a result of [`Self::bind`] call, then bound plan is returned. + /// Otherwise, an initial plan is returned. + pub fn plan(&self) -> Arc { + self.bound_plan + .clone() + .unwrap_or_else(|| Arc::clone(&self.binder.plan)) + } +} + +impl From> for ReusableExecutionPlan { + fn from(plan: Arc) -> Self { + Self::new(plan) + } +} + +#[derive(Debug, Clone)] +struct NodeWithPlaceholders { + /// The index of the node in the tree traversal. + idx: usize, + /// Positions of the placeholders among plan physical expressions. + placeholder_idx: Vec, +} + +impl NodeWithPlaceholders { + /// Returns [`Some`] if passed `node` contains placeholders and must + /// be resolved on binding stage. + fn new(node: &Arc, idx: usize) -> Option { + let placeholder_idx = if let Some(iter) = node.physical_expressions() { + iter.enumerate() + .filter_map(|(i, expr)| { + if has_placeholders(&expr) { + Some(i) + } else { + None + } + }) + .collect() + } else { + vec![] + }; + + if placeholder_idx.is_empty() { + None + } else { + Some(Self { + idx, + placeholder_idx, + }) + } + } + + fn resolve( + &self, + node: &Arc, + params: Option<&ParamValues>, + ) -> Result> { + let Some(expr) = node.physical_expressions() else { + return exec_err!("node {} does not support expressions export", node.name()); + }; + let mut exprs: Vec> = expr.collect(); + for idx in self.placeholder_idx.iter() { + exprs[*idx] = resolve_expr_placeholders(Arc::clone(&exprs[*idx]), params)?; + } + let Some(resolved_node) = + node.with_physical_expressions(ReplacePhysicalExpr { exprs })? + else { + return exec_err!( + "node {} does not support expressions replace", + node.name() + ); + }; + Ok(resolved_node) + } +} + +/// Helper to bound placeholders and reset plan nodes state. +#[derive(Clone)] +struct Binder { + /// Created during [`Binder`] construction. + /// This way we avoid runtime rebuild for expressions without placeholders. + nodes_to_resolve: Arc<[NodeWithPlaceholders]>, + plan: Arc, +} + +impl Binder { + fn new(plan: Arc) -> Self { + let mut nodes_to_resolve = vec![]; + let mut cursor = 0; + + plan.apply(|node| { + let idx = cursor; + cursor += 1; + if let Some(node) = NodeWithPlaceholders::new(node, idx) { + nodes_to_resolve.push(node); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + + Self { + plan, + nodes_to_resolve: nodes_to_resolve.into(), + } + } + + fn bind(&self, params: Option<&ParamValues>) -> Result> { + let mut cursor = 0; + let mut resolve_node_idx = 0; + Arc::clone(&self.plan) + .transform_down(|node| { + let idx = cursor; + cursor += 1; + if resolve_node_idx < self.nodes_to_resolve.len() + && self.nodes_to_resolve[resolve_node_idx].idx == idx + { + // Note: `resolve` replaces plan expressions, which also resets a plan state. + let resolved_node = + self.nodes_to_resolve[resolve_node_idx].resolve(&node, params)?; + resolve_node_idx += 1; + Ok(Transformed::yes(resolved_node)) + } else { + // Reset state. + Ok(Transformed::yes(node.reset_state()?)) + } + }) + .map(|tnr| tnr.data) + } +} diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index e12aefd3003e1..f094258c8a9f5 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -57,12 +57,13 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use crate::execution_plan::ReplacePhysicalExpr; +use crate::execution_plan::{ReplacePhysicalExpr, has_same_children_properties}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::sorts::sort::sort_batch; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use arrow::compute::concat_batches; @@ -94,7 +95,7 @@ pub struct PartialSortExec { /// Fetch highest/lowest n results fetch: Option, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl PartialSortExec { @@ -115,7 +116,7 @@ impl PartialSortExec { metrics_set: ExecutionPlanMetricsSet::new(), preserve_partitioning, fetch: None, - cache, + cache: Arc::new(cache), } } @@ -133,12 +134,8 @@ impl PartialSortExec { /// input partitions producing a single, sorted partition. pub fn with_preserve_partitioning(mut self, preserve_partitioning: bool) -> Self { self.preserve_partitioning = preserve_partitioning; - self.cache = self - .cache - .with_partitioning(Self::output_partitioning_helper( - &self.input, - self.preserve_partitioning, - )); + Arc::make_mut(&mut self.cache).partitioning = + Self::output_partitioning_helper(&self.input, self.preserve_partitioning); self } @@ -208,6 +205,16 @@ impl PartialSortExec { input.boundedness(), )) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + ..Self::clone(self) + } + } } impl DisplayAs for PartialSortExec { @@ -256,7 +263,7 @@ impl ExecutionPlan for PartialSortExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -284,6 +291,7 @@ impl ExecutionPlan for PartialSortExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); let new_partial_sort = PartialSortExec::new( self.expr.clone(), Arc::clone(&children[0]), diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 565e482592958..6e54721b0e01e 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -29,6 +29,7 @@ use parking_lot::RwLock; use crate::common::spawn_buffered; use crate::execution_plan::{ Boundedness, CardinalityEffect, EmissionType, ReplacePhysicalExpr, + has_same_children_properties, }; use crate::expressions::PhysicalSortExpr; use crate::filter_pushdown::{ @@ -953,7 +954,7 @@ pub struct SortExec { /// Normalized common sort prefix between the input and the sort expressions (only used with fetch) common_sort_prefix: Vec, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// Filter matching the state of the sort for dynamic filter pushdown. /// If `fetch` is `Some`, this will also be set and a TopK operator may be used. /// If `fetch` is `None`, this will be `None`. @@ -975,7 +976,7 @@ impl SortExec { preserve_partitioning, fetch: None, common_sort_prefix: sort_prefix, - cache, + cache: Arc::new(cache), filter: None, } } @@ -994,12 +995,8 @@ impl SortExec { /// input partitions producing a single, sorted partition. pub fn with_preserve_partitioning(mut self, preserve_partitioning: bool) -> Self { self.preserve_partitioning = preserve_partitioning; - self.cache = self - .cache - .with_partitioning(Self::output_partitioning_helper( - &self.input, - self.preserve_partitioning, - )); + Arc::make_mut(&mut self.cache).partitioning = + Self::output_partitioning_helper(&self.input, self.preserve_partitioning); self } @@ -1023,7 +1020,7 @@ impl SortExec { preserve_partitioning: self.preserve_partitioning, common_sort_prefix: self.common_sort_prefix.clone(), fetch: self.fetch, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), filter: self.filter.clone(), } } @@ -1036,12 +1033,12 @@ impl SortExec { /// operation since rows that are not going to be included /// can be dropped. pub fn with_fetch(&self, fetch: Option) -> Self { - let mut cache = self.cache.clone(); + let mut cache = PlanProperties::clone(&self.cache); // If the SortExec can emit incrementally (that means the sort requirements // and properties of the input match), the SortExec can generate its result // without scanning the entire input when a fetch value exists. let is_pipeline_friendly = matches!( - self.cache.emission_type, + cache.emission_type, EmissionType::Incremental | EmissionType::Both ); if fetch.is_some() && is_pipeline_friendly { @@ -1053,7 +1050,7 @@ impl SortExec { }); let mut new_sort = self.cloned(); new_sort.fetch = fetch; - new_sort.cache = cache; + new_sort.cache = cache.into(); new_sort.filter = filter; new_sort } @@ -1208,7 +1205,7 @@ impl ExecutionPlan for SortExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -1237,14 +1234,17 @@ impl ExecutionPlan for SortExec { let mut new_sort = self.cloned(); assert_eq!(children.len(), 1, "SortExec should have exactly one child"); new_sort.input = Arc::clone(&children[0]); - // Recompute the properties based on the new input since they may have changed - let (cache, sort_prefix) = Self::compute_properties( - &new_sort.input, - new_sort.expr.clone(), - new_sort.preserve_partitioning, - )?; - new_sort.cache = cache; - new_sort.common_sort_prefix = sort_prefix; + + if !has_same_children_properties(&self, &children)? { + // Recompute the properties based on the new input since they may have changed + let (cache, sort_prefix) = Self::compute_properties( + &new_sort.input, + new_sort.expr.clone(), + new_sort.preserve_partitioning, + )?; + new_sort.cache = Arc::new(cache); + new_sort.common_sort_prefix = sort_prefix; + } Ok(Arc::new(new_sort)) } @@ -1509,7 +1509,7 @@ mod tests { pub struct SortedUnboundedExec { schema: Schema, batch_size: u64, - cache: PlanProperties, + cache: Arc, } impl DisplayAs for SortedUnboundedExec { @@ -1549,7 +1549,7 @@ mod tests { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -2302,7 +2302,9 @@ mod tests { let source = SortedUnboundedExec { schema: schema.clone(), batch_size: 2, - cache: SortedUnboundedExec::compute_properties(Arc::new(schema.clone())), + cache: Arc::new(SortedUnboundedExec::compute_properties(Arc::new( + schema.clone(), + ))), }; let mut plan = SortExec::new( [PhysicalSortExpr::new_default(Arc::new(Column::new( diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 91ab8906dec4d..b2889cc55c0e2 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -28,6 +28,7 @@ use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; @@ -36,7 +37,9 @@ use datafusion_execution::memory_pool::MemoryConsumer; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; -use crate::execution_plan::{EvaluationType, ReplacePhysicalExpr, SchedulingType}; +use crate::execution_plan::{ + EvaluationType, ReplacePhysicalExpr, SchedulingType, has_same_children_properties, +}; use log::{debug, trace}; /// Sort preserving merge execution plan @@ -94,7 +97,7 @@ pub struct SortPreservingMergeExec { /// Optional number of rows to fetch. Stops producing rows after this fetch fetch: Option, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// Use round-robin selection of tied winners of loser tree /// /// See [`Self::with_round_robin_repartition`] for more information. @@ -110,7 +113,7 @@ impl SortPreservingMergeExec { expr, metrics: ExecutionPlanMetricsSet::new(), fetch: None, - cache, + cache: Arc::new(cache), enable_round_robin_repartition: true, } } @@ -181,6 +184,16 @@ impl SortPreservingMergeExec { .with_evaluation_type(drive) .with_scheduling_type(scheduling) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + ..Self::clone(self) + } + } } impl DisplayAs for SortPreservingMergeExec { @@ -226,7 +239,7 @@ impl ExecutionPlan for SortPreservingMergeExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -241,7 +254,7 @@ impl ExecutionPlan for SortPreservingMergeExec { expr: self.expr.clone(), metrics: self.metrics.clone(), fetch: limit, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), enable_round_robin_repartition: true, })) } @@ -281,10 +294,11 @@ impl ExecutionPlan for SortPreservingMergeExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new( - SortPreservingMergeExec::new(self.expr.clone(), Arc::clone(&children[0])) + SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) .with_fetch(self.fetch), )) } @@ -1393,7 +1407,7 @@ mod tests { #[derive(Debug, Clone)] struct CongestedExec { schema: Schema, - cache: PlanProperties, + cache: Arc, congestion: Arc, } @@ -1429,7 +1443,7 @@ mod tests { fn as_any(&self) -> &dyn Any { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } fn children(&self) -> Vec<&Arc> { @@ -1522,7 +1536,7 @@ mod tests { }; let source = CongestedExec { schema: schema.clone(), - cache: properties, + cache: Arc::new(properties), congestion: Arc::new(Congestion::new(partition_count)), }; let spm = SortPreservingMergeExec::new( diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index c8b8d95718cb8..1535482374110 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -67,7 +67,7 @@ pub struct StreamingTableExec { projected_output_ordering: Vec, infinite: bool, limit: Option, - cache: PlanProperties, + cache: Arc, metrics: ExecutionPlanMetricsSet, } @@ -111,7 +111,7 @@ impl StreamingTableExec { projected_output_ordering, infinite, limit, - cache, + cache: Arc::new(cache), metrics: ExecutionPlanMetricsSet::new(), }) } @@ -236,7 +236,7 @@ impl ExecutionPlan for StreamingTableExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -335,7 +335,7 @@ impl ExecutionPlan for StreamingTableExec { projected_output_ordering: self.projected_output_ordering.clone(), infinite: self.infinite, limit, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), metrics: self.metrics.clone(), })) } diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index c6d0940c35480..aa079e73dd086 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -75,7 +75,7 @@ pub struct TestMemoryExec { /// The maximum number of records to read from this plan. If `None`, /// all records after filtering are returned. fetch: Option, - cache: PlanProperties, + cache: Arc, } impl DisplayAs for TestMemoryExec { @@ -134,7 +134,7 @@ impl ExecutionPlan for TestMemoryExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -239,7 +239,7 @@ impl TestMemoryExec { Ok(Self { partitions: partitions.to_vec(), schema, - cache: PlanProperties::new( + cache: Arc::new(PlanProperties::new( EquivalenceProperties::new_with_orderings( Arc::clone(&projected_schema), Vec::::new(), @@ -247,7 +247,7 @@ impl TestMemoryExec { Partitioning::UnknownPartitioning(partitions.len()), EmissionType::Incremental, Boundedness::Bounded, - ), + )), projected_schema, projection, sort_information: vec![], @@ -265,7 +265,7 @@ impl TestMemoryExec { ) -> Result> { let mut source = Self::try_new(partitions, schema, projection)?; let cache = source.compute_properties(); - source.cache = cache; + source.cache = Arc::new(cache); Ok(Arc::new(source)) } @@ -273,7 +273,7 @@ impl TestMemoryExec { pub fn update_cache(source: &Arc) -> TestMemoryExec { let cache = source.compute_properties(); let mut source = (**source).clone(); - source.cache = cache; + source.cache = Arc::new(cache); source } @@ -342,7 +342,7 @@ impl TestMemoryExec { } self.sort_information = sort_information; - self.cache = self.compute_properties(); + self.cache = Arc::new(self.compute_properties()); Ok(self) } diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index 4507cccba05a9..a8b21f70f7760 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -125,7 +125,7 @@ pub struct MockExec { /// if true (the default), sends data using a separate task to ensure the /// batches are not available without this stream yielding first use_task: bool, - cache: PlanProperties, + cache: Arc, } impl MockExec { @@ -142,7 +142,7 @@ impl MockExec { data, schema, use_task: true, - cache, + cache: Arc::new(cache), } } @@ -192,7 +192,7 @@ impl ExecutionPlan for MockExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -299,7 +299,7 @@ pub struct BarrierExec { /// all streams wait on this barrier to produce barrier: Arc, - cache: PlanProperties, + cache: Arc, } impl BarrierExec { @@ -312,7 +312,7 @@ impl BarrierExec { data, schema, barrier, - cache, + cache: Arc::new(cache), } } @@ -364,7 +364,7 @@ impl ExecutionPlan for BarrierExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -429,7 +429,7 @@ impl ExecutionPlan for BarrierExec { /// A mock execution plan that errors on a call to execute #[derive(Debug)] pub struct ErrorExec { - cache: PlanProperties, + cache: Arc, } impl Default for ErrorExec { @@ -446,7 +446,9 @@ impl ErrorExec { true, )])); let cache = Self::compute_properties(schema); - Self { cache } + Self { + cache: Arc::new(cache), + } } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. @@ -487,7 +489,7 @@ impl ExecutionPlan for ErrorExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -517,7 +519,7 @@ impl ExecutionPlan for ErrorExec { pub struct StatisticsExec { stats: Statistics, schema: Arc, - cache: PlanProperties, + cache: Arc, } impl StatisticsExec { pub fn new(stats: Statistics, schema: Schema) -> Self { @@ -530,7 +532,7 @@ impl StatisticsExec { Self { stats, schema: Arc::new(schema), - cache, + cache: Arc::new(cache), } } @@ -577,7 +579,7 @@ impl ExecutionPlan for StatisticsExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -623,7 +625,7 @@ pub struct BlockingExec { /// Ref-counting helper to check if the plan and the produced stream are still in memory. refs: Arc<()>, - cache: PlanProperties, + cache: Arc, } impl BlockingExec { @@ -633,7 +635,7 @@ impl BlockingExec { Self { schema, refs: Default::default(), - cache, + cache: Arc::new(cache), } } @@ -684,7 +686,7 @@ impl ExecutionPlan for BlockingExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -766,7 +768,7 @@ pub struct PanicExec { /// Number of output partitions. Each partition will produce this /// many empty output record batches prior to panicking batches_until_panics: Vec, - cache: PlanProperties, + cache: Arc, } impl PanicExec { @@ -778,7 +780,7 @@ impl PanicExec { Self { schema, batches_until_panics, - cache, + cache: Arc::new(cache), } } @@ -830,7 +832,7 @@ impl ExecutionPlan for PanicExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index b6f943886e309..7a3fa8eb3a3fe 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -32,9 +32,10 @@ use super::{ SendableRecordBatchStream, Statistics, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; +use crate::check_if_same_properties; use crate::execution_plan::{ InvariantLevel, boundedness_from_children, check_default_invariants, - emission_type_from_children, + emission_type_from_children, has_same_children_properties, }; use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase}; use crate::metrics::BaselineMetrics; @@ -100,7 +101,7 @@ pub struct UnionExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl UnionExec { @@ -118,7 +119,7 @@ impl UnionExec { UnionExec { inputs, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), } } @@ -147,7 +148,7 @@ impl UnionExec { Ok(Arc::new(UnionExec { inputs, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), })) } } @@ -183,6 +184,17 @@ impl UnionExec { boundedness_from_children(inputs), )) } + + fn with_new_children_and_same_properties( + &self, + children: Vec>, + ) -> Self { + Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for UnionExec { @@ -210,7 +222,7 @@ impl ExecutionPlan for UnionExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -259,6 +271,7 @@ impl ExecutionPlan for UnionExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); UnionExec::try_new(children) } @@ -411,7 +424,7 @@ pub struct InterleaveExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl InterleaveExec { @@ -425,7 +438,7 @@ impl InterleaveExec { Ok(InterleaveExec { inputs, metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -447,6 +460,17 @@ impl InterleaveExec { boundedness_from_children(inputs), )) } + + fn with_new_children_and_same_properties( + &self, + children: Vec>, + ) -> Self { + Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for InterleaveExec { @@ -474,7 +498,7 @@ impl ExecutionPlan for InterleaveExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -495,6 +519,7 @@ impl ExecutionPlan for InterleaveExec { can_interleave(children.iter()), "Can not create InterleaveExec: new children can not be interleaved" ); + check_if_same_properties!(self, children); Ok(Arc::new(InterleaveExec::try_new(children)?)) } diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 5fef754e80780..d8a8937c46b5e 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -26,9 +26,10 @@ use super::metrics::{ RecordOutput, }; use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; +use crate::execution_plan::has_same_children_properties; use crate::{ DisplayFormatType, Distribution, ExecutionPlan, RecordBatchStream, - SendableRecordBatchStream, + SendableRecordBatchStream, check_if_same_properties, }; use arrow::array::{ @@ -74,7 +75,7 @@ pub struct UnnestExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl UnnestExec { @@ -100,7 +101,7 @@ impl UnnestExec { struct_column_indices, options, metrics: Default::default(), - cache, + cache: Arc::new(cache), }) } @@ -193,6 +194,17 @@ impl UnnestExec { pub fn options(&self) -> &UnnestOptions { &self.options } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for UnnestExec { @@ -221,7 +233,7 @@ impl ExecutionPlan for UnnestExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -231,10 +243,11 @@ impl ExecutionPlan for UnnestExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(UnnestExec::new( - Arc::clone(&children[0]), + children.swap_remove(0), self.list_column_indices.clone(), self.struct_column_indices.clone(), Arc::clone(&self.schema), diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index bf4c6a4a256b4..7daf42b71c747 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -28,7 +28,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use super::utils::create_schema; -use crate::execution_plan::ReplacePhysicalExpr; +use crate::execution_plan::{ReplacePhysicalExpr, has_same_children_properties}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::windows::{ calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, @@ -37,7 +37,7 @@ use crate::windows::{ use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, + SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, }; use arrow::compute::take_record_batch; @@ -95,7 +95,7 @@ pub struct BoundedWindowAggExec { // See `get_ordered_partition_by_indices` for more details. ordered_partition_by_indices: Vec, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// If `can_rerepartition` is false, partition_keys is always empty. can_repartition: bool, } @@ -136,7 +136,7 @@ impl BoundedWindowAggExec { metrics: ExecutionPlanMetricsSet::new(), input_order_mode, ordered_partition_by_indices, - cache, + cache: Arc::new(cache), can_repartition, }) } @@ -250,6 +250,17 @@ impl BoundedWindowAggExec { total_byte_size: Precision::Absent, }) } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for BoundedWindowAggExec { @@ -306,7 +317,7 @@ impl ExecutionPlan for BoundedWindowAggExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -341,6 +352,7 @@ impl ExecutionPlan for BoundedWindowAggExec { self: Arc, children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(BoundedWindowAggExec::try_new( self.window_expr.clone(), Arc::clone(&children[0]), diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index d1a87ac594ee9..865ed666e5cf5 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -23,7 +23,9 @@ use std::sync::Arc; use std::task::{Context, Poll}; use super::utils::create_schema; -use crate::execution_plan::{EmissionType, ReplacePhysicalExpr}; +use crate::execution_plan::{ + EmissionType, ReplacePhysicalExpr, has_same_children_properties, +}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::windows::{ calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, @@ -32,7 +34,7 @@ use crate::windows::{ use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PhysicalExpr, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, + SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, }; use arrow::array::ArrayRef; @@ -65,7 +67,7 @@ pub struct WindowAggExec { // see `get_ordered_partition_by_indices` for more details. ordered_partition_by_indices: Vec, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, /// If `can_partition` is false, partition_keys is always empty. can_repartition: bool, } @@ -89,7 +91,7 @@ impl WindowAggExec { schema, metrics: ExecutionPlanMetricsSet::new(), ordered_partition_by_indices, - cache, + cache: Arc::new(cache), can_repartition, }) } @@ -158,6 +160,17 @@ impl WindowAggExec { .unwrap_or_else(Vec::new) } } + + fn with_new_children_and_same_properties( + &self, + mut children: Vec>, + ) -> Self { + Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(self) + } + } } impl DisplayAs for WindowAggExec { @@ -206,7 +219,7 @@ impl ExecutionPlan for WindowAggExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -242,11 +255,12 @@ impl ExecutionPlan for WindowAggExec { fn with_new_children( self: Arc, - children: Vec>, + mut children: Vec>, ) -> Result> { + check_if_same_properties!(self, children); Ok(Arc::new(WindowAggExec::try_new( self.window_expr.clone(), - Arc::clone(&children[0]), + children.swap_remove(0), true, )?)) } diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index 1313909adbba2..b01ea513a8ede 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -109,7 +109,7 @@ pub struct WorkTableExec { /// Execution metrics metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. - cache: PlanProperties, + cache: Arc, } impl WorkTableExec { @@ -129,7 +129,7 @@ impl WorkTableExec { projection, work_table: Arc::new(WorkTable::new(name)), metrics: ExecutionPlanMetricsSet::new(), - cache, + cache: Arc::new(cache), }) } @@ -181,7 +181,7 @@ impl ExecutionPlan for WorkTableExec { self } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { &self.cache } @@ -263,7 +263,7 @@ impl ExecutionPlan for WorkTableExec { projection: self.projection.clone(), metrics: ExecutionPlanMetricsSet::new(), work_table, - cache: self.cache.clone(), + cache: Arc::clone(&self.cache), })) } } diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 8fe68a264c3cc..9a8c25831051f 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -751,8 +751,7 @@ message PhysicalPlanNode { MemoryScanExecNode memory_scan = 35; AsyncFuncExecNode async_func = 36; BufferExecNode buffer = 37; - TransformPlanExecNode transform_plan = 38; - ValuesExecNode values_scan = 39; + ValuesExecNode values_scan = 38; } } @@ -1458,17 +1457,3 @@ message BufferExecNode { PhysicalPlanNode input = 1; uint64 capacity = 2; } - -message TransformPlanExecNode { - PhysicalPlanNode input = 1; - repeated TransformationRule rules = 2; -} - -message TransformationRule { - oneof rule_type { - bytes extension = 1; - ResolvePlaceholdersRule resolve_placeholders = 2; - } -} - -message ResolvePlaceholdersRule {} diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index f579478f976e3..3644140ed0538 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -17905,9 +17905,6 @@ impl serde::Serialize for PhysicalPlanNode { physical_plan_node::PhysicalPlanType::Buffer(v) => { struct_ser.serialize_field("buffer", v)?; } - physical_plan_node::PhysicalPlanType::TransformPlan(v) => { - struct_ser.serialize_field("transformPlan", v)?; - } physical_plan_node::PhysicalPlanType::ValuesScan(v) => { struct_ser.serialize_field("valuesScan", v)?; } @@ -17979,8 +17976,6 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "async_func", "asyncFunc", "buffer", - "transform_plan", - "transformPlan", "values_scan", "valuesScan", ]; @@ -18023,7 +18018,6 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { MemoryScan, AsyncFunc, Buffer, - TransformPlan, ValuesScan, } impl<'de> serde::Deserialize<'de> for GeneratedField { @@ -18082,7 +18076,6 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { "memoryScan" | "memory_scan" => Ok(GeneratedField::MemoryScan), "asyncFunc" | "async_func" => Ok(GeneratedField::AsyncFunc), "buffer" => Ok(GeneratedField::Buffer), - "transformPlan" | "transform_plan" => Ok(GeneratedField::TransformPlan), "valuesScan" | "values_scan" => Ok(GeneratedField::ValuesScan), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } @@ -18356,13 +18349,6 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { return Err(serde::de::Error::duplicate_field("buffer")); } physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::Buffer) -; - } - GeneratedField::TransformPlan => { - if physical_plan_type__.is_some() { - return Err(serde::de::Error::duplicate_field("transformPlan")); - } - physical_plan_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_plan_node::PhysicalPlanType::TransformPlan) ; } GeneratedField::ValuesScan => { @@ -20942,77 +20928,6 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { deserializer.deserialize_struct("datafusion.RepartitionNode", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for ResolvePlaceholdersRule { - #[allow(deprecated)] - fn serialize(&self, serializer: S) -> std::result::Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let len = 0; - let struct_ser = serializer.serialize_struct("datafusion.ResolvePlaceholdersRule", len)?; - struct_ser.end() - } -} -impl<'de> serde::Deserialize<'de> for ResolvePlaceholdersRule { - #[allow(deprecated)] - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - const FIELDS: &[&str] = &[ - ]; - - #[allow(clippy::enum_variant_names)] - enum GeneratedField { - } - impl<'de> serde::Deserialize<'de> for GeneratedField { - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - struct GeneratedVisitor; - - impl serde::de::Visitor<'_> for GeneratedVisitor { - type Value = GeneratedField; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(formatter, "expected one of: {:?}", &FIELDS) - } - - #[allow(unused_variables)] - fn visit_str(self, value: &str) -> std::result::Result - where - E: serde::de::Error, - { - Err(serde::de::Error::unknown_field(value, FIELDS)) - } - } - deserializer.deserialize_identifier(GeneratedVisitor) - } - } - struct GeneratedVisitor; - impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = ResolvePlaceholdersRule; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.ResolvePlaceholdersRule") - } - - fn visit_map(self, mut map_: V) -> std::result::Result - where - V: serde::de::MapAccess<'de>, - { - while map_.next_key::()?.is_some() { - let _ = map_.next_value::()?; - } - Ok(ResolvePlaceholdersRule { - }) - } - } - deserializer.deserialize_struct("datafusion.ResolvePlaceholdersRule", FIELDS, GeneratedVisitor) - } -} impl serde::Serialize for RollupNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -23141,225 +23056,6 @@ impl<'de> serde::Deserialize<'de> for TableReference { deserializer.deserialize_struct("datafusion.TableReference", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for TransformPlanExecNode { - #[allow(deprecated)] - fn serialize(&self, serializer: S) -> std::result::Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let mut len = 0; - if self.input.is_some() { - len += 1; - } - if !self.rules.is_empty() { - len += 1; - } - let mut struct_ser = serializer.serialize_struct("datafusion.TransformPlanExecNode", len)?; - if let Some(v) = self.input.as_ref() { - struct_ser.serialize_field("input", v)?; - } - if !self.rules.is_empty() { - struct_ser.serialize_field("rules", &self.rules)?; - } - struct_ser.end() - } -} -impl<'de> serde::Deserialize<'de> for TransformPlanExecNode { - #[allow(deprecated)] - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - const FIELDS: &[&str] = &[ - "input", - "rules", - ]; - - #[allow(clippy::enum_variant_names)] - enum GeneratedField { - Input, - Rules, - } - impl<'de> serde::Deserialize<'de> for GeneratedField { - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - struct GeneratedVisitor; - - impl serde::de::Visitor<'_> for GeneratedVisitor { - type Value = GeneratedField; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(formatter, "expected one of: {:?}", &FIELDS) - } - - #[allow(unused_variables)] - fn visit_str(self, value: &str) -> std::result::Result - where - E: serde::de::Error, - { - match value { - "input" => Ok(GeneratedField::Input), - "rules" => Ok(GeneratedField::Rules), - _ => Err(serde::de::Error::unknown_field(value, FIELDS)), - } - } - } - deserializer.deserialize_identifier(GeneratedVisitor) - } - } - struct GeneratedVisitor; - impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = TransformPlanExecNode; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.TransformPlanExecNode") - } - - fn visit_map(self, mut map_: V) -> std::result::Result - where - V: serde::de::MapAccess<'de>, - { - let mut input__ = None; - let mut rules__ = None; - while let Some(k) = map_.next_key()? { - match k { - GeneratedField::Input => { - if input__.is_some() { - return Err(serde::de::Error::duplicate_field("input")); - } - input__ = map_.next_value()?; - } - GeneratedField::Rules => { - if rules__.is_some() { - return Err(serde::de::Error::duplicate_field("rules")); - } - rules__ = Some(map_.next_value()?); - } - } - } - Ok(TransformPlanExecNode { - input: input__, - rules: rules__.unwrap_or_default(), - }) - } - } - deserializer.deserialize_struct("datafusion.TransformPlanExecNode", FIELDS, GeneratedVisitor) - } -} -impl serde::Serialize for TransformationRule { - #[allow(deprecated)] - fn serialize(&self, serializer: S) -> std::result::Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let mut len = 0; - if self.rule_type.is_some() { - len += 1; - } - let mut struct_ser = serializer.serialize_struct("datafusion.TransformationRule", len)?; - if let Some(v) = self.rule_type.as_ref() { - match v { - transformation_rule::RuleType::Extension(v) => { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("extension", pbjson::private::base64::encode(&v).as_str())?; - } - transformation_rule::RuleType::ResolvePlaceholders(v) => { - struct_ser.serialize_field("resolvePlaceholders", v)?; - } - } - } - struct_ser.end() - } -} -impl<'de> serde::Deserialize<'de> for TransformationRule { - #[allow(deprecated)] - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - const FIELDS: &[&str] = &[ - "extension", - "resolve_placeholders", - "resolvePlaceholders", - ]; - - #[allow(clippy::enum_variant_names)] - enum GeneratedField { - Extension, - ResolvePlaceholders, - } - impl<'de> serde::Deserialize<'de> for GeneratedField { - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - struct GeneratedVisitor; - - impl serde::de::Visitor<'_> for GeneratedVisitor { - type Value = GeneratedField; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(formatter, "expected one of: {:?}", &FIELDS) - } - - #[allow(unused_variables)] - fn visit_str(self, value: &str) -> std::result::Result - where - E: serde::de::Error, - { - match value { - "extension" => Ok(GeneratedField::Extension), - "resolvePlaceholders" | "resolve_placeholders" => Ok(GeneratedField::ResolvePlaceholders), - _ => Err(serde::de::Error::unknown_field(value, FIELDS)), - } - } - } - deserializer.deserialize_identifier(GeneratedVisitor) - } - } - struct GeneratedVisitor; - impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = TransformationRule; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.TransformationRule") - } - - fn visit_map(self, mut map_: V) -> std::result::Result - where - V: serde::de::MapAccess<'de>, - { - let mut rule_type__ = None; - while let Some(k) = map_.next_key()? { - match k { - GeneratedField::Extension => { - if rule_type__.is_some() { - return Err(serde::de::Error::duplicate_field("extension")); - } - rule_type__ = map_.next_value::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| transformation_rule::RuleType::Extension(x.0)); - } - GeneratedField::ResolvePlaceholders => { - if rule_type__.is_some() { - return Err(serde::de::Error::duplicate_field("resolvePlaceholders")); - } - rule_type__ = map_.next_value::<::std::option::Option<_>>()?.map(transformation_rule::RuleType::ResolvePlaceholders) -; - } - } - } - Ok(TransformationRule { - rule_type: rule_type__, - }) - } - } - deserializer.deserialize_struct("datafusion.TransformationRule", FIELDS, GeneratedVisitor) - } -} impl serde::Serialize for TryCastNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 755d04c68111f..dd7bbcc2c9163 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -1079,7 +1079,7 @@ pub mod table_reference { pub struct PhysicalPlanNode { #[prost( oneof = "physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38" )] pub physical_plan_type: ::core::option::Option, } @@ -1162,8 +1162,6 @@ pub mod physical_plan_node { #[prost(message, tag = "37")] Buffer(::prost::alloc::boxed::Box), #[prost(message, tag = "38")] - TransformPlan(::prost::alloc::boxed::Box), - #[prost(message, tag = "39")] ValuesScan(super::ValuesExecNode), } } @@ -2184,30 +2182,6 @@ pub struct BufferExecNode { #[prost(uint64, tag = "2")] pub capacity: u64, } -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct TransformPlanExecNode { - #[prost(message, optional, boxed, tag = "1")] - pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - #[prost(message, repeated, tag = "2")] - pub rules: ::prost::alloc::vec::Vec, -} -#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] -pub struct TransformationRule { - #[prost(oneof = "transformation_rule::RuleType", tags = "1, 2")] - pub rule_type: ::core::option::Option, -} -/// Nested message and enum types in `TransformationRule`. -pub mod transformation_rule { - #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] - pub enum RuleType { - #[prost(bytes, tag = "1")] - Extension(::prost::alloc::vec::Vec), - #[prost(message, tag = "2")] - ResolvePlaceholders(super::ResolvePlaceholdersRule), - } -} -#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] -pub struct ResolvePlaceholdersRule {} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum WindowFrameUnits { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 4d0b3645d5025..55646f937cf1d 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -76,9 +76,6 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; use datafusion_physical_plan::metrics::MetricType; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; -use datafusion_physical_plan::plan_transformer::{ - ExecutionTransformationRule, ResolvePlaceholdersRule, TransformPlanExec, -}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; @@ -106,7 +103,6 @@ use crate::physical_plan::to_proto::{ use crate::protobuf::physical_aggregate_expr_node::AggregateFunction; use crate::protobuf::physical_expr_node::ExprType; use crate::protobuf::physical_plan_node::PhysicalPlanType; -use crate::protobuf::transformation_rule::RuleType; use crate::protobuf::{ self, ListUnnest as ProtoListUnnest, SortExprNode, SortMergeJoinExecNode, proto_error, window_agg_exec_node, @@ -317,13 +313,6 @@ impl protobuf::PhysicalPlanNode { PhysicalPlanType::Buffer(buffer) => { self.try_into_buffer_physical_plan(buffer, ctx, codec, proto_converter) } - PhysicalPlanType::TransformPlan(transform_plan) => self - .try_into_transform_plan_exec( - transform_plan, - ctx, - codec, - proto_converter, - ), PhysicalPlanType::ValuesScan(scan) => { self.try_into_values_scan_exec(scan, ctx, codec, proto_converter) } @@ -575,14 +564,6 @@ impl protobuf::PhysicalPlanNode { ); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_transform_plan_exec( - exec, - codec, - proto_converter, - ); - } - let mut buf: Vec = vec![]; match codec.try_encode(Arc::clone(&plan_clone), &mut buf) { Ok(_) => { @@ -2221,37 +2202,6 @@ impl protobuf::PhysicalPlanNode { Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) } - fn try_into_transform_plan_exec( - &self, - transform_plan: &protobuf::TransformPlanExecNode, - ctx: &TaskContext, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&transform_plan.input, ctx, codec, proto_converter)?; - - let mut rules: Vec> = - Vec::with_capacity(transform_plan.rules.len()); - - for rule in transform_plan.rules.iter() { - match &rule.rule_type { - Some(RuleType::ResolvePlaceholders(_)) => { - rules.push(Box::new(ResolvePlaceholdersRule::new())) - } - Some(RuleType::Extension(ext)) => { - rules.push(codec.try_decode_transformation_rule(ext)?) - } - None => { - return internal_err!("Missing rule_type in TransformationRule"); - } - } - } - - let transformer = TransformPlanExec::try_new(input, rules)?; - Ok(Arc::new(transformer)) - } - fn try_into_values_scan_exec( &self, scan: &protobuf::ValuesExecNode, @@ -3264,7 +3214,7 @@ impl protobuf::PhysicalPlanNode { right: Some(Box::new(right)), join_type: join_type.into(), filter, - projection: exec.projection().map_or_else(Vec::new, |v| { + projection: exec.projection().as_ref().map_or_else(Vec::new, |v| { v.iter().map(|x| *x as u32).collect::>() }), }, @@ -3667,42 +3617,6 @@ impl protobuf::PhysicalPlanNode { }) } - fn try_from_transform_plan_exec( - exec: &TransformPlanExec, - extension_codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - extension_codec, - proto_converter, - )?; - - let mut rules = Vec::with_capacity(exec.rules().len()); - for rule in exec.rules() { - let rule_type = if rule.as_any().is::() { - Some(RuleType::ResolvePlaceholders( - protobuf::ResolvePlaceholdersRule {}, - )) - } else { - let mut buf = vec![]; - extension_codec - .try_encode_transformation_rule(rule.as_ref(), &mut buf)?; - Some(RuleType::Extension(buf)) - }; - rules.push(protobuf::TransformationRule { rule_type }); - } - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::TransformPlan(Box::new( - protobuf::TransformPlanExecNode { - input: Some(Box::new(input)), - rules, - }, - ))), - }) - } - fn try_from_values_source( source: &ValuesSource, extension_codec: &dyn PhysicalExtensionCodec, @@ -3796,21 +3710,6 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync { Ok(()) } - fn try_decode_transformation_rule( - &self, - _buf: &[u8], - ) -> Result> { - not_impl_err!("PhysicalExtensionCodec is not provided") - } - - fn try_encode_transformation_rule( - &self, - _node: &dyn ExecutionTransformationRule, - _buf: &mut Vec, - ) -> Result<()> { - not_impl_err!("PhysicalExtensionCodec is not provided") - } - fn try_decode_udwf(&self, name: &str, _buf: &[u8]) -> Result> { not_impl_err!("PhysicalExtensionCodec is not provided for window function {name}") } @@ -4244,25 +4143,6 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { self.encode_protobuf(buf, |codec, data| codec.try_encode_udaf(node, data)) } - - fn try_decode_transformation_rule( - &self, - buf: &[u8], - ) -> Result> { - self.decode_protobuf(buf, |codec, data| { - codec.try_decode_transformation_rule(data) - }) - } - - fn try_encode_transformation_rule( - &self, - node: &dyn ExecutionTransformationRule, - buf: &mut Vec, - ) -> Result<()> { - self.encode_protobuf(buf, |codec, data| { - codec.try_encode_transformation_rule(node, data) - }) - } } fn into_physical_plan( diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 0f911f77ea2e0..7938c5af66b38 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -227,7 +227,6 @@ initial_physical_plan_with_schema DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true -physical_plan after PhysicalExprResolver SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -244,7 +243,6 @@ physical_plan after LimitPushdown SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after PhysicalExprResolver(Post) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true @@ -307,7 +305,6 @@ physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] 02)--GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] -physical_plan after PhysicalExprResolver SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -326,7 +323,6 @@ physical_plan after LimitPushdown DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after PhysicalExprResolver(Post) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] @@ -353,7 +349,6 @@ physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--GlobalLimitExec: skip=0, fetch=10 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet -physical_plan after PhysicalExprResolver SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -372,7 +367,6 @@ physical_plan after LimitPushdown DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after PhysicalExprResolver(Post) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet @@ -431,7 +425,7 @@ query TT explain select a from t1 where exists (select count(*) from t2); ---- logical_plan -01)LeftSemi Join: +01)LeftSemi Join: 02)--TableScan: t1 projection=[a] 03)--SubqueryAlias: __correlated_sq_1 04)----EmptyRelation: rows=1 @@ -594,7 +588,6 @@ initial_physical_plan_with_schema DataSourceExec: file_groups={1 group: [[WORKSP physical_plan after OutputRequirements 01)OutputRequirementExec: order_by=[], dist_by=Unspecified 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true -physical_plan after PhysicalExprResolver SAME TEXT AS ABOVE physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE @@ -611,7 +604,6 @@ physical_plan after LimitPushdown SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after PushdownSort SAME TEXT AS ABOVE physical_plan after EnsureCooperative SAME TEXT AS ABOVE -physical_plan after PhysicalExprResolver(Post) SAME TEXT AS ABOVE physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true diff --git a/datafusion/sqllogictest/test_files/placeholders.slt b/datafusion/sqllogictest/test_files/placeholders.slt index 3dfa7eb572dbc..2ff1837206e2f 100644 --- a/datafusion/sqllogictest/test_files/placeholders.slt +++ b/datafusion/sqllogictest/test_files/placeholders.slt @@ -17,15 +17,6 @@ ########## ## Test physical plans with placeholders. -## -## These tests verify that the physical planner correctly produces a -## `TransformPlanExec` node with the expected number of plans with -## placeholders. -## -## A test is considered successful when the physical plan contains -## `TransformPlanExec`, indicating that the plan can be traversed to find -## placeholders and successfully replace them with actual parameter values -## during execution. ########## statement ok @@ -46,10 +37,9 @@ logical_plan 02)--Projection: column1 AS id, column2 AS name 03)----Values: ($1, Utf8View("Samanta") AS Utf8("Samanta")), (Int32(5) AS Int64(5), $2) physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--DataSinkExec: sink=MemoryTable (partitions=1) -03)----ProjectionExec: expr=[column1@0 as id, column2@1 as name] -04)------DataSourceExec: placeholders=2 +01)DataSinkExec: sink=MemoryTable (partitions=1) +02)--ProjectionExec: expr=[column1@0 as id, column2@1 as name] +03)----DataSourceExec: placeholders=2 # Filter with multiple placeholders query TT @@ -60,9 +50,8 @@ logical_plan 02)--Filter: t1.name = $1 OR t1.id = $2 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: name@1 = $1 OR id@0 = $2, projection=[id@0] -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: name@1 = $1 OR id@0 = $2, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] # Projection with placeholder query TT @@ -72,9 +61,8 @@ logical_plan 01)Projection: t1.id + $1 02)--TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 + $1 as t1.id + $1] -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 + $1 as t1.id + $1] +02)--DataSourceExec: partitions=1, partition_sizes=[1] # Projection and filter with placeholders query TT @@ -85,11 +73,10 @@ logical_plan 02)--Filter: t1.name = $2 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--ProjectionExec: expr=[id@0 + $1 as t1.id + $1] -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)------FilterExec: name@1 = $2 -05)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 + $1 as t1.id + $1] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----FilterExec: name@1 = $2 +04)------DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1 @@ -112,13 +99,12 @@ logical_plan 01)Aggregate: groupBy=[[]], aggr=[[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]]] 02)--TableScan: agg_order projection=[c1, c2, c3] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--AggregateExec: mode=Final, gby=[], aggr=[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] -03)----CoalescePartitionsExec -04)------AggregateExec: mode=Partial, gby=[], aggr=[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] -05)--------SortExec: expr=[c2@1 DESC, c3@2 ASC NULLS LAST], preserve_partitioning=[true] -06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], file_type=csv, has_header=true +01)AggregateExec: mode=Final, gby=[], aggr=[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=Partial, gby=[], aggr=[array_agg(agg_order.c1 + $1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] +04)------SortExec: expr=[c2@1 DESC, c3@2 ASC NULLS LAST], preserve_partitioning=[true] +05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], file_type=csv, has_header=true statement ok DROP TABLE agg_order @@ -140,9 +126,7 @@ logical_plan 01)Projection: alltypes_plain.smallint_col 02)--Filter: alltypes_plain.int_col = $1 03)----TableScan: alltypes_plain projection=[smallint_col, int_col], partial_filters=[alltypes_plain.int_col = $1] -physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col], file_type=parquet, predicate=int_col@4 = $1, pruning_predicate=int_col_null_count@2 != row_count@3 AND int_col_min@0 <= $1 AND $1 <= int_col_max@1, required_guarantees=[] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col], file_type=parquet, predicate=int_col@4 = $1, pruning_predicate=int_col_null_count@2 != row_count@3 AND int_col_min@0 <= $1 AND $1 <= int_col_max@1, required_guarantees=[] # Projection with placeholder on parquet table query TT @@ -151,9 +135,7 @@ EXPLAIN SELECT smallint_col + $1 FROM alltypes_plain; logical_plan 01)Projection: alltypes_plain.smallint_col + $1 02)--TableScan: alltypes_plain projection=[smallint_col] -physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col@3 + $1 as alltypes_plain.smallint_col + $1], file_type=parquet +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col@3 + $1 as alltypes_plain.smallint_col + $1], file_type=parquet # Projection and filter with placeholders on parquet table query TT @@ -164,10 +146,9 @@ logical_plan 02)--Filter: alltypes_plain.int_col = $2 03)----TableScan: alltypes_plain projection=[smallint_col, int_col], partial_filters=[alltypes_plain.int_col = $2] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--ProjectionExec: expr=[smallint_col@0 + $1 as alltypes_plain.smallint_col + $1] -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col, int_col], file_type=parquet, predicate=int_col@4 = $2, pruning_predicate=int_col_null_count@2 != row_count@3 AND int_col_min@0 <= $2 AND $2 <= int_col_max@1, required_guarantees=[] +01)ProjectionExec: expr=[smallint_col@0 + $1 as alltypes_plain.smallint_col + $1] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[smallint_col, int_col], file_type=parquet, predicate=int_col@4 = $2, pruning_predicate=int_col_null_count@2 != row_count@3 AND int_col_min@0 <= $2 AND $2 <= int_col_max@1, required_guarantees=[] statement ok DROP TABLE alltypes_plain; @@ -192,12 +173,11 @@ logical_plan 03)----TableScan: t1 projection=[id, name] 04)----TableScan: t2 projection=[id, age] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[name@1 as name, age@0 as age] -03)----HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, t1.id + $1@2)], projection=[age@1, name@3] -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)------ProjectionExec: expr=[id@0 as id, name@1 as name, id@0 + $1 as t1.id + $1] -06)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[name@1 as name, age@0 as age] +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, t1.id + $1@2)], projection=[age@1, name@3] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)----ProjectionExec: expr=[id@0 as id, name@1 as name, id@0 + $1 as t1.id + $1] +05)------DataSourceExec: partitions=1, partition_sizes=[1] # Join with placeholder in filter query TT @@ -210,12 +190,11 @@ logical_plan 04)----Filter: t2.age > $1 05)------TableScan: t2 projection=[id, age] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[name@1 as name, age@0 as age] -03)----HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], projection=[age@1, name@3] -04)------FilterExec: age@1 > $1 -05)--------DataSourceExec: partitions=1, partition_sizes=[1] -06)------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[name@1 as name, age@0 as age] +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], projection=[age@1, name@3] +03)----FilterExec: age@1 > $1 +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] # Join with placeholder in projection query TT @@ -227,13 +206,12 @@ logical_plan 03)----TableScan: t1 projection=[id, name] 04)----TableScan: t2 projection=[id, age] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[name@1 as name, age@3 + $1 as t2.age + $1] -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)------ProjectionExec: expr=[id@2 as id, name@3 as name, id@0 as id, age@1 as age] -05)--------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)] -06)----------DataSourceExec: partitions=1, partition_sizes=[1] -07)----------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[name@1 as name, age@3 + $1 as t2.age + $1] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----ProjectionExec: expr=[id@2 as id, name@3 as name, id@0 as id, age@1 as age] +04)------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)--------DataSourceExec: partitions=1, partition_sizes=[1] # Join with placeholder in ON statement query TT @@ -245,11 +223,10 @@ logical_plan 03)----TableScan: t1 projection=[id, name] 04)----TableScan: t2 projection=[id, age] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[name@1 as name, age@0 as age] -03)----NestedLoopJoinExec: join_type=Inner, filter=id@0 + id@1 = $1, projection=[age@1, name@3] -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[name@1 as name, age@0 as age] +02)--NestedLoopJoinExec: join_type=Inner, filter=id@0 + id@1 = $1, projection=[age@1, name@3] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)----DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; @@ -270,11 +247,10 @@ logical_plan 02)--WindowAggr: windowExpr=[[sum(CAST(t1.id AS Int64)) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 as id, sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 + $1 as sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + $1] -03)----BoundedWindowAggExec: wdw=[sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -04)------SortExec: expr=[name@1 ASC NULLS LAST, id@0 ASC NULLS LAST], preserve_partitioning=[false] -05)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 as id, sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 + $1 as sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + $1] +02)--BoundedWindowAggExec: wdw=[sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t1.id) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[name@1 ASC NULLS LAST, id@0 ASC NULLS LAST], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] # Window function with placeholder # Here we only resolve BoundedWindowAggExec. @@ -286,11 +262,10 @@ logical_plan 02)--WindowAggr: windowExpr=[[sum(CAST(t1.id + $1 AS Int64)) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 as id, sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] -03)----BoundedWindowAggExec: wdw=[sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -04)------SortExec: expr=[name@1 ASC NULLS LAST, id@0 ASC NULLS LAST], preserve_partitioning=[false] -05)--------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 as id, sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(t1.id + $1) PARTITION BY [t1.name] ORDER BY [t1.id ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[name@1 ASC NULLS LAST, id@0 ASC NULLS LAST], preserve_partitioning=[false] +04)------DataSourceExec: partitions=1, partition_sizes=[1] # UNION with placeholder query TT @@ -303,12 +278,11 @@ logical_plan 04)--Filter: t1.id = $2 05)----TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--UnionExec -03)----FilterExec: id@0 = $1 -04)------DataSourceExec: partitions=1, partition_sizes=[1] -05)----FilterExec: id@0 = $2 -06)------DataSourceExec: partitions=1, partition_sizes=[1] +01)UnionExec +02)--FilterExec: id@0 = $1 +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--FilterExec: id@0 = $2 +05)----DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; @@ -333,12 +307,11 @@ logical_plan 06)--------Filter: t1.id = $1 07)----------TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@1 as id, name@0 as name] -03)----NestedLoopJoinExec: join_type=Right -04)------FilterExec: id@0 = $1, projection=[name@1] -05)--------DataSourceExec: partitions=1, partition_sizes=[1] -06)------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@1 as id, name@0 as name] +02)--NestedLoopJoinExec: join_type=Right +03)----FilterExec: id@0 = $1, projection=[name@1] +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)----DataSourceExec: partitions=1, partition_sizes=[1] # CTE with placeholder query TT @@ -349,9 +322,8 @@ logical_plan 02)--Filter: t1.id = $1 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: id@0 = $1 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: id@0 = $1 +02)--DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; @@ -372,12 +344,11 @@ logical_plan 02)--Aggregate: groupBy=[[t1.id]], aggr=[[count(Int64(1))]] 03)----TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 + $1 as t1.id + $1, count(Int64(1))@1 as count(*)] -03)----AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))] -04)------RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 -05)--------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[count(Int64(1))] -06)----------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 + $1 as t1.id + $1, count(Int64(1))@1 as count(*)] +02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))] +03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 +04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[count(Int64(1))] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] # Group by with placeholder in HAVING query TT @@ -389,13 +360,12 @@ logical_plan 03)----Aggregate: groupBy=[[t1.id]], aggr=[[count(Int64(1))]] 04)------TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[id@0 as id, count(Int64(1))@1 as count(*)] -03)----FilterExec: count(Int64(1))@1 > $1 -04)------AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))] -05)--------RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 -06)----------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[count(Int64(1))] -07)------------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@0 as id, count(Int64(1))@1 as count(*)] +02)--FilterExec: count(Int64(1))@1 > $1 +03)----AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))] +04)------RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 +05)--------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[count(Int64(1))] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] # Order by with placeholder query TT @@ -405,9 +375,8 @@ logical_plan 01)Sort: t1.id + CAST($1 AS Int32) ASC NULLS LAST 02)--TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--SortExec: expr=[id@0 + $1 ASC NULLS LAST], preserve_partitioning=[false] -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)SortExec: expr=[id@0 + $1 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: partitions=1, partition_sizes=[1] # Group by and Order by with placeholders query TT @@ -418,13 +387,12 @@ logical_plan 02)--Aggregate: groupBy=[[t1.name]], aggr=[[sum(CAST(t1.id AS Int64))]] 03)----TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=2] -02)--SortPreservingMergeExec: [sum(t1.id)@1 + $1 ASC NULLS LAST] -03)----SortExec: expr=[sum(t1.id)@1 + $1 ASC NULLS LAST], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[name@0 as name], aggr=[sum(t1.id)] -05)--------RepartitionExec: partitioning=Hash([name@0], 4), input_partitions=1 -06)----------AggregateExec: mode=Partial, gby=[name@1 as name], aggr=[sum(t1.id)] -07)------------DataSourceExec: partitions=1, partition_sizes=[1] +01)SortPreservingMergeExec: [sum(t1.id)@1 + $1 ASC NULLS LAST] +02)--SortExec: expr=[sum(t1.id)@1 + $1 ASC NULLS LAST], preserve_partitioning=[true] +03)----AggregateExec: mode=FinalPartitioned, gby=[name@0 as name], aggr=[sum(t1.id)] +04)------RepartitionExec: partitioning=Hash([name@0], 4), input_partitions=1 +05)--------AggregateExec: mode=Partial, gby=[name@1 as name], aggr=[sum(t1.id)] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; @@ -441,9 +409,8 @@ logical_plan 01)Projection: CAST($1 AS Int32) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[$1] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[$1] +02)--PlaceholderRowExec # CAST with placeholder query TT @@ -453,9 +420,8 @@ logical_plan 01)Projection: CAST($1 AS Int32) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[$1] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[$1] +02)--PlaceholderRowExec # TRY_CAST with placeholder query TT @@ -465,9 +431,8 @@ logical_plan 01)Projection: TRY_CAST($1 AS Int32) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[TRY_CAST($1 AS Int32) as $1] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[TRY_CAST($1 AS Int32) as $1] +02)--PlaceholderRowExec ########## ## IN and BETWEEN with placeholders @@ -484,9 +449,8 @@ logical_plan 01)Filter: t1.id = $1 OR t1.id = $2 OR t1.id = $3 02)--TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: id@0 = $1 OR id@0 = $2 OR id@0 = $3 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: id@0 = $1 OR id@0 = $2 OR id@0 = $3 +02)--DataSourceExec: partitions=1, partition_sizes=[1] # BETWEEN with placeholders query TT @@ -496,9 +460,8 @@ logical_plan 01)Filter: t1.id >= $1 AND t1.id <= $2 02)--TableScan: t1 projection=[id] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: id@0 >= $1 AND id@0 <= $2 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: id@0 >= $1 AND id@0 <= $2 +02)--DataSourceExec: partitions=1, partition_sizes=[1] ########## ## String and Arithmetic operations with placeholders @@ -512,9 +475,8 @@ logical_plan 01)Projection: CAST($1 AS Utf8View) || CAST($2 AS Utf8View) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[$1 || $2] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[$1 || $2] +02)--PlaceholderRowExec # Arithmetic with placeholders query TT @@ -524,9 +486,8 @@ logical_plan 01)Projection: $1 + CAST($2 AS Int64) * CAST($3 AS Int64) 02)--EmptyRelation: rows=1 physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--ProjectionExec: expr=[$1 + CAST($2 AS Int64) * CAST($3 AS Int64) as $1 + $2 * $3] -03)----PlaceholderRowExec +01)ProjectionExec: expr=[$1 + CAST($2 AS Int64) * CAST($3 AS Int64) as $1 + $2 * $3] +02)--PlaceholderRowExec ########## ## LIKE and Regex with placeholders @@ -540,9 +501,8 @@ logical_plan 01)Filter: t1.name LIKE $1 02)--TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: name@1 LIKE $1 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: name@1 LIKE $1 +02)--DataSourceExec: partitions=1, partition_sizes=[1] # Regex with placeholder query TT @@ -552,9 +512,8 @@ logical_plan 01)Filter: t1.name ~ $1 02)--TableScan: t1 projection=[id, name] physical_plan -01)TransformPlanExec: rules=[ResolvePlaceholders: plans_to_modify=1] -02)--FilterExec: name@1 ~ $1 -03)----DataSourceExec: partitions=1, partition_sizes=[1] +01)FilterExec: name@1 ~ $1 +02)--DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE t1; diff --git a/docs/source/library-user-guide/custom-table-providers.md b/docs/source/library-user-guide/custom-table-providers.md index 8e1dee9e843ac..50005a7527da0 100644 --- a/docs/source/library-user-guide/custom-table-providers.md +++ b/docs/source/library-user-guide/custom-table-providers.md @@ -108,7 +108,7 @@ impl ExecutionPlan for CustomExec { } - fn properties(&self) -> &PlanProperties { + fn properties(&self) -> &Arc { unreachable!() } @@ -232,7 +232,7 @@ The `scan` method of the `TableProvider` returns a `Result &PlanProperties { +# fn properties(&self) -> &Arc { # unreachable!() # } # @@ -424,7 +424,7 @@ This will allow you to use the custom table provider in DataFusion. For example, # } # # -# fn properties(&self) -> &PlanProperties { +# fn properties(&self) -> &Arc { # unreachable!() # } # diff --git a/docs/source/library-user-guide/upgrading.md b/docs/source/library-user-guide/upgrading.md index 182f2f0ef9f92..f62d3d8f948ac 100644 --- a/docs/source/library-user-guide/upgrading.md +++ b/docs/source/library-user-guide/upgrading.md @@ -61,6 +61,61 @@ FileSinkConfig { } ``` +### `ExecutionPlan` properties method return type + +Now `ExecutionPlan::properties()` must return `&Arc` instead of a reference. This was done to enable the comparison of properties and to determine that they have not changed within the +`with_new_children` method. To migrate, in all `ExecutionPlan` implementations, you need to wrap stored `PlanProperties` in an `Arc`: + +```diff +- cache: PlanProperties, ++ cache: Arc, + +... + +- fn properties(&self) -> &PlanProperties { ++ fn properties(&self) -> &Arc { + &self.cache + } +``` + +Note: The optimization for `with_new_children` can be implemented for any `ExecutionPlan`. This can reduce planning time as well as the time for resetting plan states. +To support it, you can use the macro: `check_if_same_properties`. For it to work, you need to implement the function: `with_new_children_and_same_properties` with semantics +identical to `with_new_children`, but operating under the assumption that the properties of the children plans have not changed. + +An example of supporting this optimization for `ProjectionExec`: + +```diff + impl ProjectionExec { ++ fn with_new_children_and_same_properties( ++ &self, ++ mut children: Vec>, ++ ) -> Self { ++ Self { ++ input: children.swap_remove(0), ++ metrics: ExecutionPlanMetricsSet::new(), ++ ..Self::clone(self) ++ } ++ } + } + + impl ExecutionPlan for ProjectionExec { + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { ++ check_if_same_properties!(self, children); + ProjectionExec::try_new( + self.projector.projection().into_iter().cloned(), + children.swap_remove(0), + ) + .map(|p| Arc::new(p) as _) + } + } + +... + +``` + ### `SimplifyInfo` trait removed, `SimplifyContext` now uses builder-style API The `SimplifyInfo` trait has been removed and replaced with the concrete `SimplifyContext` struct. This simplifies the expression simplification API and removes the need for trait objects.