datafusion_physical_plan/
buffer.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`BufferExec`] decouples production and consumption on messages by buffering the input in the
19//! background up to a certain capacity.
20
21use crate::execution_plan::{
22    CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary,
23};
24use crate::filter_pushdown::{
25    ChildPushdownResult, FilterDescription, FilterPushdownPhase,
26    FilterPushdownPropagation,
27};
28use crate::projection::ProjectionExec;
29use crate::statistics::{ChildStats, StatisticsArgs};
30use crate::stream::RecordBatchStreamAdapter;
31use crate::{
32    ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
33    ReplaceChildrenOptions, SortOrderPushdownResult, validate_child_count,
34};
35use arrow::array::RecordBatch;
36use datafusion_common::config::ConfigOptions;
37use datafusion_common::tree_node::TreeNodeRecursion;
38use datafusion_common::{Result, Statistics, internal_err};
39use datafusion_common_runtime::SpawnedTask;
40use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
41use datafusion_execution::{SendableRecordBatchStream, TaskContext};
42use datafusion_physical_expr_common::metrics::{
43    ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricsSet,
44};
45use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
46use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
47use futures::{FutureExt, Stream, StreamExt, TryStreamExt};
48use pin_project_lite::pin_project;
49use std::fmt;
50use std::panic::AssertUnwindSafe;
51use std::pin::Pin;
52use std::sync::Arc;
53use std::sync::atomic::{AtomicUsize, Ordering};
54use std::task::{Context, Poll};
55use tokio::sync::mpsc::UnboundedReceiver;
56use tokio::sync::{OwnedSemaphorePermit, Semaphore};
57
58/// WARNING: EXPERIMENTAL
59///
60/// Decouples production and consumption of record batches with an internal queue per partition,
61/// eagerly filling up the capacity of the queues even before any message is requested.
62///
63/// ```text
64///             ┌───────────────────────────┐
65///             │        BufferExec         │
66///             │                           │
67///             │┌────── Partition 0 ──────┐│
68///             ││            ┌────┐ ┌────┐││       ┌────┐
69/// ──background poll────────▶│    │ │    ├┼┼───────▶    │
70///             ││            └────┘ └────┘││       └────┘
71///             │└─────────────────────────┘│
72///             │┌────── Partition 1 ──────┐│
73///             ││     ┌────┐ ┌────┐ ┌────┐││       ┌────┐
74/// ──background poll─▶│    │ │    │ │    ├┼┼───────▶    │
75///             ││     └────┘ └────┘ └────┘││       └────┘
76///             │└─────────────────────────┘│
77///             │                           │
78///             │           ...             │
79///             │                           │
80///             │┌────── Partition N ──────┐│
81///             ││                   ┌────┐││       ┌────┐
82/// ──background poll───────────────▶│    ├┼┼───────▶    │
83///             ││                   └────┘││       └────┘
84///             │└─────────────────────────┘│
85///             └───────────────────────────┘
86/// ```
87///
88/// The capacity is provided in bytes, and for each buffered record batch it will take into account
89/// the size reported by [RecordBatch::get_array_memory_size].
90///
91/// If a single record batch exceeds the maximum capacity set in the `capacity` argument, it's still
92/// allowed to pass in order to not deadlock the buffer.
93///
94/// This is useful for operators that conditionally start polling one of their children only after
95/// other child has finished, allowing to perform some early work and accumulating batches in
96/// memory so that they can be served immediately when requested.
97#[derive(Debug, Clone)]
98pub struct BufferExec {
99    input: Arc<dyn ExecutionPlan>,
100    properties: Arc<PlanProperties>,
101    capacity: usize,
102    metrics: ExecutionPlanMetricsSet,
103}
104
105impl BufferExec {
106    /// Builds a new [BufferExec] with the provided capacity in bytes.
107    pub fn new(input: Arc<dyn ExecutionPlan>, capacity: usize) -> Self {
108        let properties = PlanProperties::clone(input.properties())
109            .with_scheduling_type(SchedulingType::Cooperative)
110            .with_evaluation_type(EvaluationType::Eager);
111
112        Self {
113            input,
114            properties: Arc::new(properties),
115            capacity,
116            metrics: ExecutionPlanMetricsSet::new(),
117        }
118    }
119
120    /// Returns the input [ExecutionPlan] of this [BufferExec].
121    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
122        &self.input
123    }
124
125    /// Returns the per-partition capacity in bytes for this [BufferExec].
126    pub fn capacity(&self) -> usize {
127        self.capacity
128    }
129}
130
131impl DisplayAs for BufferExec {
132    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
133        match t {
134            DisplayFormatType::Default | DisplayFormatType::Verbose => {
135                write!(f, "BufferExec: capacity={}", self.capacity)
136            }
137            DisplayFormatType::TreeRender => {
138                writeln!(f, "target_batch_size={}", self.capacity)
139            }
140        }
141    }
142}
143
144impl ExecutionPlan for BufferExec {
145    fn name(&self) -> &str {
146        "BufferExec"
147    }
148
149    fn properties(&self) -> &Arc<PlanProperties> {
150        &self.properties
151    }
152
153    fn maintains_input_order(&self) -> Vec<bool> {
154        vec![true]
155    }
156
157    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
158        vec![false]
159    }
160
161    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
162        vec![&self.input]
163    }
164
165    fn apply_expressions(
166        &self,
167        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
168    ) -> Result<TreeNodeRecursion> {
169        Ok(TreeNodeRecursion::Continue)
170    }
171
172    fn replace_children(
173        self: Arc<Self>,
174        mut children: Vec<Arc<dyn ExecutionPlan>>,
175        options: ReplaceChildrenOptions,
176    ) -> Result<Arc<dyn ExecutionPlan>> {
177        validate_child_count!(self, children);
178        match options.children_properties {
179            ChildrenPropertiesMode::Keep => Ok(Arc::new(Self {
180                input: children.swap_remove(0),
181                metrics: ExecutionPlanMetricsSet::new(),
182                ..Self::clone(&*self)
183            })),
184            ChildrenPropertiesMode::Recompute => {
185                Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity)))
186            }
187        }
188    }
189
190    fn with_new_children(
191        self: Arc<Self>,
192        children: Vec<Arc<dyn ExecutionPlan>>,
193    ) -> Result<Arc<dyn ExecutionPlan>> {
194        self.replace_children(
195            children,
196            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
197        )
198    }
199
200    fn with_new_children_and_same_properties(
201        self: Arc<Self>,
202        children: Vec<Arc<dyn ExecutionPlan>>,
203    ) -> Result<Arc<dyn ExecutionPlan>> {
204        self.replace_children(
205            children,
206            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
207        )
208    }
209
210    fn execute(
211        &self,
212        partition: usize,
213        context: Arc<TaskContext>,
214    ) -> Result<SendableRecordBatchStream> {
215        let mem_reservation = MemoryConsumer::new(format!("BufferExec[{partition}]"))
216            .register(context.memory_pool());
217        let in_stream = self.input.execute(partition, context)?;
218
219        // Set up the metrics for the stream.
220        let curr_mem_in = Arc::new(AtomicUsize::new(0));
221        let curr_mem_out = Arc::clone(&curr_mem_in);
222        let mut max_mem_in = 0;
223        let max_mem = MetricBuilder::new(&self.metrics)
224            .peak_memory_usage("max_mem_used", partition);
225
226        let curr_queued_in = Arc::new(AtomicUsize::new(0));
227        let curr_queued_out = Arc::clone(&curr_queued_in);
228        let mut max_queued_in = 0;
229        let max_queued = MetricBuilder::new(&self.metrics)
230            .with_category(MetricCategory::Rows)
231            .gauge("max_queued", partition);
232
233        // Capture metrics when an element is queued on the stream.
234        let in_stream = in_stream.inspect_ok(move |v| {
235            let size = v.get_array_memory_size();
236            let curr_size = curr_mem_in.fetch_add(size, Ordering::Relaxed) + size;
237            if curr_size > max_mem_in {
238                max_mem_in = curr_size;
239                max_mem.set(max_mem_in);
240            }
241
242            let curr_queued = curr_queued_in.fetch_add(1, Ordering::Relaxed) + 1;
243            if curr_queued > max_queued_in {
244                max_queued_in = curr_queued;
245                max_queued.set(max_queued_in);
246            }
247        });
248        // Buffer the input.
249        let out_stream =
250            MemoryBufferedStream::new(in_stream, self.capacity, mem_reservation);
251        // Update in the metrics that when an element gets out, some memory gets freed.
252        let out_stream = out_stream.inspect_ok(move |v| {
253            curr_mem_out.fetch_sub(v.get_array_memory_size(), Ordering::Relaxed);
254            curr_queued_out.fetch_sub(1, Ordering::Relaxed);
255        });
256
257        Ok(Box::pin(RecordBatchStreamAdapter::new(
258            self.schema(),
259            out_stream,
260        )))
261    }
262
263    fn metrics(&self) -> Option<MetricsSet> {
264        Some(self.metrics.clone_inner())
265    }
266
267    fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {
268        vec![ChildStats::At(partition)]
269    }
270
271    fn statistics_from_inputs(
272        &self,
273        input_stats: &[Arc<Statistics>],
274        _args: &StatisticsArgs,
275    ) -> Result<Arc<Statistics>> {
276        Ok(Arc::clone(&input_stats[0]))
277    }
278
279    fn supports_limit_pushdown(&self) -> bool {
280        self.input.supports_limit_pushdown()
281    }
282
283    fn cardinality_effect(&self) -> CardinalityEffect {
284        CardinalityEffect::Equal
285    }
286
287    fn try_swapping_with_projection(
288        &self,
289        projection: &ProjectionExec,
290    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
291        match self.input.try_swapping_with_projection(projection)? {
292            Some(new_input) => Ok(Some(replace_children_if_necessary(
293                Arc::new(self.clone()),
294                vec![new_input],
295            )?)),
296            None => Ok(None),
297        }
298    }
299
300    fn gather_filters_for_pushdown(
301        &self,
302        _phase: FilterPushdownPhase,
303        parent_filters: Vec<Arc<dyn PhysicalExpr>>,
304        _config: &ConfigOptions,
305    ) -> Result<FilterDescription> {
306        FilterDescription::from_children(parent_filters, &self.children())
307    }
308
309    fn handle_child_pushdown_result(
310        &self,
311        _phase: FilterPushdownPhase,
312        child_pushdown_result: ChildPushdownResult,
313        _config: &ConfigOptions,
314    ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
315        Ok(FilterPushdownPropagation::if_all(child_pushdown_result))
316    }
317
318    fn try_pushdown_sort(
319        &self,
320        order: &[PhysicalSortExpr],
321    ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
322        // CoalesceBatchesExec is transparent for sort ordering - it preserves order
323        // Delegate to the child and wrap with a new CoalesceBatchesExec
324        self.input.try_pushdown_sort(order)?.try_map(|new_input| {
325            Ok(Arc::new(Self::new(new_input, self.capacity)) as Arc<dyn ExecutionPlan>)
326        })
327    }
328
329    #[cfg(feature = "proto")]
330    fn try_to_proto(
331        &self,
332        ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
333    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
334        use datafusion_proto_models::protobuf;
335        let input = ctx.encode_child(self.input())?;
336        Ok(Some(protobuf::PhysicalPlanNode {
337            physical_plan_type: Some(
338                protobuf::physical_plan_node::PhysicalPlanType::Buffer(Box::new(
339                    protobuf::BufferExecNode {
340                        input: Some(Box::new(input)),
341                        capacity: self.capacity() as u64,
342                    },
343                )),
344            ),
345        }))
346    }
347}
348
349#[cfg(feature = "proto")]
350impl BufferExec {
351    /// Reconstruct a [`BufferExec`] from its protobuf representation.
352    ///
353    /// The exact inverse of [`ExecutionPlan::try_to_proto`].
354    ///
355    /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto
356    pub fn try_from_proto(
357        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
358        ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
359    ) -> Result<Arc<dyn ExecutionPlan>> {
360        use datafusion_proto_models::protobuf;
361        let buffer = crate::expect_plan_variant!(
362            node,
363            protobuf::physical_plan_node::PhysicalPlanType::Buffer,
364            "BufferExec",
365        );
366        let input =
367            ctx.decode_required_child(buffer.input.as_deref(), "BufferExec", "input")?;
368        Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize)))
369    }
370}
371
372/// Represents anything that occupies a capacity in a [MemoryBufferedStream].
373pub trait SizedMessage {
374    fn size(&self) -> usize;
375}
376
377impl SizedMessage for RecordBatch {
378    fn size(&self) -> usize {
379        self.get_array_memory_size()
380    }
381}
382
383pin_project! {
384/// Decouples production and consumption of messages in a stream with an internal queue, eagerly
385/// filling it up to the specified maximum capacity even before any message is requested.
386///
387/// Allows each message to have a different size, which is taken into account for determining if
388/// the queue is full or not.
389pub struct MemoryBufferedStream<T: SizedMessage> {
390    task: SpawnedTask<()>,
391    batch_rx: UnboundedReceiver<Result<(T, OwnedSemaphorePermit)>>,
392    memory_reservation: Arc<MemoryReservation>,
393}}
394
395impl<T: Send + SizedMessage + 'static> MemoryBufferedStream<T> {
396    /// Builds a new [MemoryBufferedStream] with the provided capacity and event handler.
397    ///
398    /// This immediately spawns a Tokio task that will start consumption of the input stream.
399    pub fn new(
400        mut input: impl Stream<Item = Result<T>> + Unpin + Send + 'static,
401        capacity: usize,
402        memory_reservation: MemoryReservation,
403    ) -> Self {
404        let semaphore = Arc::new(Semaphore::new(capacity));
405        let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel();
406
407        let memory_reservation = Arc::new(memory_reservation);
408        let memory_reservation_clone = Arc::clone(&memory_reservation);
409        let task = SpawnedTask::spawn(async move {
410            loop {
411                // Select on both the input stream and the channel being closed.
412                // By down this, we abort polling the input as soon as the consumer channel is
413                // closed. Otherwise, we would need to wait for a full new message to be available
414                // in order to consider aborting the stream
415                let item_or_err = tokio::select! {
416                    biased;
417                    _ = batch_tx.closed() => break,
418                    // Catch a panic in the input poll so it surfaces as a stream error
419                    // instead of dropping `batch_tx` and looking like a clean EOF.
420                    polled = AssertUnwindSafe(input.next()).catch_unwind() => {
421                        match polled {
422                            Ok(Some(item_or_err)) => item_or_err,
423                            Ok(None) => break, // stream finished
424                            Err(panic) => {
425                                let msg = panic
426                                    .downcast_ref::<&str>()
427                                    .map(|s| s.to_string())
428                                    .or_else(|| panic.downcast_ref::<String>().cloned())
429                                    .unwrap_or_else(|| "unknown panic".to_string());
430                                let _ = batch_tx.send(internal_err!(
431                                    "BufferExec input stream panicked: {msg}"
432                                ));
433                                break;
434                            }
435                        }
436                    }
437                };
438
439                let item = match item_or_err {
440                    Ok(batch) => batch,
441                    Err(err) => {
442                        let _ = batch_tx.send(Err(err)); // If there's an error it means the channel was closed, which is fine.
443                        break;
444                    }
445                };
446
447                let size = item.size();
448                if let Err(err) = memory_reservation.try_grow(size) {
449                    let _ = batch_tx.send(Err(err)); // If there's an error it means the channel was closed, which is fine.
450                    break;
451                }
452
453                // We need to cap the minimum between amount of permits and the actual size of the
454                // message. If at any point we try to acquire more permits than the capacity of the
455                // semaphore, the stream will deadlock.
456                let capped_size = size.min(capacity) as u32;
457
458                let semaphore = Arc::clone(&semaphore);
459                let Ok(permit) = semaphore.acquire_many_owned(capped_size).await else {
460                    let _ = batch_tx.send(internal_err!("Closed semaphore in MemoryBufferedStream. This is a bug in DataFusion, please report it!"));
461                    break;
462                };
463
464                if batch_tx.send(Ok((item, permit))).is_err() {
465                    break; // stream was closed
466                };
467            }
468        });
469
470        Self {
471            task,
472            batch_rx,
473            memory_reservation: memory_reservation_clone,
474        }
475    }
476
477    /// Returns the number of queued messages.
478    pub fn messages_queued(&self) -> usize {
479        self.batch_rx.len()
480    }
481}
482
483impl<T: SizedMessage> Stream for MemoryBufferedStream<T> {
484    type Item = Result<T>;
485
486    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
487        let self_project = self.project();
488        match self_project.batch_rx.poll_recv(cx) {
489            Poll::Ready(Some(Ok((item, _semaphore_permit)))) => {
490                self_project.memory_reservation.shrink(item.size());
491                Poll::Ready(Some(Ok(item)))
492            }
493            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
494            Poll::Ready(None) => Poll::Ready(None),
495            Poll::Pending => Poll::Pending,
496        }
497    }
498
499    fn size_hint(&self) -> (usize, Option<usize>) {
500        if self.batch_rx.is_closed() {
501            let len = self.batch_rx.len();
502            (len, Some(len))
503        } else {
504            (self.batch_rx.len(), None)
505        }
506    }
507}
508
509#[cfg(test)]
510mod tests {
511    use super::*;
512    use datafusion_common::{DataFusionError, assert_contains};
513    use datafusion_execution::memory_pool::{
514        GreedyMemoryPool, MemoryPool, UnboundedMemoryPool,
515    };
516    use std::error::Error;
517    use std::fmt::Debug;
518    use std::time::Duration;
519    use tokio::time::timeout;
520
521    #[tokio::test]
522    async fn buffers_only_some_messages() -> Result<(), Box<dyn Error>> {
523        let input = futures::stream::iter([1, 2, 3, 4]).map(Ok);
524        let (_, res) = memory_pool_and_reservation();
525
526        let buffered = MemoryBufferedStream::new(input, 4, res);
527        wait_for_buffering().await;
528        assert_eq!(buffered.messages_queued(), 2);
529        Ok(())
530    }
531
532    #[tokio::test]
533    async fn yields_all_messages() -> Result<(), Box<dyn Error>> {
534        let input = futures::stream::iter([1, 2, 3, 4]).map(Ok);
535        let (_, res) = memory_pool_and_reservation();
536
537        let mut buffered = MemoryBufferedStream::new(input, 10, res);
538        wait_for_buffering().await;
539        assert_eq!(buffered.messages_queued(), 4);
540
541        pull_ok_msg(&mut buffered).await?;
542        pull_ok_msg(&mut buffered).await?;
543        pull_ok_msg(&mut buffered).await?;
544        pull_ok_msg(&mut buffered).await?;
545        finished(&mut buffered).await?;
546        Ok(())
547    }
548
549    #[tokio::test]
550    async fn yields_first_msg_even_if_big() -> Result<(), Box<dyn Error>> {
551        let input = futures::stream::iter([25, 1, 2, 3]).map(Ok);
552        let (_, res) = memory_pool_and_reservation();
553
554        let mut buffered = MemoryBufferedStream::new(input, 10, res);
555        wait_for_buffering().await;
556        assert_eq!(buffered.messages_queued(), 1);
557        pull_ok_msg(&mut buffered).await?;
558        Ok(())
559    }
560
561    #[tokio::test]
562    async fn memory_pool_kills_stream() -> Result<(), Box<dyn Error>> {
563        let input = futures::stream::iter([1, 2, 3, 4]).map(Ok);
564        let (_, res) = bounded_memory_pool_and_reservation(7);
565
566        let mut buffered = MemoryBufferedStream::new(input, 10, res);
567        wait_for_buffering().await;
568
569        pull_ok_msg(&mut buffered).await?;
570        pull_ok_msg(&mut buffered).await?;
571        pull_ok_msg(&mut buffered).await?;
572        let msg = pull_err_msg(&mut buffered).await?;
573
574        assert_contains!(msg.to_string(), "Failed to allocate additional 4.0 B");
575        Ok(())
576    }
577
578    #[tokio::test]
579    async fn memory_pool_does_not_kill_stream() -> Result<(), Box<dyn Error>> {
580        let input = futures::stream::iter([1, 2, 3, 4]).map(Ok);
581        let (_, res) = bounded_memory_pool_and_reservation(7);
582
583        let mut buffered = MemoryBufferedStream::new(input, 3, res);
584        wait_for_buffering().await;
585        pull_ok_msg(&mut buffered).await?;
586
587        wait_for_buffering().await;
588        pull_ok_msg(&mut buffered).await?;
589
590        wait_for_buffering().await;
591        pull_ok_msg(&mut buffered).await?;
592
593        wait_for_buffering().await;
594        pull_ok_msg(&mut buffered).await?;
595
596        wait_for_buffering().await;
597        finished(&mut buffered).await?;
598        Ok(())
599    }
600
601    #[tokio::test]
602    async fn messages_pass_even_if_all_exceed_limit() -> Result<(), Box<dyn Error>> {
603        let input = futures::stream::iter([3, 3, 3, 3]).map(Ok);
604        let (_, res) = memory_pool_and_reservation();
605
606        let mut buffered = MemoryBufferedStream::new(input, 2, res);
607        wait_for_buffering().await;
608        assert_eq!(buffered.messages_queued(), 1);
609        pull_ok_msg(&mut buffered).await?;
610
611        wait_for_buffering().await;
612        assert_eq!(buffered.messages_queued(), 1);
613        pull_ok_msg(&mut buffered).await?;
614
615        wait_for_buffering().await;
616        assert_eq!(buffered.messages_queued(), 1);
617        pull_ok_msg(&mut buffered).await?;
618
619        wait_for_buffering().await;
620        assert_eq!(buffered.messages_queued(), 1);
621        pull_ok_msg(&mut buffered).await?;
622
623        wait_for_buffering().await;
624        finished(&mut buffered).await?;
625        Ok(())
626    }
627
628    #[tokio::test]
629    async fn errors_get_propagated() -> Result<(), Box<dyn Error>> {
630        let input = futures::stream::iter([1, 2, 3, 4]).map(|v| {
631            if v == 3 {
632                return internal_err!("Error on 3");
633            }
634            Ok(v)
635        });
636        let (_, res) = memory_pool_and_reservation();
637
638        let mut buffered = MemoryBufferedStream::new(input, 10, res);
639        wait_for_buffering().await;
640
641        pull_ok_msg(&mut buffered).await?;
642        pull_ok_msg(&mut buffered).await?;
643        pull_err_msg(&mut buffered).await?;
644
645        Ok(())
646    }
647
648    #[tokio::test]
649    async fn panic_in_input_is_propagated() -> Result<(), Box<dyn Error>> {
650        // A panic while polling the input must surface as a stream error, not a
651        // silent end-of-stream that drops the rest of the partition's output.
652        let input = futures::stream::iter([1, 2, 3, 4]).map(|v| {
653            if v == 3 {
654                panic!("boom on 3");
655            }
656            Ok(v)
657        });
658        let (_, res) = memory_pool_and_reservation();
659
660        let mut buffered = MemoryBufferedStream::new(input, 10, res);
661        wait_for_buffering().await;
662
663        pull_ok_msg(&mut buffered).await?;
664        pull_ok_msg(&mut buffered).await?;
665        let err = pull_err_msg(&mut buffered).await?;
666        assert_contains!(err.to_string(), "panicked");
667
668        Ok(())
669    }
670
671    #[tokio::test]
672    async fn memory_gets_released_if_stream_drops() -> Result<(), Box<dyn Error>> {
673        let input = futures::stream::iter([1, 2, 3, 4]).map(Ok);
674        let (pool, res) = memory_pool_and_reservation();
675
676        let mut buffered = MemoryBufferedStream::new(input, 10, res);
677        wait_for_buffering().await;
678        assert_eq!(buffered.messages_queued(), 4);
679        assert_eq!(pool.reserved(), 10);
680
681        pull_ok_msg(&mut buffered).await?;
682        assert_eq!(buffered.messages_queued(), 3);
683        assert_eq!(pool.reserved(), 9);
684
685        pull_ok_msg(&mut buffered).await?;
686        assert_eq!(buffered.messages_queued(), 2);
687        assert_eq!(pool.reserved(), 7);
688
689        drop(buffered);
690        assert_eq!(pool.reserved(), 0);
691        Ok(())
692    }
693
694    fn memory_pool_and_reservation() -> (Arc<dyn MemoryPool>, MemoryReservation) {
695        let pool = Arc::new(UnboundedMemoryPool::default()) as _;
696        let reservation = MemoryConsumer::new("test").register(&pool);
697        (pool, reservation)
698    }
699
700    fn bounded_memory_pool_and_reservation(
701        size: usize,
702    ) -> (Arc<dyn MemoryPool>, MemoryReservation) {
703        let pool = Arc::new(GreedyMemoryPool::new(size)) as _;
704        let reservation = MemoryConsumer::new("test").register(&pool);
705        (pool, reservation)
706    }
707
708    async fn wait_for_buffering() {
709        // We do not have control over the spawned task, so the best we can do is to yield some
710        // cycles to the tokio runtime and let the task make progress on its own.
711        tokio::time::sleep(Duration::from_millis(1)).await;
712    }
713
714    async fn pull_ok_msg<T: SizedMessage>(
715        buffered: &mut MemoryBufferedStream<T>,
716    ) -> Result<T, Box<dyn Error>> {
717        Ok(timeout(Duration::from_millis(1), buffered.next())
718            .await?
719            .unwrap_or_else(|| internal_err!("Stream should not have finished"))?)
720    }
721
722    async fn pull_err_msg<T: SizedMessage + Debug>(
723        buffered: &mut MemoryBufferedStream<T>,
724    ) -> Result<DataFusionError, Box<dyn Error>> {
725        Ok(timeout(Duration::from_millis(1), buffered.next())
726            .await?
727            .map(|v| match v {
728                Ok(v) => internal_err!(
729                    "Stream should not have failed, but succeeded with {v:?}"
730                ),
731                Err(err) => Ok(err),
732            })
733            .unwrap_or_else(|| internal_err!("Stream should not have finished"))?)
734    }
735
736    async fn finished<T: SizedMessage>(
737        buffered: &mut MemoryBufferedStream<T>,
738    ) -> Result<(), Box<dyn Error>> {
739        match timeout(Duration::from_millis(1), buffered.next())
740            .await?
741            .is_none()
742        {
743            true => Ok(()),
744            false => internal_err!("Stream should have finished")?,
745        }
746    }
747
748    impl SizedMessage for usize {
749        fn size(&self) -> usize {
750            *self
751        }
752    }
753}
datafusion_physical_plan/buffer.rs

datafusion_physical_plan/
buffer.rs
