Skip to content

Commit

Permalink
Seal ArrowPrimitiveType (#3882)
Browse files Browse the repository at this point in the history
* Seal ArrowPrimitiveType

* Fix doc

* Review feedback
  • Loading branch information
tustvold authored Mar 17, 2023
1 parent 7f5f6b8 commit b08490e
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 40 deletions.
24 changes: 2 additions & 22 deletions arrow-array/src/array/primitive_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use crate::temporal_conversions::{
};
use crate::timezone::Tz;
use crate::trusted_len::trusted_len_unzip;
use crate::{types::*, ArrowNativeTypeOp};
use crate::types::*;
use crate::{Array, ArrayAccessor};
use arrow_buffer::{i256, ArrowNativeType, Buffer, ScalarBuffer};
use arrow_data::bit_iterator::try_for_each_valid_idx;
Expand Down Expand Up @@ -229,27 +229,7 @@ pub type Decimal128Array = PrimitiveArray<Decimal128Type>;
/// scale less or equal to 76.
pub type Decimal256Array = PrimitiveArray<Decimal256Type>;

/// Trait bridging the dynamic-typed nature of Arrow (via [`DataType`]) with the
/// static-typed nature of rust types ([`ArrowNativeType`]) for all types that implement [`ArrowNativeType`].
pub trait ArrowPrimitiveType: 'static {
/// Corresponding Rust native type for the primitive type.
type Native: ArrowNativeTypeOp;

/// the corresponding Arrow data type of this primitive type.
const DATA_TYPE: DataType;

/// Returns the byte width of this primitive type.
fn get_byte_width() -> usize {
std::mem::size_of::<Self::Native>()
}

/// Returns a default value of this primitive type.
///
/// This is useful for aggregate array ops like `sum()`, `mean()`.
fn default_value() -> Self::Native {
Default::default()
}
}
pub use crate::types::ArrowPrimitiveType;

/// Array whose elements are of primitive types.
///
Expand Down
55 changes: 37 additions & 18 deletions arrow-array/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,8 @@

//! Zero-sized types used to parameterize generic array implementations
use crate::array::ArrowPrimitiveType;
use crate::delta::shift_months;
use crate::OffsetSizeTrait;
use crate::{ArrowNativeTypeOp, OffsetSizeTrait};
use arrow_buffer::i256;
use arrow_data::decimal::{validate_decimal256_precision, validate_decimal_precision};
use arrow_schema::{
Expand All @@ -39,10 +38,38 @@ use std::ops::{Add, Sub};
pub struct BooleanType {}

impl BooleanType {
/// Type represetings is arrow [`DataType`]
/// The corresponding Arrow data type
pub const DATA_TYPE: DataType = DataType::Boolean;
}

/// Trait bridging the dynamic-typed nature of Arrow (via [`DataType`]) with the
/// static-typed nature of rust types ([`ArrowNativeType`]) for all types that implement [`ArrowNativeType`].
///
/// [`ArrowNativeType`]: arrow_buffer::ArrowNativeType
pub trait ArrowPrimitiveType: primitive::PrimitiveTypeSealed + 'static {
/// Corresponding Rust native type for the primitive type.
type Native: ArrowNativeTypeOp;

/// the corresponding Arrow data type of this primitive type.
const DATA_TYPE: DataType;

/// Returns the byte width of this primitive type.
fn get_byte_width() -> usize {
std::mem::size_of::<Self::Native>()
}

/// Returns a default value of this primitive type.
///
/// This is useful for aggregate array ops like `sum()`, `mean()`.
fn default_value() -> Self::Native {
Default::default()
}
}

mod primitive {
pub trait PrimitiveTypeSealed {}
}

macro_rules! make_type {
($name:ident, $native_ty:ty, $data_ty:expr, $doc_string: literal) => {
#[derive(Debug)]
Expand All @@ -53,6 +80,8 @@ macro_rules! make_type {
type Native = $native_ty;
const DATA_TYPE: DataType = $data_ty;
}

impl primitive::PrimitiveTypeSealed for $name {}
};
}

Expand Down Expand Up @@ -240,24 +269,10 @@ impl ArrowDictionaryKeyType for UInt32Type {}

impl ArrowDictionaryKeyType for UInt64Type {}

mod run {
use super::*;

pub trait RunEndTypeSealed {}

impl RunEndTypeSealed for Int16Type {}

impl RunEndTypeSealed for Int32Type {}

impl RunEndTypeSealed for Int64Type {}
}

/// A subtype of primitive type that is used as run-ends index
/// in `RunArray`.
/// See <https://arrow.apache.org/docs/format/Columnar.html>
///
/// Note: The implementation of this trait is sealed to avoid accidental misuse.
pub trait RunEndIndexType: ArrowPrimitiveType + run::RunEndTypeSealed {}
pub trait RunEndIndexType: ArrowPrimitiveType {}

impl RunEndIndexType for Int16Type {}

Expand Down Expand Up @@ -646,6 +661,8 @@ impl ArrowPrimitiveType for Decimal128Type {
const DATA_TYPE: DataType = <Self as DecimalType>::DEFAULT_TYPE;
}

impl primitive::PrimitiveTypeSealed for Decimal128Type {}

/// The decimal type for a Decimal256Array
#[derive(Debug)]
pub struct Decimal256Type {}
Expand Down Expand Up @@ -674,6 +691,8 @@ impl ArrowPrimitiveType for Decimal256Type {
const DATA_TYPE: DataType = <Self as DecimalType>::DEFAULT_TYPE;
}

impl primitive::PrimitiveTypeSealed for Decimal256Type {}

fn format_decimal_str(value_str: &str, precision: usize, scale: i8) -> String {
let (sign, rest) = match value_str.strip_prefix('-') {
Some(stripped) => ("-", stripped),
Expand Down

0 comments on commit b08490e

Please sign in to comment.