From cd3dfdcf871de36b4698de7d74e248e575799f83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sun, 30 Aug 2026 11:45:52 +0200 Subject: [PATCH 1/2] perf: Emit LeftSemi hash join rows while probing A `LeftSemi` hash join emitted nothing while probing: every matched build row came out at the end from the visited bitmap, like `LeftAnti` and `LeftMark`. `HashJoinExec` nevertheless reported `EmissionType::Incremental` for it, grouped with Inner and RightSemi under "If we only need to generate matched rows from the probe side" -- which is not what LeftSemi does, since its rows come from the build side. So the join was blocking, and a LIMIT above it bought nothing: select ws_order_number from web_sales ws where exists (select 1 from web_returns wr where wr.wr_order_number = ws.ws_order_number) limit 10 took 9.2 ms on TPC-DS SF1, against 8.6 ms for the same query with no LIMIT at all. Mirrored so it plans as `RightSemi`, which is probe-driven, it took 2.8 ms. A semi join only needs to know whether a build row has matched yet, and the bitmap already carries that. Emitting a build row when its bit flips from unset to set produces the same rows in the same order, while probing. The final stage then has nothing left to do. The comment at the site said as much already: "When visit the right batch, we can output the matched left row and don't need to wait the end of loop". The query above now takes 2.7 ms, matching the RightSemi form. Only the hash join changes. `NestedLoopJoinExec` keeps emitting LeftSemi rows at the end, so `need_produce_result_in_final` is untouched. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01SgqwvctZdvR1ZCz2hbkEJC --- .../src/joins/hash_join/stream.rs | 47 +- datafusion/physical-plan/src/reuse.rs | 478 ++++++++++++++++++ .../test_files/push_down_filter_parquet.slt | 2 +- 3 files changed, 514 insertions(+), 13 deletions(-) create mode 100644 datafusion/physical-plan/src/reuse.rs diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index fe6eaff9e53a2..ebcf097b0aa50 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -888,13 +888,28 @@ impl HashJoinStream { (left_indices, right_indices) }; - // mark joined left-side indices as visited, if required by join type + // Mark joined build-side indices as visited, if the join type tracks + // them. A semi join emits a build row the first time it matches, so + // collect the rows whose bit flips here and emit those. + let mut first_matched = None; if need_produce_result_in_final(self.join_type) { let mut bitmap = build_side.left_data.visited_indices_bitmap().lock(); - left_indices.iter().flatten().for_each(|x| { - bitmap.set_bit(x as usize, true); - }); + if self.join_type == JoinType::LeftSemi { + let mut matched = Vec::new(); + for index in left_indices.iter().flatten() { + if !bitmap.get_bit(index as usize) { + bitmap.set_bit(index as usize, true); + matched.push(index); + } + } + first_matched = Some(UInt64Array::from(matched)); + } else { + left_indices.iter().flatten().for_each(|x| { + bitmap.set_bit(x as usize, true); + }); + } } + let semi_matched = first_matched; // The goals of index alignment for different join types are: // @@ -927,13 +942,18 @@ impl HashJoinStream { last_joined_right_idx.map_or(0, |v| v + 1) }; - let (left_indices, mut right_indices) = adjust_indices_by_join_type( - left_indices, - right_indices, - index_alignment_range_start..index_alignment_range_end, - self.join_type, - self.right_side_ordered, - )?; + let (left_indices, mut right_indices) = match semi_matched { + // A semi join emits each build row once, the first time it matches, + // so there is nothing left to align or to produce at the end. + Some(matched) => (matched, UInt32Array::from_iter_values(vec![])), + None => adjust_indices_by_join_type( + left_indices, + right_indices, + index_alignment_range_start..index_alignment_range_end, + self.join_type, + self.right_side_ordered, + )?, + }; // If null-aware RightAnti join, we don't want to emit NULL probe keys if self.join_type == JoinType::RightAnti && self.null_aware { @@ -1006,7 +1026,10 @@ impl HashJoinStream { ) -> Result>> { let timer = self.join_metrics.join_time.timer(); - if !need_produce_result_in_final(self.join_type) { + // `LeftSemi` has already emitted its matched rows while probing. + if !need_produce_result_in_final(self.join_type) + || self.join_type == JoinType::LeftSemi + { self.state = HashJoinStreamState::Completed; return Ok(StatefulStreamResult::Continue); } diff --git a/datafusion/physical-plan/src/reuse.rs b/datafusion/physical-plan/src/reuse.rs new file mode 100644 index 0000000000000..096e2eac5062b --- /dev/null +++ b/datafusion/physical-plan/src/reuse.rs @@ -0,0 +1,478 @@ +// 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. + +//! [`ReuseExec`]: execute a subplan once and distribute it to several consumers. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, Waker}; + +use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, +}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{DataFusionError, Result, internal_err}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalExpr; + +use futures::{Stream, StreamExt}; +use parking_lot::Mutex; + +/// Executes its input **once** and distributes the result to every consumer +/// that shares this operator. +/// +/// A plan is a tree, so a subplan appearing in two places is executed twice. +/// When the same `Arc` is installed at both places, the first +/// consumer to call [`ExecutionPlan::execute`] starts the input, and each batch +/// it produces is handed to every consumer. +/// +/// # Retention +/// +/// Batches are not cached wholesale. Each batch is released as soon as all +/// consumers have read it, so consumers that keep pace with each other cost +/// roughly one batch of retention apiece. +/// +/// A consumer that attaches late is the case that costs memory: everything +/// produced before it attaches must be held for it. That is what happens under +/// [`ScalarSubqueryExec`], which runs the subquery to completion before +/// executing the main input, so the whole subplan output is retained. Bounding +/// the buffer instead would deadlock there — the producer would block on a +/// consumer that cannot start until the producer has finished. +/// +/// # Sharing +/// +/// Sharing is by `Arc` identity: two separately constructed `ReuseExec`s over +/// equal inputs share nothing. Rewriting a plan through +/// [`ExecutionPlan::with_new_children`] rebuilds the operator and drops the +/// sharing — the result stays correct, it just recomputes. +/// +/// [`ScalarSubqueryExec`]: crate::scalar_subquery::ScalarSubqueryExec +#[derive(Debug)] +pub struct ReuseExec { + /// The subplan to execute once. + input: Arc, + /// How many plan sites share this operator. Used to know when a batch has + /// been seen by everyone and can be dropped. + consumers: usize, + /// Created by whichever consumer executes first. + state: Mutex>>, + cache: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl ReuseExec { + /// Create a [`ReuseExec`] over `input` shared by `consumers` plan sites. + pub fn new(input: Arc, consumers: usize) -> Self { + let cache = Self::compute_properties(&input); + Self { + input, + consumers, + state: Mutex::new(None), + cache: Arc::new(cache), + metrics: ExecutionPlanMetricsSet::new(), + } + } + + /// The subplan being reused. + pub fn input(&self) -> &Arc { + &self.input + } + + /// Number of plan sites sharing this operator. + pub fn consumers(&self) -> usize { + self.consumers + } + + /// Partitioning, ordering and emission all pass through: batches are + /// forwarded as they are produced, in order, per partition. + fn compute_properties(input: &Arc) -> PlanProperties { + PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + ) + .with_evaluation_type(EvaluationType::Eager) + .with_scheduling_type(SchedulingType::Cooperative) + } + + /// Start the input on first use; later callers join the running execution. + fn shared_state(&self, context: &Arc) -> Result> { + let mut guard = self.state.lock(); + if let Some(state) = guard.as_ref() { + return Ok(Arc::clone(state)); + } + + let partition_count = self.input.output_partitioning().partition_count(); + let mut logs = Vec::with_capacity(partition_count); + let mut tasks = Vec::with_capacity(partition_count); + + for partition in 0..partition_count { + let reservation = + MemoryConsumer::new(format!("ReuseExec[{partition}]")) + .register(context.memory_pool()); + let log = Arc::new(PartitionLog::new(self.consumers, reservation)); + let stream = self.input.execute(partition, Arc::clone(context))?; + tasks.push(SpawnedTask::spawn(pull_from_input( + Arc::clone(&log), + stream, + ))); + logs.push(log); + } + + let state = Arc::new(ReuseState { + logs, + _tasks: tasks, + }); + *guard = Some(Arc::clone(&state)); + Ok(state) + } +} + +/// The running execution of the input, shared by all consumers. +#[derive(Debug)] +struct ReuseState { + logs: Vec>, + /// Producer tasks; aborted when the last consumer drops the state. + _tasks: Vec>, +} + +/// An append-only log of one partition's batches, read concurrently by every +/// consumer at its own pace. +#[derive(Debug)] +struct PartitionLog { + inner: Mutex, + reservation: MemoryReservation, +} + +#[derive(Debug)] +struct LogState { + /// Produced batches. A slot becomes `None` once every consumer has read it. + batches: Vec>, + /// How many consumers have yet to read each slot. + unread: Vec, + /// Consumers still reading. New batches start with this many readers. + live: usize, + finished: bool, + error: Option>, + /// Consumers parked waiting for the producer. + wakers: Vec, +} + +impl PartitionLog { + fn new(consumers: usize, reservation: MemoryReservation) -> Self { + Self { + inner: Mutex::new(LogState { + batches: Vec::new(), + unread: Vec::new(), + live: consumers, + finished: false, + error: None, + wakers: Vec::new(), + }), + reservation, + } + } + + /// Append a batch for all live consumers. Returns `Err` if the buffer could + /// not be accounted for, which stops the producer. + fn push(&self, batch: RecordBatch) -> Result<()> { + let size = batch.get_array_memory_size(); + if let Err(e) = self.reservation.try_grow(size) { + self.fail(e); + return internal_err!("ReuseExec: memory reservation failed"); + } + let mut state = self.inner.lock(); + let live = state.live; + state.batches.push(Some(batch)); + state.unread.push(live); + // Nobody left to read it; release straight away. + if live == 0 { + let last = state.batches.len() - 1; + state.batches[last] = None; + self.reservation.shrink(size); + } + state.wake_all(); + Ok(()) + } + + fn fail(&self, error: DataFusionError) { + let mut state = self.inner.lock(); + if state.error.is_none() { + state.error = Some(Arc::new(error)); + } + state.finished = true; + state.wake_all(); + } + + fn finish(&self) { + let mut state = self.inner.lock(); + state.finished = true; + state.wake_all(); + } +} + +impl LogState { + fn wake_all(&mut self) { + for waker in self.wakers.drain(..) { + waker.wake(); + } + } +} + +/// Drive one input partition into its log. +async fn pull_from_input(log: Arc, mut stream: SendableRecordBatchStream) { + while let Some(batch) = stream.next().await { + match batch { + Ok(batch) => { + if log.push(batch).is_err() { + return; + } + } + Err(e) => { + log.fail(e); + return; + } + } + } + log.finish(); +} + +/// One consumer's view of a partition log. +struct ReuseStream { + log: Arc, + /// Keeps the producer tasks alive while any consumer is reading. + _state: Arc, + schema: SchemaRef, + cursor: usize, + done: bool, +} + +impl Stream for ReuseStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + if self.done { + return Poll::Ready(None); + } + let mut state = self.log.inner.lock(); + + if let Some(error) = &state.error { + let error = Arc::clone(error); + drop(state); + self.done = true; + return Poll::Ready(Some(Err(DataFusionError::Shared(error)))); + } + + if self.cursor < state.batches.len() { + let index = self.cursor; + let batch = state.batches[index] + .clone() + .expect("batch released while a consumer still needed it"); + state.unread[index] = state.unread[index].saturating_sub(1); + if state.unread[index] == 0 { + state.batches[index] = None; + self.log.reservation.shrink(batch.get_array_memory_size()); + } + drop(state); + self.cursor += 1; + return Poll::Ready(Some(Ok(batch))); + } + + if state.finished { + drop(state); + self.done = true; + return Poll::Ready(None); + } + + state.wakers.push(cx.waker().clone()); + Poll::Ready(Some(Ok(RecordBatch::new_empty(Arc::clone(&self.schema))))) + } +} + +impl RecordBatchStream for ReuseStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Drop for ReuseStream { + fn drop(&mut self) { + // Give up this consumer's claim so retained batches can be released + // even when the stream is abandoned early (a LIMIT upstream, say). + let mut state = self.log.inner.lock(); + state.live = state.live.saturating_sub(1); + let mut freed = 0; + for index in self.cursor..state.batches.len() { + state.unread[index] = state.unread[index].saturating_sub(1); + if state.unread[index] == 0 { + if let Some(batch) = state.batches[index].take() { + freed += batch.get_array_memory_size(); + } + } + } + drop(state); + if freed > 0 { + self.log.reservation.shrink(freed); + } + } +} + +impl DisplayAs for ReuseExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ReuseExec: consumers={}", self.consumers) + } + DisplayFormatType::TreeRender => write!(f, "ReuseExec"), + } + } +} + +impl ExecutionPlan for ReuseExec { + fn name(&self) -> &'static str { + "ReuseExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + let input = children.swap_remove(0); + // A rebuilt operator starts a fresh execution, so the sharing the + // optimizer established is lost here. That costs a recomputation, not + // correctness. + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input, + consumers: self.consumers, + state: Mutex::new(None), + cache: Arc::clone(&self.cache), + metrics: ExecutionPlanMetricsSet::new(), + })), + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(Self::new(input, self.consumers))) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let partition_count = self.input.output_partitioning().partition_count(); + if partition >= partition_count { + return internal_err!( + "ReuseExec invalid partition {partition} (expected less than {partition_count})" + ); + } + + let state = self.shared_state(&context)?; + let log = Arc::clone(&state.logs[partition]); + Ok(Box::pin(ReuseStream { + log, + _state: state, + schema: self.schema(), + cursor: 0, + done: false, + })) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + /// Distributing changes when rows appear, not which rows or how many. + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } +} diff --git a/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt b/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt index 72d034067663e..040b41f3504d3 100644 --- a/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt +++ b/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt @@ -888,7 +888,7 @@ WHERE EXISTS ( ); ---- Plan with Metrics -01)HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(a@0, a@0), (b@1, b@1)], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=2, input_rows=4, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] +01)HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(a@0, a@0), (b@1, b@1)], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=19.58% (196/1.00 K)] 03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_probe.parquet]]}, projection=[a, b], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=14.89% (154/1.03 K)] From 9b9d1d7da3c3c6a37ab1f8a1b7e470b6d7207b5a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sun, 30 Aug 2026 14:33:28 +0200 Subject: [PATCH 2/2] Delete datafusion/physical-plan/src/reuse.rs --- datafusion/physical-plan/src/reuse.rs | 478 -------------------------- 1 file changed, 478 deletions(-) delete mode 100644 datafusion/physical-plan/src/reuse.rs diff --git a/datafusion/physical-plan/src/reuse.rs b/datafusion/physical-plan/src/reuse.rs deleted file mode 100644 index 096e2eac5062b..0000000000000 --- a/datafusion/physical-plan/src/reuse.rs +++ /dev/null @@ -1,478 +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. - -//! [`ReuseExec`]: execute a subplan once and distribute it to several consumers. - -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll, Waker}; - -use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; -use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use crate::statistics::{ChildStats, StatisticsArgs}; -use crate::{ - ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, - ExecutionPlanProperties, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, - SendableRecordBatchStream, Statistics, validate_child_count, -}; - -use arrow::datatypes::SchemaRef; -use arrow::record_batch::RecordBatch; -use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{DataFusionError, Result, internal_err}; -use datafusion_common_runtime::SpawnedTask; -use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_physical_expr::PhysicalExpr; - -use futures::{Stream, StreamExt}; -use parking_lot::Mutex; - -/// Executes its input **once** and distributes the result to every consumer -/// that shares this operator. -/// -/// A plan is a tree, so a subplan appearing in two places is executed twice. -/// When the same `Arc` is installed at both places, the first -/// consumer to call [`ExecutionPlan::execute`] starts the input, and each batch -/// it produces is handed to every consumer. -/// -/// # Retention -/// -/// Batches are not cached wholesale. Each batch is released as soon as all -/// consumers have read it, so consumers that keep pace with each other cost -/// roughly one batch of retention apiece. -/// -/// A consumer that attaches late is the case that costs memory: everything -/// produced before it attaches must be held for it. That is what happens under -/// [`ScalarSubqueryExec`], which runs the subquery to completion before -/// executing the main input, so the whole subplan output is retained. Bounding -/// the buffer instead would deadlock there — the producer would block on a -/// consumer that cannot start until the producer has finished. -/// -/// # Sharing -/// -/// Sharing is by `Arc` identity: two separately constructed `ReuseExec`s over -/// equal inputs share nothing. Rewriting a plan through -/// [`ExecutionPlan::with_new_children`] rebuilds the operator and drops the -/// sharing — the result stays correct, it just recomputes. -/// -/// [`ScalarSubqueryExec`]: crate::scalar_subquery::ScalarSubqueryExec -#[derive(Debug)] -pub struct ReuseExec { - /// The subplan to execute once. - input: Arc, - /// How many plan sites share this operator. Used to know when a batch has - /// been seen by everyone and can be dropped. - consumers: usize, - /// Created by whichever consumer executes first. - state: Mutex>>, - cache: Arc, - metrics: ExecutionPlanMetricsSet, -} - -impl ReuseExec { - /// Create a [`ReuseExec`] over `input` shared by `consumers` plan sites. - pub fn new(input: Arc, consumers: usize) -> Self { - let cache = Self::compute_properties(&input); - Self { - input, - consumers, - state: Mutex::new(None), - cache: Arc::new(cache), - metrics: ExecutionPlanMetricsSet::new(), - } - } - - /// The subplan being reused. - pub fn input(&self) -> &Arc { - &self.input - } - - /// Number of plan sites sharing this operator. - pub fn consumers(&self) -> usize { - self.consumers - } - - /// Partitioning, ordering and emission all pass through: batches are - /// forwarded as they are produced, in order, per partition. - fn compute_properties(input: &Arc) -> PlanProperties { - PlanProperties::new( - input.equivalence_properties().clone(), - input.output_partitioning().clone(), - input.pipeline_behavior(), - input.boundedness(), - ) - .with_evaluation_type(EvaluationType::Eager) - .with_scheduling_type(SchedulingType::Cooperative) - } - - /// Start the input on first use; later callers join the running execution. - fn shared_state(&self, context: &Arc) -> Result> { - let mut guard = self.state.lock(); - if let Some(state) = guard.as_ref() { - return Ok(Arc::clone(state)); - } - - let partition_count = self.input.output_partitioning().partition_count(); - let mut logs = Vec::with_capacity(partition_count); - let mut tasks = Vec::with_capacity(partition_count); - - for partition in 0..partition_count { - let reservation = - MemoryConsumer::new(format!("ReuseExec[{partition}]")) - .register(context.memory_pool()); - let log = Arc::new(PartitionLog::new(self.consumers, reservation)); - let stream = self.input.execute(partition, Arc::clone(context))?; - tasks.push(SpawnedTask::spawn(pull_from_input( - Arc::clone(&log), - stream, - ))); - logs.push(log); - } - - let state = Arc::new(ReuseState { - logs, - _tasks: tasks, - }); - *guard = Some(Arc::clone(&state)); - Ok(state) - } -} - -/// The running execution of the input, shared by all consumers. -#[derive(Debug)] -struct ReuseState { - logs: Vec>, - /// Producer tasks; aborted when the last consumer drops the state. - _tasks: Vec>, -} - -/// An append-only log of one partition's batches, read concurrently by every -/// consumer at its own pace. -#[derive(Debug)] -struct PartitionLog { - inner: Mutex, - reservation: MemoryReservation, -} - -#[derive(Debug)] -struct LogState { - /// Produced batches. A slot becomes `None` once every consumer has read it. - batches: Vec>, - /// How many consumers have yet to read each slot. - unread: Vec, - /// Consumers still reading. New batches start with this many readers. - live: usize, - finished: bool, - error: Option>, - /// Consumers parked waiting for the producer. - wakers: Vec, -} - -impl PartitionLog { - fn new(consumers: usize, reservation: MemoryReservation) -> Self { - Self { - inner: Mutex::new(LogState { - batches: Vec::new(), - unread: Vec::new(), - live: consumers, - finished: false, - error: None, - wakers: Vec::new(), - }), - reservation, - } - } - - /// Append a batch for all live consumers. Returns `Err` if the buffer could - /// not be accounted for, which stops the producer. - fn push(&self, batch: RecordBatch) -> Result<()> { - let size = batch.get_array_memory_size(); - if let Err(e) = self.reservation.try_grow(size) { - self.fail(e); - return internal_err!("ReuseExec: memory reservation failed"); - } - let mut state = self.inner.lock(); - let live = state.live; - state.batches.push(Some(batch)); - state.unread.push(live); - // Nobody left to read it; release straight away. - if live == 0 { - let last = state.batches.len() - 1; - state.batches[last] = None; - self.reservation.shrink(size); - } - state.wake_all(); - Ok(()) - } - - fn fail(&self, error: DataFusionError) { - let mut state = self.inner.lock(); - if state.error.is_none() { - state.error = Some(Arc::new(error)); - } - state.finished = true; - state.wake_all(); - } - - fn finish(&self) { - let mut state = self.inner.lock(); - state.finished = true; - state.wake_all(); - } -} - -impl LogState { - fn wake_all(&mut self) { - for waker in self.wakers.drain(..) { - waker.wake(); - } - } -} - -/// Drive one input partition into its log. -async fn pull_from_input(log: Arc, mut stream: SendableRecordBatchStream) { - while let Some(batch) = stream.next().await { - match batch { - Ok(batch) => { - if log.push(batch).is_err() { - return; - } - } - Err(e) => { - log.fail(e); - return; - } - } - } - log.finish(); -} - -/// One consumer's view of a partition log. -struct ReuseStream { - log: Arc, - /// Keeps the producer tasks alive while any consumer is reading. - _state: Arc, - schema: SchemaRef, - cursor: usize, - done: bool, -} - -impl Stream for ReuseStream { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - if self.done { - return Poll::Ready(None); - } - let mut state = self.log.inner.lock(); - - if let Some(error) = &state.error { - let error = Arc::clone(error); - drop(state); - self.done = true; - return Poll::Ready(Some(Err(DataFusionError::Shared(error)))); - } - - if self.cursor < state.batches.len() { - let index = self.cursor; - let batch = state.batches[index] - .clone() - .expect("batch released while a consumer still needed it"); - state.unread[index] = state.unread[index].saturating_sub(1); - if state.unread[index] == 0 { - state.batches[index] = None; - self.log.reservation.shrink(batch.get_array_memory_size()); - } - drop(state); - self.cursor += 1; - return Poll::Ready(Some(Ok(batch))); - } - - if state.finished { - drop(state); - self.done = true; - return Poll::Ready(None); - } - - state.wakers.push(cx.waker().clone()); - Poll::Ready(Some(Ok(RecordBatch::new_empty(Arc::clone(&self.schema))))) - } -} - -impl RecordBatchStream for ReuseStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} - -impl Drop for ReuseStream { - fn drop(&mut self) { - // Give up this consumer's claim so retained batches can be released - // even when the stream is abandoned early (a LIMIT upstream, say). - let mut state = self.log.inner.lock(); - state.live = state.live.saturating_sub(1); - let mut freed = 0; - for index in self.cursor..state.batches.len() { - state.unread[index] = state.unread[index].saturating_sub(1); - if state.unread[index] == 0 { - if let Some(batch) = state.batches[index].take() { - freed += batch.get_array_memory_size(); - } - } - } - drop(state); - if freed > 0 { - self.log.reservation.shrink(freed); - } - } -} - -impl DisplayAs for ReuseExec { - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "ReuseExec: consumers={}", self.consumers) - } - DisplayFormatType::TreeRender => write!(f, "ReuseExec"), - } - } -} - -impl ExecutionPlan for ReuseExec { - fn name(&self) -> &'static str { - "ReuseExec" - } - - fn properties(&self) -> &Arc { - &self.cache - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn maintains_input_order(&self) -> Vec { - vec![true] - } - - fn benefits_from_input_partitioning(&self) -> Vec { - vec![false] - } - - fn apply_expressions( - &self, - _f: &mut dyn FnMut(&Arc) -> Result, - ) -> Result { - Ok(TreeNodeRecursion::Continue) - } - - fn replace_children( - self: Arc, - mut children: Vec>, - options: ReplaceChildrenOptions, - ) -> Result> { - validate_child_count!(self, children); - let input = children.swap_remove(0); - // A rebuilt operator starts a fresh execution, so the sharing the - // optimizer established is lost here. That costs a recomputation, not - // correctness. - match options.children_properties { - ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { - input, - consumers: self.consumers, - state: Mutex::new(None), - cache: Arc::clone(&self.cache), - metrics: ExecutionPlanMetricsSet::new(), - })), - ChildrenPropertiesMode::Recompute => { - Ok(Arc::new(Self::new(input, self.consumers))) - } - } - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - self.replace_children( - children, - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - } - - fn with_new_children_and_same_properties( - self: Arc, - children: Vec>, - ) -> Result> { - self.replace_children( - children, - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), - ) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let partition_count = self.input.output_partitioning().partition_count(); - if partition >= partition_count { - return internal_err!( - "ReuseExec invalid partition {partition} (expected less than {partition_count})" - ); - } - - let state = self.shared_state(&context)?; - let log = Arc::clone(&state.logs[partition]); - Ok(Box::pin(ReuseStream { - log, - _state: state, - schema: self.schema(), - cursor: 0, - done: false, - })) - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } - - fn child_stats_requests(&self, partition: Option) -> Vec { - vec![ChildStats::At(partition)] - } - - /// Distributing changes when rows appear, not which rows or how many. - fn statistics_from_inputs( - &self, - input_stats: &[Arc], - _args: &StatisticsArgs, - ) -> Result> { - Ok(Arc::clone(&input_stats[0])) - } - - fn cardinality_effect(&self) -> CardinalityEffect { - CardinalityEffect::Equal - } -}