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    /// Compression level for the output file.
35    pub compression_level: Option<u32>,
36}
37
38impl CsvWriterOptions {
39    pub fn new(
40        writer_options: WriterBuilder,
41        compression: CompressionTypeVariant,
42    ) -> Self {
43        Self {
44            writer_options,
45            compression,
46            compression_level: None,
47        }
48    }
49
50    /// Create a new `CsvWriterOptions` with the specified compression level.
51    pub fn new_with_level(
52        writer_options: WriterBuilder,
53        compression: CompressionTypeVariant,
54        compression_level: u32,
55    ) -> Self {
56        Self {
57            writer_options,
58            compression,
59            compression_level: Some(compression_level),
60        }
61    }
62}
63
64impl TryFrom<&CsvOptions> for CsvWriterOptions {
65    type Error = DataFusionError;
66
67    fn try_from(value: &CsvOptions) -> Result<Self> {
68        let mut builder = WriterBuilder::default()
69            .with_header(value.has_header.unwrap_or(true))
70            .with_quote(value.quote)
71            .with_delimiter(value.delimiter);
72
73        if let Some(v) = &value.date_format {
74            builder = builder.with_date_format(v.into())
75        }
76        if let Some(v) = &value.datetime_format {
77            builder = builder.with_datetime_format(v.into())
78        }
79        if let Some(v) = &value.timestamp_format {
80            builder = builder.with_timestamp_format(v.into())
81        }
82        if let Some(v) = &value.timestamp_tz_format {
83            builder = builder.with_timestamp_tz_format(v.into())
84        }
85        if let Some(v) = &value.time_format {
86            builder = builder.with_time_format(v.into())
87        }
88        if let Some(v) = &value.null_value {
89            builder = builder.with_null(v.into())
90        }
91        if let Some(v) = &value.escape {
92            builder = builder.with_escape(*v)
93        }
94        if let Some(v) = &value.double_quote {
95            builder = builder.with_double_quote(*v)
96        }
97        Ok(CsvWriterOptions {
98            writer_options: builder,
99            compression: value.compression,
100            compression_level: value.compression_level,
101        })
102    }
103}
datafusion_common/file_options/csv_writer.rs

datafusion_common/file_options/
csv_writer.rs
