fix agg-last GroupsProxy::slice

pola-rs · Feb 2, 2022 · b6b1f7b · b6b1f7b
1 parent 7933b2e
commit b6b1f7b
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 16 deletions.
diff --git a/polars/polars-arrow/src/kernels/take.rs b/polars/polars-arrow/src/kernels/take.rs
@@ -175,9 +175,12 @@ pub unsafe fn take_no_null_primitive_opt_iter_unchecked<
 ) -> Arc<PrimitiveArray<T>> {
     let array_values = arr.values().as_slice();
 
-    let iter = indices
-        .into_iter()
-        .map(|opt_idx| opt_idx.map(|idx| *array_values.get_unchecked(idx)));
+    let iter = indices.into_iter().map(|opt_idx| {
+        opt_idx.map(|idx| {
+            debug_assert!(idx < array_values.len());
+            *array_values.get_unchecked(idx)
+        })
+    });
     let arr = PrimitiveArray::from_trusted_len_iter_unchecked(iter).to(T::PRIMITIVE.into());
 
     Arc::new(arr)
@@ -203,6 +206,7 @@ pub unsafe fn take_primitive_opt_iter_unchecked<
     let iter = indices.into_iter().map(|opt_idx| {
         opt_idx.and_then(|idx| {
             if validity.get_bit_unchecked(idx) {
+                debug_assert!(idx < array_values.len());
                 Some(*array_values.get_unchecked(idx))
             } else {
                 None
@@ -225,10 +229,14 @@ pub unsafe fn take_no_null_bool_iter_unchecked<I: IntoIterator<Item = usize>>(
     indices: I,
 ) -> Arc<BooleanArray> {
     debug_assert!(!arr.has_validity());
-    let iter = indices
-        .into_iter()
-        .map(|idx| Some(arr.values().get_bit_unchecked(idx)));
+    let values = arr.values();
+
+    let iter = indices.into_iter().map(|idx| {
+        debug_assert!(idx < values.len());
+        Some(values.get_bit_unchecked(idx))
+    });
 
+    // TODO: use values_iter. Need to add unchecked version for arrow
     Arc::new(BooleanArray::from_trusted_len_iter_unchecked(iter))
 }
 

diff --git a/polars/polars-core/src/frame/groupby/aggregations.rs b/polars/polars-core/src/frame/groupby/aggregations.rs
@@ -175,16 +175,13 @@ impl Series {
                 unsafe { self.take_opt_iter_unchecked(&mut iter) }
             }
             GroupsProxy::Slice(groups) => {
-                let mut iter =
-                    groups.iter().map(
-                        |&[first, len]| {
-                            if len == 0 {
-                                None
-                            } else {
-                                Some(first as usize)
-                            }
-                        },
-                    );
+                let mut iter = groups.iter().map(|&[first, len]| {
+                    if len == 0 {
+                        None
+                    } else {
+                        Some((first + len - 1) as usize)
+                    }
+                });
                 unsafe { self.take_opt_iter_unchecked(&mut iter) }
             }
         };

diff --git a/polars/polars-lazy/src/tests/aggregations.rs b/polars/polars-lazy/src/tests/aggregations.rs
@@ -352,6 +352,23 @@ fn test_binary_agg_context_2() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_binary_agg_context_3() -> Result<()> {
+    let df = fruits_cars();
+
+    let out = df
+        .lazy()
+        .groupby_stable([col("cars")])
+        .agg([(col("A") - col("A").first()).last().alias("last")])
+        .collect()?;
+
+    let out = out.column("last")?;
+    assert_eq!(out.get(0), AnyValue::Int32(4));
+    assert_eq!(out.get(1), AnyValue::Int32(0));
+
+    Ok(())
+}
+
 #[test]
 fn test_shift_elementwise_issue_2509() -> Result<()> {
     let df = df![