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

tpcdsgen_arrow/tables/
ship_mode.rs
