tpcdsgen_arrow/tables/
date_dim.rs

1use crate::conversions::{
2    bool_to_yn, date_to_date32, PrimitiveArrayRecycler, StringViewArrayRecycler,
3};
4use crate::DEFAULT_BATCH_SIZE;
5use arrow::array::RecordBatch;
6use arrow::datatypes::{DataType, Date32Type, Field, Int32Type, Schema, SchemaRef};
7use arrow::error::ArrowError;
8use arrow::record_batch::RecordBatchReader;
9use std::sync::{Arc, LazyLock};
10use tpcdsgen::config::{Session, Table};
11use tpcdsgen::row::{DateDimRow, DateDimRowGenerator};
12
13/// Generate [`DateDimRow`]s in [`RecordBatch`] format
14///
15/// Each column's value and null buffers are recycled from one batch to the
16/// next if possible.
17///
18/// Buffers can be recycled if the consumer drops each [`RecordBatch`] before
19/// generating the next one.
20pub struct DateDimArrow {
21    inner: DateDimRowGenerator,
22    batch_size: usize,
23    // reused allocations across batches
24    scratch: Vec<DateDimRow>,
25    d_date_sk: PrimitiveArrayRecycler<Int32Type>,
26    d_date_id: StringViewArrayRecycler,
27    d_date: PrimitiveArrayRecycler<Date32Type>,
28    d_month_seq: PrimitiveArrayRecycler<Int32Type>,
29    d_week_seq: PrimitiveArrayRecycler<Int32Type>,
30    d_quarter_seq: PrimitiveArrayRecycler<Int32Type>,
31    d_year: PrimitiveArrayRecycler<Int32Type>,
32    d_dow: PrimitiveArrayRecycler<Int32Type>,
33    d_moy: PrimitiveArrayRecycler<Int32Type>,
34    d_dom: PrimitiveArrayRecycler<Int32Type>,
35    d_qoy: PrimitiveArrayRecycler<Int32Type>,
36    d_fy_year: PrimitiveArrayRecycler<Int32Type>,
37    d_fy_quarter_seq: PrimitiveArrayRecycler<Int32Type>,
38    d_fy_week_seq: PrimitiveArrayRecycler<Int32Type>,
39    d_day_name: StringViewArrayRecycler,
40    d_quarter_name: StringViewArrayRecycler,
41    d_holiday: StringViewArrayRecycler,
42    d_weekend: StringViewArrayRecycler,
43    d_following_holiday: StringViewArrayRecycler,
44    d_first_dom: PrimitiveArrayRecycler<Int32Type>,
45    d_last_dom: PrimitiveArrayRecycler<Int32Type>,
46    d_same_day_ly: PrimitiveArrayRecycler<Int32Type>,
47    d_same_day_lq: PrimitiveArrayRecycler<Int32Type>,
48    d_current_day: StringViewArrayRecycler,
49    d_current_week: StringViewArrayRecycler,
50    d_current_month: StringViewArrayRecycler,
51    d_current_quarter: StringViewArrayRecycler,
52    d_current_year: StringViewArrayRecycler,
53}
54
55impl DateDimArrow {
56    /// Return the schema without initializing a data generator.
57    pub fn schema_ref() -> SchemaRef {
58        Arc::clone(&SCHEMA)
59    }
60
61    pub fn new(session: Session) -> Self {
62        let row_count = session.get_scaling().get_row_count(Table::DateDim);
63        Self {
64            inner: DateDimRowGenerator::new(session, row_count),
65            batch_size: DEFAULT_BATCH_SIZE,
66            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
67            d_date_sk: PrimitiveArrayRecycler::new(),
68            d_date_id: StringViewArrayRecycler::new(),
69            d_date: PrimitiveArrayRecycler::new(),
70            d_month_seq: PrimitiveArrayRecycler::new(),
71            d_week_seq: PrimitiveArrayRecycler::new(),
72            d_quarter_seq: PrimitiveArrayRecycler::new(),
73            d_year: PrimitiveArrayRecycler::new(),
74            d_dow: PrimitiveArrayRecycler::new(),
75            d_moy: PrimitiveArrayRecycler::new(),
76            d_dom: PrimitiveArrayRecycler::new(),
77            d_qoy: PrimitiveArrayRecycler::new(),
78            d_fy_year: PrimitiveArrayRecycler::new(),
79            d_fy_quarter_seq: PrimitiveArrayRecycler::new(),
80            d_fy_week_seq: PrimitiveArrayRecycler::new(),
81            d_day_name: StringViewArrayRecycler::new(),
82            d_quarter_name: StringViewArrayRecycler::new(),
83            d_holiday: StringViewArrayRecycler::new(),
84            d_weekend: StringViewArrayRecycler::new(),
85            d_following_holiday: StringViewArrayRecycler::new(),
86            d_first_dom: PrimitiveArrayRecycler::new(),
87            d_last_dom: PrimitiveArrayRecycler::new(),
88            d_same_day_ly: PrimitiveArrayRecycler::new(),
89            d_same_day_lq: PrimitiveArrayRecycler::new(),
90            d_current_day: StringViewArrayRecycler::new(),
91            d_current_week: StringViewArrayRecycler::new(),
92            d_current_month: StringViewArrayRecycler::new(),
93            d_current_quarter: StringViewArrayRecycler::new(),
94            d_current_year: StringViewArrayRecycler::new(),
95        }
96    }
97    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
98        self.inner
99            .skip_rows_until_starting_row_number(starting_row_number);
100    }
101
102    /// Generate only source rows `starting_row_number..=ending_row_number`
103    /// (1-based, inclusive). The ending row number is clamped to the table's
104    /// row count.
105    pub fn with_source_row_range(
106        mut self,
107        starting_row_number: u64,
108        ending_row_number: u64,
109    ) -> Self {
110        self.inner
111            .set_source_row_range(starting_row_number, ending_row_number);
112        self
113    }
114
115    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
116        self.batch_size = batch_size;
117        self
118    }
119}
120
121impl RecordBatchReader for DateDimArrow {
122    fn schema(&self) -> SchemaRef {
123        Self::schema_ref()
124    }
125}
126
127impl Iterator for DateDimArrow {
128    type Item = Result<RecordBatch, ArrowError>;
129
130    fn next(&mut self) -> Option<Self::Item> {
131        self.scratch.clear();
132        self.scratch
133            .extend(self.inner.by_ref().take(self.batch_size));
134        let rows = &self.scratch;
135        if rows.is_empty() {
136            return None;
137        }
138
139        let schema = self.schema();
140        let Self {
141            d_date_sk,
142            d_date_id,
143            d_date,
144            d_month_seq,
145            d_week_seq,
146            d_quarter_seq,
147            d_year,
148            d_dow,
149            d_moy,
150            d_dom,
151            d_qoy,
152            d_fy_year,
153            d_fy_quarter_seq,
154            d_fy_week_seq,
155            d_day_name,
156            d_quarter_name,
157            d_holiday,
158            d_weekend,
159            d_following_holiday,
160            d_first_dom,
161            d_last_dom,
162            d_same_day_ly,
163            d_same_day_lq,
164            d_current_day,
165            d_current_week,
166            d_current_month,
167            d_current_quarter,
168            d_current_year,
169            ..
170        } = self;
171        let nbm = DateDimRow::null_bit_map;
172        // Only d_date_sk consults the null bitmap; every other column is non-null.
173        let int = |rec: &mut PrimitiveArrayRecycler<Int32Type>, f: fn(&DateDimRow) -> i32| {
174            rec.build(rows, f, |_| true)
175        };
176        let str = |rec: &mut StringViewArrayRecycler, f: fn(&DateDimRow) -> &str| {
177            rec.build(rows, |r| Some(f(r)))
178        };
179
180        let d_date_sk = d_date_sk.int32_sk(rows, 0, nbm, |r| r.d_date_sk);
181        let d_date_id = str(d_date_id, |r| r.d_date_id.as_str());
182        let d_date = d_date.build(rows, |r| date_to_date32(&r.d_date), |_| true);
183        let d_month_seq = int(d_month_seq, |r| r.d_month_seq);
184        let d_week_seq = int(d_week_seq, |r| r.d_week_seq);
185        let d_quarter_seq = int(d_quarter_seq, |r| r.d_quarter_seq);
186        let d_year = int(d_year, |r| r.d_year);
187        let d_dow = int(d_dow, |r| r.d_dow);
188        let d_moy = int(d_moy, |r| r.d_moy);
189        let d_dom = int(d_dom, |r| r.d_dom);
190        let d_qoy = int(d_qoy, |r| r.d_qoy);
191        let d_fy_year = int(d_fy_year, |r| r.d_fy_year);
192        let d_fy_quarter_seq = int(d_fy_quarter_seq, |r| r.d_fy_quarter_seq);
193        let d_fy_week_seq = int(d_fy_week_seq, |r| r.d_fy_week_seq);
194        let d_day_name = str(d_day_name, |r| r.d_day_name.as_str());
195        let d_quarter_name = str(d_quarter_name, |r| r.d_quarter_name.as_str());
196        let d_holiday = str(d_holiday, |r| bool_to_yn(r.d_holiday));
197        let d_weekend = str(d_weekend, |r| bool_to_yn(r.d_weekend));
198        let d_following_holiday = str(d_following_holiday, |r| bool_to_yn(r.d_following_holiday));
199        let d_first_dom = int(d_first_dom, |r| r.d_first_dom);
200        let d_last_dom = int(d_last_dom, |r| r.d_last_dom);
201        let d_same_day_ly = int(d_same_day_ly, |r| r.d_same_day_ly);
202        let d_same_day_lq = int(d_same_day_lq, |r| r.d_same_day_lq);
203        let d_current_day = str(d_current_day, |r| bool_to_yn(r.d_current_day));
204        let d_current_week = str(d_current_week, |r| bool_to_yn(r.d_current_week));
205        let d_current_month = str(d_current_month, |r| bool_to_yn(r.d_current_month));
206        let d_current_quarter = str(d_current_quarter, |r| bool_to_yn(r.d_current_quarter));
207        let d_current_year = str(d_current_year, |r| bool_to_yn(r.d_current_year));
208
209        let batch = RecordBatch::try_new(
210            schema,
211            vec![
212                Arc::new(d_date_sk),
213                Arc::new(d_date_id),
214                Arc::new(d_date),
215                Arc::new(d_month_seq),
216                Arc::new(d_week_seq),
217                Arc::new(d_quarter_seq),
218                Arc::new(d_year),
219                Arc::new(d_dow),
220                Arc::new(d_moy),
221                Arc::new(d_dom),
222                Arc::new(d_qoy),
223                Arc::new(d_fy_year),
224                Arc::new(d_fy_quarter_seq),
225                Arc::new(d_fy_week_seq),
226                Arc::new(d_day_name),
227                Arc::new(d_quarter_name),
228                Arc::new(d_holiday),
229                Arc::new(d_weekend),
230                Arc::new(d_following_holiday),
231                Arc::new(d_first_dom),
232                Arc::new(d_last_dom),
233                Arc::new(d_same_day_ly),
234                Arc::new(d_same_day_lq),
235                Arc::new(d_current_day),
236                Arc::new(d_current_week),
237                Arc::new(d_current_month),
238                Arc::new(d_current_quarter),
239                Arc::new(d_current_year),
240            ],
241        );
242        Some(batch)
243    }
244}
245
246static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
247
248fn make_schema() -> SchemaRef {
249    Arc::new(Schema::new(vec![
250        Field::new("d_date_sk", DataType::Int32, false),
251        Field::new("d_date_id", DataType::Utf8View, false),
252        Field::new("d_date", DataType::Date32, true),
253        Field::new("d_month_seq", DataType::Int32, true),
254        Field::new("d_week_seq", DataType::Int32, true),
255        Field::new("d_quarter_seq", DataType::Int32, true),
256        Field::new("d_year", DataType::Int32, true),
257        Field::new("d_dow", DataType::Int32, true),
258        Field::new("d_moy", DataType::Int32, true),
259        Field::new("d_dom", DataType::Int32, true),
260        Field::new("d_qoy", DataType::Int32, true),
261        Field::new("d_fy_year", DataType::Int32, true),
262        Field::new("d_fy_quarter_seq", DataType::Int32, true),
263        Field::new("d_fy_week_seq", DataType::Int32, true),
264        Field::new("d_day_name", DataType::Utf8View, true),
265        Field::new("d_quarter_name", DataType::Utf8View, true),
266        Field::new("d_holiday", DataType::Utf8View, true),
267        Field::new("d_weekend", DataType::Utf8View, true),
268        Field::new("d_following_holiday", DataType::Utf8View, true),
269        Field::new("d_first_dom", DataType::Int32, true),
270        Field::new("d_last_dom", DataType::Int32, true),
271        Field::new("d_same_day_ly", DataType::Int32, true),
272        Field::new("d_same_day_lq", DataType::Int32, true),
273        Field::new("d_current_day", DataType::Utf8View, true),
274        Field::new("d_current_week", DataType::Utf8View, true),
275        Field::new("d_current_month", DataType::Utf8View, true),
276        Field::new("d_current_quarter", DataType::Utf8View, true),
277        Field::new("d_current_year", DataType::Utf8View, true),
278    ]))
279}
tpcdsgen_arrow/tables/date_dim.rs

tpcdsgen_arrow/tables/
date_dim.rs
