tpcdsgen_arrow/tables/
inventory.rs

1use crate::conversions::PrimitiveArrayRecycler;
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::{InventoryRow, InventoryRowGenerator};
10
11/// Generate [`InventoryRow`]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 InventoryArrow {
19    inner: InventoryRowGenerator,
20    batch_size: usize,
21    // reused allocations across batches
22    scratch: Vec<InventoryRow>,
23    inv_date_sk: PrimitiveArrayRecycler<Int32Type>,
24    inv_item_sk: PrimitiveArrayRecycler<Int32Type>,
25    inv_warehouse_sk: PrimitiveArrayRecycler<Int32Type>,
26    inv_quantity_on_hand: PrimitiveArrayRecycler<Int32Type>,
27}
28
29impl InventoryArrow {
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::Inventory);
37        Self {
38            inner: InventoryRowGenerator::new(session, row_count),
39            batch_size: DEFAULT_BATCH_SIZE,
40            scratch: Vec::with_capacity(DEFAULT_BATCH_SIZE),
41            inv_date_sk: PrimitiveArrayRecycler::new(),
42            inv_item_sk: PrimitiveArrayRecycler::new(),
43            inv_warehouse_sk: PrimitiveArrayRecycler::new(),
44            inv_quantity_on_hand: PrimitiveArrayRecycler::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 InventoryArrow {
72    fn schema(&self) -> SchemaRef {
73        Self::schema_ref()
74    }
75}
76
77impl Iterator for InventoryArrow {
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            inv_date_sk,
92            inv_item_sk,
93            inv_warehouse_sk,
94            inv_quantity_on_hand,
95            ..
96        } = self;
97        let nbm = InventoryRow::null_bit_map;
98        let sk32 =
99            |rec: &mut PrimitiveArrayRecycler<Int32Type>, pos, f: fn(&InventoryRow) -> i64| {
100                rec.int32_sk(rows, pos, nbm, f)
101            };
102
103        let inv_date_sk = sk32(inv_date_sk, 0, InventoryRow::get_inv_date_sk);
104        let inv_item_sk = sk32(inv_item_sk, 1, InventoryRow::get_inv_item_sk);
105        let inv_warehouse_sk = sk32(inv_warehouse_sk, 2, InventoryRow::get_inv_warehouse_sk);
106        let inv_quantity_on_hand =
107            inv_quantity_on_hand.int32(rows, 3, nbm, InventoryRow::get_inv_quantity_on_hand);
108
109        let batch = RecordBatch::try_new(
110            schema,
111            vec![
112                Arc::new(inv_date_sk),
113                Arc::new(inv_item_sk),
114                Arc::new(inv_warehouse_sk),
115                Arc::new(inv_quantity_on_hand),
116            ],
117        );
118        Some(batch)
119    }
120}
121
122static SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_schema);
123
124fn make_schema() -> SchemaRef {
125    Arc::new(Schema::new(vec![
126        Field::new("inv_date_sk", DataType::Int32, false),
127        Field::new("inv_item_sk", DataType::Int32, false),
128        Field::new("inv_warehouse_sk", DataType::Int32, false),
129        Field::new("inv_quantity_on_hand", DataType::Int32, true),
130    ]))
131}
tpcdsgen_arrow/tables/inventory.rs

tpcdsgen_arrow/tables/
inventory.rs
