diff --git a/datafusion/functions/src/strings.rs b/datafusion/functions/src/strings.rs index d2fb5d58519e..caafbae6ba5f 100644 --- a/datafusion/functions/src/strings.rs +++ b/datafusion/functions/src/strings.rs @@ -185,9 +185,24 @@ impl StringArrayBuilder { unsafe { self.offsets_buffer.push_unchecked(next_offset) }; } + /// Finalise the builder into a concrete [`StringArray`]. + /// + /// # Panics + /// + /// This method can panic when: + /// + /// - the provided `null_buffer` is not the same length as the `offsets_buffer`. pub fn finish(self, null_buffer: Option) -> StringArray { + let row_count = self.offsets_buffer.len() / size_of::() - 1; + if let Some(ref null_buffer) = null_buffer { + assert_eq!( + null_buffer.len(), + row_count, + "Null buffer and offsets buffer must be the same length" + ); + } let array_builder = ArrayDataBuilder::new(DataType::Utf8) - .len(self.offsets_buffer.len() / size_of::() - 1) + .len(row_count) .add_buffer(self.offsets_buffer.into()) .add_buffer(self.value_buffer.into()) .nulls(null_buffer); @@ -335,9 +350,24 @@ impl LargeStringArrayBuilder { unsafe { self.offsets_buffer.push_unchecked(next_offset) }; } + /// Finalise the builder into a concrete [`LargeStringArray`]. + /// + /// # Panics + /// + /// This method can panic when: + /// + /// - the provided `null_buffer` is not the same length as the `offsets_buffer`. pub fn finish(self, null_buffer: Option) -> LargeStringArray { + let row_count = self.offsets_buffer.len() / size_of::() - 1; + if let Some(ref null_buffer) = null_buffer { + assert_eq!( + null_buffer.len(), + row_count, + "Null buffer and offsets buffer must be the same length" + ); + } let array_builder = ArrayDataBuilder::new(DataType::LargeUtf8) - .len(self.offsets_buffer.len() / size_of::() - 1) + .len(row_count) .add_buffer(self.offsets_buffer.into()) .add_buffer(self.value_buffer.into()) .nulls(null_buffer);