From b75ec66edd9912927f1032f33b339ae57b351669 Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Sun, 20 Sep 2026 22:57:54 +0100 Subject: [PATCH] Add Arrow type to mid-level arrays --- src/enums/array.rs | 48 ++------------------ src/enums/collections/numeric_array.rs | 58 +++++++++++++------------ src/enums/collections/temporal_array.rs | 27 ++++++------ src/enums/collections/text_array.rs | 46 ++++++++++---------- 4 files changed, 71 insertions(+), 108 deletions(-) diff --git a/src/enums/array.rs b/src/enums/array.rs index dc1c46f..fa4da29 100644 --- a/src/enums/array.rs +++ b/src/enums/array.rs @@ -2030,52 +2030,10 @@ impl Array { /// Arrow physical type for this array. pub fn arrow_type(&self) -> ArrowType { match self { - Array::NumericArray(inner) => match inner { - #[cfg(feature = "extended_numeric_types")] - NumericArray::Int8(_) => ArrowType::Int8, - #[cfg(feature = "extended_numeric_types")] - NumericArray::Int16(_) => ArrowType::Int16, - NumericArray::Int32(_) => ArrowType::Int32, - NumericArray::Int64(_) => ArrowType::Int64, - #[cfg(feature = "extended_numeric_types")] - NumericArray::UInt8(_) => ArrowType::UInt8, - #[cfg(feature = "extended_numeric_types")] - NumericArray::UInt16(_) => ArrowType::UInt16, - NumericArray::UInt32(_) => ArrowType::UInt32, - NumericArray::UInt64(_) => ArrowType::UInt64, - NumericArray::Float32(_) => ArrowType::Float32, - NumericArray::Float64(_) => ArrowType::Float64, - #[cfg(feature = "decimal")] - NumericArray::Decimal32(a) => a.arrow_type(), - #[cfg(feature = "decimal")] - NumericArray::Decimal64(a) => a.arrow_type(), - #[cfg(feature = "decimal")] - NumericArray::Decimal128(a) => a.arrow_type(), - NumericArray::Null => ArrowType::Null, - }, - Array::TextArray(inner) => match inner { - TextArray::String32(_) => ArrowType::String, - #[cfg(feature = "large_string")] - TextArray::String64(_) => ArrowType::LargeString, - #[cfg(feature = "default_categorical_8")] - TextArray::Categorical8(_) => ArrowType::Dictionary(CategoricalIndexType::UInt8), - #[cfg(feature = "extended_categorical")] - TextArray::Categorical16(_) => ArrowType::Dictionary(CategoricalIndexType::UInt16), - #[cfg(any( - not(feature = "default_categorical_8"), - feature = "extended_categorical" - ))] - TextArray::Categorical32(_) => ArrowType::Dictionary(CategoricalIndexType::UInt32), - #[cfg(feature = "extended_categorical")] - TextArray::Categorical64(_) => ArrowType::Dictionary(CategoricalIndexType::UInt64), - TextArray::Null => ArrowType::Null, - }, + Array::NumericArray(inner) => inner.arrow_type(), + Array::TextArray(inner) => inner.arrow_type(), #[cfg(feature = "datetime")] - Array::TemporalArray(inner) => match inner { - TemporalArray::Datetime32(_) => ArrowType::Date32, - TemporalArray::Datetime64(_) => ArrowType::Date64, - TemporalArray::Null => ArrowType::Null, - }, + Array::TemporalArray(inner) => inner.arrow_type(), Array::BooleanArray(_) => ArrowType::Boolean, Array::Null => ArrowType::Null, } diff --git a/src/enums/collections/numeric_array.rs b/src/enums/collections/numeric_array.rs index c3da3d4..0c73c14 100644 --- a/src/enums/collections/numeric_array.rs +++ b/src/enums/collections/numeric_array.rs @@ -30,6 +30,7 @@ use std::{ sync::Arc, }; +use crate::ffi::arrow_dtype::ArrowType; use crate::{Bitmask, FloatArray, IntegerArray, MaskedArray, Vec64}; use crate::{BooleanArray, StringArray}; #[cfg(feature = "decimal")] @@ -189,6 +190,33 @@ macro_rules! decimal_to_str { } impl NumericArray { + /// Returns the Arrow physical type for this numeric array. + pub fn arrow_type(&self) -> ArrowType { + match self { + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int8(_) => ArrowType::Int8, + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int16(_) => ArrowType::Int16, + NumericArray::Int32(_) => ArrowType::Int32, + NumericArray::Int64(_) => ArrowType::Int64, + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt8(_) => ArrowType::UInt8, + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt16(_) => ArrowType::UInt16, + NumericArray::UInt32(_) => ArrowType::UInt32, + NumericArray::UInt64(_) => ArrowType::UInt64, + NumericArray::Float32(_) => ArrowType::Float32, + NumericArray::Float64(_) => ArrowType::Float64, + #[cfg(feature = "decimal")] + NumericArray::Decimal32(a) => a.arrow_type(), + #[cfg(feature = "decimal")] + NumericArray::Decimal64(a) => a.arrow_type(), + #[cfg(feature = "decimal")] + NumericArray::Decimal128(a) => a.arrow_type(), + NumericArray::Null => ArrowType::Null, + } + } + /// Returns the logical length of the numeric array. #[inline] pub fn len(&self) -> usize { @@ -1246,40 +1274,14 @@ impl Concatenate for NumericArray { to: "NumericArray", message: Some(format!( "Cannot concatenate mismatched NumericArray variants: {:?} and {:?}", - variant_name(&lhs), - variant_name(&rhs) + lhs.arrow_type(), + rhs.arrow_type() )), }), } } } -/// Helper function to get the variant name for error messages -fn variant_name(arr: &NumericArray) -> &'static str { - match arr { - #[cfg(feature = "extended_numeric_types")] - NumericArray::Int8(_) => "Int8", - #[cfg(feature = "extended_numeric_types")] - NumericArray::Int16(_) => "Int16", - NumericArray::Int32(_) => "Int32", - NumericArray::Int64(_) => "Int64", - #[cfg(feature = "extended_numeric_types")] - NumericArray::UInt8(_) => "UInt8", - #[cfg(feature = "extended_numeric_types")] - NumericArray::UInt16(_) => "UInt16", - NumericArray::UInt32(_) => "UInt32", - NumericArray::UInt64(_) => "UInt64", - NumericArray::Float32(_) => "Float32", - NumericArray::Float64(_) => "Float64", - #[cfg(feature = "decimal")] - NumericArray::Decimal32(_) => "Decimal32", - #[cfg(feature = "decimal")] - NumericArray::Decimal64(_) => "Decimal64", - #[cfg(feature = "decimal")] - NumericArray::Decimal128(_) => "Decimal128", - NumericArray::Null => "Null", - } -} // --------------------------------------------------------------------------- // From impls - DecimalArray -> NumericArray diff --git a/src/enums/collections/temporal_array.rs b/src/enums/collections/temporal_array.rs index edeb241..f8c523c 100644 --- a/src/enums/collections/temporal_array.rs +++ b/src/enums/collections/temporal_array.rs @@ -30,6 +30,7 @@ use std::{ sync::Arc, }; +use crate::ffi::arrow_dtype::ArrowType; use crate::{Bitmask, DatetimeArray, MaskedArray, TimeUnit}; use crate::{ enums::{error::MinarrowError, shape_dim::ShapeDim}, @@ -92,6 +93,15 @@ pub enum TemporalArray { } impl TemporalArray { + /// Returns the Arrow physical type for this temporal array. + pub fn arrow_type(&self) -> ArrowType { + match self { + TemporalArray::Datetime32(_) => ArrowType::Date32, + TemporalArray::Datetime64(_) => ArrowType::Date64, + TemporalArray::Null => ArrowType::Null, + } + } + /// Returns the logical length of the temporal array. #[inline] pub fn len(&self) -> usize { @@ -225,8 +235,8 @@ impl TemporalArray { to: "TemporalArray", message: Some(format!( "Cannot insert {} into {}: incompatible types", - temporal_variant_name(rhs), - temporal_variant_name(lhs) + rhs.arrow_type(), + lhs.arrow_type() )), }), } @@ -327,8 +337,8 @@ impl Concatenate for TemporalArray { to: "TemporalArray", message: Some(format!( "Cannot concatenate mismatched TemporalArray variants: {:?} and {:?}", - temporal_variant_name(&lhs), - temporal_variant_name(&rhs) + lhs.arrow_type(), + rhs.arrow_type() )), }), } @@ -672,15 +682,6 @@ impl DatetimeOps for TemporalArray { } } -/// Helper function to get the variant name for error messages -fn temporal_variant_name(arr: &TemporalArray) -> &'static str { - match arr { - TemporalArray::Datetime32(_) => "Datetime32", - TemporalArray::Datetime64(_) => "Datetime64", - TemporalArray::Null => "Null", - } -} - impl Display for TemporalArray { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { diff --git a/src/enums/collections/text_array.rs b/src/enums/collections/text_array.rs index b19af45..7a83aa4 100644 --- a/src/enums/collections/text_array.rs +++ b/src/enums/collections/text_array.rs @@ -30,6 +30,7 @@ use std::sync::Arc; use crate::enums::error::MinarrowError; use crate::enums::shape_dim::ShapeDim; +use crate::ffi::arrow_dtype::{ArrowType, CategoricalIndexType}; use crate::traits::{concatenate::Concatenate, shape::Shape}; use crate::{Bitmask, CategoricalArray, MaskedArray, StringArray}; @@ -92,6 +93,27 @@ pub enum TextArray { } impl TextArray { + /// Returns the Arrow physical type for this text array. + pub fn arrow_type(&self) -> ArrowType { + match self { + TextArray::String32(_) => ArrowType::String, + #[cfg(feature = "large_string")] + TextArray::String64(_) => ArrowType::LargeString, + #[cfg(feature = "default_categorical_8")] + TextArray::Categorical8(_) => ArrowType::Dictionary(CategoricalIndexType::UInt8), + #[cfg(feature = "extended_categorical")] + TextArray::Categorical16(_) => ArrowType::Dictionary(CategoricalIndexType::UInt16), + #[cfg(any( + not(feature = "default_categorical_8"), + feature = "extended_categorical" + ))] + TextArray::Categorical32(_) => ArrowType::Dictionary(CategoricalIndexType::UInt32), + #[cfg(feature = "extended_categorical")] + TextArray::Categorical64(_) => ArrowType::Dictionary(CategoricalIndexType::UInt64), + TextArray::Null => ArrowType::Null, + } + } + /// Returns the logical length of the text array. #[inline] pub fn len(&self) -> usize { @@ -711,31 +733,11 @@ impl Concatenate for TextArray { to: "TextArray", message: Some(format!( "Cannot concatenate mismatched TextArray variants: {:?} and {:?}", - text_variant_name(&lhs), - text_variant_name(&rhs) + lhs.arrow_type(), + rhs.arrow_type() )), }), } } } -/// Helper function to get the variant name for error messages -fn text_variant_name(arr: &TextArray) -> &'static str { - match arr { - TextArray::String32(_) => "String32", - #[cfg(feature = "large_string")] - TextArray::String64(_) => "String64", - #[cfg(feature = "default_categorical_8")] - TextArray::Categorical8(_) => "Categorical8", - #[cfg(feature = "extended_categorical")] - TextArray::Categorical16(_) => "Categorical16", - #[cfg(any( - not(feature = "default_categorical_8"), - feature = "extended_categorical" - ))] - TextArray::Categorical32(_) => "Categorical32", - #[cfg(feature = "extended_categorical")] - TextArray::Categorical64(_) => "Categorical64", - TextArray::Null => "Null", - } -}