tpcdsgen_arrow/tables/
time_dim.rs

1use crate::conversions::{PrimitiveArrayRecycler, StringViewArrayRecycler};
2use crate::DEFAULT_BATCH_SIZE;
3use arrow::array::RecordBatch;
4use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef};
5use arrow::error::ArrowError;
6use arrow::record_batch::RecordBatchReader;
7use std::sync::{Arc, LazyLock};
8use tpcdsgen::config::{Session, Table};
9use tpcdsgen::row::{TimeDimRow, TimeDimRowGenerator};
10
11/// Generate [`TimeDimRow`]s in [`RecordBatch`] format
12///
13/// Each column's value and null buffers are recycled from one batch to the
14/// next if possible.
15///
16/// Buffers can be recycled if the consumer drops each [`RecordBatch`] before
17/// generating the next one.
18pub struct TimeDimArrow {
19    inner: TimeDimRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<TimeDimRow>,
23    t_time_sk: PrimitiveArrayRecycler<Int32Type>,
24    t_time_id: StringViewArrayRecycler,
25    t_time: PrimitiveArrayRecycler<Int32Type>,
26    t_hour: PrimitiveArrayRecycler<Int32Type>,
27    t_minute: PrimitiveArrayRecycler<Int32Type>,
28    t_second: PrimitiveArrayRecycler<Int32Type>,
29    t_am_pm: StringViewArrayRecycler,
30    t_shift: StringViewArrayRecycler,
31    t_sub_shift: StringViewArrayRecycler,
32    t_meal_time: StringViewArrayRecycler,
33}
34
35impl TimeDimArrow {
36    /// Return the schema without initializing a data generator.
37    pub fn schema_ref() -> SchemaRef {
38        Arc::clone(&SCHEMA)
39    }
40
41    pub fn new(session: Session) -> Self {
42        let row_count = session.get_scaling().get_row_count(Table::TimeDim);
43        Self {
44            inner: TimeDimRowGenerator::new(session, row_count),
45            batch_size: DEFAULT_BATCH_SIZE,
46            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
47            t_time_sk: PrimitiveArrayRecycler::new(),
48            t_time_id: StringViewArrayRecycler::new(),
49            t_time: PrimitiveArrayRecycler::new(),
50            t_hour: PrimitiveArrayRecycler::new(),
51            t_minute: PrimitiveArrayRecycler::new(),
52            t_second: PrimitiveArrayRecycler::new(),
53            t_am_pm: StringViewArrayRecycler::new(),
54            t_shift: StringViewArrayRecycler::new(),
55            t_sub_shift: StringViewArrayRecycler::new(),
56            t_meal_time: StringViewArrayRecycler::new(),
57        }
58    }
59    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
60        self.inner
61            .skip_rows_until_starting_row_number(starting_row_number);
62    }
63
64    /// Generate only source rows `starting_row_number..=ending_row_number`
65    /// (1-based, inclusive). The ending row number is clamped to the table's
66    /// row count.
67    pub fn with_source_row_range(
68        mut self,
69        starting_row_number: u64,
70        ending_row_number: u64,
71    ) -> Self {
72        self.inner
73            .set_source_row_range(starting_row_number, ending_row_number);
74        self
75    }
76
77    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
78        self.batch_size = batch_size;
79        self
80    }
81}
82
83impl RecordBatchReader for TimeDimArrow {
84    fn schema(&self) -> SchemaRef {
85        Self::schema_ref()
86    }
87}
88
89impl Iterator for TimeDimArrow {
90    type Item = Result<RecordBatch, ArrowError>;
91
92    fn next(&mut self) -> Option<Self::Item> {
93        self.scratch.clear();
94        self.scratch
95            .extend(self.inner.by_ref().take(self.batch_size));
96        let rows = &self.scratch;
97        if rows.is_empty() {
98            return None;
99        }
100
101        let schema = self.schema();
102        let Self {
103            t_time_sk,
104            t_time_id,
105            t_time,
106            t_hour,
107            t_minute,
108            t_second,
109            t_am_pm,
110            t_shift,
111            t_sub_shift,
112            t_meal_time,
113            ..
114        } = self;
115        let nbm = TimeDimRow::null_bit_map;
116        let sk32 = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&TimeDimRow) -> i64| {
117            rec.int32_sk(rows, pos, nbm, f)
118        };
119        let int = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&TimeDimRow) -> i32| {
120            rec.int32(rows, pos, nbm, f)
121        };
122        let str = |rec: &mut StringViewArrayRecycler, pos, f: fn(&TimeDimRow) -> &str| {
123            rec.str(rows, pos, nbm, f)
124        };
125
126        let t_time_sk = sk32(t_time_sk, 0, |r| r.t_time_sk);
127        let t_time_id = str(t_time_id, 1, |r| r.t_time_id.as_str());
128        let t_time = int(t_time, 2, |r| r.t_time);
129        let t_hour = int(t_hour, 3, |r| r.t_hour);
130        let t_minute = int(t_minute, 4, |r| r.t_minute);
131        let t_second = int(t_second, 5, |r| r.t_second);
132        let t_am_pm = str(t_am_pm, 6, |r| r.t_am_pm.as_str());
133        let t_shift = str(t_shift, 7, |r| r.t_shift.as_str());
134        let t_sub_shift = str(t_sub_shift, 8, |r| r.t_sub_shift.as_str());
135        // t_meal_time is an empty string (not null) for hours with no meal,
136        // but the pipe-delimited format can't distinguish empty from null,
137        // so we map empty -> None to match the .dat file convention.
138        let t_meal_time = t_meal_time.build(rows, |r| {
139            (!r.t_meal_time.is_empty()).then_some(r.t_meal_time.as_str())
140        });
141
142        let batch = RecordBatch::try_new(
143            schema,
144            vec![
145                Arc::new(t_time_sk),
146                Arc::new(t_time_id),
147                Arc::new(t_time),
148                Arc::new(t_hour),
149                Arc::new(t_minute),
150                Arc::new(t_second),
151                Arc::new(t_am_pm),
152                Arc::new(t_shift),
153                Arc::new(t_sub_shift),
154                Arc::new(t_meal_time),
155            ],
156        );
157        Some(batch)
158    }
159}
160
161static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
162
163fn make_schema() -> SchemaRef {
164    Arc::new(Schema::new(vec![
165        Field::new("t_time_sk", DataType::Int32, false),
166        Field::new("t_time_id", DataType::Utf8View, false),
167        Field::new("t_time", DataType::Int32, true),
168        Field::new("t_hour", DataType::Int32, true),
169        Field::new("t_minute", DataType::Int32, true),
170        Field::new("t_second", DataType::Int32, true),
171        Field::new("t_am_pm", DataType::Utf8View, true),
172        Field::new("t_shift", DataType::Utf8View, true),
173        Field::new("t_sub_shift", DataType::Utf8View, true),
174        Field::new("t_meal_time", DataType::Utf8View, true),
175    ]))
176}
tpcdsgen_arrow/tables/time_dim.rs

tpcdsgen_arrow/tables/
time_dim.rs
