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}