use array_ref in favor of chunks (#3368)

pola-rs · May 11, 2022 · 7166080 · 7166080
1 parent 272368b
commit 7166080
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 4 deletions.
diff --git a/polars/polars-core/src/chunked_array/kernels/take.rs b/polars/polars-core/src/chunked_array/kernels/take.rs
@@ -50,7 +50,8 @@ pub(crate) unsafe fn take_list_unchecked(
             vec![Arc::new(list_indices) as ArrayRef],
         ))
         .unwrap();
-    let taken = taken.chunks()[0].clone();
+
+    let taken = taken.array_ref(0).clone();
 
     let validity =
         // if null count > 0

diff --git a/polars/polars-core/src/chunked_array/ops/rolling_window.rs b/polars/polars-core/src/chunked_array/ops/rolling_window.rs
@@ -306,7 +306,7 @@ mod inner_mod {
             let arr = ca.downcast_iter().next().unwrap();
             let series_container =
                 ChunkedArray::<T>::from_slice("", &[T::Native::zero()]).into_series();
-            let array_ptr = &series_container.chunks()[0];
+            let array_ptr = &series_container.array_ref(0);
             let ptr = Arc::as_ptr(array_ptr) as *mut dyn Array as *mut PrimitiveArray<T::Native>;
             let mut builder = PrimitiveChunkedBuilder::<T>::new(self.name(), self.len());
 

diff --git a/polars/polars-core/src/series/ops/to_list.rs b/polars/polars-core/src/series/ops/to_list.rs
@@ -21,7 +21,7 @@ impl Series {
     /// `[1, 2, 3]` becomes `[[1, 2, 3]]`
     pub fn to_list(&self) -> Result<ListChunked> {
         let s = self.rechunk();
-        let values = &s.chunks()[0];
+        let values = s.array_ref(0);
 
         let offsets = vec![0i64, values.len() as i64];
         let inner_type = self.dtype();

diff --git a/polars/polars-core/src/series/unstable.rs b/polars/polars-core/src/series/unstable.rs
@@ -25,7 +25,7 @@ pub type ArrayBox = Box<dyn Array>;
 
 impl<'a> UnstableSeries<'a> {
     pub fn new(series: &'a Series) -> Self {
-        let inner_chunk = &series.chunks()[0];
+        let inner_chunk = series.array_ref(0);
         UnstableSeries {
             container: series,
             inner: NonNull::new(inner_chunk as *const ArrayRef as *mut ArrayRef).unwrap(),

diff --git a/py-polars/tests/test_struct.py b/py-polars/tests/test_struct.py
@@ -271,3 +271,10 @@ def test_list_to_struct() -> None:
         {"field_0": 1, "field_1": 2, "field_2": None},
         {"field_0": 1, "field_1": 2, "field_2": 1},
     ]
+
+
+def test_sort_df_with_list_struct() -> None:
+    assert pl.DataFrame([{"a": 1, "b": [{"c": 1}]}]).sort("a").to_dict(False) == {
+        "a": [1],
+        "b": [[{"c": 1}]],
+    }