Skip to content

Commit

Permalink
fix(rust): index_to_chunked_index's fast path is not correct (#11710)
Browse files Browse the repository at this point in the history
  • Loading branch information
reswqa authored Oct 14, 2023
1 parent d30b33c commit 0e0daa9
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 10 deletions.
8 changes: 4 additions & 4 deletions crates/polars-core/src/chunked_array/ops/any_value.rs
Original file line number Diff line number Diff line change
Expand Up @@ -187,12 +187,12 @@ macro_rules! get_any_value_unchecked {

macro_rules! get_any_value {
($self:ident, $index:expr) => {{
let (chunk_idx, idx) = $self.index_to_chunked_index($index);
let arr = &*$self.chunks[chunk_idx];
polars_ensure!(idx < arr.len(), oob = idx, arr.len());
if $index >= $self.len() {
polars_bail!(oob = $index, $self.len());
}
// SAFETY
// bounds are checked
Ok(unsafe { arr_to_any_value(arr, idx, $self.dtype()) })
Ok(unsafe { $self.get_any_value_unchecked($index) })
}};
}

Expand Down
8 changes: 7 additions & 1 deletion crates/polars-core/src/chunked_array/ops/downcast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,13 @@ impl<T: PolarsDataType> ChunkedArray<T> {
#[inline]
pub(crate) fn index_to_chunked_index(&self, index: usize) -> (usize, usize) {
if self.chunks.len() == 1 {
return (0, index);
// SAFETY: chunks.len() == 1 guarantees this is correct.
let len = unsafe { self.chunks.get_unchecked(0).len() };
return if index < len {
(0, index)
} else {
(1, index - len)
};
}
index_to_chunked_index(self.downcast_iter().map(|arr| arr.len()), index)
}
Expand Down
5 changes: 0 additions & 5 deletions crates/polars-ops/src/chunked_array/strings/concat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,11 +65,6 @@ pub fn hor_str_concat(cas: &[&Utf8Chunked], delimiter: &str) -> PolarsResult<Utf
ComputeError: "all series in `hor_str_concat` should have equal or unit length"
);

let has_empty_ca = cas.iter().any(|ca| ca.is_empty());
if has_empty_ca {
return Ok(Utf8Chunked::full_null(cas[0].name(), 0));
}

// Calculate total capacity needed.
let tot_strings_bytes: usize = cas
.iter()
Expand Down

0 comments on commit 0e0daa9

Please sign in to comment.