tpcdsgen_arrow/tables/
promotion.rs

1use crate::conversions::{bool_to_yn, PrimitiveArrayRecycler, StringViewArrayRecycler};
2use crate::DEFAULT_BATCH_SIZE;
3use arrow::array::RecordBatch;
4use arrow::datatypes::{DataType, Decimal128Type, Field, Int32Type, Schema, SchemaRef};
5use arrow::error::ArrowError;
6use arrow::record_batch::RecordBatchReader;
7use std::sync::{Arc, LazyLock};
8use tpcdsgen::config::{Session, Table};
9use tpcdsgen::row::{PromotionRow, PromotionRowGenerator};
10
11/// Generate [`PromotionRow`]s in [`RecordBatch`] format
12///
13/// Each column's value and null buffers are recycled from one batch to the
14/// next if possible.
15///
16/// Buffers can be recycled if the consumer drops each [`RecordBatch`] before
17/// generating the next one.
18pub struct PromotionArrow {
19    inner: PromotionRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<PromotionRow>,
23    p_promo_sk: PrimitiveArrayRecycler<Int32Type>,
24    p_promo_id: StringViewArrayRecycler,
25    p_start_date_sk: PrimitiveArrayRecycler<Int32Type>,
26    p_end_date_sk: PrimitiveArrayRecycler<Int32Type>,
27    p_item_sk: PrimitiveArrayRecycler<Int32Type>,
28    p_cost: PrimitiveArrayRecycler<Decimal128Type>,
29    p_response_target: PrimitiveArrayRecycler<Int32Type>,
30    p_promo_name: StringViewArrayRecycler,
31    p_channel_dmail: StringViewArrayRecycler,
32    p_channel_email: StringViewArrayRecycler,
33    p_channel_catalog: StringViewArrayRecycler,
34    p_channel_tv: StringViewArrayRecycler,
35    p_channel_radio: StringViewArrayRecycler,
36    p_channel_press: StringViewArrayRecycler,
37    p_channel_event: StringViewArrayRecycler,
38    p_channel_demo: StringViewArrayRecycler,
39    p_channel_details: StringViewArrayRecycler,
40    p_purpose: StringViewArrayRecycler,
41    p_discount_active: StringViewArrayRecycler,
42}
43
44impl PromotionArrow {
45    /// Return the schema without initializing a data generator.
46    pub fn schema_ref() -> SchemaRef {
47        Arc::clone(&SCHEMA)
48    }
49
50    pub fn new(session: Session) -> Self {
51        let row_count = session.get_scaling().get_row_count(Table::Promotion);
52        Self {
53            inner: PromotionRowGenerator::new(session, row_count),
54            batch_size: DEFAULT_BATCH_SIZE,
55            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
56            p_promo_sk: PrimitiveArrayRecycler::new(),
57            p_promo_id: StringViewArrayRecycler::new(),
58            p_start_date_sk: PrimitiveArrayRecycler::new(),
59            p_end_date_sk: PrimitiveArrayRecycler::new(),
60            p_item_sk: PrimitiveArrayRecycler::new(),
61            p_cost: PrimitiveArrayRecycler::new(),
62            p_response_target: PrimitiveArrayRecycler::new(),
63            p_promo_name: StringViewArrayRecycler::new(),
64            p_channel_dmail: StringViewArrayRecycler::new(),
65            p_channel_email: StringViewArrayRecycler::new(),
66            p_channel_catalog: StringViewArrayRecycler::new(),
67            p_channel_tv: StringViewArrayRecycler::new(),
68            p_channel_radio: StringViewArrayRecycler::new(),
69            p_channel_press: StringViewArrayRecycler::new(),
70            p_channel_event: StringViewArrayRecycler::new(),
71            p_channel_demo: StringViewArrayRecycler::new(),
72            p_channel_details: StringViewArrayRecycler::new(),
73            p_purpose: StringViewArrayRecycler::new(),
74            p_discount_active: StringViewArrayRecycler::new(),
75        }
76    }
77    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
78        self.inner
79            .skip_rows_until_starting_row_number(starting_row_number);
80    }
81
82    /// Generate only source rows `starting_row_number..=ending_row_number`
83    /// (1-based, inclusive). The ending row number is clamped to the table's
84    /// row count.
85    pub fn with_source_row_range(
86        mut self,
87        starting_row_number: u64,
88        ending_row_number: u64,
89    ) -> Self {
90        self.inner
91            .set_source_row_range(starting_row_number, ending_row_number);
92        self
93    }
94
95    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
96        self.batch_size = batch_size;
97        self
98    }
99}
100
101impl RecordBatchReader for PromotionArrow {
102    fn schema(&self) -> SchemaRef {
103        Self::schema_ref()
104    }
105}
106
107impl Iterator for PromotionArrow {
108    type Item = Result<RecordBatch, ArrowError>;
109
110    fn next(&mut self) -> Option<Self::Item> {
111        self.scratch.clear();
112        self.scratch
113            .extend(self.inner.by_ref().take(self.batch_size));
114        let rows = &self.scratch;
115        if rows.is_empty() {
116            return None;
117        }
118
119        let schema = self.schema();
120        let Self {
121            p_promo_sk,
122            p_promo_id,
123            p_start_date_sk,
124            p_end_date_sk,
125            p_item_sk,
126            p_cost,
127            p_response_target,
128            p_promo_name,
129            p_channel_dmail,
130            p_channel_email,
131            p_channel_catalog,
132            p_channel_tv,
133            p_channel_radio,
134            p_channel_press,
135            p_channel_event,
136            p_channel_demo,
137            p_channel_details,
138            p_purpose,
139            p_discount_active,
140            ..
141        } = self;
142        let nbm = PromotionRow::null_bit_map;
143        let sk32 =
144            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&PromotionRow) -> i64| {
145                rec.int32_sk(rows, pos, nbm, f)
146            };
147        let int = |rec: &mut PrimitiveArrayRecycler<Int32Type>,
148                   pos,
149                   f: fn(&PromotionRow) -> i32| { rec.int32(rows, pos, nbm, f) };
150        let str = |rec: &mut StringViewArrayRecycler, pos, f: fn(&PromotionRow) -> &str| {
151            rec.str(rows, pos, nbm, f)
152        };
153
154        let p_promo_sk = sk32(p_promo_sk, 0, PromotionRow::get_p_promo_sk);
155        let p_promo_id = str(p_promo_id, 1, PromotionRow::get_p_promo_id);
156        let p_start_date_sk = sk32(p_start_date_sk, 2, PromotionRow::get_p_start_date_id);
157        let p_end_date_sk = sk32(p_end_date_sk, 3, PromotionRow::get_p_end_date_id);
158        let p_item_sk = sk32(p_item_sk, 4, PromotionRow::get_p_item_sk);
159        let p_cost = p_cost.decimal_15_2(rows, 5, nbm, PromotionRow::get_p_cost);
160        let p_response_target = int(p_response_target, 6, PromotionRow::get_p_response_target);
161        let p_promo_name = str(p_promo_name, 7, PromotionRow::get_p_promo_name);
162        let p_channel_dmail = str(p_channel_dmail, 8, |r| bool_to_yn(r.get_p_channel_dmail()));
163        let p_channel_email = str(p_channel_email, 9, |r| bool_to_yn(r.get_p_channel_email()));
164        let p_channel_catalog = str(p_channel_catalog, 10, |r| {
165            bool_to_yn(r.get_p_channel_catalog())
166        });
167        let p_channel_tv = str(p_channel_tv, 11, |r| bool_to_yn(r.get_p_channel_tv()));
168        let p_channel_radio = str(p_channel_radio, 12, |r| bool_to_yn(r.get_p_channel_radio()));
169        let p_channel_press = str(p_channel_press, 13, |r| bool_to_yn(r.get_p_channel_press()));
170        let p_channel_event = str(p_channel_event, 14, |r| bool_to_yn(r.get_p_channel_event()));
171        let p_channel_demo = str(p_channel_demo, 15, |r| bool_to_yn(r.get_p_channel_demo()));
172        let p_channel_details = str(p_channel_details, 16, PromotionRow::get_p_channel_details);
173        let p_purpose = str(p_purpose, 17, PromotionRow::get_p_purpose);
174        let p_discount_active = str(p_discount_active, 18, |r| {
175            bool_to_yn(r.get_p_discount_active())
176        });
177
178        let batch = RecordBatch::try_new(
179            schema,
180            vec![
181                Arc::new(p_promo_sk),
182                Arc::new(p_promo_id),
183                Arc::new(p_start_date_sk),
184                Arc::new(p_end_date_sk),
185                Arc::new(p_item_sk),
186                Arc::new(p_cost),
187                Arc::new(p_response_target),
188                Arc::new(p_promo_name),
189                Arc::new(p_channel_dmail),
190                Arc::new(p_channel_email),
191                Arc::new(p_channel_catalog),
192                Arc::new(p_channel_tv),
193                Arc::new(p_channel_radio),
194                Arc::new(p_channel_press),
195                Arc::new(p_channel_event),
196                Arc::new(p_channel_demo),
197                Arc::new(p_channel_details),
198                Arc::new(p_purpose),
199                Arc::new(p_discount_active),
200            ],
201        );
202        Some(batch)
203    }
204}
205
206static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
207
208fn make_schema() -> SchemaRef {
209    Arc::new(Schema::new(vec![
210        Field::new("p_promo_sk", DataType::Int32, false),
211        Field::new("p_promo_id", DataType::Utf8View, false),
212        Field::new("p_start_date_sk", DataType::Int32, true),
213        Field::new("p_end_date_sk", DataType::Int32, true),
214        Field::new("p_item_sk", DataType::Int32, true),
215        Field::new("p_cost", DataType::Decimal128(15, 2), true),
216        Field::new("p_response_target", DataType::Int32, true),
217        Field::new("p_promo_name", DataType::Utf8View, true),
218        Field::new("p_channel_dmail", DataType::Utf8View, true),
219        Field::new("p_channel_email", DataType::Utf8View, true),
220        Field::new("p_channel_catalog", DataType::Utf8View, true),
221        Field::new("p_channel_tv", DataType::Utf8View, true),
222        Field::new("p_channel_radio", DataType::Utf8View, true),
223        Field::new("p_channel_press", DataType::Utf8View, true),
224        Field::new("p_channel_event", DataType::Utf8View, true),
225        Field::new("p_channel_demo", DataType::Utf8View, true),
226        Field::new("p_channel_details", DataType::Utf8View, true),
227        Field::new("p_purpose", DataType::Utf8View, true),
228        Field::new("p_discount_active", DataType::Utf8View, true),
229    ]))
230}
tpcdsgen_arrow/tables/promotion.rs

tpcdsgen_arrow/tables/
promotion.rs
