tpcdsgen_arrow/tables/
call_center.rs

1use crate::conversions::{AddressArrayRecycler, PrimitiveArrayRecycler, StringViewArrayRecycler};
2use crate::DEFAULT_BATCH_SIZE;
3use arrow::array::RecordBatch;
4use arrow::datatypes::{DataType, Date32Type, Decimal128Type, 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::{CallCenterRow, CallCenterRowGenerator};
10
11/// Generate [`CallCenterRow`]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 CallCenterArrow {
19    inner: CallCenterRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<CallCenterRow>,
23    cc_call_center_sk: PrimitiveArrayRecycler<Int32Type>,
24    cc_call_center_id: StringViewArrayRecycler,
25    cc_rec_start_date: PrimitiveArrayRecycler<Date32Type>,
26    cc_rec_end_date: PrimitiveArrayRecycler<Date32Type>,
27    cc_closed_date_sk: PrimitiveArrayRecycler<Int32Type>,
28    cc_open_date_sk: PrimitiveArrayRecycler<Int32Type>,
29    cc_name: StringViewArrayRecycler,
30    cc_class: StringViewArrayRecycler,
31    cc_employees: PrimitiveArrayRecycler<Int32Type>,
32    cc_sq_ft: PrimitiveArrayRecycler<Int32Type>,
33    cc_hours: StringViewArrayRecycler,
34    cc_manager: StringViewArrayRecycler,
35    cc_mkt_id: PrimitiveArrayRecycler<Int32Type>,
36    cc_mkt_class: StringViewArrayRecycler,
37    cc_mkt_desc: StringViewArrayRecycler,
38    cc_market_manager: StringViewArrayRecycler,
39    cc_division: PrimitiveArrayRecycler<Int32Type>,
40    cc_division_name: StringViewArrayRecycler,
41    cc_company: PrimitiveArrayRecycler<Int32Type>,
42    cc_company_name: StringViewArrayRecycler,
43    address: AddressArrayRecycler,
44    cc_tax_percentage: PrimitiveArrayRecycler<Decimal128Type>,
45}
46
47impl CallCenterArrow {
48    /// Return the schema without initializing a data generator.
49    pub fn schema_ref() -> SchemaRef {
50        Arc::clone(&SCHEMA)
51    }
52
53    pub fn new(session: Session) -> Self {
54        let row_count = session.get_scaling().get_row_count(Table::CallCenter);
55        Self {
56            inner: CallCenterRowGenerator::new(session, row_count),
57            batch_size: DEFAULT_BATCH_SIZE,
58            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
59            cc_call_center_sk: PrimitiveArrayRecycler::new(),
60            cc_call_center_id: StringViewArrayRecycler::new(),
61            cc_rec_start_date: PrimitiveArrayRecycler::new(),
62            cc_rec_end_date: PrimitiveArrayRecycler::new(),
63            cc_closed_date_sk: PrimitiveArrayRecycler::new(),
64            cc_open_date_sk: PrimitiveArrayRecycler::new(),
65            cc_name: StringViewArrayRecycler::new(),
66            cc_class: StringViewArrayRecycler::new(),
67            cc_employees: PrimitiveArrayRecycler::new(),
68            cc_sq_ft: PrimitiveArrayRecycler::new(),
69            cc_hours: StringViewArrayRecycler::new(),
70            cc_manager: StringViewArrayRecycler::new(),
71            cc_mkt_id: PrimitiveArrayRecycler::new(),
72            cc_mkt_class: StringViewArrayRecycler::new(),
73            cc_mkt_desc: StringViewArrayRecycler::new(),
74            cc_market_manager: StringViewArrayRecycler::new(),
75            cc_division: PrimitiveArrayRecycler::new(),
76            cc_division_name: StringViewArrayRecycler::new(),
77            cc_company: PrimitiveArrayRecycler::new(),
78            cc_company_name: StringViewArrayRecycler::new(),
79            address: AddressArrayRecycler::new(),
80            cc_tax_percentage: PrimitiveArrayRecycler::new(),
81        }
82    }
83    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
84        self.inner
85            .skip_rows_until_starting_row_number(starting_row_number);
86    }
87
88    /// Generate only source rows `starting_row_number..=ending_row_number`
89    /// (1-based, inclusive). The ending row number is clamped to the table's
90    /// row count.
91    pub fn with_source_row_range(
92        mut self,
93        starting_row_number: u64,
94        ending_row_number: u64,
95    ) -> Self {
96        self.inner
97            .set_source_row_range(starting_row_number, ending_row_number);
98        self
99    }
100
101    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
102        self.batch_size = batch_size;
103        self
104    }
105}
106
107impl RecordBatchReader for CallCenterArrow {
108    fn schema(&self) -> SchemaRef {
109        Self::schema_ref()
110    }
111}
112
113impl Iterator for CallCenterArrow {
114    type Item = Result<RecordBatch, ArrowError>;
115
116    fn next(&mut self) -> Option<Self::Item> {
117        self.scratch.clear();
118        self.scratch
119            .extend(self.inner.by_ref().take(self.batch_size));
120        let rows = &self.scratch;
121        if rows.is_empty() {
122            return None;
123        }
124
125        let schema = self.schema();
126        let Self {
127            cc_call_center_sk,
128            cc_call_center_id,
129            cc_rec_start_date,
130            cc_rec_end_date,
131            cc_closed_date_sk,
132            cc_open_date_sk,
133            cc_name,
134            cc_class,
135            cc_employees,
136            cc_sq_ft,
137            cc_hours,
138            cc_manager,
139            cc_mkt_id,
140            cc_mkt_class,
141            cc_mkt_desc,
142            cc_market_manager,
143            cc_division,
144            cc_division_name,
145            cc_company,
146            cc_company_name,
147            address,
148            cc_tax_percentage,
149            ..
150        } = self;
151        let nbm = CallCenterRow::get_null_bit_map;
152        let sk32 =
153            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&CallCenterRow) -> i64| {
154                rec.int32_sk(rows, pos, nbm, f)
155            };
156        let int = |rec: &mut PrimitiveArrayRecycler<Int32Type>,
157                   pos,
158                   f: fn(&CallCenterRow) -> i32| { rec.int32(rows, pos, nbm, f) };
159        let str = |rec: &mut StringViewArrayRecycler, pos, f: fn(&CallCenterRow) -> &str| {
160            rec.str(rows, pos, nbm, f)
161        };
162        let date =
163            |rec: &mut PrimitiveArrayRecycler<Date32Type>, pos, f: fn(&CallCenterRow) -> i64| {
164                rec.date32_julian(rows, pos, nbm, f)
165            };
166
167        let cc_call_center_sk = sk32(cc_call_center_sk, 0, CallCenterRow::get_cc_call_center_sk);
168        let cc_call_center_id = str(cc_call_center_id, 1, CallCenterRow::get_cc_call_center_id);
169        let cc_rec_start_date = date(
170            cc_rec_start_date,
171            2,
172            CallCenterRow::get_cc_rec_start_date_id,
173        );
174        let cc_rec_end_date = date(cc_rec_end_date, 3, CallCenterRow::get_cc_rec_end_date_id);
175        let cc_closed_date_sk = sk32(cc_closed_date_sk, 4, CallCenterRow::get_cc_closed_date_id);
176        let cc_open_date_sk = sk32(cc_open_date_sk, 5, CallCenterRow::get_cc_open_date_id);
177        let cc_name = str(cc_name, 6, CallCenterRow::get_cc_name);
178        let cc_class = str(cc_class, 7, CallCenterRow::get_cc_class);
179        let cc_employees = int(cc_employees, 8, CallCenterRow::get_cc_employees);
180        let cc_sq_ft = int(cc_sq_ft, 9, CallCenterRow::get_cc_sq_ft);
181        let cc_hours = str(cc_hours, 10, CallCenterRow::get_cc_hours);
182        let cc_manager = str(cc_manager, 11, CallCenterRow::get_cc_manager);
183        let cc_mkt_id = int(cc_mkt_id, 12, CallCenterRow::get_cc_market_id);
184        let cc_mkt_class = str(cc_mkt_class, 13, CallCenterRow::get_cc_market_class);
185        let cc_mkt_desc = str(cc_mkt_desc, 14, CallCenterRow::get_cc_market_desc);
186        let cc_market_manager = str(cc_market_manager, 15, CallCenterRow::get_cc_market_manager);
187        let cc_division = int(cc_division, 16, CallCenterRow::get_cc_division_id);
188        let cc_division_name = str(cc_division_name, 17, CallCenterRow::get_cc_division_name);
189        let cc_company = int(cc_company, 18, CallCenterRow::get_cc_company);
190        let cc_company_name = str(cc_company_name, 19, CallCenterRow::get_cc_company_name);
191        let address = address.build(rows, 20, nbm, CallCenterRow::get_cc_address);
192        let cc_tax_percentage =
193            cc_tax_percentage.decimal_5_2(rows, 30, nbm, |r| *r.get_cc_tax_percentage());
194
195        let batch = RecordBatch::try_new(
196            schema,
197            vec![
198                Arc::new(cc_call_center_sk),
199                Arc::new(cc_call_center_id),
200                Arc::new(cc_rec_start_date),
201                Arc::new(cc_rec_end_date),
202                Arc::new(cc_closed_date_sk),
203                Arc::new(cc_open_date_sk),
204                Arc::new(cc_name),
205                Arc::new(cc_class),
206                Arc::new(cc_employees),
207                Arc::new(cc_sq_ft),
208                Arc::new(cc_hours),
209                Arc::new(cc_manager),
210                Arc::new(cc_mkt_id),
211                Arc::new(cc_mkt_class),
212                Arc::new(cc_mkt_desc),
213                Arc::new(cc_market_manager),
214                Arc::new(cc_division),
215                Arc::new(cc_division_name),
216                Arc::new(cc_company),
217                Arc::new(cc_company_name),
218                Arc::new(address.street_number),
219                Arc::new(address.street_name),
220                Arc::new(address.street_type),
221                Arc::new(address.suite_number),
222                Arc::new(address.city),
223                Arc::new(address.county),
224                Arc::new(address.state),
225                Arc::new(address.zip),
226                Arc::new(address.country),
227                Arc::new(address.gmt_offset),
228                Arc::new(cc_tax_percentage),
229            ],
230        );
231        Some(batch)
232    }
233}
234
235static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
236
237fn make_schema() -> SchemaRef {
238    Arc::new(Schema::new(vec![
239        Field::new("cc_call_center_sk", DataType::Int32, false),
240        Field::new("cc_call_center_id", DataType::Utf8View, false),
241        Field::new("cc_rec_start_date", DataType::Date32, true),
242        Field::new("cc_rec_end_date", DataType::Date32, true),
243        Field::new("cc_closed_date_sk", DataType::Int32, true),
244        Field::new("cc_open_date_sk", DataType::Int32, true),
245        Field::new("cc_name", DataType::Utf8View, true),
246        Field::new("cc_class", DataType::Utf8View, true),
247        Field::new("cc_employees", DataType::Int32, true),
248        Field::new("cc_sq_ft", DataType::Int32, true),
249        Field::new("cc_hours", DataType::Utf8View, true),
250        Field::new("cc_manager", DataType::Utf8View, true),
251        Field::new("cc_mkt_id", DataType::Int32, true),
252        Field::new("cc_mkt_class", DataType::Utf8View, true),
253        Field::new("cc_mkt_desc", DataType::Utf8View, true),
254        Field::new("cc_market_manager", DataType::Utf8View, true),
255        Field::new("cc_division", DataType::Int32, true),
256        Field::new("cc_division_name", DataType::Utf8View, true),
257        Field::new("cc_company", DataType::Int32, true),
258        Field::new("cc_company_name", DataType::Utf8View, true),
259        Field::new("cc_street_number", DataType::Utf8View, true),
260        Field::new("cc_street_name", DataType::Utf8View, true),
261        Field::new("cc_street_type", DataType::Utf8View, true),
262        Field::new("cc_suite_number", DataType::Utf8View, true),
263        Field::new("cc_city", DataType::Utf8View, true),
264        Field::new("cc_county", DataType::Utf8View, true),
265        Field::new("cc_state", DataType::Utf8View, true),
266        Field::new("cc_zip", DataType::Utf8View, true),
267        Field::new("cc_country", DataType::Utf8View, true),
268        Field::new("cc_gmt_offset", DataType::Decimal128(5, 2), true),
269        Field::new("cc_tax_percentage", DataType::Decimal128(5, 2), true),
270    ]))
271}
tpcdsgen_arrow/tables/call_center.rs

tpcdsgen_arrow/tables/
call_center.rs
