fixes .list.to_array() with sliced df (#16733)

pola-rs · Jun 5, 2024 · a2a4157 · a2a4157
1 parent ddf8126
commit a2a4157
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/crates/polars-arrow/src/compute/cast/mod.rs b/crates/polars-arrow/src/compute/cast/mod.rs
@@ -175,7 +175,8 @@ fn cast_list_to_fixed_size_list<O: Offset>(
     let null_cnt = list.null_count();
     let new_values = if null_cnt == 0 {
         let offsets = list.offsets().buffer().iter();
-        let expected = (0..list.len()).map(|ix| O::from_as_usize(ix * size));
+        let expected =
+            (list.offsets().first().to_usize()..list.len()).map(|ix| O::from_as_usize(ix * size));
 
         match offsets
             .zip(expected)

diff --git a/py-polars/tests/unit/operations/namespaces/list/test_list.py b/py-polars/tests/unit/operations/namespaces/list/test_list.py
@@ -735,10 +735,14 @@ def test_list_to_array() -> None:
     s = pl.Series(data, dtype=pl.List(pl.Float32))
 
     result = s.list.to_array(2)
+    result_slice = s[1:].list.to_array(2)
 
     expected = pl.Series(data, dtype=pl.Array(pl.Float32, 2))
     assert_series_equal(result, expected)
 
+    expected_slice = pl.Series([data[1]], dtype=pl.Array(pl.Float32, 2))
+    assert_series_equal(result_slice, expected_slice)
+
     # test logical type
     df = pl.DataFrame(
         data={"duration": [[1000, 2000], None]},