|
| 1 | +use polars_core::frame::DataFrame; |
| 2 | +use polars_core::schema::SchemaRef; |
| 3 | +use polars_error::{PolarsResult, polars_err}; |
| 4 | +use polars_utils::pl_str::PlSmallStr; |
| 5 | +use polars_utils::{IdxSize, format_pl_smallstr}; |
| 6 | +use tokio::sync::oneshot; |
| 7 | + |
| 8 | +use super::{ |
| 9 | + JoinHandle, Morsel, MorselSeq, SourceNode, SourceOutput, StreamingExecutionState, TaskPriority, |
| 10 | +}; |
| 11 | +use crate::async_executor::spawn; |
| 12 | +use crate::async_primitives::connector::Receiver; |
| 13 | +use crate::async_primitives::wait_group::WaitGroup; |
| 14 | +use crate::morsel::SourceToken; |
| 15 | + |
| 16 | +type GetBatchFn = |
| 17 | + Box<dyn Fn(&StreamingExecutionState) -> PolarsResult<Option<DataFrame>> + Send + Sync>; |
| 18 | + |
| 19 | +pub struct BatchSourceNode { |
| 20 | + pub name: PlSmallStr, |
| 21 | + pub output_schema: SchemaRef, |
| 22 | + pub get_batch_fn: Option<GetBatchFn>, |
| 23 | +} |
| 24 | + |
| 25 | +impl BatchSourceNode { |
| 26 | + pub fn new(name: &str, output_schema: SchemaRef, get_batch_fn: Option<GetBatchFn>) -> Self { |
| 27 | + let name = format_pl_smallstr!("batch_source[{name}]"); |
| 28 | + Self { |
| 29 | + name, |
| 30 | + output_schema, |
| 31 | + get_batch_fn, |
| 32 | + } |
| 33 | + } |
| 34 | +} |
| 35 | + |
| 36 | +impl SourceNode for BatchSourceNode { |
| 37 | + fn name(&self) -> &str { |
| 38 | + self.name.as_str() |
| 39 | + } |
| 40 | + |
| 41 | + fn is_source_output_parallel(&self, _is_receiver_serial: bool) -> bool { |
| 42 | + false |
| 43 | + } |
| 44 | + |
| 45 | + fn spawn_source( |
| 46 | + &mut self, |
| 47 | + mut output_recv: Receiver<SourceOutput>, |
| 48 | + state: &StreamingExecutionState, |
| 49 | + join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>, |
| 50 | + unrestricted_row_count: Option<oneshot::Sender<polars_utils::IdxSize>>, |
| 51 | + ) { |
| 52 | + // We only spawn this once, so this is all fine. |
| 53 | + let output_schema = self.output_schema.clone(); |
| 54 | + let get_batch_fn = self.get_batch_fn.take().unwrap(); |
| 55 | + let state = state.clone(); |
| 56 | + join_handles.push(spawn(TaskPriority::Low, async move { |
| 57 | + let mut seq = MorselSeq::default(); |
| 58 | + let mut n_rows_seen = 0; |
| 59 | + |
| 60 | + 'phase_loop: while let Ok(phase_output) = output_recv.recv().await { |
| 61 | + let mut sender = phase_output.port.serial(); |
| 62 | + let source_token = SourceToken::new(); |
| 63 | + let wait_group = WaitGroup::default(); |
| 64 | + |
| 65 | + loop { |
| 66 | + let df = (get_batch_fn)(&state)?; |
| 67 | + let Some(df) = df else { |
| 68 | + if let Some(unrestricted_row_count) = unrestricted_row_count { |
| 69 | + if unrestricted_row_count.send(n_rows_seen).is_err() { |
| 70 | + return Ok(()); |
| 71 | + } |
| 72 | + } |
| 73 | + |
| 74 | + if n_rows_seen == 0 { |
| 75 | + let morsel = Morsel::new( |
| 76 | + DataFrame::empty_with_schema(output_schema.as_ref()), |
| 77 | + seq, |
| 78 | + source_token.clone(), |
| 79 | + ); |
| 80 | + if sender.send(morsel).await.is_err() { |
| 81 | + return Ok(()); |
| 82 | + } |
| 83 | + } |
| 84 | + |
| 85 | + break 'phase_loop; |
| 86 | + }; |
| 87 | + |
| 88 | + let num_rows = IdxSize::try_from(df.height()).map_err(|_| { |
| 89 | + polars_err!(bigidx, ctx = "batch source", size = df.height()) |
| 90 | + })?; |
| 91 | + n_rows_seen = n_rows_seen.checked_add(num_rows).ok_or_else(|| { |
| 92 | + polars_err!( |
| 93 | + bigidx, |
| 94 | + ctx = "batch source", |
| 95 | + size = n_rows_seen as usize + num_rows as usize |
| 96 | + ) |
| 97 | + })?; |
| 98 | + |
| 99 | + let mut morsel = Morsel::new(df, seq, source_token.clone()); |
| 100 | + morsel.set_consume_token(wait_group.token()); |
| 101 | + seq = seq.successor(); |
| 102 | + |
| 103 | + if sender.send(morsel).await.is_err() { |
| 104 | + return Ok(()); |
| 105 | + } |
| 106 | + |
| 107 | + wait_group.wait().await; |
| 108 | + if source_token.stop_requested() { |
| 109 | + phase_output.outcome.stop(); |
| 110 | + continue 'phase_loop; |
| 111 | + } |
| 112 | + } |
| 113 | + } |
| 114 | + |
| 115 | + Ok(()) |
| 116 | + })); |
| 117 | + } |
| 118 | +} |
0 commit comments