more collect_trusted

pola-rs · Sep 23, 2021 · 269247f · 269247f
1 parent a14187d
commit 269247f
Show file tree

Hide file tree

Showing 7 changed files with 65 additions and 33 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/cum_agg.rs b/polars/polars-core/src/chunked_array/ops/cum_agg.rs
@@ -1,4 +1,5 @@
 use crate::prelude::*;
+use crate::utils::CustomIterTools;
 use itertools::__std_iter::FromIterator;
 use num::Bounded;
 use std::ops::{Add, AddAssign};
@@ -74,8 +75,17 @@ where
     fn cummin(&self, reverse: bool) -> ChunkedArray<T> {
         let init = Bounded::max_value();
         let mut ca: Self = match reverse {
-            false => self.into_iter().scan(init, det_min).collect(),
-            true => self.into_iter().rev().scan(init, det_min).collect(),
+            false => self
+                .into_iter()
+                .scan(init, det_min)
+                .trust_my_length(self.len())
+                .collect_trusted(),
+            true => self
+                .into_iter()
+                .rev()
+                .scan(init, det_min)
+                .trust_my_length(self.len())
+                .collect_trusted(),
         };
 
         ca.rename(self.name());
@@ -89,8 +99,17 @@ where
     fn cumsum(&self, reverse: bool) -> ChunkedArray<T> {
         let init = None;
         let mut ca: Self = match reverse {
-            false => self.into_iter().scan(init, det_sum).collect(),
-            true => self.into_iter().rev().scan(init, det_sum).collect(),
+            false => self
+                .into_iter()
+                .scan(init, det_sum)
+                .trust_my_length(self.len())
+                .collect_trusted(),
+            true => self
+                .into_iter()
+                .rev()
+                .scan(init, det_sum)
+                .trust_my_length(self.len())
+                .collect_trusted(),
         };
 
         ca.rename(self.name());

diff --git a/polars/polars-core/src/chunked_array/ops/fill_null.rs b/polars/polars-core/src/chunked_array/ops/fill_null.rs
@@ -19,7 +19,8 @@ where
             }
             None => Some(*previous),
         })
-        .collect()
+        .trust_my_length(ca.len())
+        .collect_trusted()
 }
 
 macro_rules! impl_fill_forward {

diff --git a/polars/polars-core/src/chunked_array/ops/is_in.rs b/polars/polars-core/src/chunked_array/ops/is_in.rs
@@ -1,5 +1,5 @@
 use crate::prelude::*;
-use crate::utils::get_supertype;
+use crate::utils::{get_supertype, CustomIterTools};
 use hashbrown::hash_set::HashSet;
 use num::NumCast;
 use std::hash::Hash;
@@ -33,7 +33,7 @@ where
             let opt_val = *ptr;
             set.contains(&opt_val)
         })
-        .collect();
+        .collect_trusted();
     ca.rename(name);
     Ok(ca)
 }
@@ -64,7 +64,7 @@ where
                         }
                         _ => false,
                     })
-                    .collect();
+                    .collect_trusted();
                 Ok(ca)
             }
             _ => {
@@ -117,7 +117,7 @@ impl IsIn for Utf8Chunked {
                         }
                         _ => false,
                     })
-                    .collect();
+                    .collect_trusted();
                 Ok(ca)
             }
             DataType::Utf8 => {
@@ -132,7 +132,7 @@ impl IsIn for Utf8Chunked {
                 let mut ca: BooleanChunked = self
                     .into_iter()
                     .map(|opt_val| set.contains(&opt_val))
-                    .collect();
+                    .collect_trusted();
                 ca.rename(self.name());
                 Ok(ca)
             }
@@ -162,7 +162,7 @@ impl IsIn for BooleanChunked {
                         }
                         _ => false,
                     })
-                    .collect();
+                    .collect_trusted();
                 Ok(ca)
             }
             _ => Err(PolarsError::DataTypeMisMatch(

diff --git a/polars/polars-core/src/chunked_array/ops/reverse.rs b/polars/polars-core/src/chunked_array/ops/reverse.rs
@@ -1,5 +1,5 @@
 use crate::prelude::*;
-use crate::utils::NoNull;
+use crate::utils::{CustomIterTools, NoNull};
 
 impl<T> ChunkReverse<T> for ChunkedArray<T>
 where
@@ -8,12 +8,12 @@ where
 {
     fn reverse(&self) -> ChunkedArray<T> {
         if let Ok(slice) = self.cont_slice() {
-            let ca: NoNull<ChunkedArray<T>> = slice.iter().rev().copied().collect();
+            let ca: NoNull<ChunkedArray<T>> = slice.iter().rev().copied().collect_trusted();
             let mut ca = ca.into_inner();
             ca.rename(self.name());
             ca
         } else {
-            self.into_iter().rev().collect()
+            self.into_iter().rev().collect_trusted()
         }
     }
 }
@@ -29,7 +29,7 @@ macro_rules! impl_reverse {
     ($arrow_type:ident, $ca_type:ident) => {
         impl ChunkReverse<$arrow_type> for $ca_type {
             fn reverse(&self) -> Self {
-                unsafe { self.take_unchecked((0..self.len()).rev().into()) }
+                self.into_iter().rev().collect_trusted()
             }
         }
     };

diff --git a/polars/polars-core/src/chunked_array/ops/set.rs b/polars/polars-core/src/chunked_array/ops/set.rs
@@ -1,5 +1,5 @@
 use crate::prelude::*;
-use crate::utils::align_chunks_binary;
+use crate::utils::{align_chunks_binary, CustomIterTools};
 use arrow::array::ArrayRef;
 use polars_arrow::array::ValueSize;
 use polars_arrow::kernels::set::{set_at_idx_no_null, set_with_mask};
@@ -128,7 +128,7 @@ where
                     Some(true) => value,
                     _ => opt_val,
                 })
-                .collect();
+                .collect_trusted();
             Ok(ca)
         }
     }
