tpcdsgen_arrow/tables/
catalog_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::{CatalogSalesRow, CatalogSalesRowGenerator};
10use tpcdsgen::types::{Decimal, Pricing};
11
12/// Generate [`CatalogSalesRow`]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 CatalogSalesArrow {
20    inner: CatalogSalesRowGenerator,
21    batch_size: usize,
22    // reused allocations across batches
23    scratch: Vec<CatalogSalesRow>,
24    cs_sold_date_sk: PrimitiveArrayRecycler<Int32Type>,
25    cs_sold_time_sk: PrimitiveArrayRecycler<Int32Type>,
26    cs_ship_date_sk: PrimitiveArrayRecycler<Int32Type>,
27    cs_bill_customer_sk: PrimitiveArrayRecycler<Int32Type>,
28    cs_bill_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
29    cs_bill_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
30    cs_bill_addr_sk: PrimitiveArrayRecycler<Int32Type>,
31    cs_ship_customer_sk: PrimitiveArrayRecycler<Int32Type>,
32    cs_ship_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
33    cs_ship_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
34    cs_ship_addr_sk: PrimitiveArrayRecycler<Int32Type>,
35    cs_call_center_sk: PrimitiveArrayRecycler<Int32Type>,
36    cs_catalog_page_sk: PrimitiveArrayRecycler<Int32Type>,
37    cs_ship_mode_sk: PrimitiveArrayRecycler<Int32Type>,
38    cs_warehouse_sk: PrimitiveArrayRecycler<Int32Type>,
39    cs_item_sk: PrimitiveArrayRecycler<Int32Type>,
40    cs_promo_sk: PrimitiveArrayRecycler<Int32Type>,
41    cs_order_number: PrimitiveArrayRecycler<Int64Type>,
42    cs_quantity: PrimitiveArrayRecycler<Int32Type>,
43    cs_wholesale_cost: PrimitiveArrayRecycler<Decimal128Type>,
44    cs_list_price: PrimitiveArrayRecycler<Decimal128Type>,
45    cs_sales_price: PrimitiveArrayRecycler<Decimal128Type>,
46    cs_ext_discount_amt: PrimitiveArrayRecycler<Decimal128Type>,
47    cs_ext_sales_price: PrimitiveArrayRecycler<Decimal128Type>,
48    cs_ext_wholesale_cost: PrimitiveArrayRecycler<Decimal128Type>,
49    cs_ext_list_price: PrimitiveArrayRecycler<Decimal128Type>,
50    cs_ext_tax: PrimitiveArrayRecycler<Decimal128Type>,
51    cs_coupon_amt: PrimitiveArrayRecycler<Decimal128Type>,
52    cs_ext_ship_cost: PrimitiveArrayRecycler<Decimal128Type>,
53    cs_net_paid: PrimitiveArrayRecycler<Decimal128Type>,
54    cs_net_paid_inc_tax: PrimitiveArrayRecycler<Decimal128Type>,
55    cs_net_paid_inc_ship: PrimitiveArrayRecycler<Decimal128Type>,
56    cs_net_paid_inc_ship_tax: PrimitiveArrayRecycler<Decimal128Type>,
57    cs_net_profit: PrimitiveArrayRecycler<Decimal128Type>,
58}
59
60impl CatalogSalesArrow {
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::CatalogSales);
68        Self {
69            inner: CatalogSalesRowGenerator::new(session, row_count),
70            batch_size: DEFAULT_BATCH_SIZE,
71            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
72            cs_sold_date_sk: PrimitiveArrayRecycler::new(),
73            cs_sold_time_sk: PrimitiveArrayRecycler::new(),
74            cs_ship_date_sk: PrimitiveArrayRecycler::new(),
75            cs_bill_customer_sk: PrimitiveArrayRecycler::new(),
76            cs_bill_cdemo_sk: PrimitiveArrayRecycler::new(),
77            cs_bill_hdemo_sk: PrimitiveArrayRecycler::new(),
78            cs_bill_addr_sk: PrimitiveArrayRecycler::new(),
79            cs_ship_customer_sk: PrimitiveArrayRecycler::new(),
80            cs_ship_cdemo_sk: PrimitiveArrayRecycler::new(),
81            cs_ship_hdemo_sk: PrimitiveArrayRecycler::new(),
82            cs_ship_addr_sk: PrimitiveArrayRecycler::new(),
83            cs_call_center_sk: PrimitiveArrayRecycler::new(),
84            cs_catalog_page_sk: PrimitiveArrayRecycler::new(),
85            cs_ship_mode_sk: PrimitiveArrayRecycler::new(),
86            cs_warehouse_sk: PrimitiveArrayRecycler::new(),
87            cs_item_sk: PrimitiveArrayRecycler::new(),
88            cs_promo_sk: PrimitiveArrayRecycler::new(),
89            cs_order_number: PrimitiveArrayRecycler::new(),
90            cs_quantity: PrimitiveArrayRecycler::new(),
91            cs_wholesale_cost: PrimitiveArrayRecycler::new(),
92            cs_list_price: PrimitiveArrayRecycler::new(),
93            cs_sales_price: PrimitiveArrayRecycler::new(),
94            cs_ext_discount_amt: PrimitiveArrayRecycler::new(),
95            cs_ext_sales_price: PrimitiveArrayRecycler::new(),
96            cs_ext_wholesale_cost: PrimitiveArrayRecycler::new(),
97            cs_ext_list_price: PrimitiveArrayRecycler::new(),
98            cs_ext_tax: PrimitiveArrayRecycler::new(),
99            cs_coupon_amt: PrimitiveArrayRecycler::new(),
100            cs_ext_ship_cost: PrimitiveArrayRecycler::new(),
101            cs_net_paid: PrimitiveArrayRecycler::new(),
102            cs_net_paid_inc_tax: PrimitiveArrayRecycler::new(),
103            cs_net_paid_inc_ship: PrimitiveArrayRecycler::new(),
104            cs_net_paid_inc_ship_tax: PrimitiveArrayRecycler::new(),
105            cs_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 CatalogSalesArrow {
133    fn schema(&self) -> SchemaRef {
134        Self::schema_ref()
135    }
136}
137
138impl Iterator for CatalogSalesArrow {
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            cs_sold_date_sk,
153            cs_sold_time_sk,
154            cs_ship_date_sk,
155            cs_bill_customer_sk,
156            cs_bill_cdemo_sk,
157            cs_bill_hdemo_sk,
158            cs_bill_addr_sk,
159            cs_ship_customer_sk,
160            cs_ship_cdemo_sk,
161            cs_ship_hdemo_sk,
162            cs_ship_addr_sk,
163            cs_call_center_sk,
164            cs_catalog_page_sk,
165            cs_ship_mode_sk,
166            cs_warehouse_sk,
167            cs_item_sk,
168            cs_promo_sk,
169            cs_order_number,
170            cs_quantity,
171            cs_wholesale_cost,
172            cs_list_price,
173            cs_sales_price,
174            cs_ext_discount_amt,
175            cs_ext_sales_price,
176            cs_ext_wholesale_cost,
177            cs_ext_list_price,
178            cs_ext_tax,
179            cs_coupon_amt,
180            cs_ext_ship_cost,
181            cs_net_paid,
182            cs_net_paid_inc_tax,
183            cs_net_paid_inc_ship,
184            cs_net_paid_inc_ship_tax,
185            cs_net_profit,
186            ..
187        } = self;
188        let nbm = CatalogSalesRow::null_bit_map;
189        let sk32 =
190            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&CatalogSalesRow) -> 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_cs_pricing()))
196            };
197
198        let cs_sold_date_sk = sk32(cs_sold_date_sk, 0, CatalogSalesRow::get_cs_sold_date_sk);
199        let cs_sold_time_sk = sk32(cs_sold_time_sk, 1, CatalogSalesRow::get_cs_sold_time_sk);
200        let cs_ship_date_sk = sk32(cs_ship_date_sk, 2, CatalogSalesRow::get_cs_ship_date_sk);
201        let cs_bill_customer_sk = sk32(
202            cs_bill_customer_sk,
203            3,
204            CatalogSalesRow::get_cs_bill_customer_sk,
205        );
206        let cs_bill_cdemo_sk = sk32(cs_bill_cdemo_sk, 4, CatalogSalesRow::get_cs_bill_cdemo_sk);
207        let cs_bill_hdemo_sk = sk32(cs_bill_hdemo_sk, 5, CatalogSalesRow::get_cs_bill_hdemo_sk);
208        let cs_bill_addr_sk = sk32(cs_bill_addr_sk, 6, CatalogSalesRow::get_cs_bill_addr_sk);
209        let cs_ship_customer_sk = sk32(
210            cs_ship_customer_sk,
211            7,
212            CatalogSalesRow::get_cs_ship_customer_sk,
213        );
214        let cs_ship_cdemo_sk = sk32(cs_ship_cdemo_sk, 8, CatalogSalesRow::get_cs_ship_cdemo_sk);
215        let cs_ship_hdemo_sk = sk32(cs_ship_hdemo_sk, 9, CatalogSalesRow::get_cs_ship_hdemo_sk);
216        let cs_ship_addr_sk = sk32(cs_ship_addr_sk, 10, CatalogSalesRow::get_cs_ship_addr_sk);
217        let cs_call_center_sk = sk32(
218            cs_call_center_sk,
219            11,
220            CatalogSalesRow::get_cs_call_center_sk,
221        );
222        let cs_catalog_page_sk = sk32(
223            cs_catalog_page_sk,
224            12,
225            CatalogSalesRow::get_cs_catalog_page_sk,
226        );
227        let cs_ship_mode_sk = sk32(cs_ship_mode_sk, 13, CatalogSalesRow::get_cs_ship_mode_sk);
228        let cs_warehouse_sk = sk32(cs_warehouse_sk, 14, CatalogSalesRow::get_cs_warehouse_sk);
229        let cs_item_sk = sk32(cs_item_sk, 15, CatalogSalesRow::get_cs_sold_item_sk);
230        let cs_promo_sk = sk32(cs_promo_sk, 16, CatalogSalesRow::get_cs_promo_sk);
231        let cs_order_number =
232            cs_order_number.int64(rows, 17, nbm, CatalogSalesRow::get_cs_order_number);
233        let cs_quantity = cs_quantity.int32(rows, 18, nbm, |r| r.get_cs_pricing().get_quantity());
234        let cs_wholesale_cost = dec(cs_wholesale_cost, 19, Pricing::get_wholesale_cost);
235        let cs_list_price = dec(cs_list_price, 20, Pricing::get_list_price);
236        let cs_sales_price = dec(cs_sales_price, 21, Pricing::get_sales_price);
237        let cs_ext_discount_amt = dec(cs_ext_discount_amt, 24, Pricing::get_ext_discount_amount);
238        let cs_ext_sales_price = dec(cs_ext_sales_price, 23, Pricing::get_ext_sales_price);
239        let cs_ext_wholesale_cost = dec(cs_ext_wholesale_cost, 25, Pricing::get_ext_wholesale_cost);
240        let cs_ext_list_price = dec(cs_ext_list_price, 26, Pricing::get_ext_list_price);
241        let cs_ext_tax = dec(cs_ext_tax, 27, Pricing::get_ext_tax);
242        let cs_coupon_amt = dec(cs_coupon_amt, 22, Pricing::get_coupon_amount);
243        let cs_ext_ship_cost = dec(cs_ext_ship_cost, 28, Pricing::get_ext_ship_cost);
244        let cs_net_paid = dec(cs_net_paid, 29, Pricing::get_net_paid);
245        let cs_net_paid_inc_tax = dec(cs_net_paid_inc_tax, 30, Pricing::get_net_paid_including_tax);
246        let cs_net_paid_inc_ship = dec(
247            cs_net_paid_inc_ship,
248            31,
249            Pricing::get_net_paid_including_shipping,
250        );
251        let cs_net_paid_inc_ship_tax = dec(
252            cs_net_paid_inc_ship_tax,
253            32,
254            Pricing::get_net_paid_including_shipping_and_tax,
255        );
256        let cs_net_profit = dec(cs_net_profit, 33, Pricing::get_net_profit);
257
258        let batch = RecordBatch::try_new(
259            schema,
260            vec![
261                Arc::new(cs_sold_date_sk),
262                Arc::new(cs_sold_time_sk),
263                Arc::new(cs_ship_date_sk),
264                Arc::new(cs_bill_customer_sk),
265                Arc::new(cs_bill_cdemo_sk),
266                Arc::new(cs_bill_hdemo_sk),
267                Arc::new(cs_bill_addr_sk),
268                Arc::new(cs_ship_customer_sk),
269                Arc::new(cs_ship_cdemo_sk),
270                Arc::new(cs_ship_hdemo_sk),
271                Arc::new(cs_ship_addr_sk),
272                Arc::new(cs_call_center_sk),
273                Arc::new(cs_catalog_page_sk),
274                Arc::new(cs_ship_mode_sk),
275                Arc::new(cs_warehouse_sk),
276                Arc::new(cs_item_sk),
277                Arc::new(cs_promo_sk),
278                Arc::new(cs_order_number),
279                Arc::new(cs_quantity),
280                Arc::new(cs_wholesale_cost),
281                Arc::new(cs_list_price),
282                Arc::new(cs_sales_price),
283                Arc::new(cs_ext_discount_amt),
284                Arc::new(cs_ext_sales_price),
285                Arc::new(cs_ext_wholesale_cost),
286                Arc::new(cs_ext_list_price),
287                Arc::new(cs_ext_tax),
288                Arc::new(cs_coupon_amt),
289                Arc::new(cs_ext_ship_cost),
290                Arc::new(cs_net_paid),
291                Arc::new(cs_net_paid_inc_tax),
292                Arc::new(cs_net_paid_inc_ship),
293                Arc::new(cs_net_paid_inc_ship_tax),
294                Arc::new(cs_net_profit),
295            ],
296        );
297        Some(batch)
298    }
299}
300
301static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
302
303fn make_schema() -> SchemaRef {
304    Arc::new(Schema::new(vec![
305        Field::new("cs_sold_date_sk", DataType::Int32, true),
306        Field::new("cs_sold_time_sk", DataType::Int32, true),
307        Field::new("cs_ship_date_sk", DataType::Int32, true),
308        Field::new("cs_bill_customer_sk", DataType::Int32, true),
309        Field::new("cs_bill_cdemo_sk", DataType::Int32, true),
310        Field::new("cs_bill_hdemo_sk", DataType::Int32, true),
311        Field::new("cs_bill_addr_sk", DataType::Int32, true),
312        Field::new("cs_ship_customer_sk", DataType::Int32, true),
313        Field::new("cs_ship_cdemo_sk", DataType::Int32, true),
314        Field::new("cs_ship_hdemo_sk", DataType::Int32, true),
315        Field::new("cs_ship_addr_sk", DataType::Int32, true),
316        Field::new("cs_call_center_sk", DataType::Int32, true),
317        Field::new("cs_catalog_page_sk", DataType::Int32, true),
318        Field::new("cs_ship_mode_sk", DataType::Int32, true),
319        Field::new("cs_warehouse_sk", DataType::Int32, true),
320        Field::new("cs_item_sk", DataType::Int32, false),
321        Field::new("cs_promo_sk", DataType::Int32, true),
322        Field::new("cs_order_number", DataType::Int64, false),
323        Field::new("cs_quantity", DataType::Int32, true),
324        Field::new("cs_wholesale_cost", DataType::Decimal128(7, 2), true),
325        Field::new("cs_list_price", DataType::Decimal128(7, 2), true),
326        Field::new("cs_sales_price", DataType::Decimal128(7, 2), true),
327        Field::new("cs_ext_discount_amt", DataType::Decimal128(7, 2), true),
328        Field::new("cs_ext_sales_price", DataType::Decimal128(7, 2), true),
329        Field::new("cs_ext_wholesale_cost", DataType::Decimal128(7, 2), true),
330        Field::new("cs_ext_list_price", DataType::Decimal128(7, 2), true),
331        Field::new("cs_ext_tax", DataType::Decimal128(7, 2), true),
332        Field::new("cs_coupon_amt", DataType::Decimal128(7, 2), true),
333        Field::new("cs_ext_ship_cost", DataType::Decimal128(7, 2), true),
334        Field::new("cs_net_paid", DataType::Decimal128(7, 2), true),
335        Field::new("cs_net_paid_inc_tax", DataType::Decimal128(7, 2), true),
336        Field::new("cs_net_paid_inc_ship", DataType::Decimal128(7, 2), true),
337        Field::new("cs_net_paid_inc_ship_tax", DataType::Decimal128(7, 2), true),
338        Field::new("cs_net_profit", DataType::Decimal128(7, 2), true),
339    ]))
340}
tpcdsgen_arrow/tables/catalog_sales.rs

tpcdsgen_arrow/tables/
catalog_sales.rs
