datafusion_common/file_options/
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//! Options related to how csv files should be written
19
20use crate::config::CsvOptions;
21use crate::error::{DataFusionError, Result};
22use crate::parsers::CompressionTypeVariant;
23
24use arrow::csv::WriterBuilder;
25
26/// Options for writing CSV files
27#[derive(Clone, Debug)]
28pub struct CsvWriterOptions {
29    /// Struct from the arrow crate which contains all csv writing related settings
30    pub writer_options: WriterBuilder,
31    /// Compression to apply after ArrowWriter serializes RecordBatches.
32    /// This compression is applied by DataFusion not the ArrowWriter itself.
33    pub compression: CompressionTypeVariant,
34}
35
36impl CsvWriterOptions {
37    pub fn new(
38        writer_options: WriterBuilder,
39        compression: CompressionTypeVariant,
40    ) -> Self {
41        Self {
42            writer_options,
43            compression,
44        }
45    }
46}
47
48impl TryFrom<&CsvOptions> for CsvWriterOptions {
49    type Error = DataFusionError;
50
51    fn try_from(value: &CsvOptions) -> Result<Self> {
52        let mut builder = WriterBuilder::default()
53            .with_header(value.has_header.unwrap_or(true))
54            .with_quote(value.quote)
55            .with_delimiter(value.delimiter);
56
57        if let Some(v) = &value.date_format {
58            builder = builder.with_date_format(v.into())
59        }
60        if let Some(v) = &value.datetime_format {
61            builder = builder.with_datetime_format(v.into())
62        }
63        if let Some(v) = &value.timestamp_format {
64            builder = builder.with_timestamp_format(v.into())
65        }
66        if let Some(v) = &value.timestamp_tz_format {
67            builder = builder.with_timestamp_tz_format(v.into())
68        }
69        if let Some(v) = &value.time_format {
70            builder = builder.with_time_format(v.into())
71        }
72        if let Some(v) = &value.null_value {
73            builder = builder.with_null(v.into())
74        }
75        if let Some(v) = &value.escape {
76            builder = builder.with_escape(*v)
77        }
78        if let Some(v) = &value.double_quote {
79            builder = builder.with_double_quote(*v)
80        }
81        Ok(CsvWriterOptions {
82            writer_options: builder,
83            compression: value.compression,
84        })
85    }
86}
datafusion_common/file_options/csv_writer.rs

datafusion_common/file_options/
csv_writer.rs
