Skip to content

Commit

Permalink
Avoid rechunking in get_inner
Browse files Browse the repository at this point in the history
  • Loading branch information
stinodego committed May 16, 2024
1 parent dda2e53 commit d68d580
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 11 deletions.
7 changes: 4 additions & 3 deletions crates/polars-core/src/chunked_array/array/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,13 +31,14 @@ impl ArrayChunked {

/// Get the inner values as `Series`
pub fn get_inner(&self) -> Series {
let ca = self.rechunk();
let field = self.inner_dtype().to_arrow_field("item", true);
let arr = ca.downcast_iter().next().unwrap();
let chunks: Vec<_> = self.downcast_iter().map(|c| c.values().clone()).collect();

// SAFETY: Data type of arrays matches because they are chunks from the same array.
unsafe {
Series::_try_from_arrow_unchecked_with_md(
self.name(),
vec![(arr.values()).clone()],
chunks,
&field.data_type,
Some(&field.metadata),
)
Expand Down
16 changes: 9 additions & 7 deletions crates/polars-core/src/chunked_array/list/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,16 +41,18 @@ impl ListChunked {

/// Get the inner values as [`Series`], ignoring the list offsets.
pub fn get_inner(&self) -> Series {
let ca = self.rechunk();
let arr = ca.downcast_iter().next().unwrap();
// SAFETY:
// Inner dtype is passed correctly
let field = self.inner_dtype().to_arrow_field("item", true);
let chunks: Vec<_> = self.downcast_iter().map(|c| c.values().clone()).collect();

// SAFETY: Data type of arrays matches because they are chunks from the same array.
unsafe {
Series::from_chunks_and_dtype_unchecked(
Series::_try_from_arrow_unchecked_with_md(
self.name(),
vec![arr.values().clone()],
&ca.inner_dtype(),
chunks,
&field.data_type,
Some(&field.metadata),
)
.unwrap()
}
}

Expand Down
2 changes: 1 addition & 1 deletion py-polars/src/series/export.rs
Original file line number Diff line number Diff line change
Expand Up @@ -357,7 +357,7 @@ where
/// Convert arrays by flattening first, converting the flat Series, and then reshaping.
fn array_series_to_numpy(py: Python, s: &Series) -> PyObject {
let ca = s.array().unwrap();
let s_inner = ca.get_inner(); // TODO: This rechunks - is there a way to avoid this?
let s_inner = ca.get_inner();
let np_array_flat = series_to_numpy_with_copy(py, &s_inner).unwrap();

// Reshape to the original shape.
Expand Down

0 comments on commit d68d580

Please sign in to comment.