tpcdsgen_arrow/tables/
dbgen_version.rs

1use crate::conversions::{is_null, PrimitiveArrayRecycler, StringViewArrayRecycler};
2use crate::DEFAULT_BATCH_SIZE;
3use arrow::array::RecordBatch;
4use arrow::datatypes::{
5    DataType, Date32Type, Field, Schema, SchemaRef, Time32SecondType, TimeUnit,
6};
7use arrow::error::ArrowError;
8use arrow::record_batch::RecordBatchReader;
9use std::sync::{Arc, LazyLock};
10use tpcdsgen::config::{Session, Table};
11use tpcdsgen::row::{DbgenVersionRow, DbgenVersionRowGenerator};
12
13/// Generate [`DbgenVersionRow`]s in [`RecordBatch`] format
14pub struct DbgenVersionArrow {
15    inner: DbgenVersionRowGenerator,
16    batch_size: usize,
17    // reused allocations across batches
18    scratch: Vec<DbgenVersionRow>,
19    dv_version: StringViewArrayRecycler,
20    dv_create_date: PrimitiveArrayRecycler<Date32Type>,
21    dv_create_time: PrimitiveArrayRecycler<Time32SecondType>,
22    dv_cmdline_args: StringViewArrayRecycler,
23}
24
25impl DbgenVersionArrow {
26    /// Return the schema without initializing a data generator.
27    pub fn schema_ref() -> SchemaRef {
28        Arc::clone(&SCHEMA)
29    }
30
31    pub fn new(session: Session) -> Self {
32        let row_count = session.get_scaling().get_row_count(Table::DbgenVersion);
33        Self {
34            inner: DbgenVersionRowGenerator::new(session, row_count),
35            batch_size: DEFAULT_BATCH_SIZE,
36            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
37            dv_version: StringViewArrayRecycler::new(),
38            dv_create_date: PrimitiveArrayRecycler::new(),
39            dv_create_time: PrimitiveArrayRecycler::new(),
40            dv_cmdline_args: StringViewArrayRecycler::new(),
41        }
42    }
43    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
44        self.inner
45            .skip_rows_until_starting_row_number(starting_row_number);
46    }
47
48    /// Generate only source rows `starting_row_number..=ending_row_number`
49    /// (1-based, inclusive). The ending row number is clamped to the table's
50    /// row count.
51    pub fn with_source_row_range(
52        mut self,
53        starting_row_number: u64,
54        ending_row_number: u64,
55    ) -> Self {
56        self.inner
57            .set_source_row_range(starting_row_number, ending_row_number);
58        self
59    }
60
61    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
62        self.batch_size = batch_size;
63        self
64    }
65}
66
67impl RecordBatchReader for DbgenVersionArrow {
68    fn schema(&self) -> SchemaRef {
69        Self::schema_ref()
70    }
71}
72
73impl Iterator for DbgenVersionArrow {
74    type Item = Result<RecordBatch, ArrowError>;
75
76    fn next(&mut self) -> Option<Self::Item> {
77        self.scratch.clear();
78        self.scratch
79            .extend(self.inner.by_ref().take(self.batch_size));
80        let rows = &self.scratch;
81        if rows.is_empty() {
82            return None;
83        }
84
85        let schema = self.schema();
86        let Self {
87            dv_version,
88            dv_create_date,
89            dv_create_time,
90            dv_cmdline_args,
91            ..
92        } = self;
93        let nbm = DbgenVersionRow::null_bit_map;
94
95        let dv_version = dv_version.str(rows, 0, nbm, DbgenVersionRow::get_dv_version);
96        let dv_create_date = dv_create_date.date32(rows, 1, nbm, |r| *r.get_dv_create_date());
97        let dv_create_time = dv_create_time.build(rows, DbgenVersionRow::get_dv_create_time, |r| {
98            !is_null(nbm(r), 2)
99        });
100        let dv_cmdline_args =
101            dv_cmdline_args.str(rows, 3, nbm, DbgenVersionRow::get_dv_cmdline_args);
102
103        let batch = RecordBatch::try_new(
104            schema,
105            vec![
106                Arc::new(dv_version),
107                Arc::new(dv_create_date),
108                Arc::new(dv_create_time),
109                Arc::new(dv_cmdline_args),
110            ],
111        );
112        Some(batch)
113    }
114}
115
116static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
117
118fn make_schema() -> SchemaRef {
119    Arc::new(Schema::new(vec![
120        Field::new("dv_version", DataType::Utf8View, true),
121        Field::new("dv_create_date", DataType::Date32, true),
122        Field::new("dv_create_time", DataType::Time32(TimeUnit::Second), true),
123        Field::new("dv_cmdline_args", DataType::Utf8View, true),
124    ]))
125}
tpcdsgen_arrow/tables/dbgen_version.rs

tpcdsgen_arrow/tables/
dbgen_version.rs
