refactor(rust): Remove FileType in favor of ReaderCapabilities for new-streaming multiscan (#21881)

Author: nameexhaustion

Cargo.lock

@@ -12,6 +12,7 @@ description = "Private crate for the streaming execution engine for the Polars D
 arrow = { workspace = true }
 async-trait = { workspace = true }
 atomic-waker = { workspace = true }
+bitflags = { workspace = true }
 crossbeam-channel = { workspace = true }
 crossbeam-deque = { workspace = true }
 crossbeam-queue = { workspace = true }

crates/polars-stream/Cargo.toml

@@ -66,8 +66,8 @@ impl ApplyExtraOps {
                     ExtraOperations {
                         row_index,
                         pre_slice,
-                        cast_columns,
-                        missing_columns,
+                        cast_columns_policy,
+                        missing_columns_policy,
                         include_file_paths,
                         predicate,
                     },
@@ -78,15 +78,11 @@ impl ApplyExtraOps {
                 // This should always be pushed to the reader, or otherwise handled separately.
                 assert!(pre_slice.is_none());
 
-                let cast_columns = if let Some(policy) = cast_columns {
-                    CastColumns::try_init_from_policy(
-                        policy,
-                        &final_output_schema,
-                        incoming_schema,
-                    )?
-                } else {
-                    None
-                };
+                let cast_columns = CastColumns::try_init_from_policy(
+                    cast_columns_policy,
+                    &final_output_schema,
+                    incoming_schema,
+                )?;
 
                 let n_expected_extra_columns = final_output_schema.len()
                     - incoming_schema.len()
@@ -95,13 +91,11 @@ impl ApplyExtraOps {
                 let mut extra_columns: Vec<ScalarColumn> =
                     Vec::with_capacity(n_expected_extra_columns);
 
-                if let Some(policy) = missing_columns {
-                    policy.initialize_policy(
-                        &projected_file_schema,
-                        incoming_schema,
-                        &mut extra_columns,
-                    )?;
-                }
+                missing_columns_policy.initialize_policy(
+                    &projected_file_schema,
+                    incoming_schema,
+                    &mut extra_columns,
+                )?;
 
                 if let Some(hive_parts) = hive_parts {
                     extra_columns.extend(hive_parts.df().get_columns().iter().map(|c| {

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

@@ -3,9 +3,10 @@ use polars_core::schema::SchemaRef;
 use polars_error::{PolarsResult, polars_bail};
 
 /// TODO: Eventually move this enum to polars-plan
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default)]
 pub enum CastColumnsPolicy {
     /// Raise an error if the datatypes do not match
+    #[default]
     ErrorOnMismatch,
 }
 

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

@@ -5,8 +5,9 @@ use polars_core::schema::Schema;
 use polars_error::{PolarsResult, polars_bail};
 
 /// TODO: Eventually move this enum to polars-plan
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default)]
 pub enum MissingColumnsPolicy {
+    #[default]
     Raise,
     Insert,
 }

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

@@ -23,8 +23,8 @@ pub struct ExtraOperations {
     // Note: These fields are ordered according to when they (should be) applied.
     pub row_index: Option<RowIndex>,
     pub pre_slice: Option<Slice>,
-    pub cast_columns: Option<CastColumnsPolicy>,
-    pub missing_columns: Option<MissingColumnsPolicy>,
+    pub cast_columns_policy: CastColumnsPolicy,
+    pub missing_columns_policy: MissingColumnsPolicy,
     pub include_file_paths: Option<PlSmallStr>,
     pub predicate: Option<ScanIOPredicate>,
 }

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

@@ -7,27 +7,12 @@ use polars_io::cloud::CloudOptions;
 use polars_plan::dsl::ScanSource;
 
 use super::FileReader;
-
-/// `FileReaderType` to avoid confusion with a `FileType` enum from polars-plan.
-#[derive(Debug, Clone, PartialEq)]
-pub enum FileReaderType {
-    #[cfg(feature = "parquet")]
-    Parquet,
-    #[cfg(feature = "ipc")]
-    #[expect(unused)]
-    Ipc,
-    #[cfg(feature = "csv")]
-    #[expect(unused)]
-    Csv,
-    #[cfg(feature = "json")]
-    NDJson,
-    /// So that we can compile when all feature flags disabled.
-    #[expect(unused)]
-    Unknown,
-}
+use super::capabilities::ReaderCapabilities;
 
 pub trait FileReaderBuilder: Debug + Send + Sync + 'static {
-    fn file_type(&self) -> FileReaderType;
+    fn reader_name(&self) -> &str;
+
+    fn reader_capabilities(&self) -> ReaderCapabilities;
 
     fn build_file_reader(
         &self,

crates/polars-stream/src/nodes/io_sources/multi_file_reader/reader_interface/builder.rs

@@ -0,0 +1,11 @@
+use bitflags::bitflags;
+
+bitflags! {
+    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+    pub struct ReaderCapabilities: u8 {
+        const ROW_INDEX          = 1 << 0;
+        const PRE_SLICE          = 1 << 1;
+        const NEGATIVE_PRE_SLICE = 1 << 2;
+        const FILTER             = 1 << 3;
+    }
+}

crates/polars-stream/src/nodes/io_sources/multi_file_reader/reader_interface/capabilities.rs

@@ -1,16 +1,20 @@
 //! Interface for single-file readers
 
 pub mod builder;
+pub mod capabilities;
 pub mod output;
 
 use async_trait::async_trait;
 use output::FileReaderOutputRecv;
 use polars_core::schema::SchemaRef;
 use polars_error::PolarsResult;
+use polars_io::RowIndex;
+use polars_io::predicates::ScanIOPredicate;
 use polars_utils::IdxSize;
 use polars_utils::slice_enum::Slice;
 
-use super::extra_ops::apply::ApplyExtraOps;
+use super::extra_ops::cast_columns::CastColumnsPolicy;
+use super::extra_ops::missing_columns::MissingColumnsPolicy;
 use crate::async_executor::JoinHandle;
 
 /// Interface to read a single file
@@ -19,24 +23,12 @@ pub trait FileReader: Send + Sync {
     /// Initialize this FileReader. Intended to allow the reader to pre-fetch metadata.
     ///
     /// This must be called before calling any other functions of the FileReader.
-    ///
-    /// Returns the schema of the morsels that this FileReader will return.
     async fn initialize(&mut self) -> PolarsResult<()>;
 
     /// Begin reading the file into morsels.
     fn begin_read(
         &self,
-        // Note: This may contain more columns that what exist in the file. The reader should project
-        // the ones that it finds. The remaining ones will be handled in post.
-        projected_schema: &SchemaRef,
-        extra_ops: ApplyExtraOps,
-
-        num_pipelines: usize,
-        callbacks: FileReaderCallbacks,
-        // TODO
-        // We could introduce dynamic `Option<Box<dyn Any>>` for the reader to use. That would help
-        // with e.g. synchronizing row group prefetches across multiple files in Parquet. Currently
-        // every reader started concurrently will prefetch up to the row group prefetch limit.
+        args: BeginReadArgs,
     ) -> PolarsResult<(FileReaderOutputRecv, JoinHandle<PolarsResult<()>>)>;
 
     /// This FileReader must be initialized before calling this.
@@ -46,15 +38,16 @@ pub trait FileReader: Send + Sync {
     async fn n_rows_in_file(&self) -> PolarsResult<IdxSize> {
         let (tx, rx) = tokio::sync::oneshot::channel();
 
-        let (morsel_receivers, handle) = self.begin_read(
-            &Default::default(), // pass empty schema
-            ApplyExtraOps::Noop,
-            1,
-            FileReaderCallbacks {
+        let (morsel_receivers, handle) = self.begin_read(BeginReadArgs {
+            // Passing 0-0 slice indicates to the reader that we want the full row count, but it can
+            // skip actually reading the data if it is able to.
+            pre_slice: Some(Slice::Positive { offset: 0, len: 0 }),
+            callbacks: FileReaderCallbacks {
                 n_rows_in_file_tx: Some(tx),
                 ..Default::default()
             },
-        )?;
+            ..Default::default()
+        })?;
 
         drop(morsel_receivers);
 
@@ -75,15 +68,14 @@ pub trait FileReader: Send + Sync {
 
         let (tx, rx) = tokio::sync::oneshot::channel();
 
-        let (mut morsel_receivers, handle) = self.begin_read(
-            &Default::default(), // pass empty schema
-            ApplyExtraOps::Noop,
-            1,
-            FileReaderCallbacks {
+        let (mut morsel_receivers, handle) = self.begin_read(BeginReadArgs {
+            pre_slice,
+            callbacks: FileReaderCallbacks {
                 row_position_on_end_tx: Some(tx),
                 ..Default::default()
             },
-        )?;
+            ..Default::default()
+        })?;
 
         // We are using the `row_position_on_end` callback, this means we must fully consume all of
         // the morsels sent by the reader.
@@ -102,7 +94,54 @@ pub trait FileReader: Send + Sync {
     }
 }
 
-#[derive(Default)]
+#[derive(Debug)]
+pub struct BeginReadArgs {
+    /// Columns to project from the file.
+    pub projected_schema: SchemaRef,
+
+    pub row_index: Option<RowIndex>,
+    pub pre_slice: Option<Slice>,
+    pub predicate: Option<ScanIOPredicate>,
+
+    /// User-configured policy for when datatypes do not match.
+    ///
+    /// A reader may wish to use this if it is applying predicates.
+    ///
+    /// This can be ignored by the reader, as the policy is also applied in post.
+    #[expect(unused)]
+    pub cast_columns_policy: CastColumnsPolicy,
+    /// User-configured policy for when columns are not found in the file.
+    ///
+    /// A reader may wish to use this if it is applying predicates.
+    ///
+    /// This can be ignored by the reader, as the policy is also applied in post.
+    pub missing_columns_policy: MissingColumnsPolicy,
+
+    pub num_pipelines: usize,
+    pub callbacks: FileReaderCallbacks,
+    // TODO
+    // We could introduce dynamic `Option<Box<dyn Any>>` for the reader to use. That would help
+    // with e.g. synchronizing row group prefetches across multiple files in Parquet. Currently
+    // every reader started concurrently will prefetch up to the row group prefetch limit.
+}
+
+impl Default for BeginReadArgs {
+    fn default() -> Self {
+        BeginReadArgs {
+            projected_schema: SchemaRef::default(),
+            row_index: None,
+            pre_slice: None,
+            predicate: None,
+            // TODO: Use less restrictive default
+            cast_columns_policy: CastColumnsPolicy::ErrorOnMismatch,
+            missing_columns_policy: MissingColumnsPolicy::Insert,
+            num_pipelines: 1,
+            callbacks: FileReaderCallbacks::default(),
+        }
+    }
+}
+
+#[derive(Debug, Default)]
 /// We have this to avoid a footgun of accidentally swapping the arguments.
 pub struct FileReaderCallbacks {
     /// Full file schema
@@ -112,10 +151,10 @@ pub struct FileReaderCallbacks {
     /// on the source. Prefer instead to use `row_position_on_end`, which can be much faster.
     ///
     /// Notes:
+    /// * All readers must ensure that this count is sent if requested, even if the output port
+    ///   closes prematurely, or a slice is sent.
     /// * Some readers will only send this after their output morsels to be fully consumed (or if
     ///   their output port is dropped), so you should not block morsel consumption on waiting for this.
-    /// * All readers must ensure that this count is sent if requested, even if the output port
-    ///   closes prematurely.
     pub n_rows_in_file_tx: Option<tokio::sync::oneshot::Sender<IdxSize>>,
 
     /// Returns the row position reached by this reader.
@@ -139,19 +178,7 @@ pub fn calc_row_position_after_slice(n_rows_in_file: IdxSize, pre_slice: Option<
 
     let out = match pre_slice {
         None => n_rows_in_file,
-
-        Some(Slice::Positive { offset, len }) => {
-            let slice_end = offset.saturating_add(len);
-            n_rows_in_file.min(slice_end)
-        },
-
-        Some(Slice::Negative {
-            offset_from_end,
-            len,
-        }) => {
-            let n_from_end = offset_from_end.saturating_sub(len);
-            n_rows_in_file.saturating_sub(n_from_end)
-        },
+        Some(v) => v.restrict_to_bounds(n_rows_in_file).end_position(),
     };
 
     IdxSize::try_from(out).unwrap_or(IdxSize::MAX)