odbc2parquet 8.1.5

Query an ODBC data source and store the result in a Parquet file.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

use anyhow::Error;
use chrono::NaiveDate;
use odbc_api::{
    buffers::{AnySlice, BufferDesc, Item},
    sys::Date as OdbcDate,
};
use parquet::{
    basic::{LogicalType, Repetition, Type as PhysicalType},
    column::writer::{get_typed_column_writer_mut, ColumnWriter},
    data_type::Int32Type,
    schema::types::Type,
};

use crate::parquet_buffer::ParquetBuffer;

use super::column_strategy::ColumnStrategy;

pub struct Date {
    repetition: Repetition,
}

impl Date {
    pub fn new(repetition: Repetition) -> Self {
        Self { repetition }
    }
}

impl ColumnStrategy for Date {
    fn parquet_type(&self, name: &str) -> Type {
        Type::primitive_type_builder(name, PhysicalType::INT32)
            .with_repetition(self.repetition)
            .with_logical_type(Some(LogicalType::Date))
            .build()
            .unwrap()
    }

    fn buffer_desc(&self) -> BufferDesc {
        BufferDesc::Date { nullable: true }
    }

    fn copy_odbc_to_parquet(
        &self,
        parquet_buffer: &mut ParquetBuffer,
        column_writer: &mut ColumnWriter,
        column_view: AnySlice,
    ) -> Result<(), Error> {
        let it = OdbcDate::as_nullable_slice(column_view).unwrap();
        let column_writer = get_typed_column_writer_mut::<Int32Type>(column_writer);
        parquet_buffer.write_optional(column_writer, it.map(|date| date.map(days_since_epoch)))?;
        Ok(())
    }
}

/// Transform date to days since unix epoch as i32
fn days_since_epoch(date: &OdbcDate) -> i32 {
    let unix_epoch = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
    let date =
        NaiveDate::from_ymd_opt(date.year as i32, date.month as u32, date.day as u32).unwrap();
    let duration = date.signed_duration_since(unix_epoch);
    duration.num_days().try_into().unwrap()
}