Implement std::fmt::Write for StringBuilder (apache#3638)

tustvold · Feb 4, 2023 · 587e985 · 587e985
1 parent 25e10dd
commit 587e985
Showing 1 changed file with 23 additions and 0 deletions.
diff --git a/arrow-array/src/builder/generic_bytes_builder.rs b/arrow-array/src/builder/generic_bytes_builder.rs
@@ -22,6 +22,7 @@ use crate::{ArrayRef, GenericByteArray, OffsetSizeTrait};
 use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer};
 use arrow_data::ArrayDataBuilder;
 use std::any::Any;
+use std::fmt::Write;
 use std::sync::Arc;
 
 ///  Array builder for [`GenericByteArray`]
@@ -169,6 +170,13 @@ impl<T: ByteArrayType> GenericByteBuilder<T> {
     }
 }
 
+impl<O: OffsetSizeTrait> Write for GenericByteBuilder<GenericStringType<O>> {
+    fn write_str(&mut self, s: &str) -> std::fmt::Result {
+        self.value_builder.append_slice(s.as_bytes());
+        Ok(())
+    }
+}
+
 impl<T: ByteArrayType> std::fmt::Debug for GenericByteBuilder<T> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}{}Builder", T::Offset::PREFIX, T::PREFIX)?;
@@ -443,4 +451,19 @@ mod tests {
         assert_eq!(array.value_offsets(), &[0, 1, 2, 3, 3, 4, 5, 6, 7, 15, 20]);
         assert_eq!(array.value_data(), b"abcabcdcupcakeshello");
     }
+
+    #[test]
+    fn test_write() {
+        let mut builder = GenericStringBuilder::<i32>::new();
+        write!(builder, "foo").unwrap();
+        builder.append_value("");
+        writeln!(builder, "bar").unwrap();
+        builder.append_value("");
+        write!(builder, "fiz").unwrap();
+        write!(builder, "buz").unwrap();
+        builder.append_value("");
+        let a = builder.finish();
+        let r: Vec<_> = a.iter().map(|x| x.unwrap()).collect();
+        assert_eq!(r, &["foo", "bar\n", "fizbuz"])
+    }
 }