Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
don't allocate string when writing to csv (#935)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Apr 10, 2022
1 parent f60bb55 commit 4df7032
Showing 1 changed file with 26 additions and 19 deletions.
45 changes: 26 additions & 19 deletions src/io/csv/write/serialize.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use super::super::super::iterator::{BufStreamingIterator, StreamingIterator};
use crate::array::{DictionaryArray, DictionaryKey, Offset};
use csv_core::WriteResult;
use std::any::Any;
use std::fmt::{Debug, Write};

/// Options to serialize logical types to CSV
/// The default is to format times and dates as `chrono` crate formats them.
Expand Down Expand Up @@ -49,6 +50,16 @@ impl Default for SerializeOptions {
}
}

/// Utility to write to `&mut Vec<u8>` buffer
struct StringWrap<'a>(pub &'a mut Vec<u8>);

impl<'a> Write for StringWrap<'a> {
fn write_str(&mut self, s: &str) -> std::fmt::Result {
self.0.extend_from_slice(s.as_bytes());
Ok(())
}
}

fn primitive_write<'a, T: NativeType + ToLexical>(
array: &'a PrimitiveArray<T>,
) -> Box<dyn StreamingIterator<Item = [u8]> + 'a> {
Expand Down Expand Up @@ -81,7 +92,8 @@ macro_rules! dyn_date {
array.iter(),
move |x, buf| {
if let Some(x) = x {
buf.extend_from_slice(($fn)(*x).format(format).to_string().as_bytes())
let dt = ($fn)(*x).format(format);
let _ = write!(StringWrap(buf), "{}", dt);
}
},
vec![],
Expand All @@ -91,7 +103,8 @@ macro_rules! dyn_date {
array.iter(),
move |x, buf| {
if let Some(x) = x {
buf.extend_from_slice(($fn)(*x).to_string().as_bytes())
let dt = ($fn)(*x);
let _ = write!(StringWrap(buf), "{}", dt);
}
},
vec![],
Expand All @@ -111,10 +124,8 @@ fn timestamp_with_tz_default<'a>(
array.iter(),
move |x, buf| {
if let Some(x) = x {
let data =
temporal_conversions::timestamp_to_datetime(*x, time_unit, &timezone)
.to_string();
buf.extend_from_slice(data.as_bytes())
let dt = temporal_conversions::timestamp_to_datetime(*x, time_unit, &timezone);
let _ = write!(StringWrap(buf), "{}", dt);
}
},
vec![],
Expand All @@ -126,10 +137,9 @@ fn timestamp_with_tz_default<'a>(
array.iter(),
move |x, buf| {
if let Some(x) = x {
let data =
temporal_conversions::timestamp_to_datetime(*x, time_unit, &timezone)
.to_string();
buf.extend_from_slice(data.as_bytes())
let dt =
temporal_conversions::timestamp_to_datetime(*x, time_unit, &timezone);
let _ = write!(StringWrap(buf), "{}", dt);
}
},
vec![],
Expand Down Expand Up @@ -157,11 +167,9 @@ fn timestamp_with_tz_with_format<'a>(
array.iter(),
move |x, buf| {
if let Some(x) = x {
let data =
temporal_conversions::timestamp_to_datetime(*x, time_unit, &timezone)
.format(format)
.to_string();
buf.extend_from_slice(data.as_bytes())
let dt = temporal_conversions::timestamp_to_datetime(*x, time_unit, &timezone)
.format(format);
let _ = write!(StringWrap(buf), "{}", dt);
}
},
vec![],
Expand All @@ -173,11 +181,10 @@ fn timestamp_with_tz_with_format<'a>(
array.iter(),
move |x, buf| {
if let Some(x) = x {
let data =
let dt =
temporal_conversions::timestamp_to_datetime(*x, time_unit, &timezone)
.format(format)
.to_string();
buf.extend_from_slice(data.as_bytes())
.format(format);
let _ = write!(StringWrap(buf), "{}", dt);
}
},
vec![],
Expand Down

0 comments on commit 4df7032

Please sign in to comment.