feat: Add support for io-plugins in new-streaming (#21870)

Author: coastalwhite

Cargo.lock

@@ -4,5 +4,7 @@ mod predicate;
 mod prelude;
 
 pub use executors::Executor;
+#[cfg(feature = "python")]
+pub use planner::python_scan_predicate;
 pub use planner::{create_multiple_physical_plans, create_physical_plan, create_scan_predicate};
 pub use predicate::ScanPredicate;

crates/polars-mem-engine/src/lib.rs

@@ -134,6 +134,58 @@ pub fn create_multiple_physical_plans(
     })
 }
 
+#[cfg(feature = "python")]
+#[allow(clippy::type_complexity)]
+pub fn python_scan_predicate(
+    options: &mut PythonOptions,
+    expr_arena: &Arena<AExpr>,
+    state: &mut ExpressionConversionState,
+) -> PolarsResult<(
+    Option<Arc<dyn polars_expr::prelude::PhysicalExpr>>,
+    Option<Vec<u8>>,
+)> {
+    let mut predicate_serialized = None;
+    let predicate = if let PythonPredicate::Polars(e) = &options.predicate {
+        // Convert to a pyarrow eval string.
+        if matches!(options.python_source, PythonScanSource::Pyarrow) {
+            if let Some(eval_str) = polars_plan::plans::python::pyarrow::predicate_to_pa(
+                e.node(),
+                expr_arena,
+                Default::default(),
+            ) {
+                options.predicate = PythonPredicate::PyArrow(eval_str);
+                // We don't have to use a physical expression as pyarrow deals with the filter.
+                None
+            } else {
+                Some(create_physical_expr(
+                    e,
+                    Context::Default,
+                    expr_arena,
+                    &options.schema,
+                    state,
+                )?)
+            }
+        }
+        // Convert to physical expression for the case the reader cannot consume the predicate.
+        else {
+            let dsl_expr = e.to_expr(expr_arena);
+            predicate_serialized = polars_plan::plans::python::predicate::serialize(&dsl_expr)?;
+
+            Some(create_physical_expr(
+                e,
+                Context::Default,
+                expr_arena,
+                &options.schema,
+                state,
+            )?)
+        }
+    } else {
+        None
+    };
+
+    Ok((predicate, predicate_serialized))
+}
+
 #[recursive]
 fn create_physical_plan_impl(
     root: Node,
@@ -160,45 +212,9 @@ fn create_physical_plan_impl(
     match logical_plan {
         #[cfg(feature = "python")]
         PythonScan { mut options } => {
-            let mut predicate_serialized = None;
-
-            let predicate = if let PythonPredicate::Polars(e) = &options.predicate {
-                let phys_expr = || {
-                    let mut state = ExpressionConversionState::new(true, state.expr_depth);
-                    create_physical_expr(
-                        e,
-                        Context::Default,
-                        expr_arena,
-                        &options.schema,
-                        &mut state,
-                    )
-                };
-
-                // Convert to a pyarrow eval string.
-                if matches!(options.python_source, PythonScanSource::Pyarrow) {
-                    if let Some(eval_str) = polars_plan::plans::python::pyarrow::predicate_to_pa(
-                        e.node(),
-                        expr_arena,
-                        Default::default(),
-                    ) {
-                        options.predicate = PythonPredicate::PyArrow(eval_str);
-                        // We don't have to use a physical expression as pyarrow deals with the filter.
-                        None
-                    } else {
-                        Some(phys_expr()?)
-                    }
-                }
-                // Convert to physical expression for the case the reader cannot consume the predicate.
-                else {
-                    let dsl_expr = e.to_expr(expr_arena);
-                    predicate_serialized =
-                        polars_plan::plans::python::predicate::serialize(&dsl_expr)?;
-
-                    Some(phys_expr()?)
-                }
-            } else {
-                None
-            };
+            let mut expr_conv_state = ExpressionConversionState::new(true, state.expr_depth);
+            let (predicate, predicate_serialized) =
+                python_scan_predicate(&mut options, expr_arena, &mut expr_conv_state)?;
             Ok(Box::new(executors::PythonScanExec {
                 options,
                 predicate,

crates/polars-mem-engine/src/planner/lp.rs

@@ -24,6 +24,7 @@ percent-encoding = { workspace = true }
 pin-project-lite = { workspace = true }
 polars-io = { workspace = true, features = ["async"] }
 polars-utils = { workspace = true }
+pyo3 = { workspace = true, optional = true }
 rand = { workspace = true }
 rayon = { workspace = true }
 recursive = { workspace = true }
@@ -54,7 +55,7 @@ json = ["polars-mem-engine/json", "polars-plan/json", "polars-io/json"]
 cloud = ["polars-mem-engine/cloud", "polars-plan/cloud", "polars-io/cloud"]
 dtype-categorical = ["polars-core/dtype-categorical"]
 object = ["polars-ops/object"]
-python = ["polars-plan/python"]
+python = ["pyo3", "polars-plan/python"]
 
 # We need to specify default features here to match workspace defaults.
 # Otherwise we get warnings with cargo check/clippy.

crates/polars-stream/Cargo.toml

@@ -0,0 +1,118 @@
+use polars_core::frame::DataFrame;
+use polars_core::schema::SchemaRef;
+use polars_error::{PolarsResult, polars_err};
+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::spawn;
+use crate::async_primitives::connector::Receiver;
+use crate::async_primitives::wait_group::WaitGroup;
+use crate::morsel::SourceToken;
+
+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>,
+}
+
+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 SourceNode for BatchSourceNode {
+    fn name(&self) -> &str {
+        self.name.as_str()
+    }
+
+    fn is_source_output_parallel(&self, _is_receiver_serial: bool) -> bool {
+        false
+    }
+
+    fn spawn_source(
+        &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;
+                    }
+                }
+            }
+
+            Ok(())
+        }));
+    }
+}

