laminar_connectors/postgres/types/

mod.rs

//! Validated Arrow-to-PostgreSQL type contract for the sink.

use arrow_schema::{DataType, TimeUnit};

use crate::error::ConnectorError;

/// SQL spellings used by COPY table DDL and typed UNNEST parameters.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct PostgresType {
    sql: &'static str,
    ddl: &'static str,
    /// UNNEST parameter type when it differs from `sql`: the Rust parameter encoder has no
    /// lossless NUMERIC, so decimals travel as exact text and the statement casts them.
    parameter: Option<&'static str>,
}

impl PostgresType {
    const fn new(sql: &'static str, ddl: &'static str) -> Self {
        Self {
            sql,
            ddl,
            parameter: None,
        }
    }

    #[must_use]
    pub(super) const fn sql(self) -> &'static str {
        self.sql
    }
}

/// Returns the single supported type mapping used by admission, DDL, COPY, and UNNEST.
///
/// The surface is intentionally the intersection of the COPY encoder and the Rust `PostgreSQL`
/// parameter encoder. Types are added only when both write paths have the same lossless contract.
pub(super) fn postgres_type(data_type: &DataType) -> Result<PostgresType, ConnectorError> {
    let mapping = match data_type {
        DataType::Boolean => PostgresType::new("bool", "BOOLEAN"),
        DataType::Int8 | DataType::UInt8 | DataType::Int16 => PostgresType::new("int2", "SMALLINT"),
        DataType::UInt16 | DataType::Int32 => PostgresType::new("int4", "INTEGER"),
        DataType::UInt32 | DataType::Int64 | DataType::UInt64 => {
            PostgresType::new("int8", "BIGINT")
        }
        DataType::Float32 => PostgresType::new("float4", "REAL"),
        DataType::Float64 => PostgresType::new("float8", "DOUBLE PRECISION"),
        DataType::Utf8 | DataType::LargeUtf8 => PostgresType::new("text", "TEXT"),
        DataType::Binary | DataType::LargeBinary => PostgresType::new("bytea", "BYTEA"),
        DataType::Date32 => PostgresType::new("date", "DATE"),
        DataType::Decimal128(_, scale) if *scale >= 0 => PostgresType {
            parameter: Some("text"),
            ..PostgresType::new("numeric", "NUMERIC")
        },
        DataType::Timestamp(
            TimeUnit::Second | TimeUnit::Millisecond | TimeUnit::Microsecond,
            None,
        ) => PostgresType::new("timestamp", "TIMESTAMP"),
        DataType::Timestamp(
            TimeUnit::Second | TimeUnit::Millisecond | TimeUnit::Microsecond,
            Some(_),
        ) => PostgresType::new("timestamptz", "TIMESTAMPTZ"),
        unsupported => {
            return Err(ConnectorError::ConfigurationError(format!(
                "PostgreSQL sink does not support Arrow type {unsupported:?}; supported types are \
                 Boolean, signed integers, UInt8/16/32, range-checked UInt64, Float32/64, \
                 Utf8/LargeUtf8, Binary/LargeBinary, Date32, Decimal128 with a non-negative \
                 scale, and second/millisecond/microsecond Timestamp"
            )));
        }
    };
    Ok(mapping)
}

/// `PostgreSQL` type used in generated CREATE TABLE DDL.
pub(super) fn arrow_to_pg_ddl_type(data_type: &DataType) -> Result<String, ConnectorError> {
    let mapping = postgres_type(data_type)?;
    Ok(match data_type {
        DataType::Decimal128(precision, scale) => format!("NUMERIC({precision}, {scale})"),
        _ => mapping.ddl.to_string(),
    })
}

/// Typed `PostgreSQL` array parameter used by UNNEST and key-delete statements.
pub(super) fn arrow_type_to_pg_array_cast(
    data_type: &DataType,
    parameter: usize,
) -> Result<String, ConnectorError> {
    let mapping = postgres_type(data_type)?;
    Ok(match mapping.parameter {
        Some(wire) => format!("${parameter}::{wire}[]::{}[]", mapping.sql),
        None => format!("${parameter}::{}[]", mapping.sql),
    })
}

