tpcdsgen_arrow/tables/
web_sales.rs

1use crate::conversions::PrimitiveArrayRecycler;
2use crate::DEFAULT_BATCH_SIZE;
3use arrow::array::RecordBatch;
4use arrow::datatypes::{DataType, Decimal128Type, Field, Int32Type, Int64Type, Schema, SchemaRef};
5use arrow::error::ArrowError;
6use arrow::record_batch::RecordBatchReader;
7use std::sync::{Arc, LazyLock};
8use tpcdsgen::config::{Session, Table};
9use tpcdsgen::row::{WebSalesRow, WebSalesRowGenerator};
10use tpcdsgen::types::{Decimal, Pricing};
11
12/// Generate [`WebSalesRow`]s in [`RecordBatch`] format
13///
14/// Each column's value and null buffers are recycled from one batch to the
15/// next if possible.
16///
17/// Buffers can be recycled if the consumer drops each [`RecordBatch`] before
18/// generating the next one.
19pub struct WebSalesArrow {
20    inner: WebSalesRowGenerator,
21    batch_size: usize,
22    // reused allocations across batches
23    scratch: Vec<WebSalesRow>,
24    ws_sold_date_sk: PrimitiveArrayRecycler<Int32Type>,
25    ws_sold_time_sk: PrimitiveArrayRecycler<Int32Type>,
26    ws_ship_date_sk: PrimitiveArrayRecycler<Int32Type>,
27    ws_item_sk: PrimitiveArrayRecycler<Int32Type>,
28    ws_bill_customer_sk: PrimitiveArrayRecycler<Int32Type>,
29    ws_bill_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
30    ws_bill_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
31    ws_bill_addr_sk: PrimitiveArrayRecycler<Int32Type>,
32    ws_ship_customer_sk: PrimitiveArrayRecycler<Int32Type>,
33    ws_ship_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
34    ws_ship_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
35    ws_ship_addr_sk: PrimitiveArrayRecycler<Int32Type>,
36    ws_web_page_sk: PrimitiveArrayRecycler<Int32Type>,
37    ws_web_site_sk: PrimitiveArrayRecycler<Int32Type>,
38    ws_ship_mode_sk: PrimitiveArrayRecycler<Int32Type>,
39    ws_warehouse_sk: PrimitiveArrayRecycler<Int32Type>,
40    ws_promo_sk: PrimitiveArrayRecycler<Int32Type>,
41    ws_order_number: PrimitiveArrayRecycler<Int64Type>,
42    ws_quantity: PrimitiveArrayRecycler<Int32Type>,
43    ws_wholesale_cost: PrimitiveArrayRecycler<Decimal128Type>,
44    ws_list_price: PrimitiveArrayRecycler<Decimal128Type>,
45    ws_sales_price: PrimitiveArrayRecycler<Decimal128Type>,
46    ws_ext_discount_amt: PrimitiveArrayRecycler<Decimal128Type>,
47    ws_ext_sales_price: PrimitiveArrayRecycler<Decimal128Type>,
48    ws_ext_wholesale_cost: PrimitiveArrayRecycler<Decimal128Type>,
49    ws_ext_list_price: PrimitiveArrayRecycler<Decimal128Type>,
50    ws_ext_tax: PrimitiveArrayRecycler<Decimal128Type>,
51    ws_coupon_amt: PrimitiveArrayRecycler<Decimal128Type>,
52    ws_ext_ship_cost: PrimitiveArrayRecycler<Decimal128Type>,
53    ws_net_paid: PrimitiveArrayRecycler<Decimal128Type>,
54    ws_net_paid_inc_tax: PrimitiveArrayRecycler<Decimal128Type>,
55    ws_net_paid_inc_ship: PrimitiveArrayRecycler<Decimal128Type>,
56    ws_net_paid_inc_ship_tax: PrimitiveArrayRecycler<Decimal128Type>,
57    ws_net_profit: PrimitiveArrayRecycler<Decimal128Type>,
58}
59
60impl WebSalesArrow {
61    /// Return the schema without initializing a data generator.
62    pub fn schema_ref() -> SchemaRef {
63        Arc::clone(&SCHEMA)
64    }
65
66    pub fn new(session: Session) -> Self {
67        let row_count = session.get_scaling().get_row_count(Table::WebSales);
68        Self {
69            inner: WebSalesRowGenerator::new(session, row_count),
70            batch_size: DEFAULT_BATCH_SIZE,
71            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
72            ws_sold_date_sk: PrimitiveArrayRecycler::new(),
73            ws_sold_time_sk: PrimitiveArrayRecycler::new(),
74            ws_ship_date_sk: PrimitiveArrayRecycler::new(),
75            ws_item_sk: PrimitiveArrayRecycler::new(),
76            ws_bill_customer_sk: PrimitiveArrayRecycler::new(),
77            ws_bill_cdemo_sk: PrimitiveArrayRecycler::new(),
78            ws_bill_hdemo_sk: PrimitiveArrayRecycler::new(),
79            ws_bill_addr_sk: PrimitiveArrayRecycler::new(),
80            ws_ship_customer_sk: PrimitiveArrayRecycler::new(),
81            ws_ship_cdemo_sk: PrimitiveArrayRecycler::new(),
82            ws_ship_hdemo_sk: PrimitiveArrayRecycler::new(),
83            ws_ship_addr_sk: PrimitiveArrayRecycler::new(),
84            ws_web_page_sk: PrimitiveArrayRecycler::new(),
85            ws_web_site_sk: PrimitiveArrayRecycler::new(),
86            ws_ship_mode_sk: PrimitiveArrayRecycler::new(),
87            ws_warehouse_sk: PrimitiveArrayRecycler::new(),
88            ws_promo_sk: PrimitiveArrayRecycler::new(),
89            ws_order_number: PrimitiveArrayRecycler::new(),
90            ws_quantity: PrimitiveArrayRecycler::new(),
91            ws_wholesale_cost: PrimitiveArrayRecycler::new(),
92            ws_list_price: PrimitiveArrayRecycler::new(),
93            ws_sales_price: PrimitiveArrayRecycler::new(),
94            ws_ext_discount_amt: PrimitiveArrayRecycler::new(),
95            ws_ext_sales_price: PrimitiveArrayRecycler::new(),
96            ws_ext_wholesale_cost: PrimitiveArrayRecycler::new(),
97            ws_ext_list_price: PrimitiveArrayRecycler::new(),
98            ws_ext_tax: PrimitiveArrayRecycler::new(),
99            ws_coupon_amt: PrimitiveArrayRecycler::new(),
100            ws_ext_ship_cost: PrimitiveArrayRecycler::new(),
101            ws_net_paid: PrimitiveArrayRecycler::new(),
102            ws_net_paid_inc_tax: PrimitiveArrayRecycler::new(),
103            ws_net_paid_inc_ship: PrimitiveArrayRecycler::new(),
104            ws_net_paid_inc_ship_tax: PrimitiveArrayRecycler::new(),
105            ws_net_profit: PrimitiveArrayRecycler::new(),
106        }
107    }
108    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
109        self.inner
110            .skip_rows_until_starting_row_number(starting_row_number);
111    }
112
113    /// Generate only source rows `starting_row_number..=ending_row_number`
114    /// (1-based, inclusive). The ending row number is clamped to the table's
115    /// row count.
116    pub fn with_source_row_range(
117        mut self,
118        starting_row_number: u64,
119        ending_row_number: u64,
120    ) -> Self {
121        self.inner
122            .set_source_row_range(starting_row_number, ending_row_number);
123        self
124    }
125
126    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
127        self.batch_size = batch_size;
128        self
129    }
130}
131
132impl RecordBatchReader for WebSalesArrow {
133    fn schema(&self) -> SchemaRef {
134        Self::schema_ref()
135    }
136}
137
138impl Iterator for WebSalesArrow {
139    type Item = Result<RecordBatch, ArrowError>;
140
141    fn next(&mut self) -> Option<Self::Item> {
142        self.scratch.clear();
143        self.scratch
144            .extend(self.inner.by_ref().take(self.batch_size));
145        let rows = &self.scratch;
146        if rows.is_empty() {
147            return None;
148        }
149
150        let schema = self.schema();
151        let Self {
152            ws_sold_date_sk,
153            ws_sold_time_sk,
154            ws_ship_date_sk,
155            ws_item_sk,
156            ws_bill_customer_sk,
157            ws_bill_cdemo_sk,
158            ws_bill_hdemo_sk,
159            ws_bill_addr_sk,
160            ws_ship_customer_sk,
161            ws_ship_cdemo_sk,
162            ws_ship_hdemo_sk,
163            ws_ship_addr_sk,
164            ws_web_page_sk,
165            ws_web_site_sk,
166            ws_ship_mode_sk,
167            ws_warehouse_sk,
168            ws_promo_sk,
169            ws_order_number,
170            ws_quantity,
171            ws_wholesale_cost,
172            ws_list_price,
173            ws_sales_price,
174            ws_ext_discount_amt,
175            ws_ext_sales_price,
176            ws_ext_wholesale_cost,
177            ws_ext_list_price,
178            ws_ext_tax,
179            ws_coupon_amt,
180            ws_ext_ship_cost,
181            ws_net_paid,
182            ws_net_paid_inc_tax,
183            ws_net_paid_inc_ship,
184            ws_net_paid_inc_ship_tax,
185            ws_net_profit,
186            ..
187        } = self;
188        let nbm = WebSalesRow::null_bit_map;
189        let sk32 =
190            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&WebSalesRow) -> i64| {
191                rec.int32_sk(rows, pos, nbm, f)
192            };
193        let dec =
194            |rec: &mut PrimitiveArrayRecycler<Decimal128Type>, pos, f: fn(&Pricing) -> Decimal| {
195                rec.decimal_7_2(rows, pos, nbm, |r| f(r.get_ws_pricing()))
196            };
197
198        let ws_sold_date_sk = sk32(ws_sold_date_sk, 0, WebSalesRow::get_ws_sold_date_sk);
199        let ws_sold_time_sk = sk32(ws_sold_time_sk, 1, WebSalesRow::get_ws_sold_time_sk);
200        let ws_ship_date_sk = sk32(ws_ship_date_sk, 2, WebSalesRow::get_ws_ship_date_sk);
201        let ws_item_sk = sk32(ws_item_sk, 3, WebSalesRow::get_ws_item_sk);
202        let ws_bill_customer_sk =
203            sk32(ws_bill_customer_sk, 4, WebSalesRow::get_ws_bill_customer_sk);
204        let ws_bill_cdemo_sk = sk32(ws_bill_cdemo_sk, 5, WebSalesRow::get_ws_bill_cdemo_sk);
205        let ws_bill_hdemo_sk = sk32(ws_bill_hdemo_sk, 6, WebSalesRow::get_ws_bill_hdemo_sk);
206        let ws_bill_addr_sk = sk32(ws_bill_addr_sk, 7, WebSalesRow::get_ws_bill_addr_sk);
207        let ws_ship_customer_sk =
208            sk32(ws_ship_customer_sk, 8, WebSalesRow::get_ws_ship_customer_sk);
209        let ws_ship_cdemo_sk = sk32(ws_ship_cdemo_sk, 9, WebSalesRow::get_ws_ship_cdemo_sk);
210        let ws_ship_hdemo_sk = sk32(ws_ship_hdemo_sk, 10, WebSalesRow::get_ws_ship_hdemo_sk);
211        let ws_ship_addr_sk = sk32(ws_ship_addr_sk, 11, WebSalesRow::get_ws_ship_addr_sk);
212        let ws_web_page_sk = sk32(ws_web_page_sk, 12, WebSalesRow::get_ws_web_page_sk);
213        let ws_web_site_sk = sk32(ws_web_site_sk, 13, WebSalesRow::get_ws_web_site_sk);
214        let ws_ship_mode_sk = sk32(ws_ship_mode_sk, 14, WebSalesRow::get_ws_ship_mode_sk);
215        let ws_warehouse_sk = sk32(ws_warehouse_sk, 15, WebSalesRow::get_ws_warehouse_sk);
216        let ws_promo_sk = sk32(ws_promo_sk, 16, WebSalesRow::get_ws_promo_sk);
217        let ws_order_number =
218            ws_order_number.int64(rows, 17, nbm, WebSalesRow::get_ws_order_number);
219        let ws_quantity = ws_quantity.int32(rows, 18, nbm, |r| r.get_ws_pricing().get_quantity());
220        let ws_wholesale_cost = dec(ws_wholesale_cost, 19, Pricing::get_wholesale_cost);
221        let ws_list_price = dec(ws_list_price, 20, Pricing::get_list_price);
222        let ws_sales_price = dec(ws_sales_price, 21, Pricing::get_sales_price);
223        let ws_ext_discount_amt = dec(ws_ext_discount_amt, 22, Pricing::get_ext_discount_amount);
224        let ws_ext_sales_price = dec(ws_ext_sales_price, 23, Pricing::get_ext_sales_price);
225        let ws_ext_wholesale_cost = dec(ws_ext_wholesale_cost, 24, Pricing::get_ext_wholesale_cost);
226        let ws_ext_list_price = dec(ws_ext_list_price, 25, Pricing::get_ext_list_price);
227        let ws_ext_tax = dec(ws_ext_tax, 26, Pricing::get_ext_tax);
228        let ws_coupon_amt = dec(ws_coupon_amt, 27, Pricing::get_coupon_amount);
229        let ws_ext_ship_cost = dec(ws_ext_ship_cost, 28, Pricing::get_ext_ship_cost);
230        let ws_net_paid = dec(ws_net_paid, 29, Pricing::get_net_paid);
231        let ws_net_paid_inc_tax = dec(ws_net_paid_inc_tax, 30, Pricing::get_net_paid_including_tax);
232        let ws_net_paid_inc_ship = dec(
233            ws_net_paid_inc_ship,
234            31,
235            Pricing::get_net_paid_including_shipping,
236        );
237        let ws_net_paid_inc_ship_tax = dec(
238            ws_net_paid_inc_ship_tax,
239            32,
240            Pricing::get_net_paid_including_shipping_and_tax,
241        );
242        let ws_net_profit = dec(ws_net_profit, 33, Pricing::get_net_profit);
243
244        let batch = RecordBatch::try_new(
245            schema,
246            vec![
247                Arc::new(ws_sold_date_sk),
248                Arc::new(ws_sold_time_sk),
249                Arc::new(ws_ship_date_sk),
250                Arc::new(ws_item_sk),
251                Arc::new(ws_bill_customer_sk),
252                Arc::new(ws_bill_cdemo_sk),
253                Arc::new(ws_bill_hdemo_sk),
254                Arc::new(ws_bill_addr_sk),
255                Arc::new(ws_ship_customer_sk),
256                Arc::new(ws_ship_cdemo_sk),
257                Arc::new(ws_ship_hdemo_sk),
258                Arc::new(ws_ship_addr_sk),
259                Arc::new(ws_web_page_sk),
260                Arc::new(ws_web_site_sk),
261                Arc::new(ws_ship_mode_sk),
262                Arc::new(ws_warehouse_sk),
263                Arc::new(ws_promo_sk),
264                Arc::new(ws_order_number),
265                Arc::new(ws_quantity),
266                Arc::new(ws_wholesale_cost),
267                Arc::new(ws_list_price),
268                Arc::new(ws_sales_price),
269                Arc::new(ws_ext_discount_amt),
270                Arc::new(ws_ext_sales_price),
271                Arc::new(ws_ext_wholesale_cost),
272                Arc::new(ws_ext_list_price),
273                Arc::new(ws_ext_tax),
274                Arc::new(ws_coupon_amt),
275                Arc::new(ws_ext_ship_cost),
276                Arc::new(ws_net_paid),
277                Arc::new(ws_net_paid_inc_tax),
278                Arc::new(ws_net_paid_inc_ship),
279                Arc::new(ws_net_paid_inc_ship_tax),
280                Arc::new(ws_net_profit),
281            ],
282        );
283        Some(batch)
284    }
285}
286
287static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
288
289fn make_schema() -> SchemaRef {
290    Arc::new(Schema::new(vec![
291        Field::new("ws_sold_date_sk", DataType::Int32, true),
292        Field::new("ws_sold_time_sk", DataType::Int32, true),
293        Field::new("ws_ship_date_sk", DataType::Int32, true),
294        Field::new("ws_item_sk", DataType::Int32, false),
295        Field::new("ws_bill_customer_sk", DataType::Int32, true),
296        Field::new("ws_bill_cdemo_sk", DataType::Int32, true),
297        Field::new("ws_bill_hdemo_sk", DataType::Int32, true),
298        Field::new("ws_bill_addr_sk", DataType::Int32, true),
299        Field::new("ws_ship_customer_sk", DataType::Int32, true),
300        Field::new("ws_ship_cdemo_sk", DataType::Int32, true),
301        Field::new("ws_ship_hdemo_sk", DataType::Int32, true),
302        Field::new("ws_ship_addr_sk", DataType::Int32, true),
303        Field::new("ws_web_page_sk", DataType::Int32, true),
304        Field::new("ws_web_site_sk", DataType::Int32, true),
305        Field::new("ws_ship_mode_sk", DataType::Int32, true),
306        Field::new("ws_warehouse_sk", DataType::Int32, true),
307        Field::new("ws_promo_sk", DataType::Int32, true),
308        Field::new("ws_order_number", DataType::Int64, false),
309        Field::new("ws_quantity", DataType::Int32, true),
310        Field::new("ws_wholesale_cost", DataType::Decimal128(7, 2), true),
311        Field::new("ws_list_price", DataType::Decimal128(7, 2), true),
312        Field::new("ws_sales_price", DataType::Decimal128(7, 2), true),
313        Field::new("ws_ext_discount_amt", DataType::Decimal128(7, 2), true),
314        Field::new("ws_ext_sales_price", DataType::Decimal128(7, 2), true),
315        Field::new("ws_ext_wholesale_cost", DataType::Decimal128(7, 2), true),
316        Field::new("ws_ext_list_price", DataType::Decimal128(7, 2), true),
317        Field::new("ws_ext_tax", DataType::Decimal128(7, 2), true),
318        Field::new("ws_coupon_amt", DataType::Decimal128(7, 2), true),
319        Field::new("ws_ext_ship_cost", DataType::Decimal128(7, 2), true),
320        Field::new("ws_net_paid", DataType::Decimal128(7, 2), true),
321        Field::new("ws_net_paid_inc_tax", DataType::Decimal128(7, 2), true),
322        Field::new("ws_net_paid_inc_ship", DataType::Decimal128(7, 2), true),
323        Field::new("ws_net_paid_inc_ship_tax", DataType::Decimal128(7, 2), true),
324        Field::new("ws_net_profit", DataType::Decimal128(7, 2), true),
325    ]))
326}
tpcdsgen_arrow/tables/web_sales.rs

tpcdsgen_arrow/tables/
web_sales.rs
