dataset_writer/
zstd.rs

1// Copyright (C) 2025  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::io::Write;
8use std::path::PathBuf;
9
10use anyhow::{Context, Result};
11
12use crate::TableWriter;
13
14#[derive(Debug, Clone)]
15pub struct PlainZstTableWriterConfig {
16    pub extension: String,
17    pub compression_level: i32,
18}
19
20impl Default for PlainZstTableWriterConfig {
21    fn default() -> Self {
22        PlainZstTableWriterConfig {
23            extension: "zst".to_owned(),
24            compression_level: 3,
25        }
26    }
27}
28
29pub type PlainZstTableWriter<'a> = zstd::stream::AutoFinishEncoder<'a, File>;
30
31impl TableWriter for PlainZstTableWriter<'_> {
32    type Schema = ();
33    type CloseResult = ();
34    type Config = PlainZstTableWriterConfig;
35
36    fn new(mut path: PathBuf, _schema: Self::Schema, config: Self::Config) -> Result<Self> {
37        path.set_extension(&config.extension);
38        let file =
39            File::create(&path).with_context(|| format!("Could not create {}", path.display()))?;
40        let encoder = zstd::stream::write::Encoder::new(file, config.compression_level)
41            .with_context(|| format!("Could not create ZSTD encoder for {}", path.display()))?
42            .auto_finish();
43        Ok(encoder)
44    }
45
46    fn flush(&mut self) -> Result<()> {
47        Write::flush(self).context("Could not flush Zst writer")
48    }
49
50    fn close(mut self) -> Result<()> {
51        Write::flush(&mut self).context("Could not close Zst writer")
52    }
53}
dataset_writer/zstd.rs

dataset_writer/
zstd.rs
