deltalake_core/kernel/snapshot/stream.rs
1//! the code in this file is hoisted from datafusion with only slight modifications
2//!
3use arrow::{datatypes::SchemaRef, record_batch::RecordBatch};
4use futures::stream::BoxStream;
5use futures::{Stream, StreamExt};
6use std::pin::Pin;
7use tokio::sync::mpsc::{Receiver, Sender};
8use tokio::task::JoinSet;
9use tracing::Span;
10use tracing::dispatcher;
11
12use crate::DeltaTableError;
13use crate::errors::DeltaResult;
14
15/// Trait for types that stream [RecordBatch]
16///
17/// See [`SendableRecordBatchStream`] for more details.
18pub trait RecordBatchStream: Stream<Item = DeltaResult<RecordBatch>> {
19 /// Returns the schema of this `RecordBatchStream`.
20 ///
21 /// Implementation of this trait should guarantee that all `RecordBatch`'s returned by this
22 /// stream should have the same schema as returned from this method.
23 fn schema(&self) -> SchemaRef;
24}
25
26/// Trait for a [`Stream`] of [`RecordBatch`]es that can be passed between threads
27///
28/// This trait is used to retrieve the results of DataFusion execution plan nodes.
29///
30/// The trait is a specialized Rust Async [`Stream`] that also knows the schema
31/// of the data it will return (even if the stream has no data). Every
32/// `RecordBatch` returned by the stream should have the same schema as returned
33/// by [`schema`](`RecordBatchStream::schema`).
34///
35/// # See Also
36///
37/// * [`RecordBatchStreamAdapter`] to convert an existing [`Stream`]
38/// to [`SendableRecordBatchStream`]
39///
40/// [`RecordBatchStreamAdapter`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/stream/struct.RecordBatchStreamAdapter.html
41///
42/// # Error Handling
43///
44/// Once a stream returns an error, it should not be polled again (the caller
45/// should stop calling `next`) and handle the error.
46///
47/// However, returning `Ready(None)` (end of stream) is likely the safest
48/// behavior after an error. Like [`Stream`]s, `RecordBatchStream`s should not
49/// be polled after end of stream or returning an error. However, also like
50/// [`Stream`]s there is no mechanism to prevent callers polling so returning
51/// `Ready(None)` is recommended.
52pub type SendableRecordBatchStream = Pin<Box<dyn RecordBatchStream + Send>>;
53
54pub type SendableRBStream = Pin<Box<dyn Stream<Item = DeltaResult<RecordBatch>> + Send>>;
55
56/// Creates a stream from a collection of producing tasks, routing panics to the stream.
57///
58/// Note that this is similar to [`ReceiverStream` from tokio-stream], with the differences being:
59///
60/// 1. Methods to bound and "detach" tasks (`spawn_blocking()`).
61///
62/// 2. Propagates panics, whereas the `tokio` version doesn't propagate panics to the receiver.
63///
64/// 3. Automatically cancels any outstanding tasks when the receiver stream is dropped.
65///
66/// [`ReceiverStream` from tokio-stream]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.ReceiverStream.html
67pub(crate) struct ReceiverStreamBuilder<O> {
68 tx: Sender<DeltaResult<O>>,
69 rx: Receiver<DeltaResult<O>>,
70 join_set: JoinSet<DeltaResult<()>>,
71}
72
73impl<O: Send + 'static> ReceiverStreamBuilder<O> {
74 /// Create new channels with the specified buffer size
75 pub fn new(capacity: usize) -> Self {
76 let (tx, rx) = tokio::sync::mpsc::channel(capacity);
77
78 Self {
79 tx,
80 rx,
81 join_set: JoinSet::new(),
82 }
83 }
84
85 /// Get a handle for sending data to the output
86 pub fn tx(&self) -> Sender<DeltaResult<O>> {
87 self.tx.clone()
88 }
89
90 /// Spawn a blocking task that will be aborted if this builder (or the stream
91 /// built from it) are dropped.
92 ///
93 /// This is often used to spawn tasks that write to the sender
94 /// retrieved from `Self::tx`.
95 pub fn spawn_blocking<F>(&mut self, f: F)
96 where
97 F: FnOnce() -> DeltaResult<()>,
98 F: Send + 'static,
99 {
100 // Capture the current dispatcher and span
101 let dispatch = dispatcher::get_default(|d| d.clone());
102 let span = Span::current();
103
104 self.join_set.spawn_blocking(move || {
105 dispatcher::with_default(&dispatch, || {
106 let _enter = span.enter();
107 f()
108 })
109 });
110 }
111
112 /// Create a stream of all data written to `tx`
113 pub fn build(self) -> BoxStream<'static, DeltaResult<O>> {
114 let Self {
115 tx,
116 rx,
117 mut join_set,
118 } = self;
119
120 // Doesn't need tx
121 drop(tx);
122
123 // future that checks the result of the join set, and propagates panic if seen
124 let check = async move {
125 while let Some(result) = join_set.join_next().await {
126 match result {
127 Ok(task_result) => {
128 match task_result {
129 // Nothing to report
130 Ok(_) => continue,
131 // This means a blocking task error
132 Err(error) => return Some(Err(error)),
133 }
134 }
135 // This means a tokio task error, likely a panic
136 Err(e) => {
137 if e.is_panic() {
138 // resume on the main thread
139 std::panic::resume_unwind(e.into_panic());
140 } else {
141 // This should only occur if the task is
142 // cancelled, which would only occur if
143 // the JoinSet were aborted, which in turn
144 // would imply that the receiver has been
145 // dropped and this code is not running
146 return Some(Err(DeltaTableError::Generic(format!(
147 "Non Panic Task error: {e}"
148 ))));
149 }
150 }
151 }
152 }
153 None
154 };
155
156 let check_stream = futures::stream::once(check)
157 // unwrap Option / only return the error
158 .filter_map(|item| async move { item });
159
160 // Convert the receiver into a stream
161 let rx_stream = futures::stream::unfold(rx, |mut rx| async move {
162 let next_item = rx.recv().await;
163 next_item.map(|next_item| (next_item, rx))
164 });
165
166 // Merge the streams together so whichever is ready first
167 // produces the batch
168 futures::stream::select(rx_stream, check_stream).boxed()
169 }
170}
171
172pub(crate) type RecordBatchReceiverStreamBuilder = ReceiverStreamBuilder<RecordBatch>;