dataset_writer/
partitioned.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::collections::HashMap;
7use std::collections::hash_map::Entry;
8use std::ffi::OsString;
9use std::num::NonZeroU16;
10use std::path::PathBuf;
11
12use anyhow::{ensure, Context, Result};
13use rayon::prelude::*;
14
15use crate::TableWriter;
16
17/// Alias of [`U16PartitionedTableWriter`] for backward compatibility
18pub type PartitionedTableWriter<PartitionWriter> = U16PartitionedTableWriter<PartitionWriter>;
19
20/// Wraps `N` [`TableWriter`] in such a way that they each write to `base/0/x.parquet`,
21/// ..., `base/N-1/x.parquet` instead of `base/x.parquet`.
22///
23/// This allows Hive partitioning while writing with multiple threads (`x` is the
24/// thread id in the example above).
25///
26/// If `num_partitions` is `None`, disables partitioning.
27pub struct U16PartitionedTableWriter<PartitionWriter: TableWriter + Send> {
28    partition_writers: Vec<PartitionWriter>,
29}
30
31impl<PartitionWriter: TableWriter + Send> TableWriter
32    for U16PartitionedTableWriter<PartitionWriter>
33{
34    /// `(partition_column, num_partitions, underlying_schema)`
35    type Schema = (String, Option<NonZeroU16>, PartitionWriter::Schema);
36    type CloseResult = Vec<PartitionWriter::CloseResult>;
37    type Config = PartitionWriter::Config;
38
39    fn new(
40        mut path: PathBuf,
41        (partition_column, num_partitions, schema): Self::Schema,
42        config: Self::Config,
43    ) -> Result<Self> {
44        // Remove the last part of the path (the thread id), so we can insert the
45        // partition number between the base path and the thread id.
46        let thread_id = path.file_name().map(|p| p.to_owned());
47        ensure!(
48            path.pop(),
49            "Unexpected root path for partitioned writer: {}",
50            path.display()
51        );
52        let thread_id = thread_id.unwrap();
53        Ok(U16PartitionedTableWriter {
54            partition_writers: (0..num_partitions.map(NonZeroU16::get).unwrap_or(1))
55                .map(|partition_id| {
56                    let partition_path = if num_partitions.is_some() {
57                        path.join(format!("{}={}", partition_column, partition_id))
58                    } else {
59                        // Partitioning disabled
60                        path.to_owned()
61                    };
62                    std::fs::create_dir_all(&partition_path).with_context(|| {
63                        format!("Could not create {}", partition_path.display())
64                    })?;
65                    PartitionWriter::new(
66                        partition_path.join(&thread_id),
67                        schema.clone(),
68                        config.clone(),
69                    )
70                })
71                .collect::<Result<_>>()?,
72        })
73    }
74
75    fn flush(&mut self) -> Result<()> {
76        self.partition_writers
77            .par_iter_mut()
78            .try_for_each(|writer| writer.flush())
79    }
80
81    fn close(self) -> Result<Self::CloseResult> {
82        self.partition_writers
83            .into_par_iter()
84            .map(|writer| writer.close())
85            .collect()
86    }
87}
88
89impl<PartitionWriter: TableWriter + Send> U16PartitionedTableWriter<PartitionWriter> {
90    pub fn partitions(&mut self) -> &mut [PartitionWriter] {
91        &mut self.partition_writers
92    }
93}
94
95/// Wraps a set of [`TableWriter`] in such a way that they each write to a different
96/// `base/<partition_key>/x.parquet` instead of `base/x.parquet`, where `<partition_key>`
97/// is a UTF8 column.
98///
99/// This allows Hive partitioning while writing with multiple threads (`x` is the
100/// thread id in the example above).
101pub struct Utf8PartitionedTableWriter<PartitionWriter: TableWriter + Send> {
102    path: PathBuf,
103    partition_column: String,
104    schema: PartitionWriter::Schema,
105    config: PartitionWriter::Config,
106    thread_id: OsString,
107    partition_writers: HashMap<String, PartitionWriter>,
108}
109
110impl<PartitionWriter: TableWriter + Send> TableWriter
111    for Utf8PartitionedTableWriter<PartitionWriter>
112{
113    /// `(partition_column, num_partitions, underlying_schema)`
114    type Schema = (String, PartitionWriter::Schema);
115    type CloseResult = Vec<PartitionWriter::CloseResult>;
116    type Config = PartitionWriter::Config;
117
118    fn new(
119        mut path: PathBuf,
120        (partition_column, schema): Self::Schema,
121        config: Self::Config,
122    ) -> Result<Self> {
123        // Remove the last part of the path (the thread id), so we can insert the
124        // partition number between the base path and the thread id.
125        let thread_id = path.file_name().map(|p| p.to_owned());
126        ensure!(
127            path.pop(),
128            "Unexpected root path for partitioned writer: {}",
129            path.display()
130        );
131        let thread_id = thread_id.unwrap();
132        Ok(Utf8PartitionedTableWriter {
133            path,
134            partition_column,
135            schema,
136            config,
137            thread_id,
138            partition_writers: HashMap::new(),
139        })
140    }
141
142    fn flush(&mut self) -> Result<()> {
143        self.partition_writers
144            .par_iter_mut()
145            .try_for_each(|(_partitiong_key, writer)| writer.flush())
146    }
147
148    fn close(self) -> Result<Self::CloseResult> {
149        self.partition_writers
150            .into_par_iter()
151            .map(|(_partitiong_key, writer)| writer.close())
152            .collect()
153    }
154}
155
156impl<PartitionWriter: TableWriter + Send> Utf8PartitionedTableWriter<PartitionWriter> {
157    pub fn partition(&mut self, partition_key: String) -> Result<&mut PartitionWriter> {
158        match self.partition_writers.entry(partition_key) {
159            Entry::Occupied(entry) => Ok(entry.into_mut()),
160            Entry::Vacant(entry) => {
161                let partition_path = self
162                    .path
163                    .join(format!("{}={}", self.partition_column, entry.key()));
164                std::fs::create_dir_all(&partition_path)
165                    .with_context(|| format!("Could not create {}", partition_path.display()))?;
166                Ok(entry.insert(PartitionWriter::new(
167                    partition_path.join(&self.thread_id),
168                    self.schema.clone(),
169                    self.config.clone(),
170                )?))
171            }
172        }
173    }
174    pub fn partitions(&mut self) -> &mut HashMap<String, PartitionWriter> {
175        &mut self.partition_writers
176    }
177}
dataset_writer/partitioned.rs

dataset_writer/
partitioned.rs
