dataset_writer/
csv.rs

1// Copyright (C) 2024  The Software Heritage developers
2// See the AUTHORS file at the top-level directory of this distribution
3// License: GNU General Public License version 3, or any later version
4// See top-level LICENSE file for more information
5
6use std::fs::File;
7use std::path::PathBuf;
8
9use anyhow::{Context, Result};
10
11use crate::TableWriter;
12
13pub type CsvZstTableWriter<'a> = csv::Writer<zstd::stream::AutoFinishEncoder<'a, File>>;
14
15impl TableWriter for CsvZstTableWriter<'_> {
16    type Schema = ();
17    type CloseResult = ();
18    type Config = ();
19
20    fn new(mut path: PathBuf, _schema: Self::Schema, _config: ()) -> Result<Self> {
21        path.set_extension("csv.zst");
22        let file =
23            File::create(&path).with_context(|| format!("Could not create {}", path.display()))?;
24        let compression_level = 3;
25        let zstd_encoder = zstd::stream::write::Encoder::new(file, compression_level)
26            .with_context(|| format!("Could not create ZSTD encoder for {}", path.display()))?
27            .auto_finish();
28        Ok(csv::WriterBuilder::new()
29            .has_headers(true)
30            .terminator(csv::Terminator::CRLF)
31            .from_writer(zstd_encoder))
32    }
33
34    fn flush(&mut self) -> Result<()> {
35        self.flush().context("Could not flush CsvZst writer")
36    }
37
38    fn close(mut self) -> Result<()> {
39        self.flush().context("Could not close CsvZst writer")
40    }
41}
dataset_writer/csv.rs

dataset_writer/
csv.rs
