tpchgen_arrow/
conversions.rs

1//! Routines to convert TPCH types to Arrow types
2
3use arrow::array::{StringViewArray, StringViewBuilder};
4use std::fmt::Write;
5use tpchgen::dates::TPCHDate;
6use tpchgen::decimal::TPCHDecimal;
7
8/// Convert a TPCHDecimal to an Arrow Decimal(15,2)
9#[inline(always)]
10pub fn to_arrow_decimal(value: TPCHDecimal) -> i128 {
11    // TPCH decimals are stored as i64 with 2 decimal places, so
12    // we can simply convert to i128 directly
13    value.into_inner() as i128
14}
15
16/// Convert a TPCH date to an Arrow Date32.
17///
18/// * Arrow `Date32` are days since the epoch (1970-01-01)
19/// * [`TPCHDate`]s are days since MIN_GENERATE_DATE (1992-01-01)
20///
21/// ```
22/// use chrono::NaiveDate;
23/// use tpchgen::dates::TPCHDate;
24/// let arrow_epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
25///  let tpch_epoch = NaiveDate::from_ymd_opt(1992, 1, 1).unwrap();
26/// // the difference between the two epochs is 8035 days
27/// let day_offset = (tpch_epoch - arrow_epoch).num_days();
28/// let day_offset: i32 = day_offset.try_into().unwrap();
29///  assert_eq!(day_offset, TPCHDate::UNIX_EPOCH_OFFSET);
30/// ```
31#[inline(always)]
32pub fn to_arrow_date32(value: TPCHDate) -> i32 {
33    value.to_unix_epoch()
34}
35
36/// Converts an iterator of TPCH decimals to an Arrow Decimal128Array
37pub fn decimal128_array_from_iter<I>(values: I) -> arrow::array::Decimal128Array
38where
39    I: Iterator<Item = TPCHDecimal>,
40{
41    let values = values.map(to_arrow_decimal);
42    arrow::array::Decimal128Array::from_iter_values(values)
43        .with_precision_and_scale(15, 2)
44        // safe to unwrap because 15,2 is within the valid range for Decimal128 (38)
45        .unwrap()
46}
47
48/// Coverts an iterator of displayable values to an Arrow StringViewArray
49///
50/// This results in an extra copy of the data, which could be avoided for some types
51pub fn string_view_array_from_display_iter<I>(values: I) -> StringViewArray
52where
53    I: Iterator<Item: std::fmt::Display>,
54{
55    let mut buffer = String::new();
56    let values = values.into_iter();
57    let size_hint = values.size_hint().0;
58    let mut builder = StringViewBuilder::with_capacity(size_hint);
59    for v in values {
60        buffer.clear();
61        write!(&mut buffer, "{v}").unwrap();
62        builder.append_value(&buffer);
63    }
64    builder.finish()
65}
66
67// test to ensure that the conversion functions are correct
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use tpchgen::dates::MIN_GENERATE_DATE;
72
73    #[test]
74    fn test_to_arrow_decimal() {
75        let value = TPCHDecimal::new(123456789);
76        assert_eq!(to_arrow_decimal(value), 123456789);
77    }
78
79    #[test]
80    fn test_to_arrow_date32() {
81        let value = TPCHDate::new(MIN_GENERATE_DATE);
82        assert_eq!(to_arrow_date32(value), 8035);
83
84        let value = TPCHDate::new(MIN_GENERATE_DATE + 100);
85        assert_eq!(to_arrow_date32(value), 8135);
86
87        let value = TPCHDate::new(MIN_GENERATE_DATE + 1234);
88        assert_eq!(to_arrow_date32(value), 9269);
89    }
90}
tpchgen_arrow/conversions.rs

tpchgen_arrow/
conversions.rs
