dataset_writer/
lib.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
6#![cfg_attr(feature = "parquet", doc = include_str!("../README.md"))]
7
8use std::cell::{RefCell, RefMut};
9use std::path::PathBuf;
10use std::sync::atomic::{AtomicU64, Ordering};
11
12use anyhow::{Context, Result};
13#[cfg(feature = "arrow")]
14use arrow::array::StructArray;
15use rayon::prelude::*;
16use thread_local::ThreadLocal;
17#[cfg(feature = "arrow")]
18pub use arrow;
19
20#[cfg(feature = "csv")]
21mod csv;
22#[cfg(feature = "csv")]
23pub use csv::*;
24
25#[cfg(feature = "arrow-ipc")]
26mod ipc;
27#[cfg(feature = "arrow-ipc")]
28pub use ipc::*;
29
30#[cfg(feature = "parquet")]
31mod parquet_;
32#[cfg(feature = "parquet")]
33pub use parquet_::*;
34
35mod partitioned;
36pub use partitioned::*;
37
38#[cfg(feature = "zstd")]
39mod zstd;
40#[cfg(feature = "zstd")]
41pub use zstd::*;
42
43#[cfg(feature = "arrow")]
44#[allow(clippy::len_without_is_empty)]
45pub trait StructArrayBuilder {
46    /// Number of rows currently in the buffer (not capacity)
47    fn len(&self) -> usize;
48    /// Number of bytes currently in the buffer (not capacity)
49    fn buffer_size(&self) -> usize;
50    fn finish(&mut self) -> Result<StructArray>;
51}
52
53/// Writes a set of files (called tables here) to a directory.
54pub struct ParallelDatasetWriter<W: TableWriter + Send> {
55    num_files: AtomicU64,
56    schema: W::Schema,
57    path: PathBuf,
58    writers: ThreadLocal<RefCell<W>>,
59    pub config: W::Config,
60}
61
62impl<W: TableWriter<Schema = ()> + Send> ParallelDatasetWriter<W>
63where
64    W::Config: Default,
65{
66    pub fn new(path: PathBuf) -> Result<Self> {
67        std::fs::create_dir_all(&path)
68            .with_context(|| format!("Could not create {}", path.display()))?;
69        Ok(ParallelDatasetWriter {
70            num_files: AtomicU64::new(0),
71            schema: (),
72            path,
73            writers: ThreadLocal::new(),
74            config: W::Config::default(),
75        })
76    }
77}
78
79impl<W: TableWriter + Send> ParallelDatasetWriter<W>
80where
81    W::Config: Default,
82{
83    pub fn with_schema(path: PathBuf, schema: W::Schema) -> Result<Self> {
84        std::fs::create_dir_all(&path)
85            .with_context(|| format!("Could not create {}", path.display()))?;
86        Ok(ParallelDatasetWriter {
87            num_files: AtomicU64::new(0),
88            schema,
89            path,
90            writers: ThreadLocal::new(),
91            config: W::Config::default(),
92        })
93    }
94
95    fn get_new_seq_writer(&self) -> Result<RefCell<W>> {
96        let path = self
97            .path
98            .join(self.num_files.fetch_add(1, Ordering::Relaxed).to_string());
99        Ok(RefCell::new(W::new(
100            path,
101            self.schema.clone(),
102            self.config.clone(),
103        )?))
104    }
105
106    /// Returns a new sequential writer.
107    ///
108    /// # Panics
109    ///
110    /// When called from a thread holding another reference to a sequential writer
111    /// of this dataset.
112    pub fn get_thread_writer(&self) -> Result<RefMut<W>> {
113        self.writers
114            .get_or_try(|| self.get_new_seq_writer())
115            .map(|writer| writer.borrow_mut())
116    }
117
118    /// Flushes all underlying writers
119    pub fn flush(&mut self) -> Result<()> {
120        self.writers
121            .iter_mut()
122            .collect::<Vec<_>>()
123            .into_par_iter()
124            .map(|writer| writer.get_mut().flush())
125            .collect::<Result<Vec<()>>>()
126            .map(|_: Vec<()>| ())
127    }
128
129    /// Closes all underlying writers
130    pub fn close(mut self) -> Result<Vec<W::CloseResult>> {
131        let mut tmp = ThreadLocal::new();
132        std::mem::swap(&mut tmp, &mut self.writers);
133        tmp.into_iter()
134            .collect::<Vec<_>>()
135            .into_par_iter()
136            .map(|writer| writer.into_inner().close())
137            .collect()
138    }
139}
140
141impl<W: TableWriter + Send> Drop for ParallelDatasetWriter<W> {
142    fn drop(&mut self) {
143        let mut tmp = ThreadLocal::new();
144        std::mem::swap(&mut tmp, &mut self.writers);
145        tmp.into_iter()
146            .collect::<Vec<_>>()
147            .into_par_iter()
148            .try_for_each(|writer| writer.into_inner().close().map(|_| ()))
149            .expect("Could not close ParallelDatasetWriter");
150    }
151}
152
153pub trait TableWriter {
154    type Schema: Clone;
155    type CloseResult: Send;
156    type Config: Clone;
157
158    fn new(path: PathBuf, schema: Self::Schema, config: Self::Config) -> Result<Self>
159    where
160        Self: Sized;
161
162    /// Calls `.into()` on the internal builder, and writes its result to disk.
163    fn flush(&mut self) -> Result<()>;
164
165    fn close(self) -> Result<Self::CloseResult>;
166}
dataset_writer/lib.rs

dataset_writer/
lib.rs
