parquet_fromcsv/
parquet-fromcsv.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Binary file to converts csv to Parquet file
6//!
7//! # Install
8//!
9//! `parquet-fromcsv` can be installed using `cargo`:
10//!
11//! ```text
12//! cargo install parquet --features=cli
13//! ```
14//!
15//! After this `parquet-fromcsv` should be available:
16//!
17//! ```text
18//! parquet-fromcsv --schema message_schema_for_parquet.txt input.csv output.parquet
19//! ```
20//!
21//! The binary can also be built from the source code and run as follows:
22//!
23//! ```text
24//! cargo run --features=cli --bin parquet-fromcsv --schema message_schema_for_parquet.txt \
25//!    \ input.csv output.parquet
26//! ```
27//!
28//! # Options
29//!
30//! ```text
31#![cfg_attr(doc, doc = include_str!("./parquet-fromcsv-help.txt"))] // Update for this file : Run test test_command_help
32//! ```
33//!
34//! ## Parquet file options
35//!
36//! ```text
37//! - `-b`, `--batch-size` : Batch size for Parquet
38//! - `-c`, `--parquet-compression` : Compression option for Parquet, default is SNAPPY
39//! - `-s`, `--schema` : Path to message schema for generated Parquet file
40//! - `-o`, `--output-file` : Path to output Parquet file
41//! - `-w`, `--writer-version` : Writer version
42//! - `-m`, `--max-row-group-size` : Max row group size
43//! -       `--enable-bloom-filter` : Enable bloom filter during writing
44//! ```
45//!
46//! ## Input file options
47//!
48//! ```text
49//! - `-i`, `--input-file` : Path to input CSV file
50//! - `-f`, `--input-format` : Dialect for input file, `csv` or `tsv`.
51//! - `-C`, `--csv-compression` : Compression option for csv, default is UNCOMPRESSED
52//! - `-d`, `--delimiter : Field delimiter for CSV file, default depends `--input-format`
53//! - `-e`, `--escape` : Escape character for input file
54//! - `-h`, `--has-header` : Input has header
55//! - `-r`, `--record-terminator` : Record terminator character for input. default is CRLF
56//! - `-q`, `--quote-char` : Input quoting character
57//! ```
58//!
59
60use std::{
61    fmt::Display,
62    fs::{File, read_to_string},
63    io::Read,
64    path::{Path, PathBuf},
65    sync::Arc,
66};
67
68use arrow_csv::ReaderBuilder;
69use arrow_schema::{ArrowError, Schema};
70use clap::{Parser, ValueEnum};
71use parquet::arrow::arrow_writer::ArrowWriterOptions;
72use parquet::{
73    arrow::{ArrowWriter, parquet_to_arrow_schema},
74    basic::Compression,
75    errors::ParquetError,
76    file::properties::{WriterProperties, WriterVersion},
77    schema::{parser::parse_message_type, types::SchemaDescriptor},
78};
79
80#[derive(Debug)]
81enum ParquetFromCsvError {
82    CommandLineParseError(clap::Error),
83    IoError(std::io::Error),
84    ArrowError(ArrowError),
85    ParquetError(ParquetError),
86    WithContext(String, Box<Self>),
87}
88
89impl From<std::io::Error> for ParquetFromCsvError {
90    fn from(e: std::io::Error) -> Self {
91        Self::IoError(e)
92    }
93}
94
95impl From<ArrowError> for ParquetFromCsvError {
96    fn from(e: ArrowError) -> Self {
97        Self::ArrowError(e)
98    }
99}
100
101impl From<ParquetError> for ParquetFromCsvError {
102    fn from(e: ParquetError) -> Self {
103        Self::ParquetError(e)
104    }
105}
106
107impl From<clap::Error> for ParquetFromCsvError {
108    fn from(e: clap::Error) -> Self {
109        Self::CommandLineParseError(e)
110    }
111}
112
113impl ParquetFromCsvError {
114    pub fn with_context<E: Into<ParquetFromCsvError>>(
115        inner_error: E,
116        context: &str,
117    ) -> ParquetFromCsvError {
118        let inner = inner_error.into();
119        ParquetFromCsvError::WithContext(context.to_string(), Box::new(inner))
120    }
121}
122
123impl Display for ParquetFromCsvError {
124    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125        match self {
126            ParquetFromCsvError::CommandLineParseError(e) => write!(f, "{e}"),
127            ParquetFromCsvError::IoError(e) => write!(f, "{e}"),
128            ParquetFromCsvError::ArrowError(e) => write!(f, "{e}"),
129            ParquetFromCsvError::ParquetError(e) => write!(f, "{e}"),
130            ParquetFromCsvError::WithContext(c, e) => {
131                writeln!(f, "{e}")?;
132                write!(f, "context: {c}")
133            }
134        }
135    }
136}
137
138#[derive(Debug, Parser)]
139#[clap(author, version, disable_help_flag=true, about("Binary to convert csv to Parquet"), long_about=None)]
140struct Args {
141    /// Path to a text file containing a parquet schema definition
142    #[clap(short, long, help("message schema for output Parquet"))]
143    schema: PathBuf,
144    /// input CSV file path
145    #[clap(short, long, help("input CSV file"))]
146    input_file: PathBuf,
147    /// output Parquet file path
148    #[clap(short, long, help("output Parquet file"))]
149    output_file: PathBuf,
150    /// input file format
151    #[clap(
152        value_enum,
153        short('f'),
154        long,
155        help("input file format"),
156        default_value_t=CsvDialect::Csv
157    )]
158    input_format: CsvDialect,
159    /// batch size
160    #[clap(
161        short,
162        long,
163        help("batch size"),
164        default_value_t = 1000,
165        env = "PARQUET_FROM_CSV_BATCHSIZE"
166    )]
167    batch_size: usize,
168    /// has header line
169    #[clap(short, long, help("has header"))]
170    has_header: bool,
171    /// field delimiter
172    ///
173    /// default value:
174    ///  when input_format==CSV: ','
175    ///  when input_format==TSV: 'TAB'
176    #[clap(short, long, help("field delimiter"))]
177    delimiter: Option<char>,
178    #[clap(value_enum, short, long, help("record terminator"))]
179    record_terminator: Option<RecordTerminator>,
180    #[clap(short, long, help("escape character"))]
181    escape_char: Option<char>,
182    #[clap(short, long, help("quote character"))]
183    quote_char: Option<char>,
184    #[clap(short('D'), long, help("double quote"))]
185    double_quote: Option<bool>,
186    #[clap(short('C'), long, help("compression mode of csv"), default_value_t=Compression::UNCOMPRESSED)]
187    #[clap(value_parser=compression_from_str)]
188    csv_compression: Compression,
189    #[clap(short('c'), long, help("compression mode of parquet"), default_value_t=Compression::SNAPPY)]
190    #[clap(value_parser=compression_from_str)]
191    parquet_compression: Compression,
192
193    #[clap(short, long, help("writer version"))]
194    #[clap(value_parser=writer_version_from_str)]
195    writer_version: Option<WriterVersion>,
196    #[clap(short, long, help("max row group size"))]
197    max_row_group_size: Option<usize>,
198    #[clap(long, help("whether to enable bloom filter writing"))]
199    enable_bloom_filter: Option<bool>,
200
201    #[clap(long, action=clap::ArgAction::Help, help("display usage help"))]
202    help: Option<bool>,
203}
204
205fn compression_from_str(cmp: &str) -> Result<Compression, String> {
206    match cmp.to_uppercase().as_str() {
207        "UNCOMPRESSED" => Ok(Compression::UNCOMPRESSED),
208        "SNAPPY" => Ok(Compression::SNAPPY),
209        "GZIP" => Ok(Compression::GZIP(Default::default())),
210        "LZO" => Ok(Compression::LZO),
211        "BROTLI" => Ok(Compression::BROTLI(Default::default())),
212        "LZ4" => Ok(Compression::LZ4),
213        "ZSTD" => Ok(Compression::ZSTD(Default::default())),
214        v => Err(format!(
215            "Unknown compression {v} : possible values UNCOMPRESSED, SNAPPY, GZIP, LZO, BROTLI, LZ4, ZSTD \n\nFor more information try --help"
216        )),
217    }
218}
219
220fn writer_version_from_str(cmp: &str) -> Result<WriterVersion, String> {
221    match cmp.to_uppercase().as_str() {
222        "1" => Ok(WriterVersion::PARQUET_1_0),
223        "2" => Ok(WriterVersion::PARQUET_2_0),
224        v => Err(format!("Unknown writer version {v} : possible values 1, 2")),
225    }
226}
227
228impl Args {
229    fn schema_path(&self) -> &Path {
230        self.schema.as_path()
231    }
232    fn get_delimiter(&self) -> u8 {
233        match self.delimiter {
234            Some(ch) => ch as u8,
235            None => match self.input_format {
236                CsvDialect::Csv => b',',
237                CsvDialect::Tsv => b'\t',
238            },
239        }
240    }
241    fn get_terminator(&self) -> Option<u8> {
242        match self.record_terminator {
243            Some(RecordTerminator::LF) => Some(0x0a),
244            Some(RecordTerminator::CR) => Some(0x0d),
245            Some(RecordTerminator::Crlf) => None,
246            None => match self.input_format {
247                CsvDialect::Csv => None,
248                CsvDialect::Tsv => Some(0x0a),
249            },
250        }
251    }
252    fn get_escape(&self) -> Option<u8> {
253        self.escape_char.map(|ch| ch as u8)
254    }
255    fn get_quote(&self) -> Option<u8> {
256        if self.quote_char.is_none() {
257            match self.input_format {
258                CsvDialect::Csv => Some(b'\"'),
259                CsvDialect::Tsv => None,
260            }
261        } else {
262            self.quote_char.map(|c| c as u8)
263        }
264    }
265}
266
267#[derive(Debug, Clone, Copy, ValueEnum, PartialEq)]
268enum CsvDialect {
269    Csv,
270    Tsv,
271}
272
273#[derive(Debug, Clone, Copy, ValueEnum, PartialEq)]
274enum RecordTerminator {
275    LF,
276    Crlf,
277    CR,
278}
279
280fn configure_writer_properties(args: &Args) -> WriterProperties {
281    let mut properties_builder =
282        WriterProperties::builder().set_compression(args.parquet_compression);
283    if let Some(writer_version) = args.writer_version {
284        properties_builder = properties_builder.set_writer_version(writer_version);
285    }
286    if let Some(max_row_group_size) = args.max_row_group_size {
287        properties_builder =
288            properties_builder.set_max_row_group_row_count(Some(max_row_group_size));
289    }
290    if let Some(enable_bloom_filter) = args.enable_bloom_filter {
291        properties_builder = properties_builder.set_bloom_filter_enabled(enable_bloom_filter);
292    }
293    properties_builder.build()
294}
295
296fn configure_reader_builder(args: &Args, arrow_schema: Arc<Schema>) -> ReaderBuilder {
297    fn configure_reader<T, F: Fn(ReaderBuilder, T) -> ReaderBuilder>(
298        builder: ReaderBuilder,
299        value: Option<T>,
300        fun: F,
301    ) -> ReaderBuilder {
302        if let Some(val) = value {
303            fun(builder, val)
304        } else {
305            builder
306        }
307    }
308
309    let mut builder = ReaderBuilder::new(arrow_schema)
310        .with_batch_size(args.batch_size)
311        .with_header(args.has_header)
312        .with_delimiter(args.get_delimiter());
313
314    builder = configure_reader(
315        builder,
316        args.get_terminator(),
317        ReaderBuilder::with_terminator,
318    );
319    builder = configure_reader(builder, args.get_escape(), ReaderBuilder::with_escape);
320    builder = configure_reader(builder, args.get_quote(), ReaderBuilder::with_quote);
321
322    builder
323}
324
325fn convert_csv_to_parquet(args: &Args) -> Result<(), ParquetFromCsvError> {
326    let schema = read_to_string(args.schema_path()).map_err(|e| {
327        ParquetFromCsvError::with_context(
328            e,
329            &format!("Failed to open schema file {:#?}", args.schema_path()),
330        )
331    })?;
332    let parquet_schema = Arc::new(parse_message_type(&schema)?);
333    let desc = SchemaDescriptor::new(parquet_schema);
334    let arrow_schema = Arc::new(parquet_to_arrow_schema(&desc, None)?);
335
336    // create output parquet writer
337    let parquet_file = File::create(&args.output_file).map_err(|e| {
338        ParquetFromCsvError::with_context(
339            e,
340            &format!("Failed to create output file {:#?}", args.output_file),
341        )
342    })?;
343
344    let options = ArrowWriterOptions::new()
345        .with_properties(configure_writer_properties(args))
346        .with_schema_root(desc.name().to_string());
347
348    let mut arrow_writer =
349        ArrowWriter::try_new_with_options(parquet_file, arrow_schema.clone(), options)
350            .map_err(|e| ParquetFromCsvError::with_context(e, "Failed to create ArrowWriter"))?;
351
352    // open input file
353    let input_file = File::open(&args.input_file).map_err(|e| {
354        ParquetFromCsvError::with_context(
355            e,
356            &format!("Failed to open input file {:#?}", args.input_file),
357        )
358    })?;
359
360    // open input file decoder
361    let input_file_decoder = match args.csv_compression {
362        Compression::UNCOMPRESSED => Box::new(input_file) as Box<dyn Read>,
363        Compression::SNAPPY => Box::new(snap::read::FrameDecoder::new(input_file)) as Box<dyn Read>,
364        Compression::GZIP(_) => {
365            Box::new(flate2::read::MultiGzDecoder::new(input_file)) as Box<dyn Read>
366        }
367        Compression::BROTLI(_) => {
368            Box::new(brotli::Decompressor::new(input_file, 0)) as Box<dyn Read>
369        }
370        Compression::LZ4 => {
371            Box::new(lz4_flex::frame::FrameDecoder::new(input_file)) as Box<dyn Read>
372        }
373        Compression::ZSTD(_) => {
374            Box::new(zstd::Decoder::new(input_file).map_err(|e| {
375                ParquetFromCsvError::with_context(e, "Failed to create zstd::Decoder")
376            })?) as Box<dyn Read>
377        }
378        d => unimplemented!("compression type {d}"),
379    };
380
381    // create input csv reader
382    let builder = configure_reader_builder(args, arrow_schema);
383    let reader = builder.build(input_file_decoder)?;
384    for batch_result in reader {
385        let batch = batch_result.map_err(|e| {
386            ParquetFromCsvError::with_context(e, "Failed to read RecordBatch from CSV")
387        })?;
388        arrow_writer.write(&batch).map_err(|e| {
389            ParquetFromCsvError::with_context(e, "Failed to write RecordBatch to parquet")
390        })?;
391    }
392    arrow_writer
393        .close()
394        .map_err(|e| ParquetFromCsvError::with_context(e, "Failed to close parquet"))?;
395    Ok(())
396}
397
398fn main() -> Result<(), ParquetFromCsvError> {
399    let args = Args::parse();
400    convert_csv_to_parquet(&args)
401}
402
403#[cfg(test)]
404mod tests {
405    use std::{
406        io::Write,
407        path::{Path, PathBuf},
408    };
409
410    use super::*;
411    use arrow::datatypes::{DataType, Field};
412    use brotli::CompressorWriter;
413    use clap::{CommandFactory, Parser};
414    use flate2::write::GzEncoder;
415    use parquet::basic::{BrotliLevel, GzipLevel, ZstdLevel};
416    use parquet::file::reader::{FileReader, SerializedFileReader};
417    use snap::write::FrameEncoder;
418    use tempfile::NamedTempFile;
419
420    #[test]
421    fn test_command_help() {
422        let mut cmd = Args::command();
423        let dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
424        let path_buf = PathBuf::from(dir).join("src/bin/parquet-fromcsv-help.txt");
425        let expected = std::fs::read_to_string(path_buf).unwrap();
426        let mut buffer_vec = Vec::new();
427        let mut buffer = std::io::Cursor::new(&mut buffer_vec);
428        cmd.write_long_help(&mut buffer).unwrap();
429        // Remove Parquet version string from the help text
430        let mut actual = String::from_utf8(buffer_vec).unwrap();
431        let pos = actual.find('\n').unwrap() + 1;
432        actual = actual[pos..].to_string();
433        assert_eq!(
434            expected, actual,
435            "help text not match. please update to \n---\n{actual}\n---\n"
436        )
437    }
438
439    fn parse_args(mut extra_args: Vec<&str>) -> Result<Args, ParquetFromCsvError> {
440        let mut args = vec![
441            "test",
442            "--schema",
443            "test.schema",
444            "--input-file",
445            "infile.csv",
446            "--output-file",
447            "out.parquet",
448        ];
449        args.append(&mut extra_args);
450        let args = Args::try_parse_from(args.iter())?;
451        Ok(args)
452    }
453
454    #[test]
455    fn test_parse_arg_minimum() -> Result<(), ParquetFromCsvError> {
456        let args = parse_args(vec![])?;
457
458        assert_eq!(args.schema, PathBuf::from(Path::new("test.schema")));
459        assert_eq!(args.input_file, PathBuf::from(Path::new("infile.csv")));
460        assert_eq!(args.output_file, PathBuf::from(Path::new("out.parquet")));
461        // test default values
462        assert_eq!(args.input_format, CsvDialect::Csv);
463        assert_eq!(args.batch_size, 1000);
464        assert!(!args.has_header);
465        assert_eq!(args.delimiter, None);
466        assert_eq!(args.get_delimiter(), b',');
467        assert_eq!(args.record_terminator, None);
468        assert_eq!(args.get_terminator(), None); // CRLF
469        assert_eq!(args.quote_char, None);
470        assert_eq!(args.get_quote(), Some(b'\"'));
471        assert_eq!(args.double_quote, None);
472        assert_eq!(args.parquet_compression, Compression::SNAPPY);
473        Ok(())
474    }
475
476    #[test]
477    fn test_parse_arg_format_variants() -> Result<(), ParquetFromCsvError> {
478        let args = parse_args(vec!["--input-format", "csv"])?;
479        assert_eq!(args.input_format, CsvDialect::Csv);
480        assert_eq!(args.get_delimiter(), b',');
481        assert_eq!(args.get_terminator(), None); // CRLF
482        assert_eq!(args.get_quote(), Some(b'\"'));
483        assert_eq!(args.get_escape(), None);
484        let args = parse_args(vec!["--input-format", "tsv"])?;
485        assert_eq!(args.input_format, CsvDialect::Tsv);
486        assert_eq!(args.get_delimiter(), b'\t');
487        assert_eq!(args.get_terminator(), Some(b'\x0a')); // LF
488        assert_eq!(args.get_quote(), None); // quote none
489        assert_eq!(args.get_escape(), None);
490
491        let args = parse_args(vec!["--input-format", "csv", "--escape-char", "\\"])?;
492        assert_eq!(args.input_format, CsvDialect::Csv);
493        assert_eq!(args.get_delimiter(), b',');
494        assert_eq!(args.get_terminator(), None); // CRLF
495        assert_eq!(args.get_quote(), Some(b'\"'));
496        assert_eq!(args.get_escape(), Some(b'\\'));
497
498        let args = parse_args(vec!["--input-format", "tsv", "--delimiter", ":"])?;
499        assert_eq!(args.input_format, CsvDialect::Tsv);
500        assert_eq!(args.get_delimiter(), b':');
501        assert_eq!(args.get_terminator(), Some(b'\x0a')); // LF
502        assert_eq!(args.get_quote(), None); // quote none
503        assert_eq!(args.get_escape(), None);
504
505        Ok(())
506    }
507
508    #[test]
509    #[should_panic(expected = "CommandLineParseError")]
510    fn test_parse_arg_format_error() {
511        parse_args(vec!["--input-format", "excel"]).unwrap();
512    }
513
514    #[test]
515    fn test_parse_arg_compression_format() {
516        let args = parse_args(vec!["--parquet-compression", "uncompressed"]).unwrap();
517        assert_eq!(args.parquet_compression, Compression::UNCOMPRESSED);
518        let args = parse_args(vec!["--parquet-compression", "snappy"]).unwrap();
519        assert_eq!(args.parquet_compression, Compression::SNAPPY);
520        let args = parse_args(vec!["--parquet-compression", "gzip"]).unwrap();
521        assert_eq!(
522            args.parquet_compression,
523            Compression::GZIP(Default::default())
524        );
525        let args = parse_args(vec!["--parquet-compression", "lzo"]).unwrap();
526        assert_eq!(args.parquet_compression, Compression::LZO);
527        let args = parse_args(vec!["--parquet-compression", "lz4"]).unwrap();
528        assert_eq!(args.parquet_compression, Compression::LZ4);
529        let args = parse_args(vec!["--parquet-compression", "brotli"]).unwrap();
530        assert_eq!(
531            args.parquet_compression,
532            Compression::BROTLI(Default::default())
533        );
534        let args = parse_args(vec!["--parquet-compression", "zstd"]).unwrap();
535        assert_eq!(
536            args.parquet_compression,
537            Compression::ZSTD(Default::default())
538        );
539    }
540
541    #[test]
542    fn test_parse_arg_compression_format_fail() {
543        match parse_args(vec!["--parquet-compression", "zip"]) {
544            Ok(_) => panic!("unexpected success"),
545            Err(e) => {
546                let err = e.to_string();
547                assert!(err.contains("error: invalid value 'zip' for '--parquet-compression <PARQUET_COMPRESSION>': Unknown compression ZIP : possible values UNCOMPRESSED, SNAPPY, GZIP, LZO, BROTLI, LZ4, ZSTD \n\nFor more information try --help"), "{err}")
548            }
549        }
550    }
551
552    fn assert_debug_text(debug_text: &str, name: &str, value: &str) {
553        let pattern = format!(" {name}: {value}");
554        assert!(
555            debug_text.contains(&pattern),
556            "\"{debug_text}\" not contains \"{pattern}\""
557        )
558    }
559
560    #[test]
561    fn test_configure_reader_builder() {
562        let args = Args {
563            schema: PathBuf::from(Path::new("schema.arvo")),
564            input_file: PathBuf::from(Path::new("test.csv")),
565            output_file: PathBuf::from(Path::new("out.parquet")),
566            batch_size: 1000,
567            input_format: CsvDialect::Csv,
568            has_header: false,
569            delimiter: None,
570            record_terminator: None,
571            escape_char: None,
572            quote_char: None,
573            double_quote: None,
574            csv_compression: Compression::UNCOMPRESSED,
575            parquet_compression: Compression::SNAPPY,
576            writer_version: None,
577            max_row_group_size: None,
578            enable_bloom_filter: None,
579            help: None,
580        };
581        let arrow_schema = Arc::new(Schema::new(vec![
582            Field::new("field1", DataType::Utf8, false),
583            Field::new("field2", DataType::Utf8, false),
584            Field::new("field3", DataType::Utf8, false),
585            Field::new("field4", DataType::Utf8, false),
586            Field::new("field5", DataType::Utf8, false),
587        ]));
588
589        let reader_builder = configure_reader_builder(&args, arrow_schema);
590        let builder_debug = format!("{reader_builder:?}");
591        assert_debug_text(&builder_debug, "header", "false");
592        assert_debug_text(&builder_debug, "delimiter", "Some(44)");
593        assert_debug_text(&builder_debug, "quote", "Some(34)");
594        assert_debug_text(&builder_debug, "terminator", "None");
595        assert_debug_text(&builder_debug, "batch_size", "1000");
596        assert_debug_text(&builder_debug, "escape", "None");
597
598        let args = Args {
599            schema: PathBuf::from(Path::new("schema.arvo")),
600            input_file: PathBuf::from(Path::new("test.csv")),
601            output_file: PathBuf::from(Path::new("out.parquet")),
602            batch_size: 2000,
603            input_format: CsvDialect::Tsv,
604            has_header: true,
605            delimiter: None,
606            record_terminator: None,
607            escape_char: Some('\\'),
608            quote_char: None,
609            double_quote: None,
610            csv_compression: Compression::UNCOMPRESSED,
611            parquet_compression: Compression::SNAPPY,
612            writer_version: None,
613            max_row_group_size: None,
614            enable_bloom_filter: None,
615            help: None,
616        };
617        let arrow_schema = Arc::new(Schema::new(vec![
618            Field::new("field1", DataType::Utf8, false),
619            Field::new("field2", DataType::Utf8, false),
620            Field::new("field3", DataType::Utf8, false),
621            Field::new("field4", DataType::Utf8, false),
622            Field::new("field5", DataType::Utf8, false),
623        ]));
624        let reader_builder = configure_reader_builder(&args, arrow_schema);
625        let builder_debug = format!("{reader_builder:?}");
626        assert_debug_text(&builder_debug, "header", "true");
627        assert_debug_text(&builder_debug, "delimiter", "Some(9)");
628        assert_debug_text(&builder_debug, "quote", "None");
629        assert_debug_text(&builder_debug, "terminator", "Some(10)");
630        assert_debug_text(&builder_debug, "batch_size", "2000");
631        assert_debug_text(&builder_debug, "escape", "Some(92)");
632    }
633
634    fn test_convert_compressed_csv_to_parquet(csv_compression: Compression) {
635        let schema = NamedTempFile::new().unwrap();
636        let schema_text = r"message my_amazing_schema {
637            optional int32 id;
638            optional binary name (STRING);
639        }";
640        schema.as_file().write_all(schema_text.as_bytes()).unwrap();
641
642        let mut input_file = NamedTempFile::new().unwrap();
643
644        fn write_tmp_file<T: Write>(w: &mut T) {
645            for index in 1..2000 {
646                write!(w, "{index},\"name_{index}\"\r\n").unwrap();
647            }
648            w.flush().unwrap();
649        }
650
651        // make sure the input_file's lifetime being long enough
652        input_file = match csv_compression {
653            Compression::UNCOMPRESSED => {
654                write_tmp_file(&mut input_file);
655                input_file
656            }
657            Compression::SNAPPY => {
658                let mut encoder = FrameEncoder::new(input_file);
659                write_tmp_file(&mut encoder);
660                encoder.into_inner().unwrap()
661            }
662            Compression::GZIP(level) => {
663                let mut encoder = GzEncoder::new(
664                    input_file,
665                    flate2::Compression::new(level.compression_level()),
666                );
667                write_tmp_file(&mut encoder);
668                encoder.finish().unwrap()
669            }
670            Compression::BROTLI(level) => {
671                let mut encoder =
672                    CompressorWriter::new(input_file, 0, level.compression_level(), 0);
673                write_tmp_file(&mut encoder);
674                encoder.into_inner()
675            }
676            Compression::LZ4 => {
677                let mut encoder = lz4_flex::frame::FrameEncoder::new(input_file);
678                write_tmp_file(&mut encoder);
679                encoder.finish().unwrap()
680            }
681
682            Compression::ZSTD(level) => {
683                let mut encoder = zstd::Encoder::new(input_file, level.compression_level())
684                    .map_err(|e| {
685                        ParquetFromCsvError::with_context(e, "Failed to create zstd::Encoder")
686                    })
687                    .unwrap();
688                write_tmp_file(&mut encoder);
689                encoder.finish().unwrap()
690            }
691            d => unimplemented!("compression type {d}"),
692        };
693
694        let output_parquet = NamedTempFile::new().unwrap();
695
696        let args = Args {
697            schema: PathBuf::from(schema.path()),
698            input_file: PathBuf::from(input_file.path()),
699            output_file: PathBuf::from(output_parquet.path()),
700            batch_size: 1000,
701            input_format: CsvDialect::Csv,
702            has_header: false,
703            delimiter: None,
704            record_terminator: None,
705            escape_char: None,
706            quote_char: None,
707            double_quote: None,
708            csv_compression,
709            parquet_compression: Compression::SNAPPY,
710            writer_version: None,
711            max_row_group_size: None,
712            // by default we shall test bloom filter writing
713            enable_bloom_filter: Some(true),
714            help: None,
715        };
716        convert_csv_to_parquet(&args).unwrap();
717
718        let file = SerializedFileReader::new(output_parquet.into_file()).unwrap();
719        let schema_name = file.metadata().file_metadata().schema().name();
720        assert_eq!(schema_name, "my_amazing_schema");
721    }
722
723    #[test]
724    fn test_convert_csv_to_parquet() {
725        test_convert_compressed_csv_to_parquet(Compression::UNCOMPRESSED);
726        test_convert_compressed_csv_to_parquet(Compression::SNAPPY);
727        test_convert_compressed_csv_to_parquet(Compression::GZIP(GzipLevel::try_new(1).unwrap()));
728        test_convert_compressed_csv_to_parquet(Compression::BROTLI(
729            BrotliLevel::try_new(2).unwrap(),
730        ));
731        test_convert_compressed_csv_to_parquet(Compression::LZ4);
732        test_convert_compressed_csv_to_parquet(Compression::ZSTD(ZstdLevel::try_new(1).unwrap()));
733    }
734}
parquet_fromcsv/parquet-fromcsv.rs

parquet_fromcsv/
parquet-fromcsv.rs