@@ -145,7 +145,7 @@ where
                 Some(true) => f(opt_val),
                 _ => opt_val,
             })
-            .collect();
+            .collect_trusted();
         Ok(ca)
     }
 }
@@ -177,7 +177,7 @@ impl<'a> ChunkSet<'a, bool, bool> for BooleanChunked {
                 Some(true) => value,
                 _ => opt_val,
             })
-            .collect();
+            .collect_trusted();
         Ok(ca)
     }
 
@@ -193,7 +193,7 @@ impl<'a> ChunkSet<'a, bool, bool> for BooleanChunked {
                 Some(true) => f(opt_val),
                 _ => opt_val,
             })
-            .collect();
+            .collect_trusted();
         Ok(ca)
     }
 }
@@ -261,7 +261,7 @@ impl<'a> ChunkSet<'a, &'a str, String> for Utf8Chunked {
                 Some(true) => value,
                 _ => opt_val,
             })
-            .collect();
+            .collect_trusted();
         Ok(ca)
     }
 

diff --git a/polars/polars-core/src/chunked_array/ops/sort.rs b/polars/polars-core/src/chunked_array/ops/sort.rs
@@ -242,7 +242,10 @@ where
 
         assert_eq!(other.len(), reverse.len() - 1);
 
-        let compare_inner: Vec<_> = other.iter().map(|s| s.into_partial_ord_inner()).collect();
+        let compare_inner: Vec<_> = other
+            .iter()
+            .map(|s| s.into_partial_ord_inner())
+            .collect_trusted();
 
         let mut count: u32 = 0;
         let mut vals: Vec<_> = self
@@ -252,7 +255,7 @@ where
                 count += 1;
                 (i, v)
             })
-            .collect();
+            .collect_trusted();
 
         vals.sort_by(
             |tpl_a, tpl_b| match (reverse[0], sort_with_nulls(&tpl_a.1, &tpl_b.1)) {
@@ -268,7 +271,7 @@ where
                 (_, ord) => ord,
             },
         );
-        let ca: NoNull<UInt32Chunked> = vals.into_iter().map(|(idx, _v)| idx).collect();
+        let ca: NoNull<UInt32Chunked> = vals.into_iter().map(|(idx, _v)| idx).collect_trusted();
         let mut ca = ca.into_inner();
         ca.set_sorted(reverse[0]);
         Ok(ca)
@@ -298,9 +301,15 @@ fn ordering_other_columns<'a>(
 macro_rules! sort {
     ($self:ident, $reverse:ident) => {{
         if $reverse {
-            $self.into_iter().sorted_by(|a, b| b.cmp(a)).collect()
+            $self
+                .into_iter()
+                .sorted_by(|a, b| b.cmp(a))
+                .collect_trusted()
         } else {
-            $self.into_iter().sorted_by(|a, b| a.cmp(b)).collect()
+            $self
+                .into_iter()
+                .sorted_by(|a, b| a.cmp(b))
+                .collect_trusted()
         }
     }};
 }
@@ -358,8 +367,11 @@ impl ChunkSort<Utf8Type> for Utf8Chunked {
                 count += 1;
                 (i, v)
             })
-            .collect();
-        let compare_inner: Vec<_> = other.iter().map(|s| s.into_partial_ord_inner()).collect();
+            .collect_trusted();
+        let compare_inner: Vec<_> = other
+            .iter()
+            .map(|s| s.into_partial_ord_inner())
+            .collect_trusted();
 
         vals.sort_by(
             |tpl_a, tpl_b| match (reverse[0], sort_with_nulls(&tpl_a.1, &tpl_b.1)) {
@@ -375,7 +387,7 @@ impl ChunkSort<Utf8Type> for Utf8Chunked {
                 (_, ord) => ord,
             },
         );
-        let ca: NoNull<UInt32Chunked> = vals.into_iter().map(|(idx, _v)| idx).collect();
+        let ca: NoNull<UInt32Chunked> = vals.into_iter().map(|(idx, _v)| idx).collect_trusted();
         let mut ca = ca.into_inner();
         ca.set_sorted(reverse[0]);
         Ok(ca)

diff --git a/polars/polars-core/src/chunked_array/ops/zip.rs b/polars/polars-core/src/chunked_array/ops/zip.rs
@@ -1,5 +1,5 @@
 use crate::prelude::*;
-use crate::utils::align_chunks_ternary;
+use crate::utils::{align_chunks_ternary, CustomIterTools};
 use arrow::compute::if_then_else::if_then_else;
 use polars_arrow::array::default_arrays::FromData;
 
@@ -34,7 +34,7 @@ macro_rules! impl_ternary_broadcast {
                 let mut val: ChunkedArray<$ty> = $mask
                     .into_no_null_iter()
                     .map(|mask_val| ternary_apply(mask_val, left, right))
-                    .collect();
+                    .collect_trusted();
                 val.rename($self.name());
                 Ok(val)
             }
@@ -44,7 +44,7 @@ macro_rules! impl_ternary_broadcast {
                     .into_no_null_iter()
                     .zip($self)
                     .map(|(mask_val, left)| ternary_apply(mask_val, left, right))
-                    .collect();
+                    .collect_trusted();
                 val.rename($self.name());
                 Ok(val)
             }
@@ -54,7 +54,7 @@ macro_rules! impl_ternary_broadcast {
                     .into_no_null_iter()
                     .zip($other)
                     .map(|(mask_val, right)| ternary_apply(mask_val, left, right))
-                    .collect();
+                    .collect_trusted();
                 val.rename($self.name());
                 Ok(val)
             }