Skip to content

Commit bee4595

Browse files
authored
Add append_nulls to MapBuilder (apache#9432)
# Which issue does this PR close? Closes apache#9431 # Rationale for this change It would be nice to add `append_nulls` to MapBuilder, similar to `append_nulls` on `GenericListBuilder`. Appending the nulls at once, instead of using a loop has some nice performance implications: ``` Benchmark results (1,000,000 nulls): ┌─────────────────────────┬─────────┐ │ Method │ Time │ ├─────────────────────────┼─────────┤ │ append(false) in a loop │ 2.36 ms │ ├─────────────────────────┼─────────┤ │ append_nulls(N) │ 50 µs │ └─────────────────────────┴─────────┘ ``` # What changes are included in this PR? A new public API. # Are these changes tested? With some fresh unit tests. # Are there any user-facing changes? A nice and convient new public API
1 parent 01d34a8 commit bee4595

1 file changed

Lines changed: 59 additions & 4 deletions

File tree

arrow-array/src/builder/map_builder.rs

Lines changed: 59 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -154,23 +154,42 @@ impl<K: ArrayBuilder, V: ArrayBuilder> MapBuilder<K, V> {
154154
(&mut self.key_builder, &mut self.value_builder)
155155
}
156156

157-
/// Finish the current map array slot
158-
///
159-
/// Returns an error if the key and values builders are in an inconsistent state.
157+
/// Validates that key and value builders have equal lengths.
160158
#[inline]
161-
pub fn append(&mut self, is_valid: bool) -> Result<(), ArrowError> {
159+
fn validate_equal_lengths(&self) -> Result<(), ArrowError> {
162160
if self.key_builder.len() != self.value_builder.len() {
163161
return Err(ArrowError::InvalidArgumentError(format!(
164162
"Cannot append to a map builder when its keys and values have unequal lengths of {} and {}",
165163
self.key_builder.len(),
166164
self.value_builder.len()
167165
)));
168166
}
167+
Ok(())
168+
}
169+
170+
/// Finish the current map array slot
171+
///
172+
/// Returns an error if the key and values builders are in an inconsistent state.
173+
#[inline]
174+
pub fn append(&mut self, is_valid: bool) -> Result<(), ArrowError> {
175+
self.validate_equal_lengths()?;
169176
self.offsets_builder.push(self.key_builder.len() as i32);
170177
self.null_buffer_builder.append(is_valid);
171178
Ok(())
172179
}
173180

181+
/// Append `n` nulls to this [`MapBuilder`]
182+
///
183+
/// Returns an error if the key and values builders are in an inconsistent state.
184+
#[inline]
185+
pub fn append_nulls(&mut self, n: usize) -> Result<(), ArrowError> {
186+
self.validate_equal_lengths()?;
187+
let offset = self.key_builder.len() as i32;
188+
self.offsets_builder.extend(std::iter::repeat_n(offset, n));
189+
self.null_buffer_builder.append_n_nulls(n);
190+
Ok(())
191+
}
192+
174193
/// Builds the [`MapArray`]
175194
pub fn finish(&mut self) -> MapArray {
176195
let len = self.len();
@@ -436,6 +455,42 @@ mod tests {
436455
);
437456
}
438457

458+
#[test]
459+
fn test_append_nulls() {
460+
let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new());
461+
462+
builder.keys().append_value(1);
463+
builder.values().append_value(100);
464+
builder.append(true).unwrap();
465+
466+
builder.append_nulls(3).unwrap();
467+
468+
builder.keys().append_value(2);
469+
builder.values().append_value(200);
470+
builder.append(true).unwrap();
471+
472+
let map = builder.finish();
473+
assert_eq!(map.len(), 5);
474+
assert_eq!(map.null_count(), 3);
475+
assert!(map.is_valid(0));
476+
assert!(map.is_null(1));
477+
assert!(map.is_null(2));
478+
assert!(map.is_null(3));
479+
assert!(map.is_valid(4));
480+
assert_eq!(map.value_offsets(), &[0, 1, 1, 1, 1, 2]);
481+
}
482+
483+
#[test]
484+
fn test_append_nulls_inconsistent_state() {
485+
let mut builder = MapBuilder::new(None, Int32Builder::new(), Int32Builder::new());
486+
// Add a key without a matching value
487+
builder.keys().append_value(1);
488+
489+
let result = builder.append_nulls(2);
490+
assert!(result.is_err());
491+
assert!(result.unwrap_err().to_string().contains("unequal lengths"));
492+
}
493+
439494
#[test]
440495
#[should_panic(expected = "Keys field must not be nullable")]
441496
fn test_with_nullable_keys_field() {

0 commit comments

Comments
 (0)