From 1feae7a8c3e8739d8370e07b8d5756fb874ad85a Mon Sep 17 00:00:00 2001 From: Luke Manley Date: Thu, 20 Feb 2025 06:17:12 -0500 Subject: [PATCH 1/3] fix: Panic when projecting only row index from IPC file --- crates/polars-io/src/ipc/ipc_file.rs | 3 +++ py-polars/tests/unit/io/test_ipc.py | 8 ++++++++ 2 files changed, 11 insertions(+) diff --git a/crates/polars-io/src/ipc/ipc_file.rs b/crates/polars-io/src/ipc/ipc_file.rs index 26f7d970eed1..9e11618e48e1 100644 --- a/crates/polars-io/src/ipc/ipc_file.rs +++ b/crates/polars-io/src/ipc/ipc_file.rs @@ -261,6 +261,9 @@ impl SerReader for IpcReader { }; let df = DataFrame::empty_with_height(row_count); + if let Some(ri) = &self.row_index { + df.with_row_index_mut(ri.name.clone(), Some(ri.offset)); + } return PolarsResult::Ok(df); } diff --git a/py-polars/tests/unit/io/test_ipc.py b/py-polars/tests/unit/io/test_ipc.py index cd014fd0cc66..1f198cf17663 100644 --- a/py-polars/tests/unit/io/test_ipc.py +++ b/py-polars/tests/unit/io/test_ipc.py @@ -441,3 +441,11 @@ def test_categorical_lexical_sort_2732() -> None: df.write_ipc(f) f.seek(0) assert_frame_equal(df, pl.read_ipc(f)) + + +def test_project_only_row_index_21165() -> None: + f = io.BytesIO() + pl.DataFrame({"a": [1]}).write_ipc(f) + result = pl.scan_ipc(f, row_index_name="ri").select("ri").collect() + expected = pl.DataFrame({"ri": [0]}, schema={"ri": pl.UInt32}) + assert_frame_equal(result, expected) From f549f2febb3df43e71e77554ec8acc4988ebe2c5 Mon Sep 17 00:00:00 2001 From: Luke Manley Date: Thu, 20 Feb 2025 06:22:32 -0500 Subject: [PATCH 2/3] mut df --- crates/polars-io/src/ipc/ipc_file.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/polars-io/src/ipc/ipc_file.rs b/crates/polars-io/src/ipc/ipc_file.rs index 9e11618e48e1..305a312fe8bd 100644 --- a/crates/polars-io/src/ipc/ipc_file.rs +++ b/crates/polars-io/src/ipc/ipc_file.rs @@ -259,7 +259,7 @@ impl SerReader for IpcReader { } else { get_row_count(&mut self.reader)? as usize }; - let df = DataFrame::empty_with_height(row_count); + let mut df = DataFrame::empty_with_height(row_count); if let Some(ri) = &self.row_index { df.with_row_index_mut(ri.name.clone(), Some(ri.offset)); From 81dc6499053b4f4f6d3d4113bf495dc9ed501d4c Mon Sep 17 00:00:00 2001 From: Luke Manley Date: Thu, 20 Feb 2025 06:53:23 -0500 Subject: [PATCH 3/3] fix tests --- py-polars/tests/unit/io/test_ipc.py | 8 -------- py-polars/tests/unit/io/test_scan.py | 10 +--------- 2 files changed, 1 insertion(+), 17 deletions(-) diff --git a/py-polars/tests/unit/io/test_ipc.py b/py-polars/tests/unit/io/test_ipc.py index 1f198cf17663..cd014fd0cc66 100644 --- a/py-polars/tests/unit/io/test_ipc.py +++ b/py-polars/tests/unit/io/test_ipc.py @@ -441,11 +441,3 @@ def test_categorical_lexical_sort_2732() -> None: df.write_ipc(f) f.seek(0) assert_frame_equal(df, pl.read_ipc(f)) - - -def test_project_only_row_index_21165() -> None: - f = io.BytesIO() - pl.DataFrame({"a": [1]}).write_ipc(f) - result = pl.scan_ipc(f, row_index_name="ri").select("ri").collect() - expected = pl.DataFrame({"ri": [0]}, schema={"ri": pl.UInt32}) - assert_frame_equal(result, expected) diff --git a/py-polars/tests/unit/io/test_scan.py b/py-polars/tests/unit/io/test_scan.py index 0ddfefcc5891..78cc33aa8c91 100644 --- a/py-polars/tests/unit/io/test_scan.py +++ b/py-polars/tests/unit/io/test_scan.py @@ -969,15 +969,7 @@ def test_scan_csv_bytesio_memory_usage( "scan_type", [ (pl.DataFrame.write_parquet, pl.scan_parquet), - pytest.param( - (pl.DataFrame.write_ipc, pl.scan_ipc), - marks=[ - pytest.mark.xfail( - reason="Bug. See https://github.com/pola-rs/polars/issues/21165" - ), - pytest.mark.may_fail_auto_streaming, - ], - ), + (pl.DataFrame.write_ipc, pl.scan_ipc), (pl.DataFrame.write_csv, pl.scan_csv), (pl.DataFrame.write_ndjson, pl.scan_ndjson), ],