fixed segfault in access of empty array

pola-rs · Oct 9, 2021 · d6d11b6 · d6d11b6
1 parent e616d11
commit d6d11b6
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 24 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/mod.rs b/polars/polars-core/src/chunked_array/ops/mod.rs
@@ -146,7 +146,8 @@ pub trait TakeRandom {
 
     /// Get a nullable value by index.
     ///
-    /// Out of bounds access doesn't Error but will return a Null value
+    /// # Panics
+    /// Panics if `index >= self.len()`
     fn get(&self, index: usize) -> Option<Self::Item>;
 
     /// Get a value by index and ignore the null bit.
@@ -167,7 +168,8 @@ pub trait TakeRandomUtf8 {
 
     /// Get a nullable value by index.
     ///
-    /// Out of bounds access doesn't Error but will return a Null value
+    /// # Panics
+    /// Panics if `index >= self.len()`
     fn get(self, index: usize) -> Option<Self::Item>;
 
     /// Get a value by index and ignore the null bit.

diff --git a/polars/polars-core/src/chunked_array/ops/take/take_single.rs b/polars/polars-core/src/chunked_array/ops/take/take_single.rs
@@ -10,6 +10,7 @@ use std::sync::Arc;
 
 macro_rules! impl_take_random_get {
     ($self:ident, $index:ident, $array_type:ty) => {{
+        assert!($index < $self.len());
         let (chunk_idx, idx) = $self.index_to_chunked_index($index);
         // Safety:
         // bounds are checked above
@@ -177,3 +178,18 @@ impl TakeRandom for ListChunked {
         })
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    #[should_panic]
+    fn test_oob() {
+        let data: Series = [1.0, 2.0, 3.0].iter().collect();
+        let data = data.f64().unwrap();
+        let matches = data.eq(5.0);
+        let matches_indexes = matches.arg_true();
+        matches_indexes.get(0);
+    }
+}
diff --git a/polars/polars-core/src/frame/groupby/aggregations.rs b/polars/polars-core/src/frame/groupby/aggregations.rs
@@ -395,21 +395,12 @@ macro_rules! impl_agg_n_unique {
                 }
                 let taker = $self.take_rand();
 
-                if $self.null_count() == 0 {
-                    let mut set = HashSet::with_hasher(RandomState::new());
-                    for i in idx {
-                        let v = unsafe { taker.get_unchecked(*i as usize) };
-                        set.insert(v);
-                    }
-                    set.len() as u32
-                } else {
-                    let mut set = HashSet::with_hasher(RandomState::new());
-                    for i in idx {
-                        let opt_v = taker.get(*i as usize);
-                        set.insert(opt_v);
-                    }
-                    set.len() as u32
+                let mut set = HashSet::with_hasher(RandomState::new());
+                for i in idx {
+                    let v = unsafe { taker.get_unchecked(*i as usize) };
+                    set.insert(v);
                 }
+                set.len() as u32
             })
             .collect::<$ca_type>()
             .into_inner()

diff --git a/polars/polars-lazy/src/test.rs b/polars/polars-lazy/src/test.rs
@@ -1635,6 +1635,7 @@ fn test_filter_after_shift_in_groups() -> Result<()> {
                 .alias("filtered"),
         ])
         .collect()?;
+    dbg!(out.column("filtered")?);
 
     assert_eq!(
         out.column("filtered")?
@@ -1656,14 +1657,7 @@ fn test_filter_after_shift_in_groups() -> Result<()> {
             .unwrap(),
         5
     );
-    assert!(out
-        .column("filtered")?
-        .list()?
-        .get(2)
-        .unwrap()
-        .i32()?
-        .get(0)
-        .is_some());
+    assert_eq!(out.column("filtered")?.list()?.get(2).unwrap().len(), 0);
 
     Ok(())
 }