refactor(rust): Add updated multiscan pipeline (#21925)

Author: nameexhaustion

crates/polars-io/src/predicates.rs

@@ -471,6 +471,10 @@ pub struct ScanIOPredicate {
 }
 impl ScanIOPredicate {
     pub fn set_external_constant_columns(&mut self, constant_columns: Vec<(PlSmallStr, Scalar)>) {
+        if constant_columns.is_empty() {
+            return;
+        }
+
         let mut live_columns = self.live_columns.as_ref().clone();
         for (c, _) in constant_columns.iter() {
             live_columns.swap_remove(c);

crates/polars-plan/src/dsl/scan_sources.rs

@@ -100,6 +100,14 @@ impl ScanSource {
     pub fn run_async(&self) -> bool {
         self.as_scan_source_ref().run_async()
     }
+
+    pub fn is_cloud_url(&self) -> bool {
+        if let ScanSource::Path(path) = self {
+            polars_io::is_cloud_url(path.as_ref())
+        } else {
+            false
+        }
+    }
 }
 
 /// An iterator for [`ScanSources`]

crates/polars-plan/src/plans/hive.rs

@@ -102,6 +102,12 @@ impl HivePartitionsDf {
     }
 }
 
+impl From<DataFrame> for HivePartitionsDf {
+    fn from(value: DataFrame) -> Self {
+        Self(value)
+    }
+}
+
 /// Note: Returned hive partitions are ordered by their position in the `reader_schema`
 ///
 /// # Safety

crates/polars-stream/src/nodes/io_sources/multi_file_reader/extra_ops/apply.rs

@@ -12,6 +12,7 @@ use polars_io::predicates::ScanIOPredicate;
 use polars_plan::dsl::ScanSource;
 use polars_plan::plans::hive::HivePartitionsDf;
 use polars_utils::IdxSize;
+use polars_utils::slice_enum::Slice;
 
 use super::ExtraOperations;
 use super::cast_columns::CastColumns;
@@ -37,6 +38,7 @@ pub enum ApplyExtraOps {
     Initialized {
         // Note: These fields are ordered according to when they (should be) applied.
         row_index: Option<RowIndex>,
+        pre_slice: Option<Slice>,
         cast_columns: Option<CastColumns>,
         /// This will have include_file_paths, hive columns, missing columns.
         extra_columns: Vec<ScalarColumn>,
@@ -75,8 +77,10 @@ impl ApplyExtraOps {
                 scan_source_idx,
                 hive_parts,
             } => {
-                // This should always be pushed to the reader, or otherwise handled separately.
-                assert!(pre_slice.is_none());
+                // Negative slice should have been resolved earlier.
+                if let Some(Slice::Negative { .. }) = pre_slice {
+                    panic!("impl error: negative pre_slice at post")
+                }
 
                 let cast_columns = CastColumns::try_init_from_policy(
                     cast_columns_policy,
@@ -126,6 +130,7 @@ impl ApplyExtraOps {
 
                 let mut slf = Self::Initialized {
                     row_index,
+                    pre_slice,
                     cast_columns,
                     extra_columns,
                     predicate,
@@ -166,6 +171,7 @@ impl ApplyExtraOps {
                 let slf = match slf {
                     Initialized {
                         row_index: None,
+                        pre_slice: None,
                         cast_columns: None,
                         extra_columns,
                         predicate: None,
@@ -191,6 +197,7 @@ impl ApplyExtraOps {
     ) -> PolarsResult<()> {
         let Self::Initialized {
             row_index,
+            pre_slice,
             cast_columns,
             extra_columns,
             predicate,
@@ -218,6 +225,17 @@ impl ApplyExtraOps {
             };
         }
 
+        if let Some(pre_slice) = pre_slice.clone() {
+            let Slice::Positive { offset, len } = pre_slice
+                .offsetted(usize::try_from(current_row_position).unwrap())
+                .restrict_to_bounds(df.height())
+            else {
+                unreachable!()
+            };
+
+            *df = df.slice(i64::try_from(offset).unwrap(), len)
+        }
+
         if let Some(cast_columns) = cast_columns {
             cast_columns.apply_cast(df)?;
         }

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

@@ -0,0 +1,115 @@
+pub mod predicate;
+pub mod slice;
+
+use std::sync::{Arc, Mutex};
+
+use polars_error::PolarsResult;
+
+use super::MultiFileReaderConfig;
+use super::bridge::BridgeState;
+use crate::async_executor::{self, AbortOnDropHandle, TaskPriority};
+use crate::async_primitives::connector::{self};
+use crate::async_primitives::wait_group::WaitToken;
+use crate::morsel::Morsel;
+use crate::nodes::io_sources::multi_file_reader::bridge::spawn_bridge;
+
+pub struct MultiScanTaskInitializer {
+    pub(super) config: Arc<MultiFileReaderConfig>,
+}
+
+impl MultiScanTaskInitializer {
+    pub fn new(config: Arc<MultiFileReaderConfig>) -> Self {
+        Self { config }
+    }
+
+    #[expect(clippy::type_complexity)]
+    pub fn spawn_background_tasks(
+        self,
+    ) -> (
+        AbortOnDropHandle<PolarsResult<()>>,
+        connector::Sender<(connector::Sender<Morsel>, WaitToken)>,
+        Arc<Mutex<BridgeState>>,
+    ) {
+        assert!(self.config.num_pipelines() > 0);
+        let verbose = self.config.verbose();
+
+        if verbose {
+            eprintln!(
+                "[MultiScanTaskInitializer]: spawn_background_tasks(), {} sources, reader name: {}, {:?}",
+                self.config.sources.len(),
+                self.config.file_reader_builder.reader_name(),
+                self.config.file_reader_builder.reader_capabilities(),
+            )
+        }
+
+        let bridge_state = Arc::new(Mutex::new(BridgeState::NotYetStarted));
+
+        let (bridge_handle, bridge_recv_port_tx, send_phase_chan_to_bridge) =
+            spawn_bridge(bridge_state.clone());
+
+        let verbose = self.config.verbose();
+
+        let background_tasks_handle = AbortOnDropHandle::new(async_executor::spawn(
+            TaskPriority::Low,
+            async move {
+                let (skip_files_mask, predicate) = self.initialize_predicate()?;
+
+                if verbose {
+                    eprintln!(
+                        "[MultiScanTaskInitializer]: \
+                        predicate: {:?}, \
+                        skip files mask: {:?}, \
+                        predicate to reader: {:?}",
+                        self.config.predicate.is_some().then_some("<predicate>"),
+                        skip_files_mask.is_some().then_some("<skip_files>"),
+                        predicate.is_some().then_some("<predicate>"),
+                    )
+                }
+
+                #[expect(clippy::never_loop)]
+                loop {
+                    if skip_files_mask
+                        .as_ref()
+                        .is_some_and(|x| x.unset_bits() == 0)
+                    {
+                        if verbose {
+                            eprintln!(
+                                "[MultiScanTaskInitializer]: early return (skip_files_mask / predicate)"
+                            )
+                        }
+                    } else if self.config.pre_slice.as_ref().is_some_and(|x| x.len() == 0) {
+                        if cfg!(debug_assertions) {
+                            panic!("should quit earlier");
+                        }
+
+                        if verbose {
+                            eprintln!(
+                                "[MultiScanTaskInitializer]: early return (pre_slice.len == 0)"
+                            )
+                        }
+                    } else {
+                        break;
+                    }
+
+                    return Ok(());
+                }
+
+                let predicate = predicate.cloned();
+
+                self.init_and_run(bridge_recv_port_tx, skip_files_mask, predicate)
+                    .await?
+                    .await?;
+
+                bridge_handle.await;
+
+                Ok(())
+            },
+        ));
+
+        (
+            background_tasks_handle,
+            send_phase_chan_to_bridge,
+            bridge_state,
+        )
+    }
+}

crates/polars-stream/src/nodes/io_sources/multi_file_reader/initialization/predicate.rs

@@ -0,0 +1,32 @@
+use arrow::bitmap::Bitmap;
+use polars_error::PolarsResult;
+use polars_io::predicates::ScanIOPredicate;
+
+use super::MultiScanTaskInitializer;
+
+impl MultiScanTaskInitializer {
+    /// # Returns
+    /// `(skip_files_mask, scan_predicate)`
+    ///
+    /// TODO: Move logic here, rename to `evaluate_on_constant_columns`.
+    pub fn initialize_predicate(&self) -> PolarsResult<(Option<Bitmap>, Option<&ScanIOPredicate>)> {
+        if let Some(predicate) = &self.config.predicate {
+            if let Some(hive_parts) = self.config.hive_parts.as_ref() {
+                let (skip_files_mask, need_pred_for_inner_readers) =
+                    crate::nodes::io_sources::multi_scan::scan_predicate_to_mask(
+                        predicate,
+                        self.config.projected_file_schema.as_ref(),
+                        hive_parts.schema(),
+                        hive_parts,
+                    )?;
+
+                return Ok((
+                    skip_files_mask,
+                    need_pred_for_inner_readers.then_some(predicate),
+                ));
+            }
+        }
+
+        Ok((None, self.config.predicate.as_ref()))
+    }
+}