tpchgen-arrow 3.0.0

TPC-H data generator into Apache Arrow format
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use crate::{DEFAULT_BATCH_SIZE, RecordBatchIterator};
use arrow::array::{Int64Array, RecordBatch, StringViewArray};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use std::sync::{Arc, LazyLock};
use tpchgen::generators::{RegionGenerator, RegionGeneratorIterator};

/// Generate  [`Region`]s in [`RecordBatch`] format
///
/// [`Region`]: tpchgen::generators::Region
///
/// # Example
/// ```
/// # use tpchgen::generators::{RegionGenerator};
/// # use tpchgen_arrow::RegionArrow;
///
/// // Create a SF=1.0 generator and wrap it in an Arrow generator
/// let generator = RegionGenerator::new(1.0, 1, 1);
/// let mut arrow_generator = RegionArrow::new(generator)
///   .with_batch_size(10);
/// // Read the first 10 batches
/// let batch = arrow_generator.next().unwrap();
/// // compare the output by pretty printing it
/// let formatted_batches = arrow::util::pretty::pretty_format_batches(&[batch])
///   .unwrap()
///   .to_string();
/// let lines = formatted_batches.lines().collect::<Vec<_>>();
/// assert_eq!(lines, vec![
///   "+-------------+-------------+---------------------------------------------------------------------------------------------------------------------+",
///   "| r_regionkey | r_name      | r_comment                                                                                                           |",
///   "+-------------+-------------+---------------------------------------------------------------------------------------------------------------------+",
///   "| 0           | AFRICA      | lar deposits. blithely final packages cajole. regular waters are final requests. regular accounts are according to  |",
///   "| 1           | AMERICA     | hs use ironic, even requests. s                                                                                     |",
///   "| 2           | ASIA        | ges. thinly even pinto beans ca                                                                                     |",
///   "| 3           | EUROPE      | ly final courts cajole furiously final excuse                                                                       |",
///   "| 4           | MIDDLE EAST | uickly special accounts cajole carefully blithely close requests. carefully final asymptotes haggle furiousl        |",
///   "+-------------+-------------+---------------------------------------------------------------------------------------------------------------------+"
/// ]);
/// ```
pub struct RegionArrow {
    inner: RegionGeneratorIterator<'static>,
    batch_size: usize,
}

impl RegionArrow {
    pub fn new(generator: RegionGenerator<'static>) -> Self {
        Self {
            inner: generator.iter(),
            batch_size: DEFAULT_BATCH_SIZE,
        }
    }

    /// Set the batch size
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }
}

impl RecordBatchIterator for RegionArrow {
    fn schema(&self) -> &SchemaRef {
        &REGION_SCHEMA
    }
}

impl Iterator for RegionArrow {
    type Item = RecordBatch;

    fn next(&mut self) -> Option<Self::Item> {
        // Get next rows to convert
        let rows: Vec<_> = self.inner.by_ref().take(self.batch_size).collect();
        if rows.is_empty() {
            return None;
        }

        let r_regionkey = Int64Array::from_iter_values(rows.iter().map(|r| r.r_regionkey));
        let r_name = StringViewArray::from_iter_values(rows.iter().map(|r| r.r_name));
        let r_comment = StringViewArray::from_iter_values(rows.iter().map(|r| r.r_comment));

        let batch = RecordBatch::try_new(
            Arc::clone(self.schema()),
            vec![Arc::new(r_regionkey), Arc::new(r_name), Arc::new(r_comment)],
        )
        .unwrap();
        Some(batch)
    }
}

/// Schema for the Region
static REGION_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_region_schema);
fn make_region_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("r_regionkey", DataType::Int64, false),
        Field::new("r_name", DataType::Utf8View, false),
        Field::new("r_comment", DataType::Utf8View, false),
    ]))
}