This is an automated email from the ASF dual-hosted git repository.

jiayu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-spatialbench.git

commit adc9774285bb21bd78a11000b5eb9b0decc4ae40
Author: Pranav Toggi <[email protected]>
AuthorDate: Thu Jun 19 18:29:02 2025 -0700

    Add Trip skeleton
---
 tpchgen-arrow/src/lineitem.rs |   8 +-
 tpchgen-arrow/src/trip.rs     | 114 ++++++++++++
 tpchgen/src/generators.rs     | 396 ++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 502 insertions(+), 16 deletions(-)

diff --git a/tpchgen-arrow/src/lineitem.rs b/tpchgen-arrow/src/lineitem.rs
index 17470ba..df597bb 100644
--- a/tpchgen-arrow/src/lineitem.rs
+++ b/tpchgen-arrow/src/lineitem.rs
@@ -29,7 +29,7 @@ use tpchgen::generators::{LineItemGenerator, 
LineItemGeneratorIterator};
 /// let lines = formatted_batches.lines().collect::<Vec<_>>();
 /// assert_eq!(lines, vec![
 ///   
"+------------+-----------+-----------+--------------+------------+-----------------+------------+-------+--------------+--------------+------------+--------------+---------------+-------------------+------------+-------------------------------------+",
-///   "| l_orderkey | l_partkey | l_suppkey | l_linenumber | l_quantity | 
l_extendedprice | l_discount | l_tax | l_returnflag | l_linestatus | l_shipdate 
| l_commitdate | l_receiptdate | l_shipinstruct    | l_shipmode | l_comment     
                      |",
+///   "| l_orderkey | l_vehiclekey | l_suppkey | l_linenumber | l_quantity | 
l_extendedprice | l_discount | l_tax | l_returnflag | l_linestatus | l_shipdate 
| l_commitdate | l_receiptdate | l_shipinstruct    | l_shipmode | l_comment     
                      |",
 ///   
"+------------+-----------+-----------+--------------+------------+-----------------+------------+-------+--------------+--------------+------------+--------------+---------------+-------------------+------------+-------------------------------------+",
 ///   "| 1          | 155190    | 7706      | 1            | 17.00      | 
21168.23        | 0.04       | 0.02  | N            | O            | 1996-03-13 
| 1996-02-12   | 1996-03-22    | DELIVER IN PERSON | TRUCK      | egular courts 
above the             |",
 ///   "| 1          | 67310     | 7311      | 2            | 36.00      | 
45983.16        | 0.09       | 0.06  | N            | O            | 1996-04-12 
| 1996-02-28   | 1996-04-20    | TAKE BACK RETURN  | MAIL       | ly final 
dependencies: slyly bold   |",
@@ -88,7 +88,7 @@ impl Iterator for LineItemArrow {
 
         // Convert column by column
         let l_orderkey = Int64Array::from_iter_values(rows.iter().map(|row| 
row.l_orderkey));
-        let l_partkey = Int64Array::from_iter_values(rows.iter().map(|row| 
row.l_partkey));
+        let l_vehiclekey = Int64Array::from_iter_values(rows.iter().map(|row| 
row.l_vehiclekey));
         let l_suppkey = Int64Array::from_iter_values(rows.iter().map(|row| 
row.l_suppkey));
         let l_linenumber = Int32Array::from_iter_values(rows.iter().map(|row| 
row.l_linenumber));
         let l_quantity = 
Decimal128Array::from_iter_values(rows.iter().map(|row| {
@@ -126,7 +126,7 @@ impl Iterator for LineItemArrow {
             Arc::clone(self.schema()),
             vec![
                 Arc::new(l_orderkey),
-                Arc::new(l_partkey),
+                Arc::new(l_vehiclekey),
                 Arc::new(l_suppkey),
                 Arc::new(l_linenumber),
                 Arc::new(l_quantity),
@@ -155,7 +155,7 @@ static LINEITEM_SCHEMA: LazyLock<SchemaRef> = 
LazyLock::new(make_lineitem_schema
 fn make_lineitem_schema() -> SchemaRef {
     Arc::new(Schema::new(vec![
         Field::new("l_orderkey", DataType::Int64, false),
-        Field::new("l_partkey", DataType::Int64, false),
+        Field::new("l_vehiclekey", DataType::Int64, false),
         Field::new("l_suppkey", DataType::Int64, false),
         Field::new("l_linenumber", DataType::Int32, false),
         Field::new("l_quantity", DataType::Decimal128(15, 2), false),
diff --git a/tpchgen-arrow/src/trip.rs b/tpchgen-arrow/src/trip.rs
new file mode 100644
index 0000000..c7eee6e
--- /dev/null
+++ b/tpchgen-arrow/src/trip.rs
@@ -0,0 +1,114 @@
+use crate::conversions::{decimal128_array_from_iter, to_arrow_date32};
+use crate::{DEFAULT_BATCH_SIZE, RecordBatchIterator};
+use arrow::array::{Date32Array, Decimal128Array, Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use std::sync::{Arc, LazyLock};
+use tpchgen::generators::{TripGenerator, TripGeneratorIterator};
+
+/// Generate [`Trip`]s in [`RecordBatch`] format
+///
+/// [`Trip`]: tpchgen::generators::Trip
+///
+/// # Example
+/// ```
+/// # use tpchgen::generators::TripGenerator;
+/// # use tpchgen_arrow::TripArrow;
+///
+/// // Create a SF=1.0 generator and wrap it in an Arrow generator
+/// let generator = TripGenerator::new(1.0, 1, 1);
+/// let mut arrow_generator = TripArrow::new(generator)
+///   .with_batch_size(10);
+/// // Read the first batch
+/// let batch = arrow_generator.next().unwrap();
+/// ```
+pub struct TripArrow {
+    inner: TripGeneratorIterator,
+    batch_size: usize,
+}
+
+impl TripArrow {
+    pub fn new(generator: TripGenerator<'static>) -> Self {
+        Self {
+            inner: generator.iter(),
+            batch_size: DEFAULT_BATCH_SIZE,
+        }
+    }
+
+    /// Set the batch size
+    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
+        self.batch_size = batch_size;
+        self
+    }
+}
+
+impl RecordBatchIterator for TripArrow {
+    fn schema(&self) -> &SchemaRef {
+        &TRIP_SCHEMA
+    }
+}
+
+impl Iterator for TripArrow {
+    type Item = RecordBatch;
+
+    /// Generate the next batch of data, if there is one
+    fn next(&mut self) -> Option<Self::Item> {
+        // Get next rows to convert
+        let rows: Vec<_> = self.inner.by_ref().take(self.batch_size).collect();
+        if rows.is_empty() {
+            return None;
+        }
+
+        // Convert column by column
+        let t_tripkey = Int64Array::from_iter_values(rows.iter().map(|row| 
row.t_tripkey));
+        let t_custkey = Int64Array::from_iter_values(rows.iter().map(|row| 
row.t_custkey));
+        let t_driverkey = Int64Array::from_iter_values(rows.iter().map(|row| 
row.t_driverkey));
+        let t_vehiclekey = Int64Array::from_iter_values(rows.iter().map(|row| 
row.t_vehiclekey));
+        let t_pickuptime = Date32Array::from_iter_values(
+            rows.iter().map(|row| row.t_pickuptime).map(to_arrow_date32),
+        );
+        let t_dropofftime = Date32Array::from_iter_values(
+            rows.iter().map(|row| row.t_dropofftime).map(to_arrow_date32),
+        );
+        let t_fare = decimal128_array_from_iter(rows.iter().map(|row| 
row.t_fare));
+        let t_tip = decimal128_array_from_iter(rows.iter().map(|row| 
row.t_tip));
+        let t_totalamount = decimal128_array_from_iter(rows.iter().map(|row| 
row.t_totalamount));
+        let t_distance = decimal128_array_from_iter(rows.iter().map(|row| 
row.t_distance));
+
+        let batch = RecordBatch::try_new(
+            Arc::clone(self.schema()),
+            vec![
+                Arc::new(t_tripkey),
+                Arc::new(t_custkey),
+                Arc::new(t_driverkey),
+                Arc::new(t_vehiclekey),
+                Arc::new(t_pickuptime),
+                Arc::new(t_dropofftime),
+                Arc::new(t_fare),
+                Arc::new(t_tip),
+                Arc::new(t_totalamount),
+                Arc::new(t_distance),
+            ],
+        )
+            .unwrap();
+
+        Some(batch)
+    }
+}
+
+/// Schema for the Trip table
+static TRIP_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(make_trip_schema);
+
+fn make_trip_schema() -> SchemaRef {
+    Arc::new(Schema::new(vec![
+        Field::new("t_tripkey", DataType::Int64, false),
+        Field::new("t_custkey", DataType::Int64, false),
+        Field::new("t_driverkey", DataType::Int64, false),
+        Field::new("t_vehiclekey", DataType::Int64, false),
+        Field::new("t_pickuptime", DataType::Date32, false),
+        Field::new("t_dropofftime", DataType::Date32, false),
+        Field::new("t_fare", DataType::Decimal128(15, 2), false),
+        Field::new("t_tip", DataType::Decimal128(15, 2), false),
+        Field::new("t_totalamount", DataType::Decimal128(15, 2), false),
+        Field::new("t_distance", DataType::Decimal128(15, 2), false),
+    ]))
+}
\ No newline at end of file
diff --git a/tpchgen/src/generators.rs b/tpchgen/src/generators.rs
index c9efbff..2af92ee 100644
--- a/tpchgen/src/generators.rs
+++ b/tpchgen/src/generators.rs
@@ -340,7 +340,7 @@ impl fmt::Display for VehicleBrandName {
     }
 }
 
-/// The PART table
+/// The VEHICLE table
 ///
 /// The Display trait is implemented to format the line item data as a string
 /// in the default TPC-H 'tbl' format.
@@ -652,6 +652,9 @@ impl<'a> DriverGenerator<'a> {
     /// Base scale for Driver generation
     const SCALE_BASE: i32 = 10_000;
 
+    /// Base scale for vehicle-driver generation
+    const DRIVERS_PER_VEHICLE: i32 = 4;
+
     // Constants for Driver generation
     const ACCOUNT_BALANCE_MIN: i32 = -99999;
     const ACCOUNT_BALANCE_MAX: i32 = 999999;
@@ -828,14 +831,16 @@ impl<'a> DriverGeneratorIterator<'a> {
     }
 
     /// Selects a driver for a vehicle, with drivers table 5x the size of 
vehicles table
-    pub fn select_driver(vehicle_key: i64, scale_factor: f64) -> i64 {
-        // Calculate driver count as 5 times the vehicle count
-        let driver_count = 5 * (VehicleGenerator::SCALE_BASE as f64 * 
scale_factor) as i64;
+    pub fn select_driver(vehicle_key: i64, driver_number: i64, scale_factor: 
f64) -> i64 {
+        // Use supplier generator's scale base
+        let driver_count = (VehicleGenerator::SCALE_BASE as f64 * 
scale_factor) as i64;
 
-        // Map each vehicle to a specific driver
-        // Using the formula that ensures the driver key is within valid range
-        // and maintains a one-to-one relationship between vehicles and drivers
-        ((vehicle_key - 1) % driver_count) + 1
+        ((vehicle_key
+            + (driver_number
+            * ((driver_count / DriverGenerator::DRIVERS_PER_VEHICLE as i64)
+            + ((vehicle_key - 1) / driver_count))))
+            % driver_count)
+            + 1
     }
 }
 
@@ -1502,7 +1507,7 @@ impl<'a> Iterator for OrderGeneratorIterator<'a> {
 pub struct LineItem<'a> {
     /// Foreign key to ORDERS
     pub l_orderkey: i64,
-    /// Foreign key to PART
+    /// Foreign key to VEHICLE
     pub l_vehiclekey: i64,
     /// Foreign key to Driver
     pub l_suppkey: i64,
@@ -1578,8 +1583,7 @@ impl<'a> LineItemGenerator<'a> {
     const TAX_MAX: TPCHDecimal = TPCHDecimal(8); // 0.08
     const DISCOUNT_MIN: TPCHDecimal = TPCHDecimal(0); // 0.00
     const DISCOUNT_MAX: TPCHDecimal = TPCHDecimal(10); // 0.10
-    const PART_KEY_MIN: i32 = 1;
-
+    const VEHICLE_KEY_MIN: i32 = 1;
     const SHIP_DATE_MIN: i32 = 1;
     const SHIP_DATE_MAX: i32 = 121;
     const COMMIT_DATE_MIN: i32 = 30;
@@ -1680,7 +1684,7 @@ impl<'a> LineItemGenerator<'a> {
         RandomBoundedLong::new_with_seeds_per_row(
             1808217256,
             scale_factor >= 30000.0,
-            Self::PART_KEY_MIN as i64,
+            Self::VEHICLE_KEY_MIN as i64,
             (VehicleGenerator::SCALE_BASE as f64 * scale_factor) as i64,
             OrderGenerator::LINE_COUNT_MAX,
         )
@@ -1864,6 +1868,7 @@ impl<'a> LineItemGeneratorIterator<'a> {
         // let driver_number = self.driver_number_random.next_value() as i64;
         let driver_key = DriverGeneratorIterator::select_driver(
             vehicle_key,
+            self.line_number as i64,
             self.scale_factor,
         );
 
@@ -1959,6 +1964,325 @@ impl<'a> Iterator for LineItemGeneratorIterator<'a> {
     }
 }
 
+/// The TRIP table (fact table)
+///
+/// The Display trait is implemented to format the trip data as a string
+/// in the default TPC-H 'tbl' format.
+///
+/// ```text
+/// 1|150|342|78|2023-04-12 08:30:15|2023-04-12 09:15:42|25.50|4.50|30.00|12.7|
+/// 2|43|129|156|2023-04-12 10:05:22|2023-04-12 10:32:18|18.75|3.25|22.00|8.3|
+/// ```
+#[derive(Debug, Clone, PartialEq)]
+pub struct Trip {
+    /// Primary key
+    pub t_tripkey: i64,
+    /// Foreign key to CUSTOMER
+    pub t_custkey: i64,
+    /// Foreign key to DRIVER
+    pub t_driverkey: i64,
+    /// Foreign key to VEHICLE
+    pub t_vehiclekey: i64,
+    /// Pickup time
+    pub t_pickuptime: TPCHDate,
+    /// Dropoff time
+    pub t_dropofftime: TPCHDate,
+    /// Trip fare amount
+    pub t_fare: TPCHDecimal,
+    /// Trip tip amount
+    pub t_tip: TPCHDecimal,
+    /// Total amount
+    pub t_totalamount: TPCHDecimal,
+    /// Trip distance
+    pub t_distance: TPCHDecimal,
+}
+
+impl fmt::Display for Trip {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "{}|{}|{}|{}|{}|{}|{}|{}|{}|{}|",
+            self.t_tripkey,
+            self.t_custkey,
+            self.t_driverkey,
+            self.t_vehiclekey,
+            self.t_pickuptime,
+            self.t_dropofftime,
+            self.t_fare,
+            self.t_tip,
+            self.t_totalamount,
+            self.t_distance
+        )
+    }
+}
+
+/// Generator for Trip table data
+#[derive(Debug, Clone)]
+pub struct TripGenerator<'a> {
+    scale_factor: f64,
+    vehicle: i32,
+    vehicle_count: i32,
+    distributions: &'a Distributions,
+    text_pool: &'a TextPool,
+}
+
+impl<'a> TripGenerator<'a> {
+    /// Base scale for trip generation
+    const SCALE_BASE: i32 = 1_500_000;
+
+    // Constants for trip generation
+    const DISTANCE_MIN: i32 = 1;   // 1.0 miles
+    const DISTANCE_MAX: i32 = 500; // 50.0 miles
+    const FARE_MIN_PER_MILE: i32 = 150; // $1.50 per mile
+    const FARE_MAX_PER_MILE: i32 = 300; // $3.00 per mile
+    const TIP_PERCENT_MIN: i32 = 0;     // 0% tip
+    const TIP_PERCENT_MAX: i32 = 30;    // 30% tip
+    const TRIP_DURATION_MIN_MINUTES: i32 = 5;  // min duration 5 minutes
+    const TRIP_DURATION_MAX_PER_MILE: i32 = 3; // max 3 minutes per mile
+
+    /// Creates a new TripGenerator with the given scale factor
+    pub fn new(scale_factor: f64, vehicle: i32, vehicle_count: i32) -> 
TripGenerator<'static> {
+        Self::new_with_distributions_and_text_pool(
+            scale_factor,
+            vehicle,
+            vehicle_count,
+            Distributions::static_default(),
+            TextPool::get_or_init_default(),
+        )
+    }
+
+    /// Creates a TripGenerator with specified distributions and text pool
+    pub fn new_with_distributions_and_text_pool<'b>(
+        scale_factor: f64,
+        vehicle: i32,
+        vehicle_count: i32,
+        distributions: &'b Distributions,
+        text_pool: &'b TextPool,
+    ) -> TripGenerator<'b> {
+        TripGenerator {
+            scale_factor,
+            vehicle,
+            vehicle_count,
+            distributions,
+            text_pool,
+        }
+    }
+
+    /// Return the row count for the given scale factor and generator vehicle 
count
+    // pub fn calculate_row_count(scale_factor: f64, vehicle: i32, 
vehicle_count: i32) -> i64 {
+    //     GenerateUtils::calculate_row_count(Self::SCALE_BASE, scale_factor, 
vehicle, vehicle_count)
+    // }
+
+    /// Returns an iterator over the trip rows
+    pub fn iter(&self) -> TripGeneratorIterator {
+        TripGeneratorIterator::new(
+            self.distributions,
+            self.text_pool,
+            self.scale_factor,
+            GenerateUtils::calculate_start_index(
+                Self::SCALE_BASE,
+                self.scale_factor,
+                self.vehicle,
+                self.vehicle_count,
+            ),
+            GenerateUtils::calculate_row_count(
+                Self::SCALE_BASE,
+                self.scale_factor,
+                self.vehicle,
+                self.vehicle_count,
+            ),
+        )
+    }
+}
+
+impl<'a> IntoIterator for &'a TripGenerator<'a> {
+    type Item = Trip;
+    type IntoIter = TripGeneratorIterator;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
+/// Iterator that generates Trip rows
+#[derive(Debug)]
+pub struct TripGeneratorIterator {
+    customer_key_random: RandomBoundedLong,
+    driver_key_random: RandomBoundedLong,
+    vehicle_key_random: RandomBoundedLong,
+    pickup_date_random: RandomBoundedInt,
+    distance_random: RandomBoundedInt,
+    fare_per_mile_random: RandomBoundedInt,
+    tip_percent_random: RandomBoundedInt,
+    trip_minutes_per_mile_random: RandomBoundedInt,
+
+    scale_factor: f64,
+    start_index: i64,
+    row_count: i64,
+    max_customer_key: i64,
+
+    index: i64,
+    trip_number: i64,
+}
+
+impl TripGeneratorIterator {
+    fn new(
+        _distributions: &Distributions,
+        _text_pool: &TextPool,
+        scale_factor: f64,
+        start_index: i64,
+        row_count: i64,
+    ) -> Self {
+        // Create all the randomizers
+        let max_customer_key = (CustomerGenerator::SCALE_BASE as f64 * 
scale_factor) as i64;
+        let max_driver_key = (DriverGenerator::SCALE_BASE as f64 * 
scale_factor) as i64;
+        let max_vehicle_key = (VehicleGenerator::SCALE_BASE as f64 * 
scale_factor) as i64;
+
+        let mut customer_key_random = RandomBoundedLong::new(921591341, 
scale_factor >= 30000.0, 1, max_customer_key);
+        let mut driver_key_random = RandomBoundedLong::new(572982913, 
scale_factor >= 30000.0, 1, max_driver_key);
+        let mut vehicle_key_random = RandomBoundedLong::new(135497281, 
scale_factor >= 30000.0, 1, max_vehicle_key);
+
+        let mut pickup_date_random = RandomBoundedInt::new(
+            831649288,
+            dates::MIN_GENERATE_DATE,
+            dates::MIN_GENERATE_DATE + dates::TOTAL_DATE_RANGE - 
TripGenerator::TRIP_DURATION_MAX_PER_MILE * TripGenerator::DISTANCE_MAX / 60 / 
24
+        );
+
+        let mut distance_random = RandomBoundedInt::new(
+            692134278,
+            TripGenerator::DISTANCE_MIN,
+            TripGenerator::DISTANCE_MAX
+        );
+
+        let mut fare_per_mile_random = RandomBoundedInt::new(
+            109837462,
+            TripGenerator::FARE_MIN_PER_MILE,
+            TripGenerator::FARE_MAX_PER_MILE
+        );
+
+        let mut tip_percent_random = RandomBoundedInt::new(
+            483912756,
+            TripGenerator::TIP_PERCENT_MIN,
+            TripGenerator::TIP_PERCENT_MAX
+        );
+
+        let mut trip_minutes_per_mile_random = RandomBoundedInt::new(
+            748219567,
+            1,
+            TripGenerator::TRIP_DURATION_MAX_PER_MILE
+        );
+
+        // Advance all generators to the starting position
+        customer_key_random.advance_rows(start_index);
+        driver_key_random.advance_rows(start_index);
+        vehicle_key_random.advance_rows(start_index);
+        pickup_date_random.advance_rows(start_index);
+        distance_random.advance_rows(start_index);
+        fare_per_mile_random.advance_rows(start_index);
+        tip_percent_random.advance_rows(start_index);
+        trip_minutes_per_mile_random.advance_rows(start_index);
+
+        TripGeneratorIterator {
+            customer_key_random,
+            driver_key_random,
+            vehicle_key_random,
+            pickup_date_random,
+            distance_random,
+            fare_per_mile_random,
+            tip_percent_random,
+            trip_minutes_per_mile_random,
+            scale_factor,
+            start_index,
+            row_count,
+            max_customer_key,
+            index: 0,
+            trip_number: 0,
+        }
+    }
+
+    /// Creates a trip with the given key
+    fn make_trip(&mut self, trip_key: i64) -> Trip {
+
+        // generate customer key, taking into account customer mortality rate
+        let mut customer_key = self.customer_key_random.next_value();
+        let mut delta = 1;
+        while customer_key % OrderGenerator::CUSTOMER_MORTALITY as i64 == 0 {
+            customer_key += delta;
+            customer_key = customer_key.min(self.max_customer_key);
+            delta *= -1;
+        }
+
+        let vehicle_key = self.vehicle_key_random.next_value();
+        let driver_key = DriverGeneratorIterator::select_driver(
+            vehicle_key,
+            self.trip_number,
+            self.scale_factor,
+        );
+
+        let pickup_date_value = self.pickup_date_random.next_value();
+        let pickup_date = TPCHDate::new(pickup_date_value);
+
+        let distance_value = self.distance_random.next_value();
+        let distance = TPCHDecimal((distance_value * 10) as i64); // Convert 
to i64
+
+        let fare_per_mile = self.fare_per_mile_random.next_value();
+        let fare_value = (distance_value * fare_per_mile) / 100;
+        let fare = TPCHDecimal((fare_value * 100) as i64); // Convert to i64
+
+        let tip_percent = self.tip_percent_random.next_value();
+        let tip_value = (fare_value * tip_percent) / 100;
+        let tip = TPCHDecimal((tip_value * 100) as i64); // Convert to i64
+
+        let total_value = fare_value + tip_value;
+        let total = TPCHDecimal((total_value * 100) as i64); // Convert to i64
+
+        // Calculate trip duration in minutes
+        let minutes_per_mile = self.trip_minutes_per_mile_random.next_value();
+        let duration_minutes = TripGenerator::TRIP_DURATION_MIN_MINUTES + 
(distance_value * minutes_per_mile);
+        let dropoff_date_value = pickup_date_value + ((duration_minutes as 
f64) / (24.0 * 60.0)) as i32;
+        let dropoff_date = TPCHDate::new(dropoff_date_value);
+
+        Trip {
+            t_tripkey: trip_key,
+            t_custkey: customer_key,
+            t_driverkey: driver_key,
+            t_vehiclekey: vehicle_key,
+            t_pickuptime: pickup_date,
+            t_dropofftime: dropoff_date,
+            t_fare: fare,
+            t_tip: tip,
+            t_totalamount: total,
+            t_distance: distance,
+        }
+    }
+}
+
+impl<'a> Iterator for TripGeneratorIterator {
+    type Item = Trip;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.row_count {
+            return None;
+        }
+
+        let trip = self.make_trip(self.start_index + self.index + 1);
+
+        // Mark all generators as finished with this row
+        self.customer_key_random.row_finished();
+        self.driver_key_random.row_finished();
+        self.vehicle_key_random.row_finished();
+        self.pickup_date_random.row_finished();
+        self.distance_random.row_finished();
+        self.fare_per_mile_random.row_finished();
+        self.tip_percent_random.row_finished();
+        self.trip_minutes_per_mile_random.row_finished();
+
+        self.index += 1;
+
+        Some(trip)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -2083,6 +2407,54 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_trip_generation() {
+        // Create a generator with a small scale factor
+        let generator = TripGenerator::new(0.01, 1, 1);
+        let trips: Vec<_> = generator.iter().collect();
+
+        // Should have 0.01 * 1,000,000 = 10,000 trips
+        assert_eq!(trips.len(), 15000);
+
+        // Check first trip
+        let first = &trips[0];
+        assert_eq!(first.t_tripkey, 1);
+        assert!(first.t_custkey > 0);
+        assert!(first.t_driverkey > 0);
+        assert!(first.t_vehiclekey > 0);
+
+        // Check that pickup date is before or equal to dropoff date
+        // TPCHDate doesn't have a .0 field, use date comparison instead
+        assert!(first.t_pickuptime <= first.t_dropofftime);
+
+        // Check that the financial values make sense
+        assert!(first.t_fare.0 > 0);
+        assert!(first.t_tip.0 >= 0); // Tip could be zero
+        assert_eq!(first.t_totalamount.0, first.t_fare.0 + first.t_tip.0);
+        assert!(first.t_distance.0 > 0);
+
+        // Verify the string format matches the expected pattern
+        let expected_pattern = format!(
+            "{}|{}|{}|{}|{}|{}|{}|{}|{}|{}|",
+            first.t_tripkey,
+            first.t_custkey,
+            first.t_driverkey,
+            first.t_vehiclekey,
+            first.t_pickuptime,
+            first.t_dropofftime,
+            first.t_fare,
+            first.t_tip,
+            first.t_totalamount,
+            first.t_distance
+        );
+        assert_eq!(first.to_string(), expected_pattern);
+
+        // Check first Trip
+        let first = &trips[1];
+        assert_eq!(first.t_tripkey, 2);
+        assert_eq!(first.to_string(), 
"2|851|1286|1285|1997-12-24|1997-12-24|37.00|6.00|43.00|1.40|")
+    }
+
     #[test]
     fn test_make_order_key() {
         // Test order key generation logic

Reply via email to