Segfault in ByteGroupValueBuilder

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs

@@ -50,6 +50,8 @@ where
     offsets: Vec<O>,
     /// Nulls
     nulls: MaybeNullBufferBuilder,
+    /// The maximum size of the buffer for `0`
+    max_buffer_size: usize,
 }
 
 impl<O> ByteGroupValueBuilder<O>
@@ -62,6 +64,11 @@ where
             buffer: BufferBuilder::new(INITIAL_BUFFER_CAPACITY),
             offsets: vec![O::default()],
             nulls: MaybeNullBufferBuilder::new(),
+            max_buffer_size: if O::IS_LARGE {
+                i64::MAX as usize
+            } else {
+                i32::MAX as usize
+            },
         }
     }
 
@@ -187,6 +194,13 @@ where
     {
         let value: &[u8] = array.value(row).as_ref();
         self.buffer.append_slice(value);
+
+        assert!(
+            self.buffer.len() <= self.max_buffer_size,
+            "offset overflow, buffer size > {}",
+            self.max_buffer_size
+        );
+
         self.offsets.push(O::usize_as(self.buffer.len()));
     }
 
@@ -318,6 +332,7 @@ where
             mut buffer,
             offsets,
             nulls,
+            ..
         } = *self;
 
         let null_buffer = nulls.build();
@@ -410,6 +425,24 @@ mod tests {
 
     use super::GroupColumn;
 
+    #[test]
+    #[should_panic]
+    fn test_byte_group_value_builder_overflow() {
+        let mut builder = ByteGroupValueBuilder::<i32>::new(OutputType::Utf8);
+
+        let large_string = "a".repeat(1024 * 1024);
+
+        let array =
+            Arc::new(StringArray::from(vec![Some(large_string.as_str())])) as ArrayRef;
+
+        // Append items until our buffer length is 1 + i32::MAX as usize
+        for _ in 0..2048 {
+            builder.append_val(&array, 0);
+        }
+
+        assert_eq!(builder.value(2047), large_string.as_bytes());
+    }
+
     #[test]
     fn test_byte_take_n() {
         let mut builder = ByteGroupValueBuilder::<i32>::new(OutputType::Utf8);