From 587e985bc1bfefdcefd57189547820d376e8c2fd Mon Sep 17 00:00:00 2001
From: Raphael Taylor-Davies <r.taylordavies@googlemail.com>
Date: Sat, 4 Feb 2023 10:47:22 +0000
Subject: [PATCH] Implement std::fmt::Write for StringBuilder (#3638)

---
 .../src/builder/generic_bytes_builder.rs      | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/arrow-array/src/builder/generic_bytes_builder.rs b/arrow-array/src/builder/generic_bytes_builder.rs
index 73600d9e0a38..edc011bba70a 100644
--- 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"])
+    }
 }