tpcdsgen_arrow/tables/
reason.rs

1use crate::conversions::{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::{ReasonRow, ReasonRowGenerator};
10
11/// Generate [`ReasonRow`]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 ReasonArrow {
19    inner: ReasonRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<ReasonRow>,
23    r_reason_sk: PrimitiveArrayRecycler<Int32Type>,
24    r_reason_id: StringViewArrayRecycler,
25    r_reason_desc: StringViewArrayRecycler,
26}
27
28impl ReasonArrow {
29    /// Return the schema without initializing a data generator.
30    pub fn schema_ref() -> SchemaRef {
31        Arc::clone(&SCHEMA)
32    }
33
34    pub fn new(session: Session) -> Self {
35        let row_count = session.get_scaling().get_row_count(Table::Reason);
36        Self {
37            inner: ReasonRowGenerator::new(session, row_count),
38            batch_size: DEFAULT_BATCH_SIZE,
39            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
40            r_reason_sk: PrimitiveArrayRecycler::new(),
41            r_reason_id: StringViewArrayRecycler::new(),
42            r_reason_desc: StringViewArrayRecycler::new(),
43        }
44    }
45    pub fn skip_rows_until_starting_row_number(&mut self, starting_row_number: u64) {
46        self.inner
47            .skip_rows_until_starting_row_number(starting_row_number);
48    }
49
50    /// Generate only source rows `starting_row_number..=ending_row_number`
51    /// (1-based, inclusive). The ending row number is clamped to the table's
52    /// row count.
53    pub fn with_source_row_range(
54        mut self,
55        starting_row_number: u64,
56        ending_row_number: u64,
57    ) -> Self {
58        self.inner
59            .set_source_row_range(starting_row_number, ending_row_number);
60        self
61    }
62
63    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
64        self.batch_size = batch_size;
65        self
66    }
67}
68
69impl RecordBatchReader for ReasonArrow {
70    fn schema(&self) -> SchemaRef {
71        Self::schema_ref()
72    }
73}
74
75impl Iterator for ReasonArrow {
76    type Item = Result<RecordBatch, ArrowError>;
77
78    fn next(&mut self) -> Option<Self::Item> {
79        self.scratch.clear();
80        self.scratch
81            .extend(self.inner.by_ref().take(self.batch_size));
82        let rows = &self.scratch;
83        if rows.is_empty() {
84            return None;
85        }
86
87        let schema = self.schema();
88        let Self {
89            r_reason_sk,
90            r_reason_id,
91            r_reason_desc,
92            ..
93        } = self;
94        let nbm = ReasonRow::null_bit_map;
95        let sk32 = |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&ReasonRow) -> i64| {
96            rec.int32_sk(rows, pos, nbm, f)
97        };
98        let str = |rec: &mut StringViewArrayRecycler, pos, f: fn(&ReasonRow) -> &str| {
99            rec.str(rows, pos, nbm, f)
100        };
101
102        let r_reason_sk = sk32(r_reason_sk, 0, ReasonRow::get_r_reason_sk);
103        let r_reason_id = str(r_reason_id, 1, ReasonRow::get_r_reason_id);
104        let r_reason_desc = str(r_reason_desc, 2, ReasonRow::get_r_reason_desc);
105
106        let batch = RecordBatch::try_new(
107            schema,
108            vec![
109                Arc::new(r_reason_sk),
110                Arc::new(r_reason_id),
111                Arc::new(r_reason_desc),
112            ],
113        );
114        Some(batch)
115    }
116}
117
118static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
119
120fn make_schema() -> SchemaRef {
121    Arc::new(Schema::new(vec![
122        Field::new("r_reason_sk", DataType::Int32, false),
123        Field::new("r_reason_id", DataType::Utf8View, false),
124        Field::new("r_reason_desc", DataType::Utf8View, true),
125    ]))
126}
tpcdsgen_arrow/tables/reason.rs

tpcdsgen_arrow/tables/
reason.rs
