tpcdsgen_arrow/tables/
web_page.rs

1use crate::conversions::{bool_to_yn, PrimitiveArrayRecycler, StringViewArrayRecycler};
2use crate::DEFAULT_BATCH_SIZE;
3use arrow::array::RecordBatch;
4use arrow::datatypes::{DataType, Date32Type, 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::{WebPageRow, WebPageRowGenerator};
10
11/// Generate [`WebPageRow`]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 WebPageArrow {
19    inner: WebPageRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<WebPageRow>,
23    wp_web_page_sk: PrimitiveArrayRecycler<Int32Type>,
24    wp_web_page_id: StringViewArrayRecycler,
25    wp_rec_start_date: PrimitiveArrayRecycler<Date32Type>,
26    wp_rec_end_date: PrimitiveArrayRecycler<Date32Type>,
27    wp_creation_date_sk: PrimitiveArrayRecycler<Int32Type>,
28    wp_access_date_sk: PrimitiveArrayRecycler<Int32Type>,
29    wp_autogen_flag: StringViewArrayRecycler,
30    wp_customer_sk: PrimitiveArrayRecycler<Int32Type>,
31    wp_url: StringViewArrayRecycler,
32    wp_type: StringViewArrayRecycler,
33    wp_char_count: PrimitiveArrayRecycler<Int32Type>,
34    wp_link_count: PrimitiveArrayRecycler<Int32Type>,
35    wp_image_count: PrimitiveArrayRecycler<Int32Type>,
36    wp_max_ad_count: PrimitiveArrayRecycler<Int32Type>,
37}
38
39impl WebPageArrow {
40    /// Return the schema without initializing a data generator.
41    pub fn schema_ref() -> SchemaRef {
42        Arc::clone(&SCHEMA)
43    }
44
45    pub fn new(session: Session) -> Self {
46        let row_count = session.get_scaling().get_row_count(Table::WebPage);
47        Self {
48            inner: WebPageRowGenerator::new(session, row_count),
49            batch_size: DEFAULT_BATCH_SIZE,
50            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
51            wp_web_page_sk: PrimitiveArrayRecycler::new(),
52            wp_web_page_id: StringViewArrayRecycler::new(),
53            wp_rec_start_date: PrimitiveArrayRecycler::new(),
54            wp_rec_end_date: PrimitiveArrayRecycler::new(),
55            wp_creation_date_sk: PrimitiveArrayRecycler::new(),
56            wp_access_date_sk: PrimitiveArrayRecycler::new(),
57            wp_autogen_flag: StringViewArrayRecycler::new(),
58            wp_customer_sk: PrimitiveArrayRecycler::new(),
59            wp_url: StringViewArrayRecycler::new(),
60            wp_type: StringViewArrayRecycler::new(),
61            wp_char_count: PrimitiveArrayRecycler::new(),
62            wp_link_count: PrimitiveArrayRecycler::new(),
63            wp_image_count: PrimitiveArrayRecycler::new(),
64            wp_max_ad_count: PrimitiveArrayRecycler::new(),
65        }
66    }
67    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
68        self.inner
69            .skip_rows_until_starting_row_number(starting_row_number);
70    }
71
72    /// Generate only source rows `starting_row_number..=ending_row_number`
73    /// (1-based, inclusive). The ending row number is clamped to the table's
74    /// row count.
75    pub fn with_source_row_range(
76        mut self,
77        starting_row_number: u64,
78        ending_row_number: u64,
79    ) -> Self {
80        self.inner
81            .set_source_row_range(starting_row_number, ending_row_number);
82        self
83    }
84
85    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
86        self.batch_size = batch_size;
87        self
88    }
89}
90
91impl RecordBatchReader for WebPageArrow {
92    fn schema(&self) -> SchemaRef {
93        Self::schema_ref()
94    }
95}
96
97impl Iterator for WebPageArrow {
98    type Item = Result<RecordBatch, ArrowError>;
99
100    fn next(&mut self) -> Option<Self::Item> {
101        self.scratch.clear();
102        self.scratch
103            .extend(self.inner.by_ref().take(self.batch_size));
104        let rows = &self.scratch;
105        if rows.is_empty() {
106            return None;
107        }
108
109        let schema = self.schema();
110        let Self {
111            wp_web_page_sk,
112            wp_web_page_id,
113            wp_rec_start_date,
114            wp_rec_end_date,
115            wp_creation_date_sk,
116            wp_access_date_sk,
117            wp_autogen_flag,
118            wp_customer_sk,
119            wp_url,
120            wp_type,
121            wp_char_count,
122            wp_link_count,
123            wp_image_count,
124            wp_max_ad_count,
125            ..
126        } = self;
127        let nbm = WebPageRow::null_bit_map;
128        let sk32 = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&WebPageRow) -> i64| {
129            rec.int32_sk(rows, pos, nbm, f)
130        };
131        let int = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&WebPageRow) -> i32| {
132            rec.int32(rows, pos, nbm, f)
133        };
134        let str = |rec: &mut StringViewArrayRecycler, pos, f: fn(&WebPageRow) -> &str| {
135            rec.str(rows, pos, nbm, f)
136        };
137
138        let wp_web_page_sk = sk32(wp_web_page_sk, 0, WebPageRow::get_wp_page_sk);
139        let wp_web_page_id = str(wp_web_page_id, 1, WebPageRow::get_wp_page_id);
140        let wp_rec_start_date =
141            wp_rec_start_date.date32_julian(rows, 2, nbm, WebPageRow::get_wp_rec_start_date_id);
142        // wp_rec_end_date ignores the null bitmap; only the negative-Julian sentinel is NULL
143        let wp_rec_end_date =
144            wp_rec_end_date.date32_julian(rows, 3, |_| 0, WebPageRow::get_wp_rec_end_date_id);
145        let wp_creation_date_sk = sk32(wp_creation_date_sk, 4, WebPageRow::get_wp_creation_date_sk);
146        let wp_access_date_sk = sk32(wp_access_date_sk, 5, WebPageRow::get_wp_access_date_sk);
147        let wp_autogen_flag = str(wp_autogen_flag, 6, |r| bool_to_yn(r.get_wp_autogen_flag()));
148        let wp_customer_sk = sk32(wp_customer_sk, 7, WebPageRow::get_wp_customer_sk);
149        let wp_url = str(wp_url, 8, WebPageRow::get_wp_url);
150        let wp_type = str(wp_type, 9, WebPageRow::get_wp_type);
151        let wp_char_count = int(wp_char_count, 10, WebPageRow::get_wp_char_count);
152        let wp_link_count = int(wp_link_count, 11, WebPageRow::get_wp_link_count);
153        let wp_image_count = int(wp_image_count, 12, WebPageRow::get_wp_image_count);
154        let wp_max_ad_count = int(wp_max_ad_count, 13, WebPageRow::get_wp_max_ad_count);
155
156        let batch = RecordBatch::try_new(
157            schema,
158            vec![
159                Arc::new(wp_web_page_sk),
160                Arc::new(wp_web_page_id),
161                Arc::new(wp_rec_start_date),
162                Arc::new(wp_rec_end_date),
163                Arc::new(wp_creation_date_sk),
164                Arc::new(wp_access_date_sk),
165                Arc::new(wp_autogen_flag),
166                Arc::new(wp_customer_sk),
167                Arc::new(wp_url),
168                Arc::new(wp_type),
169                Arc::new(wp_char_count),
170                Arc::new(wp_link_count),
171                Arc::new(wp_image_count),
172                Arc::new(wp_max_ad_count),
173            ],
174        );
175        Some(batch)
176    }
177}
178
179static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
180
181fn make_schema() -> SchemaRef {
182    Arc::new(Schema::new(vec![
183        Field::new("wp_web_page_sk", DataType::Int32, false),
184        Field::new("wp_web_page_id", DataType::Utf8View, false),
185        Field::new("wp_rec_start_date", DataType::Date32, true),
186        Field::new("wp_rec_end_date", DataType::Date32, true),
187        Field::new("wp_creation_date_sk", DataType::Int32, true),
188        Field::new("wp_access_date_sk", DataType::Int32, true),
189        Field::new("wp_autogen_flag", DataType::Utf8View, true),
190        Field::new("wp_customer_sk", DataType::Int32, true),
191        Field::new("wp_url", DataType::Utf8View, true),
192        Field::new("wp_type", DataType::Utf8View, true),
193        Field::new("wp_char_count", DataType::Int32, true),
194        Field::new("wp_link_count", DataType::Int32, true),
195        Field::new("wp_image_count", DataType::Int32, true),
196        Field::new("wp_max_ad_count", DataType::Int32, true),
197    ]))
198}
tpcdsgen_arrow/tables/web_page.rs

tpcdsgen_arrow/tables/
web_page.rs
