diff --git a/crates/polars-core/src/chunked_array/ops/gather.rs b/crates/polars-core/src/chunked_array/ops/gather.rs index 5124cd7deb49..0d4955cca163 100644 --- a/crates/polars-core/src/chunked_array/ops/gather.rs +++ b/crates/polars-core/src/chunked_array/ops/gather.rs @@ -99,13 +99,12 @@ unsafe fn gather_idx_array_unchecked( ) -> A { let it = indices.iter().copied(); if targets.len() == 1 { - let arr = targets.iter().next().unwrap(); + let target = targets.first().unwrap(); if has_nulls { - it.map(|i| arr.get_unchecked(i as usize)) + it.map(|i| target.get_unchecked(i as usize)) .collect_arr_trusted_with_dtype(dtype) } else { - it.map(|i| arr.value_unchecked(i as usize)) - .collect_arr_trusted_with_dtype(dtype) + target.gather_unchecked_trusted(indices.iter().map(|i| *i as usize), dtype) } } else { let cumlens = cumulative_lengths(targets); @@ -152,25 +151,21 @@ impl ChunkTakeUnchecked for ChunkedArray { let targets: Vec<_> = ca.downcast_iter().collect(); let chunks = indices.downcast_iter().map(|idx_arr| { + let dtype = ca.dtype().clone(); if idx_arr.null_count() == 0 { - gather_idx_array_unchecked( - ca.dtype().clone(), - &targets, - targets_have_nulls, - idx_arr.values(), - ) + gather_idx_array_unchecked(dtype, &targets, targets_have_nulls, idx_arr.values()) } else if targets.len() == 1 { let target = targets.first().unwrap(); if targets_have_nulls { idx_arr .iter() .map(|i| target.get_unchecked(*i? as usize)) - .collect_arr_trusted_with_dtype(ca.dtype().clone()) + .collect_arr_trusted_with_dtype(dtype) } else { idx_arr .iter() .map(|i| Some(target.value_unchecked(*i? as usize))) - .collect_arr_trusted_with_dtype(ca.dtype().clone()) + .collect_arr_trusted_with_dtype(dtype) } } else { let cumlens = cumulative_lengths(&targets); @@ -178,12 +173,12 @@ impl ChunkTakeUnchecked for ChunkedArray { idx_arr .iter() .map(|i| target_get_unchecked(&targets, &cumlens, *i?)) - .collect_arr_trusted_with_dtype(ca.dtype().clone()) + .collect_arr_trusted_with_dtype(dtype) } else { idx_arr .iter() .map(|i| Some(target_value_unchecked(&targets, &cumlens, *i?))) - .collect_arr_trusted_with_dtype(ca.dtype().clone()) + .collect_arr_trusted_with_dtype(dtype) } } }); diff --git a/crates/polars-core/src/datatypes/static_array.rs b/crates/polars-core/src/datatypes/static_array.rs index d4cc0aa61960..2fcfffba1180 100644 --- a/crates/polars-core/src/datatypes/static_array.rs +++ b/crates/polars-core/src/datatypes/static_array.rs @@ -1,5 +1,6 @@ use arrow::bitmap::utils::{BitmapIter, ZipValidity}; use arrow::bitmap::Bitmap; +use polars_arrow::trusted_len::TrustedLenPush; #[cfg(feature = "object")] use crate::chunked_array::object::{ObjectArray, ObjectValueIter}; @@ -54,6 +55,11 @@ pub trait StaticArray: /// # Safety /// It is the callers responsibility that the `idx < self.len()`. unsafe fn value_unchecked(&self, idx: usize) -> Self::ValueT<'_>; + + #[inline] + unsafe fn gather_unchecked_trusted + TrustedLen>(&self, it: I, dtype: DataType) -> Self { + it.map(|i| self.value_unchecked(i)).collect_arr_with_dtype(dtype) + } fn iter(&self) -> ZipValidity, Self::ValueIterT<'_>, BitmapIter>; fn values_iter(&self) -> Self::ValueIterT<'_>; @@ -77,6 +83,13 @@ impl StaticArray for PrimitiveArray { self.values_iter().copied() } + #[inline] + unsafe fn gather_unchecked_trusted + TrustedLen>(&self, it: I, _dtype: DataType) -> Self { + let arr: &[T] = self.values(); + let v = Vec::from_trusted_len_iter(it.map(|i| *arr.get_unchecked(i))); + PrimitiveArray::from_vec(v) + } + fn iter(&self) -> ZipValidity, Self::ValueIterT<'_>, BitmapIter> { ZipValidity::new_with_validity(self.values().iter().copied(), self.validity()) } diff --git a/crates/polars-core/src/datatypes/static_array_collect.rs b/crates/polars-core/src/datatypes/static_array_collect.rs index 4920ad9c4a9a..cfb24c68a41e 100644 --- a/crates/polars-core/src/datatypes/static_array_collect.rs +++ b/crates/polars-core/src/datatypes/static_array_collect.rs @@ -298,10 +298,12 @@ macro_rules! impl_trusted_collect_vec_validity { } impl ArrayFromIter for PrimitiveArray { + #[inline] fn arr_from_iter>(iter: I) -> Self { PrimitiveArray::from_vec(iter.into_iter().collect()) } + #[inline] fn arr_from_iter_trusted(iter: I) -> Self where I: IntoIterator, @@ -310,11 +312,13 @@ impl ArrayFromIter for PrimitiveArray { PrimitiveArray::from_vec(Vec::from_trusted_len_iter(iter)) } + #[inline] fn try_arr_from_iter>>(iter: I) -> Result { let v: Result, E> = iter.into_iter().collect(); Ok(PrimitiveArray::from_vec(v?)) } + #[inline] fn try_arr_from_iter_trusted(iter: I) -> Result where I: IntoIterator>,