[arrow-rs] branch master updated: Implement std::fmt::Write for StringBuilder (#3638) (#3659)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 04500e7dd Implement std::fmt::Write for StringBuilder (#3638) (#3659)
04500e7dd is described below

commit 04500e7dd77a49d925657732e42924e46a82489f
Author: Raphael Taylor-Davies <1781103+tustv...@users.noreply.github.com>
AuthorDate: Mon Feb 6 11:54:11 2023 +0000

    Implement std::fmt::Write for StringBuilder (#3638) (#3659)
    
    * Implement std::fmt::Write for StringBuilder (#3638)
    
    * Add docs
---
 arrow-array/src/builder/generic_bytes_builder.rs | 52 +++++++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/arrow-array/src/builder/generic_bytes_builder.rs 
b/arrow-array/src/builder/generic_bytes_builder.rs
index 73600d9e0..406e79c31 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`]
@@ -232,9 +233,43 @@ impl<T: ByteArrayType, V: AsRef<T::Native>> 
Extend<Option<V>> for GenericByteBui
     }
 }
 
-///  Array builder for [`GenericStringArray`][crate::GenericStringArray]
+/// Array builder for [`GenericStringArray`][crate::GenericStringArray]
+///
+/// Values can be appended using [`GenericByteBuilder::append_value`], and 
nulls with
+/// [`GenericByteBuilder::append_null`] as normal.
+///
+/// Additionally implements [`std::fmt::Write`] with any written data included 
in the next
+/// appended value. This allows use with [`std::fmt::Display`] without 
intermediate allocations
+///
+/// ```
+/// # use std::fmt::Write;
+/// # use arrow_array::builder::GenericStringBuilder;
+/// let mut builder = GenericStringBuilder::<i32>::new();
+///
+/// // Write data
+/// write!(builder, "foo").unwrap();
+/// write!(builder, "bar").unwrap();
+///
+/// // Finish value
+/// builder.append_value("baz");
+///
+/// // Write second value
+/// write!(builder, "v2").unwrap();
+/// builder.append_value("");
+///
+/// let array = builder.finish();
+/// assert_eq!(array.value(0), "foobarbaz");
+/// assert_eq!(array.value(1), "v2");
+/// ```
 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 +478,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"])
+    }
 }