From b08490eb60e9bdde5193c142005190d3ea1abb63 Mon Sep 17 00:00:00 2001
From: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
Date: Fri, 17 Mar 2023 17:39:03 +0000
Subject: [PATCH] Seal ArrowPrimitiveType (#3882)

* Seal ArrowPrimitiveType

* Fix doc

* Review feedback
---
 arrow-array/src/array/primitive_array.rs | 24 +----------
 arrow-array/src/types.rs                 | 55 ++++++++++++++++--------
 2 files changed, 39 insertions(+), 40 deletions(-)

diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
index 0e78083c4795..9b3b11c8215e 100644
--- a/arrow-array/src/array/primitive_array.rs
+++ b/arrow-array/src/array/primitive_array.rs
@@ -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;
@@ -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.
 ///
diff --git a/arrow-array/src/types.rs b/arrow-array/src/types.rs
index 48eee4f5c3dc..9f1965b77570 100644
--- a/arrow-array/src/types.rs
+++ b/arrow-array/src/types.rs
@@ -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::{
@@ -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)]
@@ -53,6 +80,8 @@ macro_rules! make_type {
             type Native = $native_ty;
             const DATA_TYPE: DataType = $data_ty;
         }
+
+        impl primitive::PrimitiveTypeSealed for $name {}
     };
 }
 
@@ -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 {}
 
@@ -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 {}
@@ -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),