crates/polars-stream/src/nodes/io_sources/batch.rs

@@ -17,6 +17,7 @@ use crate::nodes::compute_node_prelude::*;
 
 pub mod multi_file_reader;
 
+pub mod batch;
 #[cfg(feature = "csv")]
 pub mod csv;
 #[cfg(feature = "ipc")]

crates/polars-stream/src/nodes/io_sources/mod.rs

@@ -46,6 +46,8 @@ fn visualize_plan_rec(
             ),
             &[][..],
         ),
+        #[cfg(feature = "python")]
+        PhysNodeKind::PythonScan { .. } => ("python-scan".to_string(), &[][..]),
         PhysNodeKind::SinkMultiple { sinks } => {
             for sink in sinks {
                 visualize_plan_rec(*sink, phys_sm, expr_arena, visited, out);

crates/polars-stream/src/physical_plan/fmt.rs

@@ -722,7 +722,9 @@ pub fn lower_ir(
         },
 
         #[cfg(feature = "python")]
-        IR::PythonScan { .. } => todo!(),
+        IR::PythonScan { options } => PhysNodeKind::PythonScan {
+            options: options.clone(),
+        },
 
         IR::Cache {
             input,

crates/polars-stream/src/physical_plan/lower_ir.rs

@@ -227,6 +227,11 @@ pub enum PhysNodeKind {
         file_options: Box<FileScanOptions>,
     },
 
+    #[cfg(feature = "python")]
+    PythonScan {
+        options: polars_plan::plans::python::PythonOptions,
+    },
+
     GroupBy {
         input: PhysStream,
         key: Vec<ExprIR>,
@@ -284,6 +289,8 @@ fn visit_node_inputs_mut(
             | PhysNodeKind::MultiScan { .. }
             | PhysNodeKind::FileScan { .. }
             | PhysNodeKind::InputIndependentSelect { .. } => {},
+            #[cfg(feature = "python")]
+            PhysNodeKind::PythonScan { .. } => {},
             PhysNodeKind::Select { input, .. }
             | PhysNodeKind::WithRowIndex { input, .. }
             | PhysNodeKind::Reduce { input, .. }