panic on invalid when -> then -> otherwise state (#1874)

polars/polars-arrow/src/kernels/take.rs

@@ -264,9 +264,10 @@ pub unsafe fn take_no_null_utf8_iter_unchecked<I: IntoIterator<Item = usize>>(
     arr: &LargeStringArray,
     indices: I,
 ) -> Arc<LargeStringArray> {
-    let iter = indices
-        .into_iter()
-        .map(|idx| Some(arr.value_unchecked(idx)));
+    let iter = indices.into_iter().map(|idx| {
+        debug_assert!(idx < arr.len());
+        Some(arr.value_unchecked(idx))
+    });
 
     Arc::new(LargeStringArray::from_trusted_len_iter_unchecked(iter))
 }

polars/polars-lazy/src/physical_plan/expressions/mod.rs

@@ -233,6 +233,7 @@ impl<'a> AggregationContext<'a> {
                     || !self.groups.as_ref().is_empty()
                     && self.groups[0].1.len() > 1)
                 {
+                    // todo! optimize this, we don't have to call agg_list, create the list directly.
                     Cow::Owned(s.expand_at_index(0, self.groups.iter().map(|g| g.1.len()).sum()))
                 } else {
                     Cow::Borrowed(s)

polars/polars-lazy/src/physical_plan/expressions/ternary.rs

@@ -33,8 +33,18 @@ impl PhysicalExpr for TernaryExpr {
         groups: &'a GroupTuples,
         state: &ExecutionState,
     ) -> Result<AggregationContext<'a>> {
+        let required_height = df.height();
         let ac_mask = self.predicate.evaluate_on_groups(df, groups, state)?;
         let mask_s = ac_mask.flat();
+
+        assert!(
+            !(mask_s.len() != required_height),
+            "The predicate is of a different length than the groups.\
+The predicate produced {} values. Where the original DataFrame has {} values",
+            mask_s.len(),
+            required_height
+        );
+
         let mask = mask_s.bool()?;
         let mut ac_truthy = self.truthy.evaluate_on_groups(df, groups, state)?;
         let ac_falsy = self.falsy.evaluate_on_groups(df, groups, state)?;
@@ -48,6 +58,12 @@ impl PhysicalExpr for TernaryExpr {
         }
 
         let out = ac_truthy.flat().zip_with(mask, ac_falsy.flat().as_ref())?;
+
+        assert!(!(out.len() != required_height), "The output of the `when -> then -> otherwise-expr` is of a different length than the groups.\
+The expr produced {} values. Where the original DataFrame has {} values",
+out.len(),
+required_height);
+
         ac_truthy.with_series(out, false);
 
         Ok(ac_truthy)

polars/polars-lazy/src/test.rs

@@ -2242,3 +2242,37 @@ fn test_literal_window_fn() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[should_panic]
+fn test_invalid_ternary_in_agg1() {
+    let df = df![
+        "groups" => [1, 1, 2, 2, 3, 3],
+        "vals" => [1, 2, 3, 4, 5, 6]
+    ]
+    .unwrap();
+
+    df.lazy()
+        .groupby([col("groups")])
+        .agg([when(col("vals").first().neq(lit(1)))
+            .then(lit("a"))
+            .otherwise(lit("b"))])
+        .collect();
+}
+
+#[test]
+#[should_panic]
+fn test_invalid_ternary_in_agg2() {
+    let df = df![
+        "groups" => [1, 1, 2, 2, 3, 3],
+        "vals" => [1, 2, 3, 4, 5, 6]
+    ]
+    .unwrap();
+
+    df.lazy()
+        .groupby([col("groups")])
+        .agg([when(col("vals").neq(lit(1)))
+            .then(col("vals").first())
+            .otherwise(lit("b"))])
+        .collect();
+}