tpcdsgen-arrow 4.0.0

TPC-DS 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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use crate::conversions::{is_null, PrimitiveArrayRecycler, StringViewArrayRecycler};
use crate::DEFAULT_BATCH_SIZE;
use arrow::array::RecordBatch;
use arrow::datatypes::{
    DataType, Date32Type, Field, Schema, SchemaRef, Time32SecondType, TimeUnit,
};
use arrow::error::ArrowError;
use arrow::record_batch::RecordBatchReader;
use std::sync::{Arc, LazyLock};
use tpcdsgen::config::{Session, Table};
use tpcdsgen::row::{DbgenVersionRow, DbgenVersionRowGenerator};

/// Generate [`DbgenVersionRow`]s in [`RecordBatch`] format
pub struct DbgenVersionArrow {
    inner: DbgenVersionRowGenerator,
    batch_size: usize,
    // reused allocations across batches
    scratch: Vec<DbgenVersionRow>,
    dv_version: StringViewArrayRecycler,
    dv_create_date: PrimitiveArrayRecycler<Date32Type>,
    dv_create_time: PrimitiveArrayRecycler<Time32SecondType>,
    dv_cmdline_args: StringViewArrayRecycler,
}

impl DbgenVersionArrow {
    /// Return the schema without initializing a data generator.
    pub fn schema_ref() -> SchemaRef {
        Arc::clone(&SCHEMA)
    }

    pub fn new(session: Session) -> Self {
        let row_count = session.get_scaling().get_row_count(Table::DbgenVersion);
        Self {
            inner: DbgenVersionRowGenerator::new(session, row_count),
            batch_size: DEFAULT_BATCH_SIZE,
            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
            dv_version: StringViewArrayRecycler::new(),
            dv_create_date: PrimitiveArrayRecycler::new(),
            dv_create_time: PrimitiveArrayRecycler::new(),
            dv_cmdline_args: StringViewArrayRecycler::new(),
        }
    }
    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
        self.inner
            .skip_rows_until_starting_row_number(starting_row_number);
    }

    /// Generate only source rows `starting_row_number..=ending_row_number`
    /// (1-based, inclusive). The ending row number is clamped to the table's
    /// row count.
    pub fn with_source_row_range(
        mut self,
        starting_row_number: u64,
        ending_row_number: u64,
    ) -> Self {
        self.inner
            .set_source_row_range(starting_row_number, ending_row_number);
        self
    }

    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }
}

impl RecordBatchReader for DbgenVersionArrow {
    fn schema(&self) -> SchemaRef {
        Self::schema_ref()
    }
}

impl Iterator for DbgenVersionArrow {
    type Item = Result<RecordBatch, ArrowError>;

    fn next(&mut self) -> Option<Self::Item> {
        self.scratch.clear();
        self.scratch
            .extend(self.inner.by_ref().take(self.batch_size));
        let rows = &self.scratch;
        if rows.is_empty() {
            return None;
        }

        let schema = self.schema();
        let Self {
            dv_version,
            dv_create_date,
            dv_create_time,
            dv_cmdline_args,
            ..
        } = self;
        let nbm = DbgenVersionRow::null_bit_map;

        let dv_version = dv_version.str(rows, 0, nbm, DbgenVersionRow::get_dv_version);
        let dv_create_date = dv_create_date.date32(rows, 1, nbm, |r| *r.get_dv_create_date());
        let dv_create_time = dv_create_time.build(rows, DbgenVersionRow::get_dv_create_time, |r| {
            !is_null(nbm(r), 2)
        });
        let dv_cmdline_args =
            dv_cmdline_args.str(rows, 3, nbm, DbgenVersionRow::get_dv_cmdline_args);

        let batch = RecordBatch::try_new(
            schema,
            vec![
                Arc::new(dv_version),
                Arc::new(dv_create_date),
                Arc::new(dv_create_time),
                Arc::new(dv_cmdline_args),
            ],
        );
        Some(batch)
    }
}

static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);

fn make_schema() -> SchemaRef {
    Arc::new(Schema::new(vec![
        Field::new("dv_version", DataType::Utf8View, true),
        Field::new("dv_create_date", DataType::Date32, true),
        Field::new("dv_create_time", DataType::Time32(TimeUnit::Second), true),
        Field::new("dv_cmdline_args", DataType::Utf8View, true),
    ]))
}