/// Whether an existing NUMERIC column holds every value of `data_type` without rounding.
pub(super) fn numeric_column_holds(data_type: &DataType, type_modifier: i32) -> bool {
    let DataType::Decimal128(precision, scale) = data_type else {
        return true;
    };
    if type_modifier == -1 {
        return true;
    }
    // NUMERIC(p, s) stores ((p << 16) | (s & 0x7ff)) + 4, with s sign-extended from 11 bits.
    let packed = type_modifier - 4;
    let column_precision = (packed >> 16) & 0xffff;
    let column_scale = ((packed & 0x7ff) ^ 1024) - 1024;
    let (precision, scale) = (i32::from(*precision), i32::from(*scale));
    column_scale >= scale && column_precision - column_scale >= precision - scale
}

/// `PostgreSQL` NUMERIC binary wire value (without its length prefix) for a non-negative scale.
///
/// Digits are base-10000 groups aligned on the decimal point, with leading and trailing zero
/// groups removed and `weight` naming the power of 10000 of the first group.
#[cfg(feature = "postgres-sink")]
fn numeric_wire(value: i128, scale: u8, buffer: &mut Vec<u8>) {
    let scale = usize::from(scale);
    let mut digits = value.unsigned_abs().to_string();
    if digits.len() <= scale {
        digits.insert_str(0, &"0".repeat(scale + 1 - digits.len()));
    }
    let (integer, fraction) = digits.split_at(digits.len() - scale);
    let integer = format!("{}{integer}", "0".repeat((4 - integer.len() % 4) % 4));
    let fraction = format!("{fraction}{}", "0".repeat((4 - fraction.len() % 4) % 4));
    let group = |chunk: &[u8]| {
        chunk
            .iter()
            .fold(0_i16, |total, digit| total * 10 + i16::from(digit - b'0'))
    };
    let mut groups: Vec<i16> = integer
        .as_bytes()
        .chunks(4)
        .chain(fraction.as_bytes().chunks(4))
        .map(group)
        .collect();
    let mut weight = i16::try_from(integer.len() / 4).unwrap_or(i16::MAX) - 1;
    let leading = groups.iter().take_while(|group| **group == 0).count();
    groups.drain(..leading);
    weight -= i16::try_from(leading).unwrap_or(i16::MAX);
    while groups.last() == Some(&0) {
        groups.pop();
    }
    if groups.is_empty() {
        weight = 0;
    }
    let sign: i16 = if value < 0 && !groups.is_empty() {
        0x4000
    } else {
        0
    };
    // At most 38 significant digits: ten groups and a weight of a few groups either way.
    let header = [
        i16::try_from(groups.len()).unwrap_or(i16::MAX),
        weight,
        sign,
        i16::try_from(scale).unwrap_or(i16::MAX),
    ];
    for word in header.into_iter().chain(groups) {
        buffer.extend_from_slice(&word.to_be_bytes());
    }
}

#[cfg(feature = "postgres-sink")]
fn checked_u64(value: u64, row: usize) -> Result<i64, ConnectorError> {
    i64::try_from(value).map_err(|_| {
        ConnectorError::SchemaMismatch(format!(
            "PostgreSQL BIGINT cannot represent UInt64 value {value} at row {row}"
        ))
    })
}

#[cfg(feature = "postgres-sink")]
fn checked_date32(value: i32, row: usize) -> Result<chrono::NaiveDate, ConnectorError> {
    let epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).expect("valid Unix epoch");
    epoch
        .checked_add_signed(chrono::Duration::days(i64::from(value)))
        .ok_or_else(|| {
            ConnectorError::SchemaMismatch(format!(
                "PostgreSQL sink cannot represent Date32 value {value} at row {row}"
            ))
        })
}

