tpchgen 3.0.0

//! [`TPCHDate`] and date handling
use std::{
    fmt::{Display, Formatter},
    sync::LazyLock,
};

/// The value of 1970-01-01 in the date generator system
pub const GENERATED_DATE_EPOCH_OFFSET: i32 = 83966;
/// The minimum date that can be generated by the data generator,
/// corresponding to `1992-01-01`.
///
/// The maximum date that can be generated by the data generator is this value
/// plus [`TOTAL_DATE_RANGE`].
pub const MIN_GENERATE_DATE: i32 = 92001;
const CURRENT_DATE: i32 = 95168;
/// The total number of days that can be generated by the data generator
pub const TOTAL_DATE_RANGE: i32 = 2557;

/// Month boundaries for a standard (non-leap) year
const MONTH_YEAR_DAY_START: [i32; 13] =
    [0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334, 365];

/// Lookup table for converting dates from the data generator to strings
///
/// index: dates generated by the data generator
/// value: strings like 1992-01-01
static DATE_TO_STRING: LazyLock<Vec<String>> = LazyLock::new(make_date_string_index);

/// Lookup table for julian date format used to check if a given date is in
/// the past as part of the lineitem generation.
static JULIAN_DATE: LazyLock<Vec<i32>> = LazyLock::new(|| {
    (0..TOTAL_DATE_RANGE)
        .map(|date| julian(date + MIN_GENERATE_DATE))
        .collect()
});

pub struct GenerateUtils;

impl GenerateUtils {
    /// Calculates row count for a specific part of the data
    pub fn calculate_row_count(
        scale_base: i32,
        scale_factor: f64,
        part: i32,
        part_count: i32,
    ) -> i64 {
        let total_row_count = (scale_base as f64 * scale_factor) as i64;
        let mut row_count = total_row_count / part_count as i64;

        if part == part_count {
            // for the last part, add the remainder rows
            row_count += total_row_count % part_count as i64;
        }

        row_count
    }

    /// Calculates start index for a specific part of the data
    pub fn calculate_start_index(
        scale_base: i32,
        scale_factor: f64,
        part: i32,
        part_count: i32,
    ) -> i64 {
        let total_row_count = (scale_base as f64 * scale_factor) as i64;
        let rows_per_part = total_row_count / part_count as i64;
        rows_per_part * (part as i64 - 1)
    }
}

/// Represents a date (day/year)
///
/// Example display: 1992-01-01
///
/// The date is stored as an index from the [`MIN_GENERATE_DATE`]
///
/// # Example
/// ```
/// # use tpchgen::dates::{TPCHDate, MIN_GENERATE_DATE};
/// let date = TPCHDate::new(MIN_GENERATE_DATE + 41);
/// // Convert the date to y/m/d fields
/// assert_eq!((92,2,11), date.to_ymd());
/// // format as a string using the Display impl
/// assert_eq!("1992-02-11", date.to_string());
/// ```
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct TPCHDate {
    /// date index (0 based) from MIN_GENERATE_DATE
    date_index: i32,
}

impl Display for TPCHDate {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        // uses a pre-computed table to avoid recalculating the date
        write!(f, "{}", &DATE_TO_STRING[self.date_index as usize])
    }
}

impl TPCHDate {
    /// Number of days that must be added to a TPCH date to get a Unix epoch
    /// relative date.
    ///
    /// * Unix epoch relative dates are days since the epoch (1970-01-01)
    /// * [`TPCHDate`]s are days since [`MIN_GENERATE_DATE`] (1992-01-01)
    ///
    /// This value is `8035` because `1992-01-01` is `8035` days after `1970-01-01`
    pub const UNIX_EPOCH_OFFSET: i32 = 8035;

    /// Create a new TPCHDate from a generated date
    pub fn new(generated_date: i32) -> Self {
        Self {
            date_index: generated_date - MIN_GENERATE_DATE,
        }
    }

    /// Returns the (year, month, day) of this date
    ///
    /// For a date that represents `1992-02-03`, this function will return
    ///
    /// * year: `92`
    /// * month: `2`
    /// * day: `3`
    pub fn to_ymd(&self) -> (i32, i32, i32) {
        to_ymd(self.date_index + 1)
    }

