tpcdsgen_arrow/tables/
store_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::{StoreSalesRow, StoreSalesRowGenerator};
10use tpcdsgen::types::{Decimal, Pricing};
11
12/// Generate [`StoreSalesRow`]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 StoreSalesArrow {
20    inner: StoreSalesRowGenerator,
21    batch_size: usize,
22    // reused allocations across batches
23    scratch: Vec<StoreSalesRow>,
24    ss_sold_date_sk: PrimitiveArrayRecycler<Int32Type>,
25    ss_sold_time_sk: PrimitiveArrayRecycler<Int32Type>,
26    ss_item_sk: PrimitiveArrayRecycler<Int32Type>,
27    ss_customer_sk: PrimitiveArrayRecycler<Int32Type>,
28    ss_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
29    ss_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
30    ss_addr_sk: PrimitiveArrayRecycler<Int32Type>,
31    ss_store_sk: PrimitiveArrayRecycler<Int32Type>,
32    ss_promo_sk: PrimitiveArrayRecycler<Int32Type>,
33    ss_ticket_number: PrimitiveArrayRecycler<Int64Type>,
34    ss_quantity: PrimitiveArrayRecycler<Int32Type>,
35    ss_wholesale_cost: PrimitiveArrayRecycler<Decimal128Type>,
36    ss_list_price: PrimitiveArrayRecycler<Decimal128Type>,
37    ss_sales_price: PrimitiveArrayRecycler<Decimal128Type>,
38    ss_ext_discount_amt: PrimitiveArrayRecycler<Decimal128Type>,
39    ss_ext_sales_price: PrimitiveArrayRecycler<Decimal128Type>,
40    ss_ext_wholesale_cost: PrimitiveArrayRecycler<Decimal128Type>,
41    ss_ext_list_price: PrimitiveArrayRecycler<Decimal128Type>,
42    ss_ext_tax: PrimitiveArrayRecycler<Decimal128Type>,
43    ss_coupon_amt: PrimitiveArrayRecycler<Decimal128Type>,
44    ss_net_paid: PrimitiveArrayRecycler<Decimal128Type>,
45    ss_net_paid_inc_tax: PrimitiveArrayRecycler<Decimal128Type>,
46    ss_net_profit: PrimitiveArrayRecycler<Decimal128Type>,
47}
48
49impl StoreSalesArrow {
50    /// Return the schema without initializing a data generator.
51    pub fn schema_ref() -> SchemaRef {
52        Arc::clone(&SCHEMA)
53    }
54
55    pub fn new(session: Session) -> Self {
56        let row_count = session.get_scaling().get_row_count(Table::StoreSales);
57        Self {
58            inner: StoreSalesRowGenerator::new(session, row_count),
59            batch_size: DEFAULT_BATCH_SIZE,
60            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
61            ss_sold_date_sk: PrimitiveArrayRecycler::new(),
62            ss_sold_time_sk: PrimitiveArrayRecycler::new(),
63            ss_item_sk: PrimitiveArrayRecycler::new(),
64            ss_customer_sk: PrimitiveArrayRecycler::new(),
65            ss_cdemo_sk: PrimitiveArrayRecycler::new(),
66            ss_hdemo_sk: PrimitiveArrayRecycler::new(),
67            ss_addr_sk: PrimitiveArrayRecycler::new(),
68            ss_store_sk: PrimitiveArrayRecycler::new(),
69            ss_promo_sk: PrimitiveArrayRecycler::new(),
70            ss_ticket_number: PrimitiveArrayRecycler::new(),
71            ss_quantity: PrimitiveArrayRecycler::new(),
72            ss_wholesale_cost: PrimitiveArrayRecycler::new(),
73            ss_list_price: PrimitiveArrayRecycler::new(),
74            ss_sales_price: PrimitiveArrayRecycler::new(),
75            ss_ext_discount_amt: PrimitiveArrayRecycler::new(),
76            ss_ext_sales_price: PrimitiveArrayRecycler::new(),
77            ss_ext_wholesale_cost: PrimitiveArrayRecycler::new(),
78            ss_ext_list_price: PrimitiveArrayRecycler::new(),
79            ss_ext_tax: PrimitiveArrayRecycler::new(),
80            ss_coupon_amt: PrimitiveArrayRecycler::new(),
81            ss_net_paid: PrimitiveArrayRecycler::new(),
82            ss_net_paid_inc_tax: PrimitiveArrayRecycler::new(),
83            ss_net_profit: PrimitiveArrayRecycler::new(),
84        }
85    }
86    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
87        self.inner
88            .skip_rows_until_starting_row_number(starting_row_number);
89    }
90
91    /// Generate only source rows `starting_row_number..=ending_row_number`
92    /// (1-based, inclusive). The ending row number is clamped to the table's
93    /// row count.
94    pub fn with_source_row_range(
95        mut self,
96        starting_row_number: u64,
97        ending_row_number: u64,
98    ) -> Self {
99        self.inner
100            .set_source_row_range(starting_row_number, ending_row_number);
101        self
102    }
103
104    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
105        self.batch_size = batch_size;
106        self
107    }
108}
109
110impl RecordBatchReader for StoreSalesArrow {
111    fn schema(&self) -> SchemaRef {
112        Self::schema_ref()
113    }
114}
115
116impl Iterator for StoreSalesArrow {
117    type Item = Result<RecordBatch, ArrowError>;
118
119    fn next(&mut self) -> Option<Self::Item> {
120        self.scratch.clear();
121        self.scratch
122            .extend(self.inner.by_ref().take(self.batch_size));
123        let rows = &self.scratch;
124        if rows.is_empty() {
125            return None;
126        }
127
128        let schema = self.schema();
129        let Self {
130            ss_sold_date_sk,
131            ss_sold_time_sk,
132            ss_item_sk,
133            ss_customer_sk,
134            ss_cdemo_sk,
135            ss_hdemo_sk,
136            ss_addr_sk,
137            ss_store_sk,
138            ss_promo_sk,
139            ss_ticket_number,
140            ss_quantity,
141            ss_wholesale_cost,
142            ss_list_price,
143            ss_sales_price,
144            ss_ext_discount_amt,
145            ss_ext_sales_price,
146            ss_ext_wholesale_cost,
147            ss_ext_list_price,
148            ss_ext_tax,
149            ss_coupon_amt,
150            ss_net_paid,
151            ss_net_paid_inc_tax,
152            ss_net_profit,
153            ..
154        } = self;
155        let nbm = StoreSalesRow::null_bit_map;
156        let sk32 =
157            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&StoreSalesRow) -> i64| {
158                rec.int32_sk(rows, pos, nbm, f)
159            };
160        let sk64 =
161            |rec: &mut PrimitiveArrayRecycler<Int64Type>, pos, f: fn(&StoreSalesRow) -> i64| {
162                rec.int64_sk(rows, pos, nbm, f)
163            };
164        let dec =
165            |rec: &mut PrimitiveArrayRecycler<Decimal128Type>, pos, f: fn(&Pricing) -> Decimal| {
166                rec.decimal_7_2(rows, pos, nbm, |r| f(r.get_ss_pricing()))
167            };
168
169        let ss_sold_date_sk = sk32(ss_sold_date_sk, 0, StoreSalesRow::get_ss_sold_date_sk);
170        let ss_sold_time_sk = sk32(ss_sold_time_sk, 1, StoreSalesRow::get_ss_sold_time_sk);
171        let ss_item_sk = sk32(ss_item_sk, 2, StoreSalesRow::get_ss_sold_item_sk);
172        let ss_customer_sk = sk32(ss_customer_sk, 3, StoreSalesRow::get_ss_sold_customer_sk);
173        let ss_cdemo_sk = sk32(ss_cdemo_sk, 4, StoreSalesRow::get_ss_sold_cdemo_sk);
174        let ss_hdemo_sk = sk32(ss_hdemo_sk, 5, StoreSalesRow::get_ss_sold_hdemo_sk);
175        let ss_addr_sk = sk32(ss_addr_sk, 6, StoreSalesRow::get_ss_sold_addr_sk);
176        let ss_store_sk = sk32(ss_store_sk, 7, StoreSalesRow::get_ss_sold_store_sk);
177        let ss_promo_sk = sk32(ss_promo_sk, 8, StoreSalesRow::get_ss_sold_promo_sk);
178        let ss_ticket_number = sk64(ss_ticket_number, 9, StoreSalesRow::get_ss_ticket_number);
179        let ss_quantity = ss_quantity.int32(rows, 10, nbm, |r| r.get_ss_pricing().get_quantity());
180        let ss_wholesale_cost = dec(ss_wholesale_cost, 11, Pricing::get_wholesale_cost);
181        let ss_list_price = dec(ss_list_price, 12, Pricing::get_list_price);
182        let ss_sales_price = dec(ss_sales_price, 13, Pricing::get_sales_price);
183        // Java bug: coupon_amount appears at position 14 instead of ext_discount_amount
184        let ss_ext_discount_amt = dec(ss_ext_discount_amt, 14, Pricing::get_coupon_amount);
185        let ss_ext_sales_price = dec(ss_ext_sales_price, 15, Pricing::get_ext_sales_price);
186        let ss_ext_wholesale_cost = dec(ss_ext_wholesale_cost, 16, Pricing::get_ext_wholesale_cost);
187        let ss_ext_list_price = dec(ss_ext_list_price, 17, Pricing::get_ext_list_price);
188        let ss_ext_tax = dec(ss_ext_tax, 18, Pricing::get_ext_tax);
189        // Java bug: coupon_amount appears again at position 19
190        let ss_coupon_amt = dec(ss_coupon_amt, 14, Pricing::get_coupon_amount);
191        let ss_net_paid = dec(ss_net_paid, 19, Pricing::get_net_paid);
192        let ss_net_paid_inc_tax = dec(ss_net_paid_inc_tax, 20, Pricing::get_net_paid_including_tax);
193        let ss_net_profit = dec(ss_net_profit, 21, Pricing::get_net_profit);
194
195        let batch = RecordBatch::try_new(
196            schema,
197            vec![
198                Arc::new(ss_sold_date_sk),
199                Arc::new(ss_sold_time_sk),
200                Arc::new(ss_item_sk),
201                Arc::new(ss_customer_sk),
202                Arc::new(ss_cdemo_sk),
203                Arc::new(ss_hdemo_sk),
204                Arc::new(ss_addr_sk),
205                Arc::new(ss_store_sk),
206                Arc::new(ss_promo_sk),
207                Arc::new(ss_ticket_number),
208                Arc::new(ss_quantity),
209                Arc::new(ss_wholesale_cost),
210                Arc::new(ss_list_price),
211                Arc::new(ss_sales_price),
212                Arc::new(ss_ext_discount_amt),
213                Arc::new(ss_ext_sales_price),
214                Arc::new(ss_ext_wholesale_cost),
215                Arc::new(ss_ext_list_price),
216                Arc::new(ss_ext_tax),
217                Arc::new(ss_coupon_amt),
218                Arc::new(ss_net_paid),
219                Arc::new(ss_net_paid_inc_tax),
220                Arc::new(ss_net_profit),
221            ],
222        );
223        Some(batch)
224    }
225}
226
227static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
228
229fn make_schema() -> SchemaRef {
230    Arc::new(Schema::new(vec![
231        Field::new("ss_sold_date_sk", DataType::Int32, true),
232        Field::new("ss_sold_time_sk", DataType::Int32, true),
233        Field::new("ss_item_sk", DataType::Int32, false),
234        Field::new("ss_customer_sk", DataType::Int32, true),
235        Field::new("ss_cdemo_sk", DataType::Int32, true),
236        Field::new("ss_hdemo_sk", DataType::Int32, true),
237        Field::new("ss_addr_sk", DataType::Int32, true),
238        Field::new("ss_store_sk", DataType::Int32, true),
239        Field::new("ss_promo_sk", DataType::Int32, true),
240        Field::new("ss_ticket_number", DataType::Int64, false),
241        Field::new("ss_quantity", DataType::Int32, true),
242        Field::new("ss_wholesale_cost", DataType::Decimal128(7, 2), true),
243        Field::new("ss_list_price", DataType::Decimal128(7, 2), true),
244        Field::new("ss_sales_price", DataType::Decimal128(7, 2), true),
245        Field::new("ss_ext_discount_amt", DataType::Decimal128(7, 2), true),
246        Field::new("ss_ext_sales_price", DataType::Decimal128(7, 2), true),
247        Field::new("ss_ext_wholesale_cost", DataType::Decimal128(7, 2), true),
248        Field::new("ss_ext_list_price", DataType::Decimal128(7, 2), true),
249        Field::new("ss_ext_tax", DataType::Decimal128(7, 2), true),
250        Field::new("ss_coupon_amt", DataType::Decimal128(7, 2), true),
251        Field::new("ss_net_paid", DataType::Decimal128(7, 2), true),
252        Field::new("ss_net_paid_inc_tax", DataType::Decimal128(7, 2), true),
253        Field::new("ss_net_profit", DataType::Decimal128(7, 2), true),
254    ]))
255}
tpcdsgen_arrow/tables/store_sales.rs

tpcdsgen_arrow/tables/
store_sales.rs
