refactor(rust): Dispatch new-streaming IO plugin source to updated multiscan (#22009)

Author: nameexhaustion

Diff for: crates/polars-stream/src/execute.rs

@@ -18,6 +18,15 @@ pub struct StreamingExecutionState {
     pub in_memory_exec_state: ExecutionState,
 }
 
+impl Default for StreamingExecutionState {
+    fn default() -> Self {
+        Self {
+            num_pipelines: POOL.current_num_threads(),
+            in_memory_exec_state: ExecutionState::default(),
+        }
+    }
+}
+
 /// Finds all runnable pipeline blockers in the graph, that is, nodes which:
 ///  - Only have blocked output ports.
 ///  - Have at least one ready input port connected to a ready output port.

Diff for: crates/polars-stream/src/nodes/io_sources/batch.rs

@@ -1,118 +1,229 @@
+//! Reads batches from a `dyn Fn`
+
+use async_trait::async_trait;
 use polars_core::frame::DataFrame;
 use polars_core::schema::SchemaRef;
 use polars_error::{PolarsResult, polars_err};
+use polars_utils::IdxSize;
 use polars_utils::pl_str::PlSmallStr;
-use polars_utils::{IdxSize, format_pl_smallstr};
-use tokio::sync::oneshot;
 
-use super::{
-    JoinHandle, Morsel, MorselSeq, SourceNode, SourceOutput, StreamingExecutionState, TaskPriority,
+use crate::async_executor::{JoinHandle, TaskPriority, spawn};
+use crate::execute::StreamingExecutionState;
+use crate::morsel::{Morsel, MorselSeq, SourceToken};
+use crate::nodes::io_sources::multi_file_reader::reader_interface::output::{
+    FileReaderOutputRecv, FileReaderOutputSend,
+};
+use crate::nodes::io_sources::multi_file_reader::reader_interface::{
+    BeginReadArgs, FileReader, FileReaderCallbacks,
 };
-use crate::async_executor::spawn;
-use crate::async_primitives::connector::Receiver;
-use crate::async_primitives::wait_group::WaitGroup;
-use crate::morsel::SourceToken;
 
-type GetBatchFn =
+pub mod builder {
+
+    use std::sync::{Arc, Mutex};
+
+    use polars_utils::pl_str::PlSmallStr;
+
+    use super::BatchFnReader;
+    use crate::nodes::io_sources::multi_file_reader::reader_interface::FileReader;
+    use crate::nodes::io_sources::multi_file_reader::reader_interface::builder::FileReaderBuilder;
+    use crate::nodes::io_sources::multi_file_reader::reader_interface::capabilities::ReaderCapabilities;
+
+    pub struct BatchFnReaderBuilder {
+        pub name: PlSmallStr,
+        pub reader: Mutex<Option<BatchFnReader>>,
+    }
+
+    impl FileReaderBuilder for BatchFnReaderBuilder {
+        fn reader_name(&self) -> &str {
+            &self.name
+        }
+
+        fn reader_capabilities(&self) -> ReaderCapabilities {
+            ReaderCapabilities::empty()
+        }
+
+        fn build_file_reader(
+            &self,
+            _source: polars_plan::prelude::ScanSource,
+            _cloud_options: Option<Arc<polars_io::cloud::CloudOptions>>,
+            scan_source_idx: usize,
+        ) -> Box<dyn FileReader> {
+            assert_eq!(scan_source_idx, 0);
+
+            Box::new(
+                self.reader
+                    .try_lock()
+                    .unwrap()
+                    .take()
+                    .expect("BatchFnReaderBuilder called more than once"),
+            ) as Box<dyn FileReader>
+        }
+    }
+
+    impl std::fmt::Debug for BatchFnReaderBuilder {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            f.write_str("BatchFnReaderBuilder: name: ")?;
+            f.write_str(&self.name)?;
+
+            Ok(())
+        }
+    }
+}
+
+pub type GetBatchFn =
     Box<dyn Fn(&StreamingExecutionState) -> PolarsResult<Option<DataFrame>> + Send + Sync>;
 
-pub struct BatchSourceNode {
-    pub name: PlSmallStr,
-    pub output_schema: SchemaRef,
-    pub get_batch_fn: Option<GetBatchFn>,
+/// Wraps `GetBatchFn` to support peeking.
+pub struct GetBatchState {
+    func: GetBatchFn,
+    peek: Option<DataFrame>,
 }
 
-impl BatchSourceNode {
-    pub fn new(name: &str, output_schema: SchemaRef, get_batch_fn: Option<GetBatchFn>) -> Self {
-        let name = format_pl_smallstr!("batch_source[{name}]");
-        Self {
-            name,
-            output_schema,
-            get_batch_fn,
+impl GetBatchState {
+    pub fn peek(&mut self, state: &StreamingExecutionState) -> PolarsResult<Option<&DataFrame>> {
+        if self.peek.is_none() {
+            self.peek = (self.func)(state)?;
+        }
+
+        Ok(self.peek.as_ref())
+    }
+
+    pub fn next(&mut self, state: &StreamingExecutionState) -> PolarsResult<Option<DataFrame>> {
+        if let Some(df) = self.peek.take() {
+            Ok(Some(df))
+        } else {
+            (self.func)(state)
         }
     }
 }
 
-impl SourceNode for BatchSourceNode {
-    fn name(&self) -> &str {
-        self.name.as_str()
+impl From<GetBatchFn> for GetBatchState {
+    fn from(func: GetBatchFn) -> Self {
+        Self { func, peek: None }
     }
+}
+
+pub struct BatchFnReader {
+    pub name: PlSmallStr,
+    pub output_schema: Option<SchemaRef>,
+    pub get_batch_state: Option<GetBatchState>,
+    pub verbose: bool,
+}
 
-    fn is_source_output_parallel(&self, _is_receiver_serial: bool) -> bool {
-        false
+#[async_trait]
+impl FileReader for BatchFnReader {
+    async fn initialize(&mut self) -> PolarsResult<()> {
+        Ok(())
     }
 
-    fn spawn_source(
+    fn begin_read(
         &mut self,
-        mut output_recv: Receiver<SourceOutput>,
-        state: &StreamingExecutionState,
-        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
-        unrestricted_row_count: Option<oneshot::Sender<polars_utils::IdxSize>>,
-    ) {
-        // We only spawn this once, so this is all fine.
-        let output_schema = self.output_schema.clone();
-        let get_batch_fn = self.get_batch_fn.take().unwrap();
-        let state = state.clone();
-        join_handles.push(spawn(TaskPriority::Low, async move {
-            let mut seq = MorselSeq::default();
-            let mut n_rows_seen = 0;
-
-            'phase_loop: while let Ok(phase_output) = output_recv.recv().await {
-                let mut sender = phase_output.port.serial();
-                let source_token = SourceToken::new();
-                let wait_group = WaitGroup::default();
-
-                loop {
-                    let df = (get_batch_fn)(&state)?;
-                    let Some(df) = df else {
-                        if let Some(unrestricted_row_count) = unrestricted_row_count {
-                            if unrestricted_row_count.send(n_rows_seen).is_err() {
-                                return Ok(());
-                            }
-                        }
-
-                        if n_rows_seen == 0 {
-                            let morsel = Morsel::new(
-                                DataFrame::empty_with_schema(output_schema.as_ref()),
-                                seq,
-                                source_token.clone(),
-                            );
-                            if sender.send(morsel).await.is_err() {
-                                return Ok(());
-                            }
-                        }
-
-                        break 'phase_loop;
-                    };
-
-                    let num_rows = IdxSize::try_from(df.height()).map_err(|_| {
-                        polars_err!(bigidx, ctx = "batch source", size = df.height())
-                    })?;
-                    n_rows_seen = n_rows_seen.checked_add(num_rows).ok_or_else(|| {
-                        polars_err!(
-                            bigidx,
-                            ctx = "batch source",
-                            size = n_rows_seen as usize + num_rows as usize
-                        )
-                    })?;
-
-                    let mut morsel = Morsel::new(df, seq, source_token.clone());
-                    morsel.set_consume_token(wait_group.token());
-                    seq = seq.successor();
-
-                    if sender.send(morsel).await.is_err() {
-                        return Ok(());
-                    }
-
-                    wait_group.wait().await;
-                    if source_token.stop_requested() {
-                        phase_output.outcome.stop();
-                        continue 'phase_loop;
-                    }
+        args: BeginReadArgs,
+    ) -> PolarsResult<(FileReaderOutputRecv, JoinHandle<PolarsResult<()>>)> {
+        let BeginReadArgs {
+            projected_schema: _,
+            row_index: None,
+            pre_slice: None,
+            predicate: None,
+            cast_columns_policy: _,
+            num_pipelines: _,
+            callbacks:
+                FileReaderCallbacks {
+                    file_schema_tx,
+                    n_rows_in_file_tx,
+                    row_position_on_end_tx,
+                },
+        } = args
+        else {
+            panic!("unsupported args: {:?}", &args)
+        };
+
+        // Must send this first before we `take()` the GetBatchState.
+        if let Some(mut file_schema_tx) = file_schema_tx {
+            _ = file_schema_tx.try_send(self._file_schema()?);
+        }
+
+        let mut get_batch_state = self
+            .get_batch_state
+            .take()
+            // If this is ever needed we can buffer
+            .expect("unimplemented: BatchFnReader called more than once");
+
+        // FIXME: Propagate this from BeginReadArgs.
+        let exec_state = StreamingExecutionState::default();
+
+        let verbose = self.verbose;
+
+        if verbose {
+            eprintln!("[BatchFnReader]: name: {}", self.name);
+        }
+
+        let (mut morsel_sender, morsel_rx) = FileReaderOutputSend::new_serial();
+
+        let handle = spawn(TaskPriority::Low, async move {
+            let mut seq: u64 = 0;
+            // Note: We don't use this (it is handled by the bridge). But morsels require a source token.
+            let source_token = SourceToken::new();
+
+            let mut n_rows_seen: usize = 0;
+
+            while let Some(df) = get_batch_state.next(&exec_state)? {
+                n_rows_seen = n_rows_seen.saturating_add(df.height());
+
+                if morsel_sender
+                    .send_morsel(Morsel::new(df, MorselSeq::new(seq), source_token.clone()))
+                    .await
+                    .is_err()
+                {
+                    break;
+                };
+                seq = seq.saturating_add(1);
+            }
+
+            if let Some(mut row_position_on_end_tx) = row_position_on_end_tx {
+                let n_rows_seen = IdxSize::try_from(n_rows_seen)
+                    .map_err(|_| polars_err!(bigidx, ctx = "batch reader", size = n_rows_seen))?;
+
+                _ = row_position_on_end_tx.try_send(n_rows_seen)
+            }
+
+            if let Some(mut n_rows_in_file_tx) = n_rows_in_file_tx {
+                if verbose {
+                    eprintln!("[BatchFnReader]: read to end for full row count");
+                }
+
+                while let Some(df) = get_batch_state.next(&exec_state)? {
+                    n_rows_seen = n_rows_seen.saturating_add(df.height());
                 }
+
+                let n_rows_seen = IdxSize::try_from(n_rows_seen)
+                    .map_err(|_| polars_err!(bigidx, ctx = "batch reader", size = n_rows_seen))?;
+
+                _ = n_rows_in_file_tx.try_send(n_rows_seen)
             }
 
             Ok(())
-        }));
+        });
+
+        Ok((morsel_rx, handle))
+    }
+}
+
+impl BatchFnReader {
+    pub fn _file_schema(&mut self) -> PolarsResult<SchemaRef> {
+        if self.output_schema.is_none() {
+            let exec_state = StreamingExecutionState::default();
+
+            let schema =
+                if let Some(df) = self.get_batch_state.as_mut().unwrap().peek(&exec_state)? {
+                    df.schema().clone()
+                } else {
+                    SchemaRef::default()
+                };
+
+            self.output_schema = Some(schema);
+        }
+
+        Ok(self.output_schema.clone().unwrap())
     }
 }