arrow/util/string_writer.rs
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! String Writer
19//! This string writer encapsulates `std::string::String` and
20//! implements `std::io::Write` trait, which makes String as a
21//! writable object like File.
22//!
23//! Example:
24//!
25//! ```
26//! #[cfg(feature = "csv")]
27//! {
28//! use arrow::array::*;
29//! use arrow::csv;
30//! use arrow::datatypes::*;
31//! use arrow::record_batch::RecordBatch;
32//! use arrow::util::string_writer::StringWriter;
33//! use std::sync::Arc;
34//!
35//! let schema = Schema::new(vec![
36//! Field::new("c1", DataType::Utf8, false),
37//! Field::new("c2", DataType::Float64, true),
38//! Field::new("c3", DataType::UInt32, false),
39//! Field::new("c3", DataType::Boolean, true),
40//! ]);
41//! let c1 = StringArray::from(vec![
42//! "Lorem ipsum dolor sit amet",
43//! "consectetur adipiscing elit",
44//! "sed do eiusmod tempor",
45//! ]);
46//! let c2 = PrimitiveArray::<Float64Type>::from(vec![
47//! Some(123.564532),
48//! None,
49//! Some(-556132.25),
50//! ]);
51//! let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
52//! let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
53//!
54//! let batch = RecordBatch::try_new(
55//! Arc::new(schema),
56//! vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
57//! )
58//! .unwrap();
59//!
60//! let sw = StringWriter::new();
61//! let mut writer = csv::Writer::new(sw);
62//! writer.write(&batch).unwrap();
63//! }
64//! ```
65
66use core::str;
67use std::fmt::Formatter;
68use std::io::{Error, ErrorKind, Result, Write};
69
70/// A writer that allows writing to a `String`
71/// like an `std::io::Write` object.
72#[derive(Debug, Default)]
73pub struct StringWriter {
74 data: String,
75}
76
77impl StringWriter {
78 /// Create a new `StringWriter`
79 pub fn new() -> Self {
80 Self::default()
81 }
82}
83
84impl std::fmt::Display for StringWriter {
85 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
86 write!(f, "{}", self.data)
87 }
88}
89
90impl Write for StringWriter {
91 fn write(&mut self, buf: &[u8]) -> Result<usize> {
92 let string = match str::from_utf8(buf) {
93 Ok(x) => x,
94 Err(e) => {
95 return Err(Error::new(ErrorKind::InvalidData, e));
96 }
97 };
98 self.data.push_str(string);
99 Ok(string.len())
100 }
101
102 fn flush(&mut self) -> Result<()> {
103 Ok(())
104 }
105}