    /// Return the inner date index
    pub fn into_inner(self) -> i32 {
        self.date_index
    }

    /// Checks if a date is in the past
    pub fn is_in_past(date: i32) -> bool {
        Self::to_julian(date) <= CURRENT_DATE
    }

    /// Lookup if a date is in the past.
    fn to_julian(date: i32) -> i32 {
        JULIAN_DATE[(date - MIN_GENERATE_DATE) as usize]
    }

    /// Returns the number of days since the Unix epoch this date
    /// represents.
    #[inline(always)]
    pub fn to_unix_epoch(&self) -> i32 {
        self.date_index + Self::UNIX_EPOCH_OFFSET
    }
}

/// Creates a index table of formatted strings
///
/// index: dates generated by the data generator
/// value: strings like 1992-01-01
fn make_date_string_index() -> Vec<String> {
    (0..TOTAL_DATE_RANGE)
        .map(|i| {
            let (y, m, dy) = to_ymd(i + 1);
            format_ymd(y, m, dy)
        })
        .collect()
}

/// Returns the (year, month, day) for the specified date index
///
/// See [`TPCHDate::to_ymd`] for more information
fn to_ymd(index: i32) -> (i32, i32, i32) {
    let y = julian(index + MIN_GENERATE_DATE - 1) / 1000;
    let d = julian(index + MIN_GENERATE_DATE - 1) % 1000;

    let mut m = 0;
    while d > MONTH_YEAR_DAY_START[m as usize] + leap_year_adjustment(y, m) {
        m += 1;
    }

    let dy =
        d - MONTH_YEAR_DAY_START[(m - 1) as usize] - if is_leap_year(y) && m > 2 { 1 } else { 0 };

    (y, m, dy)
}

/// Formants the specified y, m, d as `yyyy-mm-dd`
fn format_ymd(y: i32, m: i32, dy: i32) -> String {
    format!("{:04}-{:02}-{:02}", 1900 + y, m, dy)
}

/// Helpers duplicated to avoid circular references
const fn julian(date: i32) -> i32 {
    let mut offset = date - MIN_GENERATE_DATE;
    let mut result = MIN_GENERATE_DATE;

    loop {
        let year = result / 1000;
        let year_end = year * 1000 + 365 + if is_leap_year(year) { 1 } else { 0 };

        if result + offset <= year_end {
            break;
        }

        offset -= year_end - result + 1;
        result += 1000;
    }

    result + offset
}

const fn is_leap_year(year: i32) -> bool {
    year % 4 == 0 && year % 100 != 0
}

const fn leap_year_adjustment(year: i32, month: i32) -> i32 {
    if is_leap_year(year) && month >= 2 {
        1
    } else {
        0
    }
}

#[cfg(test)]
mod test {
    use super::*;
    #[test]
    fn test_date_strings() {
        let date = TPCHDate::new(MIN_GENERATE_DATE + 1);
        assert_eq!(date.to_string(), "1992-01-02");

        let date = TPCHDate::new(MIN_GENERATE_DATE + 1234);
        assert_eq!(date.to_string(), "1995-05-19");

        let date = TPCHDate::new(MIN_GENERATE_DATE + TOTAL_DATE_RANGE - 1);
        assert_eq!(date.to_string(), "1998-12-31");
    }

    #[test]
    fn test_display_dates() {
        for index in [1, 23, 321, 623, 1234, 2345, 2556] {
            let date = TPCHDate::new(MIN_GENERATE_DATE + index);
            let (y, m, dy) = date.to_ymd();
            assert_eq!(format_ymd(y, m, dy), date.to_string());
        }
    }

    #[test]
    fn test_date_epoch_consistency() {
        // Check that dates are actually machine some epochs.
        let date = TPCHDate::new(MIN_GENERATE_DATE + 1);
        assert_eq!(date.to_unix_epoch(), 8036);

        let date = TPCHDate::new(MIN_GENERATE_DATE + 1234);
        // 1995-05-19 00:00:00 (12:00:00 AM)
        assert_eq!(date.to_string(), "1995-05-19");
        assert_eq!(date.to_unix_epoch(), 9269);
    }
}