Skip to main content

tpcgen_cli/tpch_cli/
csv.rs

1//! Implementations of [`Source`] for generating data in TBL format
2use super::generate::Source;
3use std::io::Write;
4use tpchgen::csv::{
5    CustomerCsv, LineItemCsv, NationCsv, OrderCsv, PartCsv, PartSuppCsv, RegionCsv, SupplierCsv,
6};
7use tpchgen::generators::{
8    CustomerGenerator, LineItemGenerator, NationGenerator, OrderGenerator, PartGenerator,
9    PartSuppGenerator, RegionGenerator, SupplierGenerator,
10};
11
12/// Define a Source that writes the table in CSV format
13macro_rules! define_csv_source {
14    ($SOURCE_NAME:ident, $GENERATOR_TYPE:ty, $FORMATTER:ty) => {
15        pub struct $SOURCE_NAME {
16            inner: $GENERATOR_TYPE,
17            delimiter: char,
18        }
19
20        impl $SOURCE_NAME {
21            pub fn new(inner: $GENERATOR_TYPE, delimiter: char) -> Self {
22                Self { inner, delimiter }
23            }
24        }
25
26        impl Source for $SOURCE_NAME {
27            fn header(&self, buffer: Vec<u8>) -> Vec<u8> {
28                let mut buffer = buffer;
29                writeln!(
30                    &mut buffer,
31                    "{}",
32                    <$FORMATTER>::header_with_delimiter(self.delimiter)
33                )
34                .expect("writing to memory is infallible");
35                buffer
36            }
37
38            fn create(self, mut buffer: Vec<u8>) -> Vec<u8> {
39                for item in self.inner.iter() {
40                    let formatter = <$FORMATTER>::with_delimiter(item, self.delimiter);
41                    writeln!(&mut buffer, "{formatter}").expect("writing to memory is infallible");
42                }
43                buffer
44            }
45        }
46    };
47}
48
49// Define .csv sources for all tables
50define_csv_source!(NationCsvSource, NationGenerator<'static>, NationCsv);
51define_csv_source!(RegionCsvSource, RegionGenerator<'static>, RegionCsv);
52define_csv_source!(PartCsvSource, PartGenerator<'static>, PartCsv);
53define_csv_source!(SupplierCsvSource, SupplierGenerator<'static>, SupplierCsv);
54define_csv_source!(PartSuppCsvSource, PartSuppGenerator<'static>, PartSuppCsv);
55define_csv_source!(CustomerCsvSource, CustomerGenerator<'static>, CustomerCsv);
56define_csv_source!(OrderCsvSource, OrderGenerator<'static>, OrderCsv);
57define_csv_source!(LineItemCsvSource, LineItemGenerator<'static>, LineItemCsv);