c

Author: nameexhaustion

crates/polars-expr/src/reduce/convert.rs

@@ -25,6 +25,7 @@ pub fn into_reduction(
     node: Node,
     expr_arena: &mut Arena<AExpr>,
     schema: &Schema,
+    is_aggregation_context: bool,
 ) -> PolarsResult<(Box<dyn GroupedReduction>, Vec<Node>)> {
     let get_dt = |node| {
         expr_arena
@@ -94,6 +95,12 @@ pub fn into_reduction(
                 //   project to the height of the DataFrame (in the PhysicalExpr impl).
                 // * This approach is not sound for `update_groups()`, but currently that case is
                 //   not hit (it would need group-by -> len on empty morsels).
+                polars_ensure!(
+                    !is_aggregation_context,
+                    ComputeError:
+                    "not implemented: len() of groups with no columns"
+                );
+
                 let out: Box<dyn GroupedReduction> = new_sum_reduction(DataType::IDX_DTYPE)?;
                 let expr = expr_arena.add(AExpr::Len);
 

crates/polars-stream/src/nodes/in_memory_source.rs

@@ -22,6 +22,17 @@ impl InMemorySourceNode {
             seq_offset,
         }
     }
+
+    pub fn new_no_morsel_split(source: Arc<DataFrame>, seq_offset: MorselSeq) -> Self {
+        let morsel_size = source.height();
+
+        InMemorySourceNode {
+            source: Some(source),
+            morsel_size,
+            seq: AtomicU64::new(0),
+            seq_offset,
+        }
+    }
 }
 
 impl ComputeNode for InMemorySourceNode {

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

@@ -180,6 +180,7 @@ impl FileReader for BatchFnReader {
             predicate: None,
             cast_columns_policy: _,
             num_pipelines: _,
+            disable_morsel_split: _,
             callbacks:
                 FileReaderCallbacks {
                     mut file_schema_tx,

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

@@ -137,6 +137,7 @@ impl FileReader for CsvFileReader {
             predicate: None,
             cast_columns_policy: _,
             num_pipelines,
+            disable_morsel_split: _,
             callbacks:
                 FileReaderCallbacks {
                     file_schema_tx,

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

@@ -162,6 +162,7 @@ impl FileReader for IpcFileReader {
             predicate: None,
             cast_columns_policy: _,
             num_pipelines,
+            disable_morsel_split,
             callbacks:
                 FileReaderCallbacks {
                     file_schema_tx,
@@ -397,6 +398,20 @@ impl FileReader for IpcFileReader {
                 if df.height() == 0 {
                     continue;
                 }
+
+                if disable_morsel_split {
+                    if morsel_send
+                        .send_morsel(Morsel::new(df, morsel_seq, source_token.clone()))
+                        .await
+                        .is_err()
+                    {
+                        return Ok(());
+                    }
+                    drop(permit);
+                    morsel_seq = morsel_seq.successor();
+                    continue;
+                }
+
                 next = Some((df, permit));
                 break;
             }

crates/polars-stream/src/nodes/io_sources/multi_scan/components/projection/mod.rs

@@ -35,6 +35,10 @@ pub enum Projection {
 }
 
 impl Projection {
+    pub fn is_empty(&self) -> bool {
+        self.projected_schema().is_empty()
+    }
+
     /// Returns the full projected schema, keyed by the output name.
     pub fn projected_schema(&self) -> &SchemaRef {
         match self {

crates/polars-stream/src/nodes/io_sources/multi_scan/components/row_deletions.rs

@@ -157,6 +157,7 @@ impl DeletionFilesProvider {
                                     predicate: None,
                                     cast_columns_policy: CastColumnsPolicy::ERROR_ON_MISMATCH,
                                     num_pipelines,
+                                    disable_morsel_split: false,
                                     callbacks: FileReaderCallbacks {
                                         file_schema_tx: None,
                                         n_rows_in_file_tx: None,

crates/polars-stream/src/nodes/io_sources/multi_scan/config.rs

@@ -49,6 +49,7 @@ pub struct MultiScanConfig {
     /// step.
     pub n_readers_pre_init: RelaxedCell<usize>,
     pub max_concurrent_scans: RelaxedCell<usize>,
+    pub disable_morsel_split: bool,
 
     pub verbose: bool,
 }

crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/initialization.rs

@@ -42,12 +42,14 @@ pub fn initialize_multi_scan_pipeline(
             reader name: {}, \
             {:?}, \
             n_readers_pre_init: {}, \
-            max_concurrent_scans: {}",
+            max_concurrent_scans: {}, \
+            disable_morsel_split: {}",
             config.sources.len(),
             config.file_reader_builder.reader_name(),
             config.reader_capabilities(),
             config.n_readers_pre_init(),
             config.max_concurrent_scans(),
+            config.disable_morsel_split,
         );
     }
 
@@ -363,6 +365,7 @@ async fn finish_initialize_multi_scan_pipeline(
     let final_output_schema = config.final_output_schema.clone();
     let file_projection_builder = config.file_projection_builder.clone();
     let max_concurrent_scans = config.max_concurrent_scans();
+    let disable_morsel_split = config.disable_morsel_split;
 
     let (started_reader_tx, started_reader_rx) =
         tokio::sync::mpsc::channel(max_concurrent_scans.max(2) - 1);
@@ -387,6 +390,7 @@ async fn finish_initialize_multi_scan_pipeline(
                 missing_columns_policy,
                 forbid_extra_columns: config.forbid_extra_columns.clone(),
                 num_pipelines,
+                disable_morsel_split,
                 verbose,
             },
             verbose,

crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/models.rs

@@ -83,6 +83,7 @@ pub(super) struct StartReaderArgsConstant {
     pub(super) missing_columns_policy: MissingColumnsPolicy,
     pub(super) forbid_extra_columns: Option<ForbidExtraColumns>,
     pub(super) num_pipelines: usize,
+    pub(super) disable_morsel_split: bool,
     pub(super) verbose: bool,
 }