Method execute

datafusion/physical-plan/src/buffer.rs:181–233

(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    )

    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        let mem_reservation = MemoryConsumer::new(format!("BufferExec[{partition}]"))
            .register(context.memory_pool());
        let in_stream = self.input.execute(partition, context)?;

        // Set up the metrics for the stream.
        let curr_mem_in = Arc::new(AtomicUsize::new(0));
        let curr_mem_out = Arc::clone(&curr_mem_in);
        let mut max_mem_in = 0;
        let max_mem = MetricBuilder::new(&self.metrics)
            .with_category(MetricCategory::Bytes)
            .gauge("max_mem_used", partition);

        let curr_queued_in = Arc::new(AtomicUsize::new(0));
        let curr_queued_out = Arc::clone(&curr_queued_in);
        let mut max_queued_in = 0;
        let max_queued = MetricBuilder::new(&self.metrics)
            .with_category(MetricCategory::Rows)
            .gauge("max_queued", partition);

        // Capture metrics when an element is queued on the stream.
        let in_stream = in_stream.inspect_ok(move |v| {
            let size = v.get_array_memory_size();
            let curr_size = curr_mem_in.fetch_add(size, Ordering::Relaxed) + size;
            if curr_size > max_mem_in {
                max_mem_in = curr_size;
                max_mem.set(max_mem_in);
            }

            let curr_queued = curr_queued_in.fetch_add(1, Ordering::Relaxed) + 1;
            if curr_queued > max_queued_in {
                max_queued_in = curr_queued;
                max_queued.set(max_queued_in);
            }
        });
        // Buffer the input.
        let out_stream =
            MemoryBufferedStream::new(in_stream, self.capacity, mem_reservation);
        // Update in the metrics that when an element gets out, some memory gets freed.
        let out_stream = out_stream.inspect_ok(move |v| {
            curr_mem_out.fetch_sub(v.get_array_memory_size(), Ordering::Relaxed);
            curr_queued_out.fetch_sub(1, Ordering::Relaxed);
        });

        Ok(Box::pin(RecordBatchStreamAdapter::new(
            self.schema(),
            out_stream,
        )))
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metrics.clone_inner())