ensure schema update of row_count arguments (#4280)

polars/polars-lazy/src/frame/ipc.rs

@@ -29,12 +29,19 @@ impl LazyFrame {
             n_rows: args.n_rows,
             cache: args.cache,
             with_columns: None,
-            row_count: args.row_count,
+            row_count: None,
             rechunk: args.rechunk,
             memmap: args.memmap,
         };
+        let row_count = args.row_count;
         let mut lf: LazyFrame = LogicalPlanBuilder::scan_ipc(path, options)?.build().into();
         lf.opt_state.file_caching = true;
+
+        // it is a bit hacky, but this row_count function updates the schema
+        if let Some(row_count) = row_count {
+            lf = lf.with_row_count(&row_count.name, Some(row_count.offset))
+        }
+
         Ok(lf)
     }
 

polars/polars-lazy/src/frame/parquet.rs

@@ -37,10 +37,16 @@ impl LazyFrame {
         low_memory: bool,
     ) -> Result<Self> {
         let mut lf: LazyFrame = LogicalPlanBuilder::scan_parquet(
-            path, n_rows, cache, parallel, row_count, rechunk, low_memory,
+            path, n_rows, cache, parallel, None, rechunk, low_memory,
         )?
         .build()
         .into();
+
+        // it is a bit hacky, but this row_count function updates the schema
+        if let Some(row_count) = row_count {
+            lf = lf.with_row_count(&row_count.name, Some(row_count.offset))
+        }
+
         lf.opt_state.file_caching = true;
         Ok(lf)
     }

py-polars/tests/io/test_lazy_ipc.py

@@ -39,3 +39,9 @@ def test_is_in_type_coercion(foods_ipc: str) -> None:
         .collect()
     )
     assert out.shape == (7, 1)
+
+
+def test_row_count_schema(foods_ipc: str) -> None:
+    assert (
+        pl.scan_ipc(foods_ipc, row_count_name="id").select(["id", "category"]).collect()
+    ).dtypes == [pl.UInt32, pl.Utf8]

py-polars/tests/io/test_lazy_parquet.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from os import path
 from pathlib import Path
 
@@ -101,3 +102,10 @@ def test_parquet_stats(io_test_dir: str) -> None:
     assert (
         pl.scan_parquet(file).filter(4 < pl.col("a")).select(pl.col("a").sum())
     ).collect()[0, "a"] == 10.0
+
+
+def test_row_count_schema(io_test_dir: str) -> None:
+    f = os.path.join(io_test_dir, "..", "files", "small.parquet")
+    assert (
+        pl.scan_parquet(f, row_count_name="id").select(["id", "b"]).collect()
+    ).dtypes == [pl.UInt32, pl.Utf8]