/// Validates range-sensitive values before a batch is admitted to the sink buffer.
#[cfg(feature = "postgres-sink")]
pub(super) fn validate_postgres_array_values(
    array: &dyn arrow_array::Array,
) -> Result<(), ConnectorError> {
    use arrow_array::{Array as _, Date32Array, UInt64Array};

    postgres_type(array.data_type())?;
    match array.data_type() {
        DataType::UInt64 => {
            let values = array
                .as_any()
                .downcast_ref::<UInt64Array>()
                .ok_or_else(|| ConnectorError::Internal("downcast to UInt64Array failed".into()))?;
            for row in 0..values.len() {
                if !values.is_null(row) {
                    checked_u64(values.value(row), row)?;
                }
            }
        }
        DataType::Date32 => {
            let values = array
                .as_any()
                .downcast_ref::<Date32Array>()
                .ok_or_else(|| ConnectorError::Internal("downcast to Date32Array failed".into()))?;
            for row in 0..values.len() {
                if !values.is_null(row) {
                    checked_date32(values.value(row), row)?;
                }
            }
        }
        DataType::Timestamp(unit, _) => validate_timestamp_values(array, *unit)?,
        _ => {}
    }
    Ok(())
}

/// Produces the batch schema expected by COPY BINARY.
///
/// `pgpq` encodes Arrow `UInt64` as `PostgreSQL` NUMERIC. The sink deliberately exposes `UInt64` as a
/// range-checked BIGINT so COPY and UNNEST have identical table types; values are therefore widened
/// to an Int64 Arrow column after validation and before the COPY encoder is constructed.
///
/// `pgpq` 0.11 drops leading zero fractional groups from decimals (`1.00005` becomes `1.5`), so
/// decimals are pre-encoded as NUMERIC wire values in a Binary column. COPY BINARY fields carry
/// no type, so the target NUMERIC column decodes the bytes.
#[cfg(feature = "postgres-sink")]
pub(super) fn postgres_copy_batch(
    batch: &arrow_array::RecordBatch,
) -> Result<arrow_array::RecordBatch, ConnectorError> {
    use std::sync::Arc;

    use arrow_array::{Array as _, ArrayRef, Int64Array, UInt64Array};
    use arrow_schema::Schema;

    let mut changed = false;
    let mut fields = Vec::with_capacity(batch.num_columns());
    let mut columns: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
    for (field, column) in batch.schema().fields().iter().zip(batch.columns()) {
        postgres_type(column.data_type())?;
        if let DataType::Decimal128(_, scale) = field.data_type() {
            let values = column
                .as_any()
                .downcast_ref::<arrow_array::Decimal128Array>()
                .ok_or_else(|| {
                    ConnectorError::Internal("downcast to Decimal128Array failed".into())
                })?;
            let scale = u8::try_from(*scale).map_err(|_| {
                ConnectorError::Internal("negative decimal scale passed type admission".into())
            })?;
            let mut builder =
                arrow_array::builder::BinaryBuilder::with_capacity(values.len(), values.len() * 16);
            let mut wire = Vec::with_capacity(32);
            for row in 0..values.len() {
                if values.is_null(row) {
                    builder.append_null();
                } else {
                    wire.clear();
                    numeric_wire(values.value(row), scale, &mut wire);
                    builder.append_value(&wire);
                }
            }
            fields.push(Arc::new(
                field.as_ref().clone().with_data_type(DataType::Binary),
            ));
            columns.push(Arc::new(builder.finish()));
            changed = true;
        } else if field.data_type() == &DataType::UInt64 {
            let values = column
                .as_any()
                .downcast_ref::<UInt64Array>()
                .ok_or_else(|| ConnectorError::Internal("downcast to UInt64Array failed".into()))?;
            let converted = (0..values.len())
                .map(|row| {
                    if values.is_null(row) {
                        Ok(None)
                    } else {
                        checked_u64(values.value(row), row).map(Some)
                    }
                })
                .collect::<Result<Vec<Option<i64>>, ConnectorError>>()?;
            fields.push(Arc::new(
                field.as_ref().clone().with_data_type(DataType::Int64),
            ));
            columns.push(Arc::new(Int64Array::from(converted)));
            changed = true;
        } else {
            fields.push(field.clone());
            columns.push(column.clone());
        }
    }

    if changed {
        arrow_array::RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).map_err(|error| {
            ConnectorError::Internal(format!("build PostgreSQL COPY batch: {error}"))
        })
    } else {
        Ok(batch.clone())
    }
}

#[cfg(feature = "postgres-sink")]
fn validate_timestamp_values(
    array: &dyn arrow_array::Array,
    unit: TimeUnit,
) -> Result<(), ConnectorError> {
    use arrow_array::{
        Array as _, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampSecondArray,
    };

    macro_rules! validate {
        ($array_type:ty) => {{
            let values = array
                .as_any()
                .downcast_ref::<$array_type>()
                .ok_or_else(|| {
                    ConnectorError::Internal(format!(
                        "timestamp array does not match declared unit {unit:?}"
                    ))
                })?;
            for row in 0..values.len() {
                if !values.is_null(row) && to_naive_datetime(values.value(row), unit).is_none() {
                    return Err(ConnectorError::SchemaMismatch(format!(
                        "PostgreSQL sink cannot represent {unit:?} timestamp value {} at row {row}",
                        values.value(row)
                    )));
                }
            }
        }};
    }

    match unit {
        TimeUnit::Second => validate!(TimestampSecondArray),
        TimeUnit::Millisecond => validate!(TimestampMillisecondArray),
        TimeUnit::Microsecond => validate!(TimestampMicrosecondArray),
        TimeUnit::Nanosecond => unreachable!("nanosecond timestamps fail type admission"),
    }
    Ok(())
}

/// Converts an Arrow column to a `PostgreSQL` array parameter for UNNEST.
#[cfg(feature = "postgres-sink")]
pub(super) fn arrow_column_to_pg_array(
    array: &dyn arrow_array::Array,
) -> Result<Box<dyn postgres_types::ToSql + Sync + Send>, ConnectorError> {
    use arrow_array::{
        Array as _, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Float32Array,
        Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray,
        LargeStringArray, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray,
        TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
    };

    postgres_type(array.data_type())?;

    // Nulls stay `None`; every non-null value is converted by `$value(values, row)`.
    macro_rules! nullable {
        ($array_type:ty, $value:expr) => {{
            let values = array
                .as_any()
                .downcast_ref::<$array_type>()
                .ok_or_else(|| {
                    ConnectorError::Internal(format!(
                        "downcast to {} failed",
                        stringify!($array_type)
                    ))
                })?;
            Ok(Box::new(
                (0..values.len())
                    .map(|row| (!values.is_null(row)).then(|| $value(values, row)))
                    .collect::<Vec<Option<_>>>(),
            ))
        }};
    }

    macro_rules! widen {
        ($array_type:ty, $target:ty) => {{
            let values = array
                .as_any()
                .downcast_ref::<$array_type>()
                .ok_or_else(|| {
                    ConnectorError::Internal(format!(
                        "downcast to {} failed",
                        stringify!($array_type)
                    ))
                })?;
            let converted: Vec<Option<$target>> = (0..values.len())
                .map(|row| (!values.is_null(row)).then(|| <$target>::from(values.value(row))))
                .collect();
            Ok(Box::new(converted))
        }};
    }

    match array.data_type() {
        DataType::Boolean => nullable!(BooleanArray, |v: &BooleanArray, row| v.value(row)),
        DataType::Int8 => widen!(Int8Array, i16),
        DataType::UInt8 => widen!(UInt8Array, i16),
        DataType::Int16 => widen!(Int16Array, i16),
        DataType::UInt16 => widen!(UInt16Array, i32),
        DataType::Int32 => widen!(Int32Array, i32),
        DataType::UInt32 => widen!(UInt32Array, i64),
        DataType::Int64 => widen!(Int64Array, i64),
        DataType::UInt64 => {
            let values = array
                .as_any()
                .downcast_ref::<UInt64Array>()
                .ok_or_else(|| ConnectorError::Internal("downcast to UInt64Array failed".into()))?;
            let converted = (0..values.len())
                .map(|row| {
                    if values.is_null(row) {
                        Ok(None)
                    } else {
                        checked_u64(values.value(row), row).map(Some)
                    }
                })
                .collect::<Result<Vec<Option<i64>>, ConnectorError>>()?;
            Ok(Box::new(converted))
        }
        DataType::Float32 => widen!(Float32Array, f32),
        DataType::Float64 => widen!(Float64Array, f64),
        // Exact decimal text; the statement casts it to NUMERIC.
        DataType::Decimal128(_, _) => nullable!(Decimal128Array, |v: &Decimal128Array, row| v
            .value_as_string(row)),
        DataType::Utf8 => nullable!(StringArray, |v: &StringArray, row| v.value(row).to_owned()),
        DataType::LargeUtf8 => nullable!(LargeStringArray, |v: &LargeStringArray, row| v
            .value(row)
            .to_owned()),
        DataType::Binary => nullable!(BinaryArray, |v: &BinaryArray, row| v.value(row).to_vec()),
        DataType::LargeBinary => nullable!(LargeBinaryArray, |v: &LargeBinaryArray, row| v
            .value(row)
            .to_vec()),
        DataType::Date32 => {
            let values = array
                .as_any()
                .downcast_ref::<Date32Array>()
                .ok_or_else(|| ConnectorError::Internal("downcast to Date32Array failed".into()))?;
            let converted = (0..values.len())
                .map(|row| {
                    if values.is_null(row) {
                        Ok(None)
                    } else {
                        checked_date32(values.value(row), row).map(Some)
                    }
                })
                .collect::<Result<Vec<Option<chrono::NaiveDate>>, ConnectorError>>()?;
            Ok(Box::new(converted))
        }
        DataType::Timestamp(unit, timezone) => {
            macro_rules! timestamps {
                ($array_type:ty) => {{
                    let values = array
                        .as_any()
                        .downcast_ref::<$array_type>()
                        .ok_or_else(|| {
                            ConnectorError::Internal(format!(
                                "timestamp array does not match declared unit {unit:?}"
                            ))
                        })?;
                    (0..values.len())
                        .map(|row| {
                            if values.is_null(row) {
                                Ok(None)
                            } else {
                                to_naive_datetime(values.value(row), *unit)
                                    .map(Some)
                                    .ok_or_else(|| {
                                        ConnectorError::SchemaMismatch(format!(
                                            "PostgreSQL sink cannot represent {unit:?} timestamp \
                                             value {} at row {row}",
                                            values.value(row)
                                        ))
                                    })
                            }
                        })
                        .collect::<Result<Vec<Option<chrono::NaiveDateTime>>, ConnectorError>>()?
                }};
            }

            let values = match unit {
                TimeUnit::Second => timestamps!(TimestampSecondArray),
                TimeUnit::Millisecond => timestamps!(TimestampMillisecondArray),
                TimeUnit::Microsecond => timestamps!(TimestampMicrosecondArray),
                TimeUnit::Nanosecond => unreachable!("nanosecond timestamps fail type admission"),
            };
            if timezone.is_some() {
                Ok(Box::new(
                    values
                        .into_iter()
                        .map(|value| value.map(|timestamp| timestamp.and_utc()))
                        .collect::<Vec<Option<chrono::DateTime<chrono::Utc>>>>(),
                ))
            } else {
                Ok(Box::new(values))
            }
        }
        unsupported => Err(ConnectorError::ConfigurationError(format!(
            "PostgreSQL sink does not support Arrow type {unsupported:?}"
        ))),
    }
}

/// Converts a Unix timestamp using Euclidean division so negative fractional values retain their
/// correct instant (for example, -1 ms is 1969-12-31 23:59:59.999, not one second late).
#[cfg(feature = "postgres-sink")]
fn to_naive_datetime(value: i64, unit: TimeUnit) -> Option<chrono::NaiveDateTime> {
    let (units_per_second, nanos_per_unit) = match unit {
        TimeUnit::Second => (1_i64, 1_000_000_000_u32),
        TimeUnit::Millisecond => (1_000, 1_000_000),
        TimeUnit::Microsecond => (1_000_000, 1_000),
        TimeUnit::Nanosecond => (1_000_000_000, 1),
    };
    let seconds = value.div_euclid(units_per_second);
    let fraction = u32::try_from(value.rem_euclid(units_per_second)).ok()?;
    let nanos = fraction.checked_mul(nanos_per_unit)?;
    chrono::DateTime::from_timestamp(seconds, nanos).map(|timestamp| timestamp.naive_utc())
}

#[cfg(test)]
mod tests;