spatialbench_arrow/
vehicle.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::conversions::string_view_array_from_display_iter;
19use crate::{DEFAULT_BATCH_SIZE, RecordBatchIterator};
20use arrow::array::{Int64Array, RecordBatch, StringViewArray};
21use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
22use spatialbench::generators::{VehicleGenerator, VehicleGeneratorIterator};
23use std::sync::{Arc, LazyLock};
24
25/// Generate [`Vehicle`]s in [`RecordBatch`] format
26///
27/// [`Vehicle`]: spatialbench::generators::Vehicle
28///
29/// # Example
30/// ```
31/// # use spatialbench::generators::{VehicleGenerator};
32/// # use spatialbench_arrow::VehicleArrow;
33///
34/// // Create a SF=1.0 generator and wrap it in an Arrow generator
35/// let generator = VehicleGenerator::new(1.0, 1, 1);
36/// let mut arrow_generator = VehicleArrow::new(generator)
37///   .with_batch_size(10);
38/// // Read the first 10 batches
39/// let batch = arrow_generator.next().unwrap();
40/// // compare the output by pretty printing it
41/// let formatted_batches = arrow::util::pretty::pretty_format_batches(&[batch])
42///   .unwrap()
43///   .to_string();
44/// let lines = formatted_batches.lines().collect::<Vec<_>>();
45/// assert_eq!(lines, vec![
46///   "+--------------+----------------+----------+-------------------------+----------------------+",
47///   "| v_vehiclekey | v_mfgr         | v_brand  | v_type                  | v_comment            |",
48///   "+--------------+----------------+----------+-------------------------+----------------------+",
49///   "| 1            | Manufacturer#1 | Brand#13 | PROMO BURNISHED COPPER  | ly. slyly ironi      |",
50///   "| 2            | Manufacturer#1 | Brand#13 | LARGE BRUSHED BRASS     | lar accounts amo     |",
51///   "| 3            | Manufacturer#4 | Brand#42 | STANDARD POLISHED BRASS | egular deposits hag  |",
52///   "| 4            | Manufacturer#3 | Brand#34 | SMALL PLATED BRASS      | p furiously r        |",
53///   "| 5            | Manufacturer#3 | Brand#32 | STANDARD POLISHED TIN   |  wake carefully      |",
54///   "| 6            | Manufacturer#2 | Brand#24 | PROMO PLATED STEEL      | sual a               |",
55///   "| 7            | Manufacturer#1 | Brand#11 | SMALL PLATED COPPER     | lyly. ex             |",
56///   "| 8            | Manufacturer#4 | Brand#44 | PROMO BURNISHED TIN     | eposi                |",
57///   "| 9            | Manufacturer#4 | Brand#43 | SMALL BURNISHED STEEL   | ironic foxe          |",
58///   "| 10           | Manufacturer#5 | Brand#54 | LARGE BURNISHED STEEL   | ithely final deposit |",
59///   "+--------------+----------------+----------+-------------------------+----------------------+"
60/// ]);
61/// ```
62pub struct VehicleArrow {
63    inner: VehicleGeneratorIterator<'static>,
64    batch_size: usize,
65}
66
67impl VehicleArrow {
68    pub fn new(generator: VehicleGenerator<'static>) -> Self {
69        Self {
70            inner: generator.iter(),
71            batch_size: DEFAULT_BATCH_SIZE,
72        }
73    }
74
75    /// Set the batch size
76    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
77        self.batch_size = batch_size;
78        self
79    }
80}
81
82impl RecordBatchIterator for VehicleArrow {
83    fn schema(&self) -> &SchemaRef {
84        &VEHICLE_SCHEMA
85    }
86}
87
88impl Iterator for VehicleArrow {
89    type Item = RecordBatch;
90
91    fn next(&mut self) -> Option<Self::Item> {
92        // Get next rows to convert
93        let rows: Vec<_> = self.inner.by_ref().take(self.batch_size).collect();
94        if rows.is_empty() {
95            return None;
96        }
97
98        let v_vehiclekey = Int64Array::from_iter_values(rows.iter().map(|r| r.v_vehiclekey));
99        let v_mfgr = string_view_array_from_display_iter(rows.iter().map(|r| r.v_mfgr));
100        let v_brand = string_view_array_from_display_iter(rows.iter().map(|r| r.v_brand));
101        let v_type = StringViewArray::from_iter_values(rows.iter().map(|r| r.v_type));
102        let v_license = StringViewArray::from_iter_values(rows.iter().map(|r| r.v_license));
103
104        let batch = RecordBatch::try_new(
105            Arc::clone(self.schema()),
106            vec![
107                Arc::new(v_vehiclekey),
108                Arc::new(v_mfgr),
109                Arc::new(v_brand),
110                Arc::new(v_type),
111                Arc::new(v_license),
112            ],
113        )
114        .unwrap();
115        Some(batch)
116    }
117}
118
119/// Schema for the Vehicle
120static VEHICLE_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_vehicle_schema);
121fn make_vehicle_schema() -> SchemaRef {
122    Arc::new(Schema::new(vec![
123        Field::new("v_vehiclekey", DataType::Int64, false),
124        Field::new("v_mfgr", DataType::Utf8View, false),
125        Field::new("v_brand", DataType::Utf8View, false),
126        Field::new("v_type", DataType::Utf8View, false),
127        Field::new("v_comment", DataType::Utf8View, false),
128    ]))
129}
spatialbench_arrow/vehicle.rs

spatialbench_arrow/
vehicle.rs
