Skip to main content

arrow_csv/
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//! CSV Writing: [`Writer`] and [`WriterBuilder`]
19//!
20//! This CSV writer allows Arrow data (in record batches) to be written as CSV files.
21//! The writer does not support writing `ListArray` and `StructArray`.
22//!
23//! # Example
24//! ```
25//! # use arrow_array::*;
26//! # use arrow_array::types::*;
27//! # use arrow_csv::Writer;
28//! # use arrow_schema::*;
29//! # use std::sync::Arc;
30//!
31//! let schema = Schema::new(vec![
32//!     Field::new("c1", DataType::Utf8, false),
33//!     Field::new("c2", DataType::Float64, true),
34//!     Field::new("c3", DataType::UInt32, false),
35//!     Field::new("c4", DataType::Boolean, true),
36//! ]);
37//! let c1 = StringArray::from(vec![
38//!     "Lorem ipsum dolor sit amet",
39//!     "consectetur adipiscing elit",
40//!     "sed do eiusmod tempor",
41//! ]);
42//! let c2 = PrimitiveArray::<Float64Type>::from(vec![
43//!     Some(123.564532),
44//!     None,
45//!     Some(-556132.25),
46//! ]);
47//! let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
48//! let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
49//!
50//! let batch = RecordBatch::try_new(
51//!     Arc::new(schema),
52//!     vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
53//! )
54//! .unwrap();
55//!
56//! let mut output = Vec::with_capacity(1024);
57//!
58//! let mut writer = Writer::new(&mut output);
59//! let batches = vec![&batch, &batch];
60//! for batch in batches {
61//!     writer.write(batch).unwrap();
62//! }
63//! ```
64//!
65//! # Whitespace Handling
66//!
67//! The writer supports trimming leading and trailing whitespace from string values,
68//! compatible with Apache Spark's CSV options `ignoreLeadingWhiteSpace` and
69//! `ignoreTrailingWhiteSpace`. This is useful when working with data that may have
70//! unwanted padding.
71//!
72//! Whitespace trimming is applied to all string data types:
73//! - `DataType::Utf8`
74//! - `DataType::LargeUtf8`
75//! - `DataType::Utf8View`
76//!
77//! ## Example: Use [`WriterBuilder`] to control whitespace handling
78//!
79//! ```
80//! # use arrow_array::*;
81//! # use arrow_csv::WriterBuilder;
82//! # use arrow_schema::*;
83//! # use std::sync::Arc;
84//! let schema = Schema::new(vec![
85//!     Field::new("name", DataType::Utf8, false),
86//!     Field::new("comment", DataType::Utf8, false),
87//! ]);
88//!
89//! let name = StringArray::from(vec![
90//!     "  Alice  ",   // Leading and trailing spaces
91//!     "Bob",         // No spaces
92//!     "  Charlie",   // Leading spaces only
93//! ]);
94//! let comment = StringArray::from(vec![
95//!     "  Great job!  ",
96//!     "Well done",
97//!     "Excellent  ",
98//! ]);
99//!
100//! let batch = RecordBatch::try_new(
101//!     Arc::new(schema),
102//!     vec![Arc::new(name), Arc::new(comment)],
103//! )
104//! .unwrap();
105//!
106//! // Trim both leading and trailing whitespace
107//! let mut output = Vec::new();
108//! WriterBuilder::new()
109//!     .with_ignore_leading_whitespace(true)
110//!     .with_ignore_trailing_whitespace(true)
111//!     .build(&mut output)
112//!     .write(&batch)
113//!     .unwrap();
114//! assert_eq!(
115//!     String::from_utf8(output).unwrap(),
116//!     "\
117//! name,comment\n\
118//! Alice,Great job!\n\
119//! Bob,Well done\n\
120//! Charlie,Excellent\n"
121//! );
122//! ```
123//!
124//! # Quoting Styles
125//!
126//! The writer supports different quoting styles for fields, compatible with Apache Spark's
127//! CSV options like `quoteAll`. You can control when fields are quoted using the
128//! [`QuoteStyle`] enum.
129//!
130//! ## Example
131//!
132//! ```
133//! # use arrow_array::*;
134//! # use arrow_csv::{WriterBuilder, QuoteStyle};
135//! # use arrow_schema::*;
136//! # use std::sync::Arc;
137//!
138//! let schema = Schema::new(vec![
139//!     Field::new("product", DataType::Utf8, false),
140//!     Field::new("price", DataType::Float64, false),
141//! ]);
142//!
143//! let product = StringArray::from(vec!["apple", "banana,organic", "cherry"]);
144//! let price = Float64Array::from(vec![1.50, 2.25, 3.00]);
145//!
146//! let batch = RecordBatch::try_new(
147//!     Arc::new(schema),
148//!     vec![Arc::new(product), Arc::new(price)],
149//! )
150//! .unwrap();
151//!
152//! // Default behavior (QuoteStyle::Necessary)
153//! let mut output = Vec::new();
154//! WriterBuilder::new()
155//!     .build(&mut output)
156//!     .write(&batch)
157//!     .unwrap();
158//! assert_eq!(
159//!     String::from_utf8(output).unwrap(),
160//!     "product,price\napple,1.5\n\"banana,organic\",2.25\ncherry,3.0\n"
161//! );
162//!
163//! // Quote all fields (Spark's quoteAll=true)
164//! let mut output = Vec::new();
165//! WriterBuilder::new()
166//!     .with_quote_style(QuoteStyle::Always)
167//!     .build(&mut output)
168//!     .write(&batch)
169//!     .unwrap();
170//! assert_eq!(
171//!     String::from_utf8(output).unwrap(),
172//!     "\"product\",\"price\"\n\"apple\",\"1.5\"\n\"banana,organic\",\"2.25\"\n\"cherry\",\"3.0\"\n"
173//! );
174//! ```
175
176use arrow_array::*;
177use arrow_cast::display::*;
178use arrow_schema::*;
179use csv::ByteRecord;
180use std::io::Write;
181
182use crate::map_csv_error;
183const DEFAULT_NULL_VALUE: &str = "";
184
185/// The quoting style to use when writing CSV files.
186///
187/// This type is re-exported from the `csv` crate and supports different
188/// strategies for quoting fields. It is compatible with Apache Spark's
189/// CSV options like `quoteAll`.
190///
191/// # Example
192///
193/// ```
194/// use arrow_csv::{WriterBuilder, QuoteStyle};
195///
196/// let builder = WriterBuilder::new()
197///     .with_quote_style(QuoteStyle::Always); // Equivalent to Spark's quoteAll=true
198/// ```
199pub use csv::QuoteStyle;
200
201/// A CSV writer
202///
203/// See the [module documentation](crate::writer) for examples.
204#[derive(Debug)]
205pub struct Writer<W: Write> {
206    /// The object to write to
207    writer: csv::Writer<W>,
208    /// Whether file should be written with headers, defaults to `true`
209    has_headers: bool,
210    /// The date format for date arrays, defaults to RFC3339
211    date_format: Option<String>,
212    /// The datetime format for datetime arrays, defaults to RFC3339
213    datetime_format: Option<String>,
214    /// The timestamp format for timestamp arrays, defaults to RFC3339
215    timestamp_format: Option<String>,
216    /// The timestamp format for timestamp (with timezone) arrays, defaults to RFC3339
217    timestamp_tz_format: Option<String>,
218    /// The time format for time arrays, defaults to RFC3339
219    time_format: Option<String>,
220    /// Is the beginning-of-writer
221    beginning: bool,
222    /// The value to represent null entries, defaults to [`DEFAULT_NULL_VALUE`]
223    null_value: Option<String>,
224    /// Whether to ignore leading whitespace in string values
225    ignore_leading_whitespace: bool,
226    /// Whether to ignore trailing whitespace in string values
227    ignore_trailing_whitespace: bool,
228}
229
230impl<W: Write> Writer<W> {
231    /// Create a new CsvWriter from a writable object, with default options
232    ///
233    /// See [`WriterBuilder`] for configure options, and the [module
234    /// documentation](crate::writer) for examples.
235    pub fn new(writer: W) -> Self {
236        let delimiter = b',';
237        WriterBuilder::new().with_delimiter(delimiter).build(writer)
238    }
239
240    /// Write a RecordBatch to the underlying writer
241    pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
242        let num_columns = batch.num_columns();
243        if self.beginning {
244            if self.has_headers {
245                let mut headers: Vec<String> = Vec::with_capacity(num_columns);
246                batch
247                    .schema()
248                    .fields()
249                    .iter()
250                    .for_each(|field| headers.push(field.name().to_string()));
251                self.writer
252                    .write_record(&headers[..])
253                    .map_err(map_csv_error)?;
254            }
255            self.beginning = false;
256        }
257
258        let options = FormatOptions::default()
259            .with_null(self.null_value.as_deref().unwrap_or(DEFAULT_NULL_VALUE))
260            .with_date_format(self.date_format.as_deref())
261            .with_datetime_format(self.datetime_format.as_deref())
262            .with_timestamp_format(self.timestamp_format.as_deref())
263            .with_timestamp_tz_format(self.timestamp_tz_format.as_deref())
264            .with_time_format(self.time_format.as_deref());
265
266        let converters = batch
267            .columns()
268            .iter()
269            .map(|a| {
270                if a.data_type().is_nested() {
271                    Err(ArrowError::CsvError(format!(
272                        "Nested type {} is not supported in CSV",
273                        a.data_type()
274                    )))
275                } else {
276                    ArrayFormatter::try_new(a.as_ref(), &options)
277                }
278            })
279            .collect::<Result<Vec<_>, ArrowError>>()?;
280
281        let mut buffer = String::with_capacity(1024);
282        let mut byte_record = ByteRecord::with_capacity(1024, converters.len());
283
284        for row_idx in 0..batch.num_rows() {
285            byte_record.clear();
286            for (col_idx, converter) in converters.iter().enumerate() {
287                buffer.clear();
288                converter.value(row_idx).write(&mut buffer).map_err(|e| {
289                    ArrowError::CsvError(format!(
290                        "Error processing row {}, col {}: {e}",
291                        row_idx + 1,
292                        col_idx + 1
293                    ))
294                })?;
295
296                let field_bytes =
297                    self.get_trimmed_field_bytes(&buffer, batch.column(col_idx).data_type());
298                byte_record.push_field(field_bytes);
299            }
300
301            self.writer
302                .write_byte_record(&byte_record)
303                .map_err(map_csv_error)?;
304        }
305        self.writer.flush()?;
306
307        Ok(())
308    }
309
310    /// Returns the bytes for a field, applying whitespace trimming if configured and applicable
311    fn get_trimmed_field_bytes<'a>(&self, buffer: &'a str, data_type: &DataType) -> &'a [u8] {
312        // Only trim string types when trimming is enabled
313        let should_trim = (self.ignore_leading_whitespace || self.ignore_trailing_whitespace)
314            && matches!(
315                data_type,
316                DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
317            );
318
319        if !should_trim {
320            return buffer.as_bytes();
321        }
322
323        let mut trimmed = buffer;
324        if self.ignore_leading_whitespace {
325            trimmed = trimmed.trim_start();
326        }
327        if self.ignore_trailing_whitespace {
328            trimmed = trimmed.trim_end();
329        }
330        trimmed.as_bytes()
331    }
332
333    /// Unwraps this `Writer<W>`, returning the underlying writer.
334    pub fn into_inner(self) -> W {
335        // Safe to call `unwrap` since `write` always flushes the writer.
336        self.writer.into_inner().unwrap()
337    }
338}
339
340impl<W: Write> RecordBatchWriter for Writer<W> {
341    fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
342        self.write(batch)
343    }
344
345    fn close(self) -> Result<(), ArrowError> {
346        Ok(())
347    }
348}
349
350/// A CSV writer builder
351#[derive(Clone, Debug)]
352pub struct WriterBuilder {
353    /// Optional column delimiter. Defaults to `b','`
354    delimiter: u8,
355    /// Whether to write column names as file headers. Defaults to `true`
356    has_header: bool,
357    /// Optional quote character. Defaults to `b'"'`
358    quote: u8,
359    /// Optional escape character. Defaults to `b'\\'`
360    escape: u8,
361    /// Optional line terminator. Defaults to `LF` (`\n`)
362    terminator: Terminator,
363    /// Enable double quote escapes. Defaults to `true`
364    double_quote: bool,
365    /// Optional date format for date arrays
366    date_format: Option<String>,
367    /// Optional datetime format for datetime arrays
368    datetime_format: Option<String>,
369    /// Optional timestamp format for timestamp arrays
370    timestamp_format: Option<String>,
371    /// Optional timestamp format for timestamp with timezone arrays
372    timestamp_tz_format: Option<String>,
373    /// Optional time format for time arrays
374    time_format: Option<String>,
375    /// Optional value to represent null
376    null_value: Option<String>,
377    /// Whether to ignore leading whitespace in string values. Defaults to `false`
378    ignore_leading_whitespace: bool,
379    /// Whether to ignore trailing whitespace in string values. Defaults to `false`
380    ignore_trailing_whitespace: bool,
381    /// The quoting style to use. Defaults to `QuoteStyle::Necessary`
382    quote_style: QuoteStyle,
383}
384
385/// The line terminator to use when writing CSV files.
386#[derive(Clone, Debug)]
387pub enum Terminator {
388    /// Use CRLF (`\r\n`) as the line terminator
389    CRLF,
390    /// Use the specified byte character as the line terminator
391    Any(u8),
392}
393
394impl Default for WriterBuilder {
395    fn default() -> Self {
396        WriterBuilder {
397            delimiter: b',',
398            has_header: true,
399            quote: b'"',
400            escape: b'\\',
401            terminator: Terminator::Any(b'\n'),
402            double_quote: true,
403            date_format: None,
404            datetime_format: None,
405            timestamp_format: None,
406            timestamp_tz_format: None,
407            time_format: None,
408            null_value: None,
409            ignore_leading_whitespace: false,
410            ignore_trailing_whitespace: false,
411            quote_style: QuoteStyle::default(),
412        }
413    }
414}
415
416impl WriterBuilder {
417    /// Create a new builder for configuring CSV [`Writer`] options.
418    ///
419    /// To convert a builder into a writer, call [`WriterBuilder::build`]. See
420    /// the [module documentation](crate::writer) for more examples.
421    ///
422    /// # Example
423    ///
424    /// ```
425    /// # use arrow_csv::{Writer, WriterBuilder};
426    /// # use std::fs::File;
427    ///
428    /// fn example() -> Writer<File> {
429    ///     let file = File::create("target/out.csv").unwrap();
430    ///
431    ///     // create a builder that doesn't write headers
432    ///     let builder = WriterBuilder::new().with_header(false);
433    ///     let writer = builder.build(file);
434    ///
435    ///     writer
436    /// }
437    /// ```
438    pub fn new() -> Self {
439        Self::default()
440    }
441
442    /// Set whether to write the CSV file with a header
443    pub fn with_header(mut self, header: bool) -> Self {
444        self.has_header = header;
445        self
446    }
447
448    /// Returns `true` if this writer is configured to write a header
449    pub fn header(&self) -> bool {
450        self.has_header
451    }
452
453    /// Set the CSV file's column delimiter as a byte character
454    pub fn with_delimiter(mut self, delimiter: u8) -> Self {
455        self.delimiter = delimiter;
456        self
457    }
458
459    /// Get the CSV file's column delimiter as a byte character
460    pub fn delimiter(&self) -> u8 {
461        self.delimiter
462    }
463
464    /// Set the CSV file's quote character as a byte character
465    pub fn with_quote(mut self, quote: u8) -> Self {
466        self.quote = quote;
467        self
468    }
469
470    /// Get the CSV file's quote character as a byte character
471    pub fn quote(&self) -> u8 {
472        self.quote
473    }
474
475    /// Set the CSV file's escape character as a byte character
476    ///
477    /// In some variants of CSV, quotes are escaped using a special escape
478    /// character like `\` (instead of escaping quotes by doubling them).
479    ///
480    /// By default, writing these idiosyncratic escapes is disabled, and is
481    /// only used when `double_quote` is disabled.
482    pub fn with_escape(mut self, escape: u8) -> Self {
483        self.escape = escape;
484        self
485    }
486
487    /// Get the CSV file's escape character as a byte character
488    pub fn escape(&self) -> u8 {
489        self.escape
490    }
491
492    /// Set whether to enable double quote escapes
493    ///
494    /// When enabled (which is the default), quotes are escaped by doubling
495    /// them. e.g., `"` escapes to `""`.
496    ///
497    /// When disabled, quotes are escaped with the escape character (which
498    /// is `\\` by default).
499    pub fn with_double_quote(mut self, double_quote: bool) -> Self {
500        self.double_quote = double_quote;
501        self
502    }
503
504    /// Get whether double quote escapes are enabled
505    pub fn double_quote(&self) -> bool {
506        self.double_quote
507    }
508
509    /// Set the CSV file's date format
510    pub fn with_date_format(mut self, format: String) -> Self {
511        self.date_format = Some(format);
512        self
513    }
514
515    /// Get the CSV file's date format if set, defaults to RFC3339
516    pub fn date_format(&self) -> Option<&str> {
517        self.date_format.as_deref()
518    }
519
520    /// Set the CSV file's datetime format
521    pub fn with_datetime_format(mut self, format: String) -> Self {
522        self.datetime_format = Some(format);
523        self
524    }
525
526    /// Get the CSV file's datetime format if set, defaults to RFC3339
527    pub fn datetime_format(&self) -> Option<&str> {
528        self.datetime_format.as_deref()
529    }
530
531    /// Set the CSV file's time format
532    pub fn with_time_format(mut self, format: String) -> Self {
533        self.time_format = Some(format);
534        self
535    }
536
537    /// Get the CSV file's datetime time if set, defaults to RFC3339
538    pub fn time_format(&self) -> Option<&str> {
539        self.time_format.as_deref()
540    }
541
542    /// Set the CSV file's timestamp format
543    pub fn with_timestamp_format(mut self, format: String) -> Self {
544        self.timestamp_format = Some(format);
545        self
546    }
547
548    /// Get the CSV file's timestamp format if set, defaults to RFC3339
549    pub fn timestamp_format(&self) -> Option<&str> {
550        self.timestamp_format.as_deref()
551    }
552
553    /// Set the CSV file's timestamp tz format
554    pub fn with_timestamp_tz_format(mut self, tz_format: String) -> Self {
555        self.timestamp_tz_format = Some(tz_format);
556        self
557    }
558
559    /// Get the CSV file's timestamp tz format if set, defaults to RFC3339
560    pub fn timestamp_tz_format(&self) -> Option<&str> {
561        self.timestamp_tz_format.as_deref()
562    }
563
564    /// Set the value to represent null in output
565    pub fn with_null(mut self, null_value: String) -> Self {
566        self.null_value = Some(null_value);
567        self
568    }
569
570    /// Get the value to represent null in output
571    pub fn null(&self) -> &str {
572        self.null_value.as_deref().unwrap_or(DEFAULT_NULL_VALUE)
573    }
574
575    /// Set whether to ignore leading whitespace in string values
576    /// For example, a string value such as "   foo" will be written as "foo"
577    pub fn with_ignore_leading_whitespace(mut self, ignore: bool) -> Self {
578        self.ignore_leading_whitespace = ignore;
579        self
580    }
581
582    /// Get whether to ignore leading whitespace in string values
583    pub fn ignore_leading_whitespace(&self) -> bool {
584        self.ignore_leading_whitespace
585    }
586
587    /// Set whether to ignore trailing whitespace in string values
588    /// For example, a string value such as "foo    " will be written as "foo"
589    pub fn with_ignore_trailing_whitespace(mut self, ignore: bool) -> Self {
590        self.ignore_trailing_whitespace = ignore;
591        self
592    }
593
594    /// Get whether to ignore trailing whitespace in string values
595    pub fn ignore_trailing_whitespace(&self) -> bool {
596        self.ignore_trailing_whitespace
597    }
598
599    /// Set the quoting style for writing CSV files
600    ///
601    /// # Example
602    ///
603    /// ```
604    /// use arrow_csv::{WriterBuilder, QuoteStyle};
605    ///
606    /// // Quote all fields (equivalent to Spark's quoteAll=true)
607    /// let builder = WriterBuilder::new()
608    ///     .with_quote_style(QuoteStyle::Always);
609    ///
610    /// // Only quote when necessary (default)
611    /// let builder = WriterBuilder::new()
612    ///     .with_quote_style(QuoteStyle::Necessary);
613    /// ```
614    pub fn with_quote_style(mut self, quote_style: QuoteStyle) -> Self {
615        self.quote_style = quote_style;
616        self
617    }
618
619    /// Get the configured quoting style
620    pub fn quote_style(&self) -> QuoteStyle {
621        self.quote_style
622    }
623
624    /// Set the CSV file's line terminator
625    pub fn with_line_terminator(mut self, terminator: Terminator) -> Self {
626        self.terminator = terminator;
627        self
628    }
629
630    /// Get the CSV file's line terminator, defaults to `LF` (`\n`)
631    pub fn line_terminator(&self) -> &Terminator {
632        &self.terminator
633    }
634
635    /// Create a new `Writer`
636    pub fn build<W: Write>(self, writer: W) -> Writer<W> {
637        let mut builder = csv::WriterBuilder::new();
638
639        let terminator = match self.terminator {
640            Terminator::CRLF => csv::Terminator::CRLF,
641            Terminator::Any(byte) => csv::Terminator::Any(byte),
642        };
643
644        let writer = builder
645            .delimiter(self.delimiter)
646            .quote(self.quote)
647            .quote_style(self.quote_style)
648            .double_quote(self.double_quote)
649            .escape(self.escape)
650            .terminator(terminator)
651            .from_writer(writer);
652        Writer {
653            writer,
654            beginning: true,
655            has_headers: self.has_header,
656            date_format: self.date_format,
657            datetime_format: self.datetime_format,
658            time_format: self.time_format,
659            timestamp_format: self.timestamp_format,
660            timestamp_tz_format: self.timestamp_tz_format,
661            null_value: self.null_value,
662            ignore_leading_whitespace: self.ignore_leading_whitespace,
663            ignore_trailing_whitespace: self.ignore_trailing_whitespace,
664        }
665    }
666}
667
668#[cfg(test)]
669#[expect(
670    clippy::verbose_file_reads,
671    reason = "`tempfile::tempfile()` has no path, so the contents can only be read back through the handle"
672)]
673mod tests {
674    use super::*;
675
676    use crate::ReaderBuilder;
677    use arrow_array::builder::{
678        BinaryBuilder, Decimal32Builder, Decimal64Builder, Decimal128Builder, Decimal256Builder,
679        FixedSizeBinaryBuilder, LargeBinaryBuilder,
680    };
681    use arrow_array::types::*;
682    use arrow_buffer::i256;
683    use core::str;
684    use std::io::{Cursor, Read, Seek};
685    use std::sync::Arc;
686
687    #[test]
688    fn test_write_csv() {
689        let schema = Schema::new(vec![
690            Field::new("c1", DataType::Utf8, false),
691            Field::new("c2", DataType::Float64, true),
692            Field::new("c3", DataType::UInt32, false),
693            Field::new("c4", DataType::Boolean, true),
694            Field::new("c5", DataType::Timestamp(TimeUnit::Millisecond, None), true),
695            Field::new("c6", DataType::Time32(TimeUnit::Second), false),
696            Field::new_dictionary("c7", DataType::Int32, DataType::Utf8, false),
697        ]);
698
699        let c1 = StringArray::from(vec![
700            "Lorem ipsum dolor sit amet",
701            "consectetur adipiscing elit",
702            "sed do eiusmod tempor",
703        ]);
704        let c2 =
705            PrimitiveArray::<Float64Type>::from(vec![Some(123.564532), None, Some(-556132.25)]);
706        let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
707        let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
708        let c5 =
709            TimestampMillisecondArray::from(vec![None, Some(1555584887378), Some(1555555555555)]);
710        let c6 = Time32SecondArray::from(vec![1234, 24680, 85563]);
711        let c7: DictionaryArray<Int32Type> =
712            vec!["cupcakes", "cupcakes", "foo"].into_iter().collect();
713
714        let batch = RecordBatch::try_new(
715            Arc::new(schema),
716            vec![
717                Arc::new(c1),
718                Arc::new(c2),
719                Arc::new(c3),
720                Arc::new(c4),
721                Arc::new(c5),
722                Arc::new(c6),
723                Arc::new(c7),
724            ],
725        )
726        .unwrap();
727
728        let mut file = tempfile::tempfile().unwrap();
729
730        let mut writer = Writer::new(&mut file);
731        let batches = vec![&batch, &batch];
732        for batch in batches {
733            writer.write(batch).unwrap();
734        }
735        drop(writer);
736
737        // check that file was written successfully
738        file.rewind().unwrap();
739        let mut buffer: Vec<u8> = vec![];
740        file.read_to_end(&mut buffer).unwrap();
741
742        let expected = r"c1,c2,c3,c4,c5,c6,c7
743Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,cupcakes
744consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378,06:51:20,cupcakes
745sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555,23:46:03,foo
746Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,cupcakes
747consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378,06:51:20,cupcakes
748sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555,23:46:03,foo
749";
750        assert_eq!(expected, str::from_utf8(&buffer).unwrap());
751    }
752
753    #[test]
754    fn test_write_csv_decimal() {
755        let schema = Schema::new(vec![
756            Field::new("c1", DataType::Decimal32(9, 6), true),
757            Field::new("c2", DataType::Decimal64(17, 6), true),
758            Field::new("c3", DataType::Decimal128(38, 6), true),
759            Field::new("c4", DataType::Decimal256(76, 6), true),
760        ]);
761
762        let mut c1_builder = Decimal32Builder::new().with_data_type(DataType::Decimal32(9, 6));
763        c1_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]);
764        let c1 = c1_builder.finish();
765
766        let mut c2_builder = Decimal64Builder::new().with_data_type(DataType::Decimal64(17, 6));
767        c2_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]);
768        let c2 = c2_builder.finish();
769
770        let mut c3_builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(38, 6));
771        c3_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]);
772        let c3 = c3_builder.finish();
773
774        let mut c4_builder = Decimal256Builder::new().with_data_type(DataType::Decimal256(76, 6));
775        c4_builder.extend(vec![
776            Some(i256::from_i128(-3335724)),
777            Some(i256::from_i128(2179404)),
778            None,
779            Some(i256::from_i128(290472)),
780        ]);
781        let c4 = c4_builder.finish();
782
783        let batch = RecordBatch::try_new(
784            Arc::new(schema),
785            vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
786        )
787        .unwrap();
788
789        let mut file = tempfile::tempfile().unwrap();
790
791        let mut writer = Writer::new(&mut file);
792        let batches = vec![&batch, &batch];
793        for batch in batches {
794            writer.write(batch).unwrap();
795        }
796        drop(writer);
797
798        // check that file was written successfully
799        file.rewind().unwrap();
800        let mut buffer: Vec<u8> = vec![];
801        file.read_to_end(&mut buffer).unwrap();
802
803        let expected = r"c1,c2,c3,c4
804-3.335724,-3.335724,-3.335724,-3.335724
8052.179404,2.179404,2.179404,2.179404
806,,,
8070.290472,0.290472,0.290472,0.290472
808-3.335724,-3.335724,-3.335724,-3.335724
8092.179404,2.179404,2.179404,2.179404
810,,,
8110.290472,0.290472,0.290472,0.290472
812";
813        assert_eq!(expected, str::from_utf8(&buffer).unwrap());
814    }
815
816    #[test]
817    fn test_write_csv_custom_options() {
818        let schema = Schema::new(vec![
819            Field::new("c1", DataType::Utf8, false),
820            Field::new("c2", DataType::Float64, true),
821            Field::new("c3", DataType::UInt32, false),
822            Field::new("c4", DataType::Boolean, true),
823            Field::new("c6", DataType::Time32(TimeUnit::Second), false),
824        ]);
825
826        let c1 = StringArray::from(vec![
827            "Lorem ipsum \ndolor sit amet",
828            "consectetur \"adipiscing\" elit",
829            "sed do eiusmod tempor",
830        ]);
831        let c2 =
832            PrimitiveArray::<Float64Type>::from(vec![Some(123.564532), None, Some(-556132.25)]);
833        let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
834        let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
835        let c6 = Time32SecondArray::from(vec![1234, 24680, 85563]);
836
837        let batch = RecordBatch::try_new(
838            Arc::new(schema),
839            vec![
840                Arc::new(c1),
841                Arc::new(c2),
842                Arc::new(c3),
843                Arc::new(c4),
844                Arc::new(c6),
845            ],
846        )
847        .unwrap();
848
849        let mut file = tempfile::tempfile().unwrap();
850
851        let builder = WriterBuilder::new()
852            .with_header(false)
853            .with_delimiter(b'|')
854            .with_quote(b'\'')
855            .with_null("NULL".to_string())
856            .with_time_format("%r".to_string());
857        let mut writer = builder.build(&mut file);
858        let batches = vec![&batch];
859        for batch in batches {
860            writer.write(batch).unwrap();
861        }
862        drop(writer);
863
864        // check that file was written successfully
865        file.rewind().unwrap();
866        let mut buffer: Vec<u8> = vec![];
867        file.read_to_end(&mut buffer).unwrap();
868
869        assert_eq!(
870            "'Lorem ipsum \ndolor sit amet'|123.564532|3|true|12:20:34 AM\nconsectetur \"adipiscing\" elit|NULL|2|false|06:51:20 AM\nsed do eiusmod tempor|-556132.25|1|NULL|11:46:03 PM\n"
871            .to_string(),
872            String::from_utf8(buffer).unwrap()
873        );
874
875        let mut file = tempfile::tempfile().unwrap();
876
877        let builder = WriterBuilder::new()
878            .with_header(true)
879            .with_double_quote(false)
880            .with_escape(b'$');
881        let mut writer = builder.build(&mut file);
882        let batches = vec![&batch];
883        for batch in batches {
884            writer.write(batch).unwrap();
885        }
886        drop(writer);
887
888        file.rewind().unwrap();
889        let mut buffer: Vec<u8> = vec![];
890        file.read_to_end(&mut buffer).unwrap();
891
892        assert_eq!(
893            "c1,c2,c3,c4,c6\n\"Lorem ipsum \ndolor sit amet\",123.564532,3,true,00:20:34\n\"consectetur $\"adipiscing$\" elit\",,2,false,06:51:20\nsed do eiusmod tempor,-556132.25,1,,23:46:03\n"
894            .to_string(),
895            String::from_utf8(buffer).unwrap()
896        );
897    }
898
899    #[test]
900    fn test_conversion_consistency() {
901        // test if we can serialize and deserialize whilst retaining the same type information/ precision
902
903        let schema = Schema::new(vec![
904            Field::new("c1", DataType::Date32, false),
905            Field::new("c2", DataType::Date64, false),
906            Field::new("c3", DataType::Timestamp(TimeUnit::Nanosecond, None), false),
907        ]);
908
909        let nanoseconds = vec![
910            1599566300000000000,
911            1599566200000000000,
912            1599566100000000000,
913        ];
914        let c1 = Date32Array::from(vec![3, 2, 1]);
915        let c2 = Date64Array::from(vec![3, 2, 1]);
916        let c3 = TimestampNanosecondArray::from(nanoseconds.clone());
917
918        let batch = RecordBatch::try_new(
919            Arc::new(schema.clone()),
920            vec![Arc::new(c1), Arc::new(c2), Arc::new(c3)],
921        )
922        .unwrap();
923
924        let builder = WriterBuilder::new().with_header(false);
925
926        let mut buf: Cursor<Vec<u8>> = Default::default();
927        // drop the writer early to release the borrow.
928        {
929            let mut writer = builder.build(&mut buf);
930            writer.write(&batch).unwrap();
931        }
932        buf.set_position(0);
933
934        let mut reader = ReaderBuilder::new(Arc::new(schema))
935            .with_batch_size(3)
936            .build_buffered(buf)
937            .unwrap();
938
939        let rb = reader.next().unwrap().unwrap();
940        let c1 = rb.column(0).as_any().downcast_ref::<Date32Array>().unwrap();
941        let c2 = rb.column(1).as_any().downcast_ref::<Date64Array>().unwrap();
942        let c3 = rb
943            .column(2)
944            .as_any()
945            .downcast_ref::<TimestampNanosecondArray>()
946            .unwrap();
947
948        let actual = c1.into_iter().collect::<Vec<_>>();
949        let expected = vec![Some(3), Some(2), Some(1)];
950        assert_eq!(actual, expected);
951        let actual = c2.into_iter().collect::<Vec<_>>();
952        let expected = vec![Some(3), Some(2), Some(1)];
953        assert_eq!(actual, expected);
954        let actual = c3.into_iter().collect::<Vec<_>>();
955        let expected = nanoseconds.into_iter().map(Some).collect::<Vec<_>>();
956        assert_eq!(actual, expected);
957    }
958
959    #[test]
960    fn test_write_csv_invalid_cast() {
961        let schema = Schema::new(vec![
962            Field::new("c0", DataType::UInt32, false),
963            Field::new("c1", DataType::Date64, false),
964        ]);
965
966        let c0 = UInt32Array::from(vec![Some(123), Some(234)]);
967        let c1 = Date64Array::from(vec![Some(1926632005177), Some(1926632005177685347)]);
968        let batch =
969            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c0), Arc::new(c1)]).unwrap();
970
971        let mut file = tempfile::tempfile().unwrap();
972        let mut writer = Writer::new(&mut file);
973        let batches = vec![&batch, &batch];
974
975        for batch in batches {
976            let err = writer.write(batch).unwrap_err().to_string();
977            assert_eq!(
978                err,
979                "Csv error: Error processing row 2, col 2: Cast error: Failed to convert 1926632005177685347 to temporal for Date64"
980            )
981        }
982        drop(writer);
983    }
984
985    #[test]
986    fn test_write_csv_using_rfc3339() {
987        let schema = Schema::new(vec![
988            Field::new(
989                "c1",
990                DataType::Timestamp(TimeUnit::Millisecond, Some("+00:00".into())),
991                true,
992            ),
993            Field::new("c2", DataType::Timestamp(TimeUnit::Millisecond, None), true),
994            Field::new("c3", DataType::Date32, false),
995            Field::new("c4", DataType::Time32(TimeUnit::Second), false),
996        ]);
997
998        let c1 = TimestampMillisecondArray::from(vec![Some(1555584887378), Some(1635577147000)])
999            .with_timezone("+00:00".to_string());
1000        let c2 = TimestampMillisecondArray::from(vec![Some(1555584887378), Some(1635577147000)]);
1001        let c3 = Date32Array::from(vec![3, 2]);
1002        let c4 = Time32SecondArray::from(vec![1234, 24680]);
1003
1004        let batch = RecordBatch::try_new(
1005            Arc::new(schema),
1006            vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
1007        )
1008        .unwrap();
1009
1010        let mut file = tempfile::tempfile().unwrap();
1011
1012        let builder = WriterBuilder::new();
1013        let mut writer = builder.build(&mut file);
1014        let batches = vec![&batch];
1015        for batch in batches {
1016            writer.write(batch).unwrap();
1017        }
1018        drop(writer);
1019
1020        file.rewind().unwrap();
1021        let mut buffer: Vec<u8> = vec![];
1022        file.read_to_end(&mut buffer).unwrap();
1023
1024        assert_eq!(
1025            "c1,c2,c3,c4
10262019-04-18T10:54:47.378Z,2019-04-18T10:54:47.378,1970-01-04,00:20:34
10272021-10-30T06:59:07Z,2021-10-30T06:59:07,1970-01-03,06:51:20\n",
1028            String::from_utf8(buffer).unwrap()
1029        );
1030    }
1031
1032    #[test]
1033    fn test_write_csv_tz_format() {
1034        let schema = Schema::new(vec![
1035            Field::new(
1036                "c1",
1037                DataType::Timestamp(TimeUnit::Millisecond, Some("+02:00".into())),
1038                true,
1039            ),
1040            Field::new(
1041                "c2",
1042                DataType::Timestamp(TimeUnit::Second, Some("+04:00".into())),
1043                true,
1044            ),
1045        ]);
1046        let c1 = TimestampMillisecondArray::from(vec![Some(1_000), Some(2_000)])
1047            .with_timezone("+02:00".to_string());
1048        let c2 = TimestampSecondArray::from(vec![Some(1_000_000), None])
1049            .with_timezone("+04:00".to_string());
1050        let batch =
1051            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
1052
1053        let mut file = tempfile::tempfile().unwrap();
1054        let mut writer = WriterBuilder::new()
1055            .with_timestamp_tz_format("%M:%H".to_string())
1056            .build(&mut file);
1057        writer.write(&batch).unwrap();
1058
1059        drop(writer);
1060        file.rewind().unwrap();
1061        let mut buffer: Vec<u8> = vec![];
1062        file.read_to_end(&mut buffer).unwrap();
1063
1064        let output = String::from_utf8(buffer).unwrap();
1065        assert_eq!(output, "c1,c2\n00:02,46:17\n00:02,\n");
1066    }
1067
1068    #[test]
1069    fn test_write_csv_with_lf_terminator() {
1070        let output = write_batch_with_terminator(Terminator::Any(b'\n'));
1071        assert_eq!(output, "c1,c2\nhello,1\nworld,2\n");
1072    }
1073
1074    #[test]
1075    fn test_write_csv_with_crlf_terminator() {
1076        let output = write_batch_with_terminator(Terminator::CRLF);
1077        assert_eq!(output, "c1,c2\r\nhello,1\r\nworld,2\r\n");
1078    }
1079
1080    #[test]
1081    fn test_write_csv_with_any_terminator() {
1082        let output = write_batch_with_terminator(Terminator::Any(b'|'));
1083        assert_eq!(output, "c1,c2|hello,1|world,2|");
1084    }
1085
1086    fn write_batch_with_terminator(terminator: Terminator) -> String {
1087        let schema = Schema::new(vec![
1088            Field::new("c1", DataType::Utf8, false),
1089            Field::new("c2", DataType::UInt32, false),
1090        ]);
1091
1092        let c1 = StringArray::from(vec!["hello", "world"]);
1093        let c2 = PrimitiveArray::<UInt32Type>::from(vec![1, 2]);
1094
1095        let batch =
1096            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
1097
1098        let mut buf = Vec::new();
1099        let mut writer = WriterBuilder::new()
1100            .with_line_terminator(terminator)
1101            .build(&mut buf);
1102        writer.write(&batch).unwrap();
1103        drop(writer);
1104
1105        String::from_utf8(buf).unwrap()
1106    }
1107
1108    #[test]
1109    fn test_write_csv_binary() {
1110        let fixed_size = 8;
1111        let schema = SchemaRef::new(Schema::new(vec![
1112            Field::new("c1", DataType::Binary, true),
1113            Field::new("c2", DataType::FixedSizeBinary(fixed_size), true),
1114            Field::new("c3", DataType::LargeBinary, true),
1115        ]));
1116        let mut c1_builder = BinaryBuilder::new();
1117        c1_builder.append_value(b"Homer");
1118        c1_builder.append_value(b"Bart");
1119        c1_builder.append_null();
1120        c1_builder.append_value(b"Ned");
1121        let mut c2_builder = FixedSizeBinaryBuilder::new(fixed_size);
1122        c2_builder.append_value(b"Simpson ").unwrap();
1123        c2_builder.append_value(b"Simpson ").unwrap();
1124        c2_builder.append_null();
1125        c2_builder.append_value(b"Flanders").unwrap();
1126        let mut c3_builder = LargeBinaryBuilder::new();
1127        c3_builder.append_null();
1128        c3_builder.append_null();
1129        c3_builder.append_value(b"Comic Book Guy");
1130        c3_builder.append_null();
1131
1132        let batch = RecordBatch::try_new(
1133            schema,
1134            vec![
1135                Arc::new(c1_builder.finish()) as ArrayRef,
1136                Arc::new(c2_builder.finish()) as ArrayRef,
1137                Arc::new(c3_builder.finish()) as ArrayRef,
1138            ],
1139        )
1140        .unwrap();
1141
1142        let mut buf = Vec::new();
1143        let builder = WriterBuilder::new();
1144        let mut writer = builder.build(&mut buf);
1145        writer.write(&batch).unwrap();
1146        drop(writer);
1147        assert_eq!(
1148            "\
1149            c1,c2,c3\n\
1150            486f6d6572,53696d70736f6e20,\n\
1151            42617274,53696d70736f6e20,\n\
1152            ,,436f6d696320426f6f6b20477579\n\
1153            4e6564,466c616e64657273,\n\
1154            ",
1155            String::from_utf8(buf).unwrap()
1156        );
1157    }
1158
1159    #[test]
1160    fn test_write_csv_whitespace_handling() {
1161        let schema = Schema::new(vec![
1162            Field::new("c1", DataType::Utf8, false),
1163            Field::new("c2", DataType::Float64, true),
1164            Field::new("c3", DataType::Utf8, true),
1165        ]);
1166
1167        let c1 = StringArray::from(vec![
1168            "  leading space",
1169            "trailing space  ",
1170            "  both spaces  ",
1171            "no spaces",
1172        ]);
1173        let c2 = PrimitiveArray::<Float64Type>::from(vec![
1174            Some(123.45),
1175            Some(678.90),
1176            None,
1177            Some(111.22),
1178        ]);
1179        let c3 = StringArray::from(vec![
1180            Some("  test  "),
1181            Some("value  "),
1182            None,
1183            Some("  another"),
1184        ]);
1185
1186        let batch = RecordBatch::try_new(
1187            Arc::new(schema),
1188            vec![Arc::new(c1), Arc::new(c2), Arc::new(c3)],
1189        )
1190        .unwrap();
1191
1192        // Test with no whitespace handling (default)
1193        let mut buf = Vec::new();
1194        let builder = WriterBuilder::new();
1195        let mut writer = builder.build(&mut buf);
1196        writer.write(&batch).unwrap();
1197        drop(writer);
1198        assert_eq!(
1199            "c1,c2,c3\n  leading space,123.45,  test  \ntrailing space  ,678.9,value  \n  both spaces  ,,\nno spaces,111.22,  another\n",
1200            String::from_utf8(buf).unwrap()
1201        );
1202
1203        // Test with ignore leading whitespace only
1204        let mut buf = Vec::new();
1205        let builder = WriterBuilder::new().with_ignore_leading_whitespace(true);
1206        let mut writer = builder.build(&mut buf);
1207        writer.write(&batch).unwrap();
1208        drop(writer);
1209        assert_eq!(
1210            "c1,c2,c3\nleading space,123.45,test  \ntrailing space  ,678.9,value  \nboth spaces  ,,\nno spaces,111.22,another\n",
1211            String::from_utf8(buf).unwrap()
1212        );
1213
1214        // Test with ignore trailing whitespace only
1215        let mut buf = Vec::new();
1216        let builder = WriterBuilder::new().with_ignore_trailing_whitespace(true);
1217        let mut writer = builder.build(&mut buf);
1218        writer.write(&batch).unwrap();
1219        drop(writer);
1220        assert_eq!(
1221            "c1,c2,c3\n  leading space,123.45,  test\ntrailing space,678.9,value\n  both spaces,,\nno spaces,111.22,  another\n",
1222            String::from_utf8(buf).unwrap()
1223        );
1224
1225        // Test with both ignore leading and trailing whitespace
1226        let mut buf = Vec::new();
1227        let builder = WriterBuilder::new()
1228            .with_ignore_leading_whitespace(true)
1229            .with_ignore_trailing_whitespace(true);
1230        let mut writer = builder.build(&mut buf);
1231        writer.write(&batch).unwrap();
1232        drop(writer);
1233        assert_eq!(
1234            "c1,c2,c3\nleading space,123.45,test\ntrailing space,678.9,value\nboth spaces,,\nno spaces,111.22,another\n",
1235            String::from_utf8(buf).unwrap()
1236        );
1237    }
1238
1239    #[test]
1240    fn test_write_csv_whitespace_with_special_chars() {
1241        let schema = Schema::new(vec![Field::new("c1", DataType::Utf8, false)]);
1242
1243        let c1 = StringArray::from(vec![
1244            "  quoted \"value\"  ",
1245            "  new\nline  ",
1246            "  comma,value  ",
1247            "\ttab\tvalue\t",
1248        ]);
1249
1250        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1)]).unwrap();
1251
1252        // Test with both ignore leading and trailing whitespace
1253        let mut buf = Vec::new();
1254        let builder = WriterBuilder::new()
1255            .with_ignore_leading_whitespace(true)
1256            .with_ignore_trailing_whitespace(true);
1257        let mut writer = builder.build(&mut buf);
1258        writer.write(&batch).unwrap();
1259        drop(writer);
1260
1261        // Note: tabs are trimmed as they are whitespace characters
1262        assert_eq!(
1263            "c1\n\"quoted \"\"value\"\"\"\n\"new\nline\"\n\"comma,value\"\ntab\tvalue\n",
1264            String::from_utf8(buf).unwrap()
1265        );
1266    }
1267
1268    #[test]
1269    fn test_write_csv_whitespace_all_string_types() {
1270        use arrow_array::{LargeStringArray, StringViewArray};
1271
1272        let schema = Schema::new(vec![
1273            Field::new("utf8", DataType::Utf8, false),
1274            Field::new("large_utf8", DataType::LargeUtf8, false),
1275            Field::new("utf8_view", DataType::Utf8View, false),
1276        ]);
1277
1278        let utf8 = StringArray::from(vec!["  leading", "trailing  ", "  both  ", "no_spaces"]);
1279
1280        let large_utf8 =
1281            LargeStringArray::from(vec!["  leading", "trailing  ", "  both  ", "no_spaces"]);
1282
1283        let utf8_view =
1284            StringViewArray::from(vec!["  leading", "trailing  ", "  both  ", "no_spaces"]);
1285
1286        let batch = RecordBatch::try_new(
1287            Arc::new(schema),
1288            vec![Arc::new(utf8), Arc::new(large_utf8), Arc::new(utf8_view)],
1289        )
1290        .unwrap();
1291
1292        // Test with no whitespace handling (default)
1293        let mut buf = Vec::new();
1294        let builder = WriterBuilder::new();
1295        let mut writer = builder.build(&mut buf);
1296        writer.write(&batch).unwrap();
1297        drop(writer);
1298        assert_eq!(
1299            "utf8,large_utf8,utf8_view\n  leading,  leading,  leading\ntrailing  ,trailing  ,trailing  \n  both  ,  both  ,  both  \nno_spaces,no_spaces,no_spaces\n",
1300            String::from_utf8(buf).unwrap()
1301        );
1302
1303        // Test with both ignore leading and trailing whitespace
1304        let mut buf = Vec::new();
1305        let builder = WriterBuilder::new()
1306            .with_ignore_leading_whitespace(true)
1307            .with_ignore_trailing_whitespace(true);
1308        let mut writer = builder.build(&mut buf);
1309        writer.write(&batch).unwrap();
1310        drop(writer);
1311        assert_eq!(
1312            "utf8,large_utf8,utf8_view\nleading,leading,leading\ntrailing,trailing,trailing\nboth,both,both\nno_spaces,no_spaces,no_spaces\n",
1313            String::from_utf8(buf).unwrap()
1314        );
1315
1316        // Test with only leading whitespace trimming
1317        let mut buf = Vec::new();
1318        let builder = WriterBuilder::new().with_ignore_leading_whitespace(true);
1319        let mut writer = builder.build(&mut buf);
1320        writer.write(&batch).unwrap();
1321        drop(writer);
1322        assert_eq!(
1323            "utf8,large_utf8,utf8_view\nleading,leading,leading\ntrailing  ,trailing  ,trailing  \nboth  ,both  ,both  \nno_spaces,no_spaces,no_spaces\n",
1324            String::from_utf8(buf).unwrap()
1325        );
1326
1327        // Test with only trailing whitespace trimming
1328        let mut buf = Vec::new();
1329        let builder = WriterBuilder::new().with_ignore_trailing_whitespace(true);
1330        let mut writer = builder.build(&mut buf);
1331        writer.write(&batch).unwrap();
1332        drop(writer);
1333        assert_eq!(
1334            "utf8,large_utf8,utf8_view\n  leading,  leading,  leading\ntrailing,trailing,trailing\n  both,  both,  both\nno_spaces,no_spaces,no_spaces\n",
1335            String::from_utf8(buf).unwrap()
1336        );
1337    }
1338
1339    fn write_quote_style(batch: &RecordBatch, quote_style: QuoteStyle) -> String {
1340        let mut buf = Vec::new();
1341        let mut writer = WriterBuilder::new()
1342            .with_quote_style(quote_style)
1343            .build(&mut buf);
1344        writer.write(batch).unwrap();
1345        drop(writer);
1346        String::from_utf8(buf).unwrap()
1347    }
1348
1349    fn write_quote_style_with_null(
1350        batch: &RecordBatch,
1351        quote_style: QuoteStyle,
1352        null_value: &str,
1353    ) -> String {
1354        let mut buf = Vec::new();
1355        let mut writer = WriterBuilder::new()
1356            .with_quote_style(quote_style)
1357            .with_null(null_value.to_string())
1358            .build(&mut buf);
1359        writer.write(batch).unwrap();
1360        drop(writer);
1361        String::from_utf8(buf).unwrap()
1362    }
1363
1364    #[test]
1365    fn test_write_csv_quote_style() {
1366        let schema = Schema::new(vec![
1367            Field::new("text", DataType::Utf8, false),
1368            Field::new("number", DataType::Int32, false),
1369            Field::new("float", DataType::Float64, false),
1370        ]);
1371
1372        let text = StringArray::from(vec!["hello", "world", "comma,value", "quote\"test"]);
1373        let number = Int32Array::from(vec![1, 2, 3, 4]);
1374        let float = Float64Array::from(vec![1.1, 2.2, 3.3, 4.4]);
1375
1376        let batch = RecordBatch::try_new(
1377            Arc::new(schema),
1378            vec![Arc::new(text), Arc::new(number), Arc::new(float)],
1379        )
1380        .unwrap();
1381
1382        // Test with QuoteStyle::Necessary (default)
1383        assert_eq!(
1384            "text,number,float\nhello,1,1.1\nworld,2,2.2\n\"comma,value\",3,3.3\n\"quote\"\"test\",4,4.4\n",
1385            write_quote_style(&batch, QuoteStyle::Necessary)
1386        );
1387
1388        // Test with QuoteStyle::Always (equivalent to Spark's quoteAll=true)
1389        assert_eq!(
1390            "\"text\",\"number\",\"float\"\n\"hello\",\"1\",\"1.1\"\n\"world\",\"2\",\"2.2\"\n\"comma,value\",\"3\",\"3.3\"\n\"quote\"\"test\",\"4\",\"4.4\"\n",
1391            write_quote_style(&batch, QuoteStyle::Always)
1392        );
1393
1394        // Test with QuoteStyle::NonNumeric
1395        assert_eq!(
1396            "\"text\",\"number\",\"float\"\n\"hello\",1,1.1\n\"world\",2,2.2\n\"comma,value\",3,3.3\n\"quote\"\"test\",4,4.4\n",
1397            write_quote_style(&batch, QuoteStyle::NonNumeric)
1398        );
1399
1400        // Test with QuoteStyle::Never (warning: can produce invalid CSV)
1401        // Note: This produces invalid CSV for fields with commas or quotes
1402        assert_eq!(
1403            "text,number,float\nhello,1,1.1\nworld,2,2.2\ncomma,value,3,3.3\nquote\"test,4,4.4\n",
1404            write_quote_style(&batch, QuoteStyle::Never)
1405        );
1406    }
1407
1408    #[test]
1409    fn test_write_csv_quote_style_with_nulls() {
1410        let schema = Schema::new(vec![
1411            Field::new("text", DataType::Utf8, true),
1412            Field::new("number", DataType::Int32, true),
1413        ]);
1414
1415        let text = StringArray::from(vec![Some("hello"), None, Some("world")]);
1416        let number = Int32Array::from(vec![Some(1), Some(2), None]);
1417
1418        let batch =
1419            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(text), Arc::new(number)]).unwrap();
1420
1421        // Test with QuoteStyle::Always
1422        assert_eq!(
1423            "\"text\",\"number\"\n\"hello\",\"1\"\n\"\",\"2\"\n\"world\",\"\"\n",
1424            write_quote_style(&batch, QuoteStyle::Always)
1425        );
1426
1427        // Test with QuoteStyle::Always and custom null value
1428        assert_eq!(
1429            "\"text\",\"number\"\n\"hello\",\"1\"\n\"NULL\",\"2\"\n\"world\",\"NULL\"\n",
1430            write_quote_style_with_null(&batch, QuoteStyle::Always, "NULL")
1431        );
1432    }
1433}