tpcdsgen_arrow/tables/
store.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::{StoreRow, StoreRowGenerator};
10
11/// Generate [`StoreRow`]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 StoreArrow {
19    inner: StoreRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<StoreRow>,
23    s_store_sk: PrimitiveArrayRecycler<Int32Type>,
24    s_store_id: StringViewArrayRecycler,
25    s_rec_start_date: PrimitiveArrayRecycler<Date32Type>,
26    s_rec_end_date: PrimitiveArrayRecycler<Date32Type>,
27    s_closed_date_sk: PrimitiveArrayRecycler<Int32Type>,
28    s_store_name: StringViewArrayRecycler,
29    s_number_employees: PrimitiveArrayRecycler<Int32Type>,
30    s_floor_space: PrimitiveArrayRecycler<Int32Type>,
31    s_hours: StringViewArrayRecycler,
32    s_manager: StringViewArrayRecycler,
33    s_market_id: PrimitiveArrayRecycler<Int32Type>,
34    s_geography_class: StringViewArrayRecycler,
35    s_market_desc: StringViewArrayRecycler,
36    s_market_manager: StringViewArrayRecycler,
37    s_division_id: PrimitiveArrayRecycler<Int32Type>,
38    s_division_name: StringViewArrayRecycler,
39    s_company_id: PrimitiveArrayRecycler<Int32Type>,
40    s_company_name: StringViewArrayRecycler,
41    address: AddressArrayRecycler,
42    s_tax_precentage: PrimitiveArrayRecycler<Decimal128Type>,
43}
44
45impl StoreArrow {
46    /// Return the schema without initializing a data generator.
47    pub fn schema_ref() -> SchemaRef {
48        Arc::clone(&SCHEMA)
49    }
50
51    pub fn new(session: Session) -> Self {
52        let row_count = session.get_scaling().get_row_count(Table::Store);
53        Self {
54            inner: StoreRowGenerator::new(session, row_count),
55            batch_size: DEFAULT_BATCH_SIZE,
56            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
57            s_store_sk: PrimitiveArrayRecycler::new(),
58            s_store_id: StringViewArrayRecycler::new(),
59            s_rec_start_date: PrimitiveArrayRecycler::new(),
60            s_rec_end_date: PrimitiveArrayRecycler::new(),
61            s_closed_date_sk: PrimitiveArrayRecycler::new(),
62            s_store_name: StringViewArrayRecycler::new(),
63            s_number_employees: PrimitiveArrayRecycler::new(),
64            s_floor_space: PrimitiveArrayRecycler::new(),
65            s_hours: StringViewArrayRecycler::new(),
66            s_manager: StringViewArrayRecycler::new(),
67            s_market_id: PrimitiveArrayRecycler::new(),
68            s_geography_class: StringViewArrayRecycler::new(),
69            s_market_desc: StringViewArrayRecycler::new(),
70            s_market_manager: StringViewArrayRecycler::new(),
71            s_division_id: PrimitiveArrayRecycler::new(),
72            s_division_name: StringViewArrayRecycler::new(),
73            s_company_id: PrimitiveArrayRecycler::new(),
74            s_company_name: StringViewArrayRecycler::new(),
75            address: AddressArrayRecycler::new(),
76            s_tax_precentage: PrimitiveArrayRecycler::new(),
77        }
78    }
79    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
80        self.inner
81            .skip_rows_until_starting_row_number(starting_row_number);
82    }
83
84    /// Generate only source rows `starting_row_number..=ending_row_number`
85    /// (1-based, inclusive). The ending row number is clamped to the table's
86    /// row count.
87    pub fn with_source_row_range(
88        mut self,
89        starting_row_number: u64,
90        ending_row_number: u64,
91    ) -> Self {
92        self.inner
93            .set_source_row_range(starting_row_number, ending_row_number);
94        self
95    }
96
97    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
98        self.batch_size = batch_size;
99        self
100    }
101}
102
103impl RecordBatchReader for StoreArrow {
104    fn schema(&self) -> SchemaRef {
105        Self::schema_ref()
106    }
107}
108
109impl Iterator for StoreArrow {
110    type Item = Result<RecordBatch, ArrowError>;
111
112    fn next(&mut self) -> Option<Self::Item> {
113        self.scratch.clear();
114        self.scratch
115            .extend(self.inner.by_ref().take(self.batch_size));
116        let rows = &self.scratch;
117        if rows.is_empty() {
118            return None;
119        }
120
121        let schema = self.schema();
122        let Self {
123            s_store_sk,
124            s_store_id,
125            s_rec_start_date,
126            s_rec_end_date,
127            s_closed_date_sk,
128            s_store_name,
129            s_number_employees,
130            s_floor_space,
131            s_hours,
132            s_manager,
133            s_market_id,
134            s_geography_class,
135            s_market_desc,
136            s_market_manager,
137            s_division_id,
138            s_division_name,
139            s_company_id,
140            s_company_name,
141            address,
142            s_tax_precentage,
143            ..
144        } = self;
145        let nbm = StoreRow::null_bit_map;
146        let sk32 = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&StoreRow) -> i64| {
147            rec.int32_sk(rows, pos, nbm, f)
148        };
149        let int = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&StoreRow) -> i32| {
150            rec.int32(rows, pos, nbm, f)
151        };
152        let int64 = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&StoreRow) -> i64| {
153            rec.int32(rows, pos, nbm, |r| {
154                i32::try_from(f(r)).expect("TPC-DS INTEGER value exceeds i32 range")
155            })
156        };
157        let str = |rec: &mut StringViewArrayRecycler, pos, f: fn(&StoreRow) -> &str| {
158            rec.str(rows, pos, nbm, f)
159        };
160        let date = |rec: &mut PrimitiveArrayRecycler<Date32Type>, pos, f: fn(&StoreRow) -> i64| {
161            rec.date32_julian(rows, pos, nbm, f)
162        };
163
164        let s_store_sk = sk32(s_store_sk, 0, StoreRow::get_store_sk);
165        let s_store_id = str(s_store_id, 1, StoreRow::get_store_id);
166        let s_rec_start_date = date(s_rec_start_date, 2, StoreRow::get_rec_start_date_id);
167        // only the negative-Julian sentinel marks null here, not the bitmap
168        let s_rec_end_date =
169            s_rec_end_date.date32_julian(rows, 3, |_| 0, StoreRow::get_rec_end_date_id);
170        let s_closed_date_sk = sk32(s_closed_date_sk, 4, StoreRow::get_closed_date_id);
171        let s_store_name = str(s_store_name, 5, StoreRow::get_store_name);
172        let s_number_employees = int(s_number_employees, 6, StoreRow::get_employees);
173        let s_floor_space = int(s_floor_space, 7, StoreRow::get_floor_space);
174        let s_hours = str(s_hours, 8, StoreRow::get_hours);
175        let s_manager = str(s_manager, 9, StoreRow::get_store_manager);
176        let s_market_id = int(s_market_id, 10, StoreRow::get_market_id);
177        // Null bits follow the generator column numbering, where the tax
178        // percentage (bit 11) comes before geography class, even though it
179        // is the last output column.
180        let s_geography_class = str(s_geography_class, 12, StoreRow::get_geography_class);
181        let s_market_desc = str(s_market_desc, 13, StoreRow::get_market_desc);
182        let s_market_manager = str(s_market_manager, 14, StoreRow::get_market_manager);
183        let s_division_id = int64(s_division_id, 15, StoreRow::get_division_id);
184        let s_division_name = str(s_division_name, 16, StoreRow::get_division_name);
185        let s_company_id = int64(s_company_id, 17, StoreRow::get_company_id);
186        let s_company_name = str(s_company_name, 18, StoreRow::get_company_name);
187        let address = address.build(rows, 19, nbm, StoreRow::get_address);
188        let s_tax_precentage =
189            s_tax_precentage.decimal_5_2(rows, 11, nbm, StoreRow::get_d_tax_percentage);
190
191        let batch = RecordBatch::try_new(
192            schema,
193            vec![
194                Arc::new(s_store_sk),
195                Arc::new(s_store_id),
196                Arc::new(s_rec_start_date),
197                Arc::new(s_rec_end_date),
198                Arc::new(s_closed_date_sk),
199                Arc::new(s_store_name),
200                Arc::new(s_number_employees),
201                Arc::new(s_floor_space),
202                Arc::new(s_hours),
203                Arc::new(s_manager),
204                Arc::new(s_market_id),
205                Arc::new(s_geography_class),
206                Arc::new(s_market_desc),
207                Arc::new(s_market_manager),
208                Arc::new(s_division_id),
209                Arc::new(s_division_name),
210                Arc::new(s_company_id),
211                Arc::new(s_company_name),
212                Arc::new(address.street_number),
213                Arc::new(address.street_name),
214                Arc::new(address.street_type),
215                Arc::new(address.suite_number),
216                Arc::new(address.city),
217                Arc::new(address.county),
218                Arc::new(address.state),
219                Arc::new(address.zip),
220                Arc::new(address.country),
221                Arc::new(address.gmt_offset),
222                Arc::new(s_tax_precentage),
223            ],
224        );
225        Some(batch)
226    }
227}
228
229static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
230
231fn make_schema() -> SchemaRef {
232    Arc::new(Schema::new(vec![
233        Field::new("s_store_sk", DataType::Int32, false),
234        Field::new("s_store_id", DataType::Utf8View, false),
235        Field::new("s_rec_start_date", DataType::Date32, true),
236        Field::new("s_rec_end_date", DataType::Date32, true),
237        Field::new("s_closed_date_sk", DataType::Int32, true),
238        Field::new("s_store_name", DataType::Utf8View, true),
239        Field::new("s_number_employees", DataType::Int32, true),
240        Field::new("s_floor_space", DataType::Int32, true),
241        Field::new("s_hours", DataType::Utf8View, true),
242        Field::new("s_manager", DataType::Utf8View, true),
243        Field::new("s_market_id", DataType::Int32, true),
244        Field::new("s_geography_class", DataType::Utf8View, true),
245        Field::new("s_market_desc", DataType::Utf8View, true),
246        Field::new("s_market_manager", DataType::Utf8View, true),
247        Field::new("s_division_id", DataType::Int32, true),
248        Field::new("s_division_name", DataType::Utf8View, true),
249        Field::new("s_company_id", DataType::Int32, true),
250        Field::new("s_company_name", DataType::Utf8View, true),
251        Field::new("s_street_number", DataType::Utf8View, true),
252        Field::new("s_street_name", DataType::Utf8View, true),
253        Field::new("s_street_type", DataType::Utf8View, true),
254        Field::new("s_suite_number", DataType::Utf8View, true),
255        Field::new("s_city", DataType::Utf8View, true),
256        Field::new("s_county", DataType::Utf8View, true),
257        Field::new("s_state", DataType::Utf8View, true),
258        Field::new("s_zip", DataType::Utf8View, true),
259        Field::new("s_country", DataType::Utf8View, true),
260        Field::new("s_gmt_offset", DataType::Decimal128(5, 2), true),
261        Field::new("s_tax_precentage", DataType::Decimal128(5, 2), true),
262    ]))
263}
tpcdsgen_arrow/tables/store.rs

tpcdsgen_arrow/tables/
store.rs
