feat: Rewrite IR::Scan to IR::DataFrameScan in expand_datasets when applicable (#25106)

nameexhaustion · web-flow · commit 7d5d966d9e14 · 2025-11-03T14:19:42.000+01:00
diff --git a/crates/polars-core/src/utils/mod.rs b/crates/polars-core/src/utils/mod.rs
@@ -1320,22 +1320,6 @@ pub fn coalesce_nulls_columns(a: &Column, b: &Column) -> (Column, Column) {
     }
 }
 
-pub fn operation_exceeded_idxsize_msg(operation: &str) -> String {
-    if size_of::<IdxSize>() == size_of::<u32>() {
-        format!(
-            "{} exceeded the maximum supported limit of {} rows. Consider installing 'polars[rt64]'.",
-            operation,
-            IdxSize::MAX,
-        )
-    } else {
-        format!(
-            "{} exceeded the maximum supported limit of {} rows.",
-            operation,
-            IdxSize::MAX,
-        )
-    }
-}
-
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/crates/polars-plan/src/dsl/file_scan/mod.rs b/crates/polars-plan/src/dsl/file_scan/mod.rs
@@ -297,7 +297,9 @@ pub struct UnifiedScanArgs {
     pub table_statistics: Option<TableStatistics>,
     /// Stores (physical, deleted) row counts of the table if known upfront (e.g. for Iceberg).
     /// This allows for row-count queries to succeed without scanning all files.
-    pub row_count: Option<(IdxSize, IdxSize)>,
+    ///
+    /// Note, intentionally store u64 instead of IdxSize to avoid erroring if it's unused.
+    pub row_count: Option<(u64, u64)>,
 }
 
 impl UnifiedScanArgs {
diff --git a/crates/polars-plan/src/plans/optimizer/expand_datasets.rs b/crates/polars-plan/src/plans/optimizer/expand_datasets.rs
@@ -3,10 +3,12 @@ use std::sync::Arc;
 
 use polars_core::config;
 use polars_core::error::{PolarsResult, polars_bail};
+use polars_core::frame::DataFrame;
 use polars_utils::arena::{Arena, Node};
 use polars_utils::pl_str::PlSmallStr;
 #[cfg(feature = "python")]
 use polars_utils::python_function::PythonObject;
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 use polars_utils::{format_pl_smallstr, unitvec};
 
@@ -330,6 +332,64 @@ pub(super) fn expand_datasets(
         }
 
         apply_scan_predicate_to_scan_ir(node, ir_arena, expr_arena)?;
+
+        let output_schema = ir_arena.get(node).schema(ir_arena);
+
+        let IR::Scan {
+            sources,
+            file_info: _,
+            hive_parts: _,
+            predicate,
+            predicate_file_skip_applied: _,
+            output_schema: _,
+            scan_type,
+            unified_scan_args,
+        } = ir_arena.get(node)
+        else {
+            unreachable!()
+        };
+
+        let df: DataFrame = if (sources.is_empty()
+            && !matches!(scan_type.as_ref(), FileScanIR::Anonymous { .. }))
+            || unified_scan_args
+                .pre_slice
+                .as_ref()
+                .is_some_and(|slice| slice.len() == 0)
+        {
+            if config::verbose() {
+                eprintln!("expand_datasets: scan IR replaced with empty DataFrameScan")
+            }
+
+            DataFrame::empty_with_schema(output_schema.as_ref())
+        } else if output_schema.is_empty()
+            && let Some((physical_rows, deleted_rows)) = unified_scan_args.row_count
+            && unified_scan_args.pre_slice.is_none()
+            && predicate.is_none()
+        {
+            let row_counter = RowCounter::new(physical_rows, deleted_rows);
+            row_counter.num_rows_idxsize()?;
+            let num_rows = row_counter.num_rows()?;
+
+            if config::verbose() {
+                eprintln!(
+                    "expand_datasets: scan IR replaced with 0-width DataFrameScan with height {} ({:?})",
+                    num_rows, &row_counter
+                )
+            }
+
+            DataFrame::empty_with_height(num_rows)
+        } else {
+            continue;
+        };
+
+        let schema = df.schema().clone();
+        let new_ir = IR::DataFrameScan {
+            df: Arc::new(df),
+            schema: schema.clone(),
+            output_schema: Some(schema),
+        };
+
+        ir_arena.replace(node, new_ir);
     }
 
     Ok(())
diff --git a/crates/polars-python/src/io/mod.rs b/crates/polars-python/src/io/mod.rs
@@ -66,7 +66,7 @@ impl PyScanOptions<'_> {
             retries: usize,
             deletion_files: Option<Wrap<DeletionFilesList>>,
             table_statistics: Option<Wrap<TableStatistics>>,
