tpcdsgen_arrow/tables/
web_returns.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::{WebReturnsRow, WebReturnsRowGenerator};
10use tpcdsgen::types::{Decimal, Pricing};
11
12/// Generate [`WebReturnsRow`]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 WebReturnsArrow {
20    inner: WebReturnsRowGenerator,
21    batch_size: usize,
22    // reused allocations across batches
23    scratch: Vec<WebReturnsRow>,
24    wr_returned_date_sk: PrimitiveArrayRecycler<Int32Type>,
25    wr_returned_time_sk: PrimitiveArrayRecycler<Int32Type>,
26    wr_item_sk: PrimitiveArrayRecycler<Int32Type>,
27    wr_refunded_customer_sk: PrimitiveArrayRecycler<Int32Type>,
28    wr_refunded_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
29    wr_refunded_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
30    wr_refunded_addr_sk: PrimitiveArrayRecycler<Int32Type>,
31    wr_returning_customer_sk: PrimitiveArrayRecycler<Int32Type>,
32    wr_returning_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
33    wr_returning_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
34    wr_returning_addr_sk: PrimitiveArrayRecycler<Int32Type>,
35    wr_web_page_sk: PrimitiveArrayRecycler<Int32Type>,
36    wr_reason_sk: PrimitiveArrayRecycler<Int32Type>,
37    wr_order_number: PrimitiveArrayRecycler<Int64Type>,
38    wr_return_quantity: PrimitiveArrayRecycler<Int32Type>,
39    wr_return_amt: PrimitiveArrayRecycler<Decimal128Type>,
40    wr_return_tax: PrimitiveArrayRecycler<Decimal128Type>,
41    wr_return_amt_inc_tax: PrimitiveArrayRecycler<Decimal128Type>,
42    wr_fee: PrimitiveArrayRecycler<Decimal128Type>,
43    wr_return_ship_cost: PrimitiveArrayRecycler<Decimal128Type>,
44    wr_refunded_cash: PrimitiveArrayRecycler<Decimal128Type>,
45    wr_reversed_charge: PrimitiveArrayRecycler<Decimal128Type>,
46    wr_account_credit: PrimitiveArrayRecycler<Decimal128Type>,
47    wr_net_loss: PrimitiveArrayRecycler<Decimal128Type>,
48}
49
50impl WebReturnsArrow {
51    /// Return the schema without initializing a data generator.
52    pub fn schema_ref() -> SchemaRef {
53        Arc::clone(&SCHEMA)
54    }
55
56    pub fn new(session: Session) -> Self {
57        let row_count = session.get_scaling().get_row_count(Table::WebReturns);
58        Self {
59            inner: WebReturnsRowGenerator::new(session, row_count),
60            batch_size: DEFAULT_BATCH_SIZE,
61            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
62            wr_returned_date_sk: PrimitiveArrayRecycler::new(),
63            wr_returned_time_sk: PrimitiveArrayRecycler::new(),
64            wr_item_sk: PrimitiveArrayRecycler::new(),
65            wr_refunded_customer_sk: PrimitiveArrayRecycler::new(),
66            wr_refunded_cdemo_sk: PrimitiveArrayRecycler::new(),
67            wr_refunded_hdemo_sk: PrimitiveArrayRecycler::new(),
68            wr_refunded_addr_sk: PrimitiveArrayRecycler::new(),
69            wr_returning_customer_sk: PrimitiveArrayRecycler::new(),
70            wr_returning_cdemo_sk: PrimitiveArrayRecycler::new(),
71            wr_returning_hdemo_sk: PrimitiveArrayRecycler::new(),
72            wr_returning_addr_sk: PrimitiveArrayRecycler::new(),
73            wr_web_page_sk: PrimitiveArrayRecycler::new(),
74            wr_reason_sk: PrimitiveArrayRecycler::new(),
75            wr_order_number: PrimitiveArrayRecycler::new(),
76            wr_return_quantity: PrimitiveArrayRecycler::new(),
77            wr_return_amt: PrimitiveArrayRecycler::new(),
78            wr_return_tax: PrimitiveArrayRecycler::new(),
79            wr_return_amt_inc_tax: PrimitiveArrayRecycler::new(),
80            wr_fee: PrimitiveArrayRecycler::new(),
81            wr_return_ship_cost: PrimitiveArrayRecycler::new(),
82            wr_refunded_cash: PrimitiveArrayRecycler::new(),
83            wr_reversed_charge: PrimitiveArrayRecycler::new(),
84            wr_account_credit: PrimitiveArrayRecycler::new(),
85            wr_net_loss: PrimitiveArrayRecycler::new(),
86        }
87    }
88    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
89        self.inner
90            .skip_rows_until_starting_row_number(starting_row_number);
91    }
92
93    /// Generate only source rows `starting_row_number..=ending_row_number`
94    /// (1-based, inclusive). The ending row number is clamped to the table's
95    /// row count.
96    pub fn with_source_row_range(
97        mut self,
98        starting_row_number: u64,
99        ending_row_number: u64,
100    ) -> Self {
101        self.inner
102            .set_source_row_range(starting_row_number, ending_row_number);
103        self
104    }
105
106    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
107        self.batch_size = batch_size;
108        self
109    }
110}
111
112impl RecordBatchReader for WebReturnsArrow {
113    fn schema(&self) -> SchemaRef {
114        Self::schema_ref()
115    }
116}
117
118impl Iterator for WebReturnsArrow {
119    type Item = Result<RecordBatch, ArrowError>;
120
121    fn next(&mut self) -> Option<Self::Item> {
122        self.scratch.clear();
123        self.scratch
124            .extend(self.inner.by_ref().take(self.batch_size));
125        let rows = &self.scratch;
126        if rows.is_empty() {
127            return None;
128        }
129
130        let schema = self.schema();
131        let Self {
132            wr_returned_date_sk,
133            wr_returned_time_sk,
134            wr_item_sk,
135            wr_refunded_customer_sk,
136            wr_refunded_cdemo_sk,
137            wr_refunded_hdemo_sk,
138            wr_refunded_addr_sk,
139            wr_returning_customer_sk,
140            wr_returning_cdemo_sk,
141            wr_returning_hdemo_sk,
142            wr_returning_addr_sk,
143            wr_web_page_sk,
144            wr_reason_sk,
145            wr_order_number,
146            wr_return_quantity,
147            wr_return_amt,
148            wr_return_tax,
149            wr_return_amt_inc_tax,
150            wr_fee,
151            wr_return_ship_cost,
152            wr_refunded_cash,
153            wr_reversed_charge,
154            wr_account_credit,
155            wr_net_loss,
156            ..
157        } = self;
158        let nbm = WebReturnsRow::null_bit_map;
159        let sk32 =
160            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&WebReturnsRow) -> i64| {
161                rec.int32_sk(rows, pos, nbm, f)
162            };
163        let dec =
164            |rec: &mut PrimitiveArrayRecycler<Decimal128Type>, pos, f: fn(&Pricing) -> Decimal| {
165                rec.decimal_7_2(rows, pos, nbm, |r| f(r.get_wr_pricing()))
166            };
167
168        let wr_returned_date_sk = sk32(
169            wr_returned_date_sk,
170            0,
171            WebReturnsRow::get_wr_returned_date_sk,
172        );
173        let wr_returned_time_sk = sk32(
174            wr_returned_time_sk,
175            1,
176            WebReturnsRow::get_wr_returned_time_sk,
177        );
178        let wr_item_sk = sk32(wr_item_sk, 2, WebReturnsRow::get_wr_item_sk);
179        let wr_refunded_customer_sk = sk32(
180            wr_refunded_customer_sk,
181            3,
182            WebReturnsRow::get_wr_refunded_customer_sk,
183        );
184        let wr_refunded_cdemo_sk = sk32(
185            wr_refunded_cdemo_sk,
186            4,
187            WebReturnsRow::get_wr_refunded_cdemo_sk,
188        );
189        let wr_refunded_hdemo_sk = sk32(
190            wr_refunded_hdemo_sk,
191            5,
192            WebReturnsRow::get_wr_refunded_hdemo_sk,
193        );
194        let wr_refunded_addr_sk = sk32(
195            wr_refunded_addr_sk,
196            6,
197            WebReturnsRow::get_wr_refunded_addr_sk,
198        );
199        let wr_returning_customer_sk = sk32(
200            wr_returning_customer_sk,
201            7,
202            WebReturnsRow::get_wr_returning_customer_sk,
203        );
204        let wr_returning_cdemo_sk = sk32(
205            wr_returning_cdemo_sk,
206            8,
207            WebReturnsRow::get_wr_returning_cdemo_sk,
208        );
209        let wr_returning_hdemo_sk = sk32(
210            wr_returning_hdemo_sk,
211            9,
212            WebReturnsRow::get_wr_returning_hdemo_sk,
213        );
214        let wr_returning_addr_sk = sk32(
215            wr_returning_addr_sk,
216            10,
217            WebReturnsRow::get_wr_returning_addr_sk,
218        );
219        let wr_web_page_sk = sk32(wr_web_page_sk, 11, WebReturnsRow::get_wr_web_page_sk);
220        let wr_reason_sk = sk32(wr_reason_sk, 12, WebReturnsRow::get_wr_reason_sk);
221        let wr_order_number =
222            wr_order_number.int64(rows, 13, nbm, WebReturnsRow::get_wr_order_number);
223        let wr_return_quantity =
224            wr_return_quantity.int32(rows, 14, nbm, |r| r.get_wr_pricing().get_quantity());
225        let wr_return_amt = dec(wr_return_amt, 15, Pricing::get_net_paid);
226        let wr_return_tax = dec(wr_return_tax, 16, Pricing::get_ext_tax);
227        let wr_return_amt_inc_tax = dec(
228            wr_return_amt_inc_tax,
229            17,
230            Pricing::get_net_paid_including_tax,
231        );
232        let wr_fee = dec(wr_fee, 18, Pricing::get_fee);
233        let wr_return_ship_cost = dec(wr_return_ship_cost, 19, Pricing::get_ext_ship_cost);
234        let wr_refunded_cash = dec(wr_refunded_cash, 20, Pricing::get_refunded_cash);
235        let wr_reversed_charge = dec(wr_reversed_charge, 21, Pricing::get_reversed_charge);
236        let wr_account_credit = dec(wr_account_credit, 22, Pricing::get_store_credit);
237        let wr_net_loss = dec(wr_net_loss, 23, Pricing::get_net_loss);
238
239        let batch = RecordBatch::try_new(
240            schema,
241            vec![
242                Arc::new(wr_returned_date_sk),
243                Arc::new(wr_returned_time_sk),
244                Arc::new(wr_item_sk),
245                Arc::new(wr_refunded_customer_sk),
246                Arc::new(wr_refunded_cdemo_sk),
247                Arc::new(wr_refunded_hdemo_sk),
248                Arc::new(wr_refunded_addr_sk),
249                Arc::new(wr_returning_customer_sk),
250                Arc::new(wr_returning_cdemo_sk),
251                Arc::new(wr_returning_hdemo_sk),
252                Arc::new(wr_returning_addr_sk),
253                Arc::new(wr_web_page_sk),
254                Arc::new(wr_reason_sk),
255                Arc::new(wr_order_number),
256                Arc::new(wr_return_quantity),
257                Arc::new(wr_return_amt),
258                Arc::new(wr_return_tax),
259                Arc::new(wr_return_amt_inc_tax),
260                Arc::new(wr_fee),
261                Arc::new(wr_return_ship_cost),
262                Arc::new(wr_refunded_cash),
263                Arc::new(wr_reversed_charge),
264                Arc::new(wr_account_credit),
265                Arc::new(wr_net_loss),
266            ],
267        );
268        Some(batch)
269    }
270}
271
272static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
273
274fn make_schema() -> SchemaRef {
275    Arc::new(Schema::new(vec![
276        Field::new("wr_returned_date_sk", DataType::Int32, true),
277        Field::new("wr_returned_time_sk", DataType::Int32, true),
278        Field::new("wr_item_sk", DataType::Int32, false),
279        Field::new("wr_refunded_customer_sk", DataType::Int32, true),
280        Field::new("wr_refunded_cdemo_sk", DataType::Int32, true),
281        Field::new("wr_refunded_hdemo_sk", DataType::Int32, true),
282        Field::new("wr_refunded_addr_sk", DataType::Int32, true),
283        Field::new("wr_returning_customer_sk", DataType::Int32, true),
284        Field::new("wr_returning_cdemo_sk", DataType::Int32, true),
285        Field::new("wr_returning_hdemo_sk", DataType::Int32, true),
286        Field::new("wr_returning_addr_sk", DataType::Int32, true),
287        Field::new("wr_web_page_sk", DataType::Int32, true),
288        Field::new("wr_reason_sk", DataType::Int32, true),
289        Field::new("wr_order_number", DataType::Int64, false),
290        Field::new("wr_return_quantity", DataType::Int32, true),
291        Field::new("wr_return_amt", DataType::Decimal128(7, 2), true),
292        Field::new("wr_return_tax", DataType::Decimal128(7, 2), true),
293        Field::new("wr_return_amt_inc_tax", DataType::Decimal128(7, 2), true),
294        Field::new("wr_fee", DataType::Decimal128(7, 2), true),
295        Field::new("wr_return_ship_cost", DataType::Decimal128(7, 2), true),
296        Field::new("wr_refunded_cash", DataType::Decimal128(7, 2), true),
297        Field::new("wr_reversed_charge", DataType::Decimal128(7, 2), true),
298        Field::new("wr_account_credit", DataType::Decimal128(7, 2), true),
299        Field::new("wr_net_loss", DataType::Decimal128(7, 2), true),
300    ]))
301}
tpcdsgen_arrow/tables/web_returns.rs

tpcdsgen_arrow/tables/
web_returns.rs
