From 6bfa27bc04c94ca69239c5b6c71670fcaa795ab2 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..3a59f71a9c91 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`]
@@ -235,6 +236,13 @@ impl<T: ByteArrayType, V: AsRef<T::Native>> Extend<Option<V>> for GenericByteBui
 ///  Array builder for [`GenericStringArray`][crate::GenericStringArray]
 pub type GenericStringBuilder<O> = GenericByteBuilder<GenericStringType<O>>;
 
+impl<O: OffsetSizeTrait> Write for GenericStringBuilder<O> {
+    fn write_str(&mut self, s: &str) -> std::fmt::Result {
+        self.value_builder.append_slice(s.as_bytes());
+        Ok(())
+    }
+}
+
 ///  Array builder for [`GenericBinaryArray`][crate::GenericBinaryArray]
 pub type GenericBinaryBuilder<O> = GenericByteBuilder<GenericBinaryType<O>>;
 
@@ -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"])
+    }
 }