tpcdsgen_arrow/tables/
customer_address.rs

1use crate::conversions::{AddressArrayRecycler, 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::{CustomerAddressRow, CustomerAddressRowGenerator};
10
11/// Generate [`CustomerAddressRow`]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 CustomerAddressArrow {
19    inner: CustomerAddressRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<CustomerAddressRow>,
23    ca_address_sk: PrimitiveArrayRecycler<Int32Type>,
24    ca_address_id: StringViewArrayRecycler,
25    ca_address: AddressArrayRecycler,
26    ca_location_type: StringViewArrayRecycler,
27}
28
29impl CustomerAddressArrow {
30    /// Return the schema without initializing a data generator.
31    pub fn schema_ref() -> SchemaRef {
32        Arc::clone(&SCHEMA)
33    }
34
35    pub fn new(session: Session) -> Self {
36        let row_count = session.get_scaling().get_row_count(Table::CustomerAddress);
37        Self {
38            inner: CustomerAddressRowGenerator::new(session, row_count),
39            batch_size: DEFAULT_BATCH_SIZE,
40            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
41            ca_address_sk: PrimitiveArrayRecycler::new(),
42            ca_address_id: StringViewArrayRecycler::new(),
43            ca_address: AddressArrayRecycler::new(),
44            ca_location_type: StringViewArrayRecycler::new(),
45        }
46    }
47    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
48        self.inner
49            .skip_rows_until_starting_row_number(starting_row_number);
50    }
51
52    /// Generate only source rows `starting_row_number..=ending_row_number`
53    /// (1-based, inclusive). The ending row number is clamped to the table's
54    /// row count.
55    pub fn with_source_row_range(
56        mut self,
57        starting_row_number: u64,
58        ending_row_number: u64,
59    ) -> Self {
60        self.inner
61            .set_source_row_range(starting_row_number, ending_row_number);
62        self
63    }
64
65    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
66        self.batch_size = batch_size;
67        self
68    }
69}
70
71impl RecordBatchReader for CustomerAddressArrow {
72    fn schema(&self) -> SchemaRef {
73        Self::schema_ref()
74    }
75}
76
77impl Iterator for CustomerAddressArrow {
78    type Item = Result<RecordBatch, ArrowError>;
79
80    fn next(&mut self) -> Option<Self::Item> {
81        self.scratch.clear();
82        self.scratch
83            .extend(self.inner.by_ref().take(self.batch_size));
84        let rows = &self.scratch;
85        if rows.is_empty() {
86            return None;
87        }
88
89        let schema = self.schema();
90        let Self {
91            ca_address_sk,
92            ca_address_id,
93            ca_address,
94            ca_location_type,
95            ..
96        } = self;
97        let nbm = CustomerAddressRow::null_bit_map;
98
99        let ca_address_sk =
100            ca_address_sk.int32_sk(rows, 0, nbm, CustomerAddressRow::get_ca_addr_sk);
101        let ca_address_id = ca_address_id.str(rows, 1, nbm, CustomerAddressRow::get_ca_addr_id);
102        let address = ca_address.build(rows, 2, nbm, CustomerAddressRow::get_ca_address);
103        let ca_location_type =
104            ca_location_type.str(rows, 12, nbm, CustomerAddressRow::get_ca_location_type);
105
106        let batch = RecordBatch::try_new(
107            schema,
108            vec![
109                Arc::new(ca_address_sk),
110                Arc::new(ca_address_id),
111                Arc::new(address.street_number),
112                Arc::new(address.street_name),
113                Arc::new(address.street_type),
114                Arc::new(address.suite_number),
115                Arc::new(address.city),
116                Arc::new(address.county),
117                Arc::new(address.state),
118                Arc::new(address.zip),
119                Arc::new(address.country),
120                Arc::new(address.gmt_offset),
121                Arc::new(ca_location_type),
122            ],
123        );
124        Some(batch)
125    }
126}
127
128static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
129
130fn make_schema() -> SchemaRef {
131    Arc::new(Schema::new(vec![
132        Field::new("ca_address_sk", DataType::Int32, false),
133        Field::new("ca_address_id", DataType::Utf8View, false),
134        Field::new("ca_street_number", DataType::Utf8View, true),
135        Field::new("ca_street_name", DataType::Utf8View, true),
136        Field::new("ca_street_type", DataType::Utf8View, true),
137        Field::new("ca_suite_number", DataType::Utf8View, true),
138        Field::new("ca_city", DataType::Utf8View, true),
139        Field::new("ca_county", DataType::Utf8View, true),
140        Field::new("ca_state", DataType::Utf8View, true),
141        Field::new("ca_zip", DataType::Utf8View, true),
142        Field::new("ca_country", DataType::Utf8View, true),
143        Field::new("ca_gmt_offset", DataType::Decimal128(5, 2), true),
144        Field::new("ca_location_type", DataType::Utf8View, true),
145    ]))
146}
tpcdsgen_arrow/tables/customer_address.rs

tpcdsgen_arrow/tables/
customer_address.rs
