use OffsetBuffer::new_unchecked in convert_to_state.

apache · Jan 27, 2025 · b5144b8 · b5144b8
1 parent d8d9a19
commit b5144b8
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs
@@ -382,7 +382,7 @@ impl<T: ArrowNumericType + Send> GroupsAccumulator for MedianGroupsAccumulator<T
         // Build offsets
         let mut offsets = Vec::with_capacity(self.group_values.len() + 1);
         offsets.push(0);
-        let mut cur_len = 0;
+        let mut cur_len = 0_i32;
         for group_value in &emit_group_values {
             cur_len += group_value.len() as i32;
             offsets.push(cur_len);
@@ -443,8 +443,9 @@ impl<T: ArrowNumericType + Send> GroupsAccumulator for MedianGroupsAccumulator<T
             .with_data_type(self.data_type.clone());
 
         // `offsets` in `ListArray`, each row as a list element
+        assert!(input_array.len() <= i32::MAX as usize);
         let offsets = (0..=input_array.len() as i32).collect::<Vec<_>>();
-        let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets));
+        let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) };
 
         // `nulls` for converted `ListArray`
         let nulls = filtered_null_mask(opt_filter, input_array);