use arrow_array::types::{
TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
TimestampSecondType,
};
use arrow_array::{Array, ArrayRef, RecordBatch};
use arrow_schema::{DataType, TimeUnit};
use oracledb_protocol::dpl::DirectPathColumnValue;
use oracledb_protocol::thin::{
ColumnMetadata, CS_FORM_NCHAR, ORA_TYPE_NUM_BINARY_DOUBLE, ORA_TYPE_NUM_BINARY_FLOAT,
ORA_TYPE_NUM_LONG, ORA_TYPE_NUM_NUMBER,
};
use super::{
arrow_type_name, check_convert_from_arrow, db_type_name, ArrowConversionError, Result,
};
pub fn record_batch_to_direct_path_rows(
batch: &RecordBatch,
column_metadata: &[ColumnMetadata],
) -> Result<Vec<Vec<DirectPathColumnValue>>> {
use arrow_array::cast::AsArray;
use arrow_array::types::{
Date32Type, Date64Type, Decimal128Type, Float32Type, Float64Type, Int16Type, Int32Type,
Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
};
if batch.num_columns() != column_metadata.len() {
return Err(ArrowConversionError::InvalidValue {
column_name: String::new(),
reason: format!(
"record batch has {} columns but the table has {}",
batch.num_columns(),
column_metadata.len()
),
});
}
let num_rows = batch.num_rows();
let mut rows: Vec<Vec<DirectPathColumnValue>> = (0..num_rows)
.map(|_| Vec::with_capacity(column_metadata.len()))
.collect();
for (array, column) in batch.columns().iter().zip(column_metadata) {
check_convert_from_arrow(array.data_type(), column)?;
let number_target = column.ora_type_num() == ORA_TYPE_NUM_NUMBER;
let utf16_target =
column.ora_type_num() == ORA_TYPE_NUM_LONG && column.csfrm() == CS_FORM_NCHAR;
for (row_index, row) in rows.iter_mut().enumerate() {
if array.is_null(row_index) {
row.push(DirectPathColumnValue::Null);
continue;
}
let value = match array.data_type() {
DataType::Null => DirectPathColumnValue::Null,
DataType::Int8 => DirectPathColumnValue::Number(
array
.as_primitive::<Int8Type>()
.value(row_index)
.to_string(),
),
DataType::Int16 => DirectPathColumnValue::Number(
array
.as_primitive::<Int16Type>()
.value(row_index)
.to_string(),
),
DataType::Int32 => DirectPathColumnValue::Number(
array
.as_primitive::<Int32Type>()
.value(row_index)
.to_string(),
),
DataType::Int64 => DirectPathColumnValue::Number(
array
.as_primitive::<Int64Type>()
.value(row_index)
.to_string(),
),
DataType::UInt8 => DirectPathColumnValue::Number(
array
.as_primitive::<UInt8Type>()
.value(row_index)
.to_string(),
),
DataType::UInt16 => DirectPathColumnValue::Number(
array
.as_primitive::<UInt16Type>()
.value(row_index)
.to_string(),
),
DataType::UInt32 => DirectPathColumnValue::Number(
array
.as_primitive::<UInt32Type>()
.value(row_index)
.to_string(),
),
DataType::UInt64 => DirectPathColumnValue::Number(
array
.as_primitive::<UInt64Type>()
.value(row_index)
.to_string(),
),
DataType::Decimal128(_, _) => {
let value = array
.as_primitive::<Decimal128Type>()
.value_as_string(row_index);
DirectPathColumnValue::Number(value)
}
DataType::Float64 => {
let value = array.as_primitive::<Float64Type>().value(row_index);
if number_target {
DirectPathColumnValue::Number(float_to_number_text(value))
} else if column.ora_type_num() == ORA_TYPE_NUM_BINARY_FLOAT {
DirectPathColumnValue::BinaryFloat(value as f32)
} else {
DirectPathColumnValue::BinaryDouble(value)
}
}
DataType::Float32 => {
let value = array.as_primitive::<Float32Type>().value(row_index);
if number_target {
DirectPathColumnValue::Number(float_to_number_text(f64::from(value)))
} else if column.ora_type_num() == ORA_TYPE_NUM_BINARY_DOUBLE {
DirectPathColumnValue::BinaryDouble(f64::from(value))
} else {
DirectPathColumnValue::BinaryFloat(value)
}
}
DataType::Boolean => {
DirectPathColumnValue::Boolean(array.as_boolean().value(row_index))
}
DataType::Utf8 => string_direct_path_value(
array.as_string::<i32>().value(row_index),
utf16_target,
),
DataType::LargeUtf8 => string_direct_path_value(
array.as_string::<i64>().value(row_index),
utf16_target,
),
DataType::Utf8View => {
string_direct_path_value(array.as_string_view().value(row_index), utf16_target)
}
DataType::Binary => {
bytes_direct_path_value(array.as_binary::<i32>().value(row_index))
}
DataType::LargeBinary => {
bytes_direct_path_value(array.as_binary::<i64>().value(row_index))
}
DataType::BinaryView => {
bytes_direct_path_value(array.as_binary_view().value(row_index))
}
DataType::FixedSizeBinary(_) => {
bytes_direct_path_value(array.as_fixed_size_binary().value(row_index))
}
DataType::Date32 => {
let days = array.as_primitive::<Date32Type>().value(row_index);
datetime_from_epoch(i64::from(days) * 86_400, 0)?
}
DataType::Date64 => {
let millis = array.as_primitive::<Date64Type>().value(row_index);
datetime_from_epoch(
millis.div_euclid(1_000),
u32::try_from(millis.rem_euclid(1_000)).unwrap_or(0) * 1_000_000,
)?
}
DataType::Timestamp(unit, None) => {
let raw = timestamp_raw_value(array, row_index)?;
let (seconds, nanos) = match unit {
TimeUnit::Second => (raw, 0u32),
TimeUnit::Millisecond => (
raw.div_euclid(1_000),
u32::try_from(raw.rem_euclid(1_000)).unwrap_or(0) * 1_000_000,
),
TimeUnit::Microsecond => (
raw.div_euclid(1_000_000),
u32::try_from(raw.rem_euclid(1_000_000)).unwrap_or(0) * 1_000,
),
TimeUnit::Nanosecond => (
raw.div_euclid(1_000_000_000),
u32::try_from(raw.rem_euclid(1_000_000_000)).unwrap_or(0),
),
};
datetime_from_epoch(seconds, nanos)?
}
other => {
return Err(ArrowConversionError::CannotConvertFromArrow {
arrow_type: arrow_type_name(other),
db_type: db_type_name(column),
})
}
};
row.push(value);
}
}
Ok(rows)
}
fn float_to_number_text(value: f64) -> String {
format!("{value}")
}
fn string_direct_path_value(text: &str, utf16_target: bool) -> DirectPathColumnValue {
if text.is_empty() {
return DirectPathColumnValue::Null;
}
if utf16_target {
let mut bytes = Vec::with_capacity(text.len() * 2);
for unit in text.encode_utf16() {
bytes.extend_from_slice(&unit.to_be_bytes());
}
DirectPathColumnValue::Bytes(bytes)
} else {
DirectPathColumnValue::Bytes(text.as_bytes().to_vec())
}
}
fn bytes_direct_path_value(bytes: &[u8]) -> DirectPathColumnValue {
if bytes.is_empty() {
DirectPathColumnValue::Null
} else {
DirectPathColumnValue::Bytes(bytes.to_vec())
}
}
fn timestamp_raw_value(array: &ArrayRef, row_index: usize) -> Result<i64> {
use arrow_array::cast::AsArray;
let value = match array.data_type() {
DataType::Timestamp(TimeUnit::Second, _) => {
array.as_primitive::<TimestampSecondType>().value(row_index)
}
DataType::Timestamp(TimeUnit::Millisecond, _) => array
.as_primitive::<TimestampMillisecondType>()
.value(row_index),
DataType::Timestamp(TimeUnit::Microsecond, _) => array
.as_primitive::<TimestampMicrosecondType>()
.value(row_index),
DataType::Timestamp(TimeUnit::Nanosecond, _) => array
.as_primitive::<TimestampNanosecondType>()
.value(row_index),
other => {
return Err(ArrowConversionError::CannotConvertFromArrow {
arrow_type: arrow_type_name(other),
db_type: "DB_TYPE_TIMESTAMP".to_string(),
})
}
};
Ok(value)
}
fn civil_from_days(days: i64) -> (i32, u8, u8) {
let days = days + 719_468;
let era = if days >= 0 { days } else { days - 146_096 } / 146_097;
let day_of_era = days - era * 146_097;
let year_of_era =
(day_of_era - day_of_era / 1_460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
let year = year_of_era + era * 400;
let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
let month_prime = (5 * day_of_year + 2) / 153;
let day = day_of_year - (153 * month_prime + 2) / 5 + 1;
let month = month_prime + if month_prime < 10 { 3 } else { -9 };
let year = year + i64::from(month <= 2);
(year as i32, month as u8, day as u8)
}
fn datetime_from_epoch(seconds: i64, nanos: u32) -> Result<DirectPathColumnValue> {
let days = seconds.div_euclid(86_400);
let seconds_of_day = seconds.rem_euclid(86_400);
let (year, month, day) = civil_from_days(days);
Ok(DirectPathColumnValue::DateTime {
year,
month,
day,
hour: (seconds_of_day / 3_600) as u8,
minute: ((seconds_of_day % 3_600) / 60) as u8,
second: (seconds_of_day % 60) as u8,
nanosecond: nanos,
})
}