improve utf8 take no null

pola-rs · Nov 25, 2021 · afa250b · afa250b
1 parent 05b9762
commit afa250b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/polars/polars-arrow/src/kernels/take.rs b/polars/polars-arrow/src/kernels/take.rs
@@ -266,10 +266,9 @@ pub unsafe fn take_no_null_utf8_iter_unchecked<I: IntoIterator<Item = usize>>(
 ) -> Arc<LargeStringArray> {
     let iter = indices.into_iter().map(|idx| {
         debug_assert!(idx < arr.len());
-        Some(arr.value_unchecked(idx))
+        arr.value_unchecked(idx)
     });
-
-    Arc::new(LargeStringArray::from_trusted_len_iter_unchecked(iter))
+    Arc::new(MutableUtf8Array::<i64>::from_trusted_len_values_iter_unchecked(iter).into())
 }
 
 /// # Safety
@@ -282,6 +281,7 @@ pub unsafe fn take_utf8_iter_unchecked<I: IntoIterator<Item = usize>>(
 ) -> Arc<LargeStringArray> {
     let validity = arr.validity().expect("should have nulls");
     let iter = indices.into_iter().map(|idx| {
+        debug_assert!(idx < arr.len());
         if validity.get_bit_unchecked(idx) {
             Some(arr.value_unchecked(idx))
         } else {