Skip to content

Commit bd2c975

Browse files
authored
fix: add null_buffer length check to StringArrayBuilder/LargeStringArrayBuilder (#13758)
* fix: add `null_buffer` check for `LargeStringArray` Add a safety check to ensure that the alignment of buffers cannot be overflowed. This introduces a panic if they are not aligned through a runtime assertion. * fix: remove value_buffer assertion These buffers can be misaligned and it is not problematic, it is the `null_buffer` which we care about being of the same length. * feat: add `null_buffer` check to `StringArray` This is in a similar vein to `LargeStringArray`, as the code is the same, except for `i32`'s instead of `i64`. * feat: use `row_count` var to avoid drift
1 parent e70319f commit bd2c975

File tree

1 file changed

+32
-2
lines changed

1 file changed

+32
-2
lines changed

datafusion/functions/src/strings.rs

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -185,9 +185,24 @@ impl StringArrayBuilder {
185185
unsafe { self.offsets_buffer.push_unchecked(next_offset) };
186186
}
187187

188+
/// Finalise the builder into a concrete [`StringArray`].
189+
///
190+
/// # Panics
191+
///
192+
/// This method can panic when:
193+
///
194+
/// - the provided `null_buffer` is not the same length as the `offsets_buffer`.
188195
pub fn finish(self, null_buffer: Option<NullBuffer>) -> StringArray {
196+
let row_count = self.offsets_buffer.len() / size_of::<i32>() - 1;
197+
if let Some(ref null_buffer) = null_buffer {
198+
assert_eq!(
199+
null_buffer.len(),
200+
row_count,
201+
"Null buffer and offsets buffer must be the same length"
202+
);
203+
}
189204
let array_builder = ArrayDataBuilder::new(DataType::Utf8)
190-
.len(self.offsets_buffer.len() / size_of::<i32>() - 1)
205+
.len(row_count)
191206
.add_buffer(self.offsets_buffer.into())
192207
.add_buffer(self.value_buffer.into())
193208
.nulls(null_buffer);
@@ -335,9 +350,24 @@ impl LargeStringArrayBuilder {
335350
unsafe { self.offsets_buffer.push_unchecked(next_offset) };
336351
}
337352

353+
/// Finalise the builder into a concrete [`LargeStringArray`].
354+
///
355+
/// # Panics
356+
///
357+
/// This method can panic when:
358+
///
359+
/// - the provided `null_buffer` is not the same length as the `offsets_buffer`.
338360
pub fn finish(self, null_buffer: Option<NullBuffer>) -> LargeStringArray {
361+
let row_count = self.offsets_buffer.len() / size_of::<i64>() - 1;
362+
if let Some(ref null_buffer) = null_buffer {
363+
assert_eq!(
364+
null_buffer.len(),
365+
row_count,
366+
"Null buffer and offsets buffer must be the same length"
367+
);
368+
}
339369
let array_builder = ArrayDataBuilder::new(DataType::LargeUtf8)
340-
.len(self.offsets_buffer.len() / size_of::<i64>() - 1)
370+
.len(row_count)
341371
.add_buffer(self.offsets_buffer.into())
342372
.add_buffer(self.value_buffer.into())
343373
.nulls(null_buffer);

0 commit comments

Comments
 (0)