tpcdsgen_arrow/tables/
store_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::{StoreReturnsRow, StoreReturnsRowGenerator};
10use tpcdsgen::types::{Decimal, Pricing};
11
12/// Generate [`StoreReturnsRow`]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 StoreReturnsArrow {
20    inner: StoreReturnsRowGenerator,
21    batch_size: usize,
22    // reused allocations across batches
23    scratch: Vec<StoreReturnsRow>,
24    sr_returned_date_sk: PrimitiveArrayRecycler<Int32Type>,
25    sr_returned_time_sk: PrimitiveArrayRecycler<Int32Type>,
26    sr_item_sk: PrimitiveArrayRecycler<Int32Type>,
27    sr_customer_sk: PrimitiveArrayRecycler<Int32Type>,
28    sr_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
29    sr_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
30    sr_addr_sk: PrimitiveArrayRecycler<Int32Type>,
31    sr_store_sk: PrimitiveArrayRecycler<Int32Type>,
32    sr_reason_sk: PrimitiveArrayRecycler<Int32Type>,
33    sr_ticket_number: PrimitiveArrayRecycler<Int64Type>,
34    sr_return_quantity: PrimitiveArrayRecycler<Int32Type>,
35    sr_return_amt: PrimitiveArrayRecycler<Decimal128Type>,
36    sr_return_tax: PrimitiveArrayRecycler<Decimal128Type>,
37    sr_return_amt_inc_tax: PrimitiveArrayRecycler<Decimal128Type>,
38    sr_fee: PrimitiveArrayRecycler<Decimal128Type>,
39    sr_return_ship_cost: PrimitiveArrayRecycler<Decimal128Type>,
40    sr_refunded_cash: PrimitiveArrayRecycler<Decimal128Type>,
41    sr_reversed_charge: PrimitiveArrayRecycler<Decimal128Type>,
42    sr_store_credit: PrimitiveArrayRecycler<Decimal128Type>,
43    sr_net_loss: PrimitiveArrayRecycler<Decimal128Type>,
44}
45
46impl StoreReturnsArrow {
47    /// Return the schema without initializing a data generator.
48    pub fn schema_ref() -> SchemaRef {
49        Arc::clone(&SCHEMA)
50    }
51
52    pub fn new(session: Session) -> Self {
53        let row_count = session.get_scaling().get_row_count(Table::StoreReturns);
54        Self {
55            inner: StoreReturnsRowGenerator::new(session, row_count),
56            batch_size: DEFAULT_BATCH_SIZE,
57            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
58            sr_returned_date_sk: PrimitiveArrayRecycler::new(),
59            sr_returned_time_sk: PrimitiveArrayRecycler::new(),
60            sr_item_sk: PrimitiveArrayRecycler::new(),
61            sr_customer_sk: PrimitiveArrayRecycler::new(),
62            sr_cdemo_sk: PrimitiveArrayRecycler::new(),
63            sr_hdemo_sk: PrimitiveArrayRecycler::new(),
64            sr_addr_sk: PrimitiveArrayRecycler::new(),
65            sr_store_sk: PrimitiveArrayRecycler::new(),
66            sr_reason_sk: PrimitiveArrayRecycler::new(),
67            sr_ticket_number: PrimitiveArrayRecycler::new(),
68            sr_return_quantity: PrimitiveArrayRecycler::new(),
69            sr_return_amt: PrimitiveArrayRecycler::new(),
70            sr_return_tax: PrimitiveArrayRecycler::new(),
71            sr_return_amt_inc_tax: PrimitiveArrayRecycler::new(),
72            sr_fee: PrimitiveArrayRecycler::new(),
73            sr_return_ship_cost: PrimitiveArrayRecycler::new(),
74            sr_refunded_cash: PrimitiveArrayRecycler::new(),
75            sr_reversed_charge: PrimitiveArrayRecycler::new(),
76            sr_store_credit: PrimitiveArrayRecycler::new(),
77            sr_net_loss: PrimitiveArrayRecycler::new(),
78        }
79    }
80    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
81        self.inner
82            .skip_rows_until_starting_row_number(starting_row_number);
83    }
84
85    /// Generate only source rows `starting_row_number..=ending_row_number`
86    /// (1-based, inclusive). The ending row number is clamped to the table's
87    /// row count.
88    pub fn with_source_row_range(
89        mut self,
90        starting_row_number: u64,
91        ending_row_number: u64,
92    ) -> Self {
93        self.inner
94            .set_source_row_range(starting_row_number, ending_row_number);
95        self
96    }
97
98    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
99        self.batch_size = batch_size;
100        self
101    }
102}
103
104impl RecordBatchReader for StoreReturnsArrow {
105    fn schema(&self) -> SchemaRef {
106        Self::schema_ref()
107    }
108}
109
110impl Iterator for StoreReturnsArrow {
111    type Item = Result<RecordBatch, ArrowError>;
112
113    fn next(&mut self) -> Option<Self::Item> {
114        self.scratch.clear();
115        self.scratch
116            .extend(self.inner.by_ref().take(self.batch_size));
117        let rows = &self.scratch;
118        if rows.is_empty() {
119            return None;
120        }
121
122        let schema = self.schema();
123        let Self {
124            sr_returned_date_sk,
125            sr_returned_time_sk,
126            sr_item_sk,
127            sr_customer_sk,
128            sr_cdemo_sk,
129            sr_hdemo_sk,
130            sr_addr_sk,
131            sr_store_sk,
132            sr_reason_sk,
133            sr_ticket_number,
134            sr_return_quantity,
135            sr_return_amt,
136            sr_return_tax,
137            sr_return_amt_inc_tax,
138            sr_fee,
139            sr_return_ship_cost,
140            sr_refunded_cash,
141            sr_reversed_charge,
142            sr_store_credit,
143            sr_net_loss,
144            ..
145        } = self;
146        let nbm = StoreReturnsRow::null_bit_map;
147        let sk32 =
148            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&StoreReturnsRow) -> i64| {
149                rec.int32_sk(rows, pos, nbm, f)
150            };
151        let dec =
152            |rec: &mut PrimitiveArrayRecycler<Decimal128Type>, pos, f: fn(&Pricing) -> Decimal| {
153                rec.decimal_7_2(rows, pos, nbm, |r| f(r.get_sr_pricing()))
154            };
155
156        let sr_returned_date_sk = sk32(
157            sr_returned_date_sk,
158            0,
159            StoreReturnsRow::get_sr_returned_date_sk,
160        );
161        let sr_returned_time_sk = sk32(
162            sr_returned_time_sk,
163            1,
164            StoreReturnsRow::get_sr_returned_time_sk,
165        );
166        let sr_item_sk = sk32(sr_item_sk, 2, StoreReturnsRow::get_sr_item_sk);
167        let sr_customer_sk = sk32(sr_customer_sk, 3, StoreReturnsRow::get_sr_customer_sk);
168        let sr_cdemo_sk = sk32(sr_cdemo_sk, 4, StoreReturnsRow::get_sr_cdemo_sk);
169        let sr_hdemo_sk = sk32(sr_hdemo_sk, 5, StoreReturnsRow::get_sr_hdemo_sk);
170        let sr_addr_sk = sk32(sr_addr_sk, 6, StoreReturnsRow::get_sr_addr_sk);
171        let sr_store_sk = sk32(sr_store_sk, 7, StoreReturnsRow::get_sr_store_sk);
172        let sr_reason_sk = sk32(sr_reason_sk, 8, StoreReturnsRow::get_sr_reason_sk);
173        let sr_ticket_number =
174            sr_ticket_number.int64_sk(rows, 9, nbm, StoreReturnsRow::get_sr_ticket_number);
175        let sr_return_quantity =
176            sr_return_quantity.int32(rows, 10, nbm, |r| r.get_sr_pricing().get_quantity());
177        let sr_return_amt = dec(sr_return_amt, 11, Pricing::get_net_paid);
178        let sr_return_tax = dec(sr_return_tax, 12, Pricing::get_ext_tax);
179        let sr_return_amt_inc_tax = dec(
180            sr_return_amt_inc_tax,
181            13,
182            Pricing::get_net_paid_including_tax,
183        );
184        let sr_fee = dec(sr_fee, 14, Pricing::get_fee);
185        let sr_return_ship_cost = dec(sr_return_ship_cost, 15, Pricing::get_ext_ship_cost);
186        let sr_refunded_cash = dec(sr_refunded_cash, 16, Pricing::get_refunded_cash);
187        let sr_reversed_charge = dec(sr_reversed_charge, 17, Pricing::get_reversed_charge);
188        let sr_store_credit = dec(sr_store_credit, 18, Pricing::get_store_credit);
189        let sr_net_loss = dec(sr_net_loss, 19, Pricing::get_net_loss);
190
191        let batch = RecordBatch::try_new(
192            schema,
193            vec![
194                Arc::new(sr_returned_date_sk),
195                Arc::new(sr_returned_time_sk),
196                Arc::new(sr_item_sk),
197                Arc::new(sr_customer_sk),
198                Arc::new(sr_cdemo_sk),
199                Arc::new(sr_hdemo_sk),
200                Arc::new(sr_addr_sk),
201                Arc::new(sr_store_sk),
202                Arc::new(sr_reason_sk),
203                Arc::new(sr_ticket_number),
204                Arc::new(sr_return_quantity),
205                Arc::new(sr_return_amt),
206                Arc::new(sr_return_tax),
207                Arc::new(sr_return_amt_inc_tax),
208                Arc::new(sr_fee),
209                Arc::new(sr_return_ship_cost),
210                Arc::new(sr_refunded_cash),
211                Arc::new(sr_reversed_charge),
212                Arc::new(sr_store_credit),
213                Arc::new(sr_net_loss),
214            ],
215        );
216        Some(batch)
217    }
218}
219
220static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
221
222fn make_schema() -> SchemaRef {
223    Arc::new(Schema::new(vec![
224        Field::new("sr_returned_date_sk", DataType::Int32, true),
225        Field::new("sr_return_time_sk", DataType::Int32, true),
226        Field::new("sr_item_sk", DataType::Int32, false),
227        Field::new("sr_customer_sk", DataType::Int32, true),
228        Field::new("sr_cdemo_sk", DataType::Int32, true),
229        Field::new("sr_hdemo_sk", DataType::Int32, true),
230        Field::new("sr_addr_sk", DataType::Int32, true),
231        Field::new("sr_store_sk", DataType::Int32, true),
232        Field::new("sr_reason_sk", DataType::Int32, true),
233        Field::new("sr_ticket_number", DataType::Int64, false),
234        Field::new("sr_return_quantity", DataType::Int32, true),
235        Field::new("sr_return_amt", DataType::Decimal128(7, 2), true),
236        Field::new("sr_return_tax", DataType::Decimal128(7, 2), true),
237        Field::new("sr_return_amt_inc_tax", DataType::Decimal128(7, 2), true),
238        Field::new("sr_fee", DataType::Decimal128(7, 2), true),
239        Field::new("sr_return_ship_cost", DataType::Decimal128(7, 2), true),
240        Field::new("sr_refunded_cash", DataType::Decimal128(7, 2), true),
241        Field::new("sr_reversed_charge", DataType::Decimal128(7, 2), true),
242        Field::new("sr_store_credit", DataType::Decimal128(7, 2), true),
243        Field::new("sr_net_loss", DataType::Decimal128(7, 2), true),
244    ]))
245}
tpcdsgen_arrow/tables/store_returns.rs

tpcdsgen_arrow/tables/
store_returns.rs