-            row_count: Option<(IdxSize, IdxSize)>,
+            row_count: Option<(u64, u64)>,
         }
 
         let Extract {
diff --git a/crates/polars-stream/src/nodes/io_sources/multi_scan/components/apply_extra_ops.rs b/crates/polars-stream/src/nodes/io_sources/multi_scan/components/apply_extra_ops.rs
@@ -10,13 +10,13 @@ use polars_io::RowIndex;
 use polars_io::predicates::ScanIOPredicate;
 use polars_plan::dsl::{CastColumnsPolicy, MissingColumnsPolicy, ScanSource};
 use polars_plan::plans::hive::HivePartitionsDf;
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 
 use crate::nodes::io_sources::multi_scan::components::column_selector::ColumnSelector;
 use crate::nodes::io_sources::multi_scan::components::column_selector::builder::ColumnSelectorBuilder;
 use crate::nodes::io_sources::multi_scan::components::errors::missing_column_err;
 use crate::nodes::io_sources::multi_scan::components::projection::Projection;
-use crate::nodes::io_sources::multi_scan::components::row_counter::RowCounter;
 use crate::nodes::io_sources::multi_scan::components::row_deletions::ExternalFilterMask;
 use crate::nodes::io_sources::multi_scan::pipeline::models::ExtraOperations;
 
diff --git a/crates/polars-stream/src/nodes/io_sources/multi_scan/components/mod.rs b/crates/polars-stream/src/nodes/io_sources/multi_scan/components/mod.rs
@@ -8,5 +8,4 @@ pub mod forbid_extra_columns;
 pub mod physical_slice;
 pub mod projection;
 pub mod reader_operation_pushdown;
-pub mod row_counter;
 pub mod row_deletions;
diff --git a/crates/polars-stream/src/nodes/io_sources/multi_scan/components/physical_slice.rs b/crates/polars-stream/src/nodes/io_sources/multi_scan/components/physical_slice.rs
@@ -1,6 +1,6 @@
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 
-use crate::nodes::io_sources::multi_scan::components::row_counter::RowCounter;
 use crate::nodes::io_sources::multi_scan::components::row_deletions::ExternalFilterMask;
 
 /// Represents a [`Slice`] that has been potentially adjusted to account for deleted rows.
diff --git a/crates/polars-stream/src/nodes/io_sources/multi_scan/functions/resolve_slice.rs b/crates/polars-stream/src/nodes/io_sources/multi_scan/functions/resolve_slice.rs
@@ -4,10 +4,10 @@ use components::row_deletions::DeletionFilesProvider;
 use futures::StreamExt;
 use polars_core::prelude::{InitHashMaps, PlHashMap};
 use polars_error::PolarsResult;
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 
 use crate::async_executor::{self, AbortOnDropHandle, TaskPriority};
-use crate::nodes::io_sources::multi_scan::components::row_counter::RowCounter;
 use crate::nodes::io_sources::multi_scan::pipeline::models::ResolvedSliceInfo;
 use crate::nodes::io_sources::multi_scan::{MultiScanConfig, components};
 
diff --git a/crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/initialization.rs b/crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/initialization.rs
@@ -5,12 +5,12 @@ use futures::StreamExt;
 use polars_core::prelude::PlHashMap;
 use polars_error::PolarsResult;
 use polars_mem_engine::scan_predicate::initialize_scan_predicate;
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 
 use crate::async_executor::{self, AbortOnDropHandle, TaskPriority};
 use crate::async_primitives::connector::{self};
 use crate::nodes::io_sources::multi_scan::components::bridge::{BridgeRecvPort, BridgeState};
-use crate::nodes::io_sources::multi_scan::components::row_counter::RowCounter;
 use crate::nodes::io_sources::multi_scan::components::row_deletions::{
     DeletionFilesProvider, ExternalFilterMask, RowDeletionsInit,
 };
diff --git a/crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/models.rs b/crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/models.rs
@@ -2,7 +2,6 @@
 use std::collections::VecDeque;
 use std::sync::{Arc, Mutex};
 
-use components::row_counter::RowCounter;
 use components::row_deletions::ExternalFilterMask;
 use polars_core::prelude::PlHashMap;
 use polars_core::schema::SchemaRef;
@@ -12,6 +11,7 @@ use polars_io::predicates::ScanIOPredicate;
 use polars_plan::dsl::{CastColumnsPolicy, MissingColumnsPolicy, ScanSource};
 use polars_plan::plans::hive::HivePartitionsDf;
 use polars_utils::pl_str::PlSmallStr;
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 
 use crate::async_executor::AbortOnDropHandle;
diff --git a/crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/tasks/post_apply_extra_ops.rs b/crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/tasks/post_apply_extra_ops.rs
@@ -1,14 +1,14 @@
 use std::sync::Arc;
 
 use polars_error::PolarsResult;
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 
 use crate::async_executor::{self, AbortOnDropHandle, TaskPriority};
 use crate::async_primitives::distributor_channel::distributor_channel;
 use crate::async_primitives::morsel_linearizer::MorselLinearizer;
 use crate::morsel::Morsel;
 use crate::nodes::io_sources::multi_scan::components::apply_extra_ops::ApplyExtraOps;
-use crate::nodes::io_sources::multi_scan::components::row_counter::RowCounter;
 use crate::nodes::io_sources::multi_scan::reader_interface::output::FileReaderOutputRecv;
 
 pub struct PostApplyExtraOps {
diff --git a/crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/tasks/reader_starter.rs b/crates/polars-stream/src/nodes/io_sources/multi_scan/pipeline/tasks/reader_starter.rs
@@ -12,6 +12,7 @@ use polars_error::PolarsResult;
 use polars_mem_engine::scan_predicate::skip_files_mask::SkipFilesMask;
 use polars_plan::dsl::{MissingColumnsPolicy, ScanSource};
 use polars_utils::IdxSize;
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 
 use crate::async_executor::{self, AbortOnDropHandle, TaskPriority};
@@ -23,7 +24,6 @@ use crate::nodes::io_sources::multi_scan::components::errors::missing_column_err
 use crate::nodes::io_sources::multi_scan::components::physical_slice::PhysicalSlice;
 use crate::nodes::io_sources::multi_scan::components::projection::builder::ProjectionBuilder;
 use crate::nodes::io_sources::multi_scan::components::reader_operation_pushdown::ReaderOperationPushdown;
-use crate::nodes::io_sources::multi_scan::components::row_counter::RowCounter;
 use crate::nodes::io_sources::multi_scan::pipeline::models::{
     ExtraOperations, StartReaderArgsConstant, StartReaderArgsPerFile, StartedReaderState,
 };
diff --git a/crates/polars-stream/src/physical_plan/lower_ir.rs b/crates/polars-stream/src/physical_plan/lower_ir.rs
@@ -22,6 +22,7 @@ use polars_plan::prelude::GroupbyOptions;
 use polars_utils::arena::{Arena, Node};
 use polars_utils::itertools::Itertools;
 use polars_utils::pl_str::PlSmallStr;
+use polars_utils::row_counter::RowCounter;
 use polars_utils::slice_enum::Slice;
 use polars_utils::unique_id::UniqueId;
 use polars_utils::{IdxSize, format_pl_smallstr, unique_column_name};
@@ -32,7 +33,6 @@ use super::{PhysNode, PhysNodeKey, PhysNodeKind, PhysStream};
 use crate::nodes::io_sources::multi_scan;
 use crate::nodes::io_sources::multi_scan::components::forbid_extra_columns::ForbidExtraColumns;
 use crate::nodes::io_sources::multi_scan::components::projection::builder::ProjectionBuilder;
-use crate::nodes::io_sources::multi_scan::components::row_counter::RowCounter;
 use crate::nodes::io_sources::multi_scan::reader_interface::builder::FileReaderBuilder;
 use crate::physical_plan::lower_expr::{ExprCache, build_select_stream, lower_exprs};
 use crate::physical_plan::lower_group_by::build_group_by_stream;
diff --git a/crates/polars-utils/src/lib.rs b/crates/polars-utils/src/lib.rs
@@ -34,6 +34,7 @@ pub mod plpath;
 pub mod priority;
 pub mod regex_cache;
 pub mod relaxed_cell;
+pub mod row_counter;
 pub mod select;
 pub mod slice;
 pub mod slice_enum;
diff --git a/crates/polars-utils/src/row_counter.rs b/crates/polars-utils/src/row_counter.rs
@@ -1,7 +1,8 @@
 use std::fmt::Debug;
 
 use polars_error::{PolarsResult, polars_err};
-use polars_utils::IdxSize;
+
+use crate::IdxSize;
 
 /// Tracker counting physical and deleted rows.
 #[derive(Debug, Clone, Copy, Default, PartialEq)]
@@ -14,7 +15,6 @@ pub struct RowCounter {
 
 impl RowCounter {
     /// `usize::MAX` physical rows, 0 deleted rows
-    #[expect(unused)]
     pub const MAX: Self = Self {
         physical_rows: usize::MAX,
         deleted_rows: 0,
@@ -75,7 +75,7 @@ impl RowCounter {
     ///
     /// # Panics
     /// Panics if there are deleted rows and addition overflows.
-    #[inline]
+    #[allow(clippy::should_implement_trait)]
     pub fn add(self, other: Self) -> Self {
         (|| {
             let physical_rows = self.physical_rows.checked_add(other.physical_rows);
@@ -99,7 +99,7 @@ impl RowCounter {
 
     /// # Panics
     /// Panics if subtraction overflows.
-    #[inline]
+    #[allow(clippy::should_implement_trait)]
     pub fn sub(self, other: Self) -> Self {
         let func = |a: usize, b: usize| {
             a.checked_sub(b)
@@ -128,13 +128,23 @@ impl RowCounter {
     }
 
     /// Returns [`RowCounter::num_rows`] as a usize.
-    ///
-    /// # Panics
-    /// Panics if `usize` to `IdxSize` conversion fails.
     #[inline]
-    #[expect(unused)]
     pub fn num_rows_idxsize(&self) -> PolarsResult<IdxSize> {
-        self.num_rows().map(|x| IdxSize::try_from(x).unwrap())
+        self.num_rows().and_then(|x| {
+            IdxSize::try_from(x).map_err(|_| {
+                let consider_installing_64 = if cfg!(feature = "bigidx") {
+                    ""
+                } else {
+                    ". Consider installing 'polars[rt64]'."
+                };
+
+                polars_err!(
+                    ComputeError:
+                    "row count ({}) exceeded maximum supported of {} (counter: {:?}){}",
+                    x, IdxSize::MAX, self, consider_installing_64
+                )
+            })
+        })
     }
 
     #[inline]
@@ -151,7 +161,6 @@ impl RowCounter {
     }
 
     #[inline]
-    #[expect(unused)]
     pub fn num_physical_rows_idxsize_saturating(&self) -> IdxSize {
         IdxSize::try_from(self.physical_rows).unwrap_or(IdxSize::MAX)
     }
diff --git a/py-polars/tests/unit/io/test_hive.py b/py-polars/tests/unit/io/test_hive.py
@@ -1175,7 +1175,7 @@ def test_hive_filter_in_ir(
         )
 
     plan = pl.scan_parquet(tmp_path).filter(pl.col("a") < 0).explain()
-    assert plan.startswith("Parquet SCAN []")
+    assert plan.startswith("DF [")
 
     assert_frame_equal(
         pl.scan_parquet(tmp_path).with_row_index().filter(pl.col("a") == 2).collect(),
diff --git a/py-polars/tests/unit/io/test_iceberg.py b/py-polars/tests/unit/io/test_iceberg.py
@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import contextlib
+import io
 import itertools
 import os
 import pickle
@@ -2054,3 +2055,30 @@ def test_scan_iceberg_fast_count(tmp_path: Path) -> None:
         pl.scan_iceberg(tbl, reader_override="native").select(pl.len()).collect().item()
         == 5
     )
+
+
+def test_scan_iceberg_idxsize_limit() -> None:
+    if isinstance(pl.get_index_type(), pl.UInt64):
+        assert (
+            pl.scan_parquet([b""], schema={}, _row_count=(1 << 32, 0))
+            .select(pl.len())
+            .collect()
+            .item()
+            == 1 << 32
+        )
+
+        return
+
+    f = io.BytesIO()
+
+    pl.DataFrame({"x": 1}).write_parquet(f)
+
+    q = pl.scan_parquet([f.getvalue()], schema={"x": pl.Int64}, _row_count=(1 << 32, 0))
+
+    assert_frame_equal(q.collect(), pl.DataFrame({"x": 1}))
+
+    with pytest.raises(
+        pl.exceptions.ComputeError,
+        match=r"row count \(4294967296\) exceeded maximum supported of 4294967295.*Consider installing 'polars\[rt64\]'.",
+    ):
+        q.select(pl.len()).collect()

Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ impl PyScanOptions<'_> {`
`66`	`66`	`retries: usize,`
`67`	`67`	`deletion_files: Option<Wrap<DeletionFilesList>>,`
`68`	`68`	`table_statistics: Option<Wrap<TableStatistics>>,`
`69`		`- row_count: Option<(IdxSize, IdxSize)>,`
	`69`	`+ row_count: Option<(u64, u64)>,`
`70`	`70`	`}`
`71`	`71`
`72`	`72`	`let Extract {`