dataset_writer/
ipc.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 arrow::datatypes::Schema;
12use arrow::ipc::writer::FileWriter;
13
14use super::{StructArrayBuilder, TableWriter};
15
16/// Writer to a .arrow file, usable with [`ParallelDatasetWriter`](super::ParallelDatasetWriter)
17///
18/// `Builder` should follow the pattern documented by
19/// [`arrow::builder`](https://docs.rs/arrow/latest/arrow/array/builder/index.html)
20pub struct ArrowTableWriter<Builder: Default + StructArrayBuilder> {
21    path: PathBuf,
22    file_writer: FileWriter<File>,
23    builder: Builder,
24    pub flush_threshold: usize,
25}
26
27impl<Builder: Default + StructArrayBuilder> TableWriter for ArrowTableWriter<Builder> {
28    type Schema = Schema;
29    type CloseResult = ();
30    type Config = Option<usize>;
31
32    fn new(mut path: PathBuf, schema: Self::Schema, config: Option<usize>) -> Result<Self> {
33        path.set_extension("arrow");
34        let file =
35            File::create(&path).with_context(|| format!("Could not create {}", path.display()))?;
36        let file_writer = FileWriter::try_new(file, &schema).with_context(|| {
37            format!(
38                "Could not create writer for {} with schema {}",
39                path.display(),
40                schema
41            )
42        })?;
43
44        Ok(ArrowTableWriter {
45            path,
46            file_writer,
47            flush_threshold: config.unwrap_or(1024 * 1024), // Arbitrary
48            builder: Builder::default(),
49        })
50    }
51
52    fn flush(&mut self) -> Result<()> {
53        let mut tmp = Builder::default();
54        std::mem::swap(&mut tmp, &mut self.builder);
55        let struct_array = tmp.finish()?;
56        self.file_writer
57            .write(&struct_array.into())
58            .with_context(|| format!("Could not write to {}", self.path.display()))
59    }
60
61    fn close(mut self) -> Result<()> {
62        self.flush()?;
63        self.file_writer
64            .finish()
65            .with_context(|| format!("Could not close {}", self.path.display()))
66    }
67}
68
69impl<Builder: Default + StructArrayBuilder> ArrowTableWriter<Builder> {
70    /// Flushes the internal buffer is too large, then returns the array builder.
71    pub fn builder(&mut self) -> Result<&mut Builder> {
72        if self.builder.len() >= self.flush_threshold {
73            self.flush()?;
74        }
75
76        Ok(&mut self.builder)
77    }
78}
79
80impl<Builder: Default + StructArrayBuilder> Drop for ArrowTableWriter<Builder> {
81    fn drop(&mut self) {
82        self.flush().unwrap();
83        self.file_writer
84            .finish()
85            .with_context(|| format!("Could not close {}", self.path.display()))
86            .unwrap();
87    }
88}
dataset_writer/ipc.rs

dataset_writer/
ipc.rs
