tpcdsgen_arrow/tables/
customer.rs

1use crate::conversions::{bool_to_yn, PrimitiveArrayRecycler, StringViewArrayRecycler};
2use crate::DEFAULT_BATCH_SIZE;
3use arrow::array::RecordBatch;
4use arrow::datatypes::{DataType, 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::{CustomerRow, CustomerRowGenerator};
10
11/// Generate [`CustomerRow`]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 CustomerArrow {
19    inner: CustomerRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<CustomerRow>,
23    c_customer_sk: PrimitiveArrayRecycler<Int32Type>,
24    c_customer_id: StringViewArrayRecycler,
25    c_current_cdemo_sk: PrimitiveArrayRecycler<Int32Type>,
26    c_current_hdemo_sk: PrimitiveArrayRecycler<Int32Type>,
27    c_current_addr_sk: PrimitiveArrayRecycler<Int32Type>,
28    c_first_shipto_date_sk: PrimitiveArrayRecycler<Int32Type>,
29    c_first_sales_date_sk: PrimitiveArrayRecycler<Int32Type>,
30    c_salutation: StringViewArrayRecycler,
31    c_first_name: StringViewArrayRecycler,
32    c_last_name: StringViewArrayRecycler,
33    c_preferred_cust_flag: StringViewArrayRecycler,
34    c_birth_day: PrimitiveArrayRecycler<Int32Type>,
35    c_birth_month: PrimitiveArrayRecycler<Int32Type>,
36    c_birth_year: PrimitiveArrayRecycler<Int32Type>,
37    c_birth_country: StringViewArrayRecycler,
38    c_login: StringViewArrayRecycler,
39    c_email_address: StringViewArrayRecycler,
40    c_last_review_date_sk: PrimitiveArrayRecycler<Int32Type>,
41}
42
43impl CustomerArrow {
44    /// Return the schema without initializing a data generator.
45    pub fn schema_ref() -> SchemaRef {
46        Arc::clone(&SCHEMA)
47    }
48
49    pub fn new(session: Session) -> Self {
50        let row_count = session.get_scaling().get_row_count(Table::Customer);
51        Self {
52            inner: CustomerRowGenerator::new(session, row_count),
53            batch_size: DEFAULT_BATCH_SIZE,
54            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
55            c_customer_sk: PrimitiveArrayRecycler::new(),
56            c_customer_id: StringViewArrayRecycler::new(),
57            c_current_cdemo_sk: PrimitiveArrayRecycler::new(),
58            c_current_hdemo_sk: PrimitiveArrayRecycler::new(),
59            c_current_addr_sk: PrimitiveArrayRecycler::new(),
60            c_first_shipto_date_sk: PrimitiveArrayRecycler::new(),
61            c_first_sales_date_sk: PrimitiveArrayRecycler::new(),
62            c_salutation: StringViewArrayRecycler::new(),
63            c_first_name: StringViewArrayRecycler::new(),
64            c_last_name: StringViewArrayRecycler::new(),
65            c_preferred_cust_flag: StringViewArrayRecycler::new(),
66            c_birth_day: PrimitiveArrayRecycler::new(),
67            c_birth_month: PrimitiveArrayRecycler::new(),
68            c_birth_year: PrimitiveArrayRecycler::new(),
69            c_birth_country: StringViewArrayRecycler::new(),
70            c_login: StringViewArrayRecycler::new(),
71            c_email_address: StringViewArrayRecycler::new(),
72            c_last_review_date_sk: PrimitiveArrayRecycler::new(),
73        }
74    }
75    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
76        self.inner
77            .skip_rows_until_starting_row_number(starting_row_number);
78    }
79
80    /// Generate only source rows `starting_row_number..=ending_row_number`
81    /// (1-based, inclusive). The ending row number is clamped to the table's
82    /// row count.
83    pub fn with_source_row_range(
84        mut self,
85        starting_row_number: u64,
86        ending_row_number: u64,
87    ) -> Self {
88        self.inner
89            .set_source_row_range(starting_row_number, ending_row_number);
90        self
91    }
92
93    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
94        self.batch_size = batch_size;
95        self
96    }
97}
98
99impl RecordBatchReader for CustomerArrow {
100    fn schema(&self) -> SchemaRef {
101        Self::schema_ref()
102    }
103}
104
105impl Iterator for CustomerArrow {
106    type Item = Result<RecordBatch, ArrowError>;
107
108    fn next(&mut self) -> Option<Self::Item> {
109        self.scratch.clear();
110        self.scratch
111            .extend(self.inner.by_ref().take(self.batch_size));
112        let rows = &self.scratch;
113        if rows.is_empty() {
114            return None;
115        }
116
117        let schema = self.schema();
118        let Self {
119            c_customer_sk,
120            c_customer_id,
121            c_current_cdemo_sk,
122            c_current_hdemo_sk,
123            c_current_addr_sk,
124            c_first_shipto_date_sk,
125            c_first_sales_date_sk,
126            c_salutation,
127            c_first_name,
128            c_last_name,
129            c_preferred_cust_flag,
130            c_birth_day,
131            c_birth_month,
132            c_birth_year,
133            c_birth_country,
134            c_login,
135            c_email_address,
136            c_last_review_date_sk,
137            ..
138        } = self;
139        let nbm = CustomerRow::null_bit_map;
140        let sk32 =
141            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&CustomerRow) -> i64| {
142                rec.int32_sk(rows, pos, nbm, f)
143            };
144        let int = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&CustomerRow) -> i32| {
145            rec.int32(rows, pos, nbm, f)
146        };
147        let str = |rec: &mut StringViewArrayRecycler, pos, f: fn(&CustomerRow) -> &str| {
148            rec.str(rows, pos, nbm, f)
149        };
150
151        let c_customer_sk = sk32(c_customer_sk, 0, CustomerRow::get_c_customer_sk);
152        let c_customer_id = str(c_customer_id, 1, CustomerRow::get_c_customer_id);
153        let c_current_cdemo_sk = sk32(c_current_cdemo_sk, 2, CustomerRow::get_c_current_cdemo_sk);
154        let c_current_hdemo_sk = sk32(c_current_hdemo_sk, 3, CustomerRow::get_c_current_hdemo_sk);
155        let c_current_addr_sk = sk32(c_current_addr_sk, 4, CustomerRow::get_c_current_addr_sk);
156        let c_first_shipto_date_sk = int(
157            c_first_shipto_date_sk,
158            5,
159            CustomerRow::get_c_first_shipto_date_id,
160        );
161        let c_first_sales_date_sk = int(
162            c_first_sales_date_sk,
163            6,
164            CustomerRow::get_c_first_sales_date_id,
165        );
166        let c_salutation = str(c_salutation, 7, CustomerRow::get_c_salutation);
167        let c_first_name = str(c_first_name, 8, CustomerRow::get_c_first_name);
168        let c_last_name = str(c_last_name, 9, CustomerRow::get_c_last_name);
169        let c_preferred_cust_flag = str(c_preferred_cust_flag, 10, |r| {
170            bool_to_yn(r.get_c_preferred_cust_flag())
171        });
172        let c_birth_day = int(c_birth_day, 11, CustomerRow::get_c_birth_day);
173        let c_birth_month = int(c_birth_month, 12, CustomerRow::get_c_birth_month);
174        let c_birth_year = int(c_birth_year, 13, CustomerRow::get_c_birth_year);
175        let c_birth_country = str(c_birth_country, 14, CustomerRow::get_c_birth_country);
176        // always null per TPC-DS spec
177        let c_login = c_login.build(rows, |_| None);
178        let c_email_address = str(c_email_address, 16, CustomerRow::get_c_email_address);
179        let c_last_review_date_sk = int(
180            c_last_review_date_sk,
181            17,
182            CustomerRow::get_c_last_review_date,
183        );
184
185        let batch = RecordBatch::try_new(
186            schema,
187            vec![
188                Arc::new(c_customer_sk),
189                Arc::new(c_customer_id),
190                Arc::new(c_current_cdemo_sk),
191                Arc::new(c_current_hdemo_sk),
192                Arc::new(c_current_addr_sk),
193                Arc::new(c_first_shipto_date_sk),
194                Arc::new(c_first_sales_date_sk),
195                Arc::new(c_salutation),
196                Arc::new(c_first_name),
197                Arc::new(c_last_name),
198                Arc::new(c_preferred_cust_flag),
199                Arc::new(c_birth_day),
200                Arc::new(c_birth_month),
201                Arc::new(c_birth_year),
202                Arc::new(c_birth_country),
203                Arc::new(c_login),
204                Arc::new(c_email_address),
205                Arc::new(c_last_review_date_sk),
206            ],
207        );
208        Some(batch)
209    }
210}
211
212static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
213
214fn make_schema() -> SchemaRef {
215    Arc::new(Schema::new(vec![
216        Field::new("c_customer_sk", DataType::Int32, false),
217        Field::new("c_customer_id", DataType::Utf8View, false),
218        Field::new("c_current_cdemo_sk", DataType::Int32, true),
219        Field::new("c_current_hdemo_sk", DataType::Int32, true),
220        Field::new("c_current_addr_sk", DataType::Int32, true),
221        Field::new("c_first_shipto_date_sk", DataType::Int32, true),
222        Field::new("c_first_sales_date_sk", DataType::Int32, true),
223        Field::new("c_salutation", DataType::Utf8View, true),
224        Field::new("c_first_name", DataType::Utf8View, true),
225        Field::new("c_last_name", DataType::Utf8View, true),
226        Field::new("c_preferred_cust_flag", DataType::Utf8View, true),
227        Field::new("c_birth_day", DataType::Int32, true),
228        Field::new("c_birth_month", DataType::Int32, true),
229        Field::new("c_birth_year", DataType::Int32, true),
230        Field::new("c_birth_country", DataType::Utf8View, true),
231        Field::new("c_login", DataType::Utf8View, true),
232        Field::new("c_email_address", DataType::Utf8View, true),
233        Field::new("c_last_review_date_sk", DataType::Int32, true),
234    ]))
235}
tpcdsgen_arrow/tables/customer.rs

tpcdsgen_arrow/tables/
customer.rs
