Returns a stream which yields data
(
&self,
partition: usize,
_context: Arc<TaskContext>,
)
| 204 | |
| 205 | /// Returns a stream which yields data |
| 206 | fn execute( |
| 207 | &self, |
| 208 | partition: usize, |
| 209 | _context: Arc<TaskContext>, |
| 210 | ) -> Result<SendableRecordBatchStream> { |
| 211 | assert_eq!(partition, 0); |
| 212 | |
| 213 | // Result doesn't implement clone, so do it ourself |
| 214 | let data: Vec<_> = self |
| 215 | .data |
| 216 | .iter() |
| 217 | .map(|r| match r { |
| 218 | Ok(batch) => Ok(batch.clone()), |
| 219 | Err(e) => Err(clone_error(e)), |
| 220 | }) |
| 221 | .collect(); |
| 222 | |
| 223 | if self.use_task { |
| 224 | let mut builder = RecordBatchReceiverStream::builder(self.schema(), 2); |
| 225 | // send data in order but in a separate task (to ensure |
| 226 | // the batches are not available without the stream |
| 227 | // yielding). |
| 228 | let tx = builder.tx(); |
| 229 | builder.spawn(async move { |
| 230 | for batch in data { |
| 231 | println!("Sending batch via delayed stream"); |
| 232 | if let Err(e) = tx.send(batch).await { |
| 233 | println!("ERROR batch via delayed stream: {e}"); |
| 234 | } |
| 235 | } |
| 236 | |
| 237 | Ok(()) |
| 238 | }); |
| 239 | // returned stream simply reads off the rx stream |
| 240 | Ok(builder.build()) |
| 241 | } else { |
| 242 | // make an input that will error |
| 243 | let stream = futures::stream::iter(data); |
| 244 | Ok(Box::pin(RecordBatchStreamAdapter::new( |
| 245 | self.schema(), |
| 246 | stream, |
| 247 | ))) |
| 248 | } |
| 249 | } |
| 250 | |
| 251 | // Panics if one of the batches is an error |
| 252 | fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> { |