Repository: jorgecarleitao/arrow2
Pull request #1246 (2022-09-11 10:50:27)

Added schema to FileReader


Patch:
diff --git a/Cargo.toml b/Cargo.toml
index e52ad35508..c22c3351ac 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,6 +22,9 @@ dyn-clone = "1"
 bytemuck = { version = "1", features = ["derive"] }
 chrono = { version = "0.4", default_features = false, features = ["std"] }
 
+# for decimal i256
+ethnum = "1"
+
 # We need to Hash values before sending them to an hasher. This
 # crate provides HashMap that assumes pre-hashed values.
 hash_hasher = "^2.0.3"
diff --git a/src/array/mod.rs b/src/array/mod.rs
index a4fd5e4362..4a5e7fceba 100644
--- a/src/array/mod.rs
+++ b/src/array/mod.rs
@@ -206,13 +206,14 @@ macro_rules! with_match_primitive_type {(
 ) => ({
     macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
     use crate::datatypes::PrimitiveType::*;
-    use crate::types::{days_ms, months_days_ns, f16};
+    use crate::types::{days_ms, months_days_ns, f16, i256};
     match $key_type {
         Int8 => __with_ty__! { i8 },
         Int16 => __with_ty__! { i16 },
         Int32 => __with_ty__! { i32 },
         Int64 => __with_ty__! { i64 },
         Int128 => __with_ty__! { i128 },
+        Int256 => __with_ty__! { i256 },
         DaysMs => __with_ty__! { days_ms },
         MonthDayNano => __with_ty__! { months_days_ns },
         UInt8 => __with_ty__! { u8 },
diff --git a/src/array/primitive/fmt.rs b/src/array/primitive/fmt.rs
index 86ad35b127..f21b095862 100644
--- a/src/array/primitive/fmt.rs
+++ b/src/array/primitive/fmt.rs
@@ -2,7 +2,7 @@ use std::fmt::{Debug, Formatter, Result, Write};
 
 use crate::array::Array;
 use crate::datatypes::{IntervalUnit, TimeUnit};
-use crate::types::{days_ms, months_days_ns};
+use crate::types::{days_ms, i256, months_days_ns};
 
 use super::PrimitiveArray;
 use crate::array::fmt::write_vec;
@@ -118,13 +118,24 @@ pub fn get_write_value<'a, T: NativeType, F: Write>(
         Decimal(_, scale) => {
             // The number 999.99 has a precision of 5 and scale of 2
             let scale = *scale as u32;
-            let display = move |x| {
-                let base = x / 10i128.pow(scale);
-                let decimals = x - base * 10i128.pow(scale);
+            let factor = 10i128.pow(scale);
+            let display = move |x: i128| {
+                let base = x / factor;
+                let decimals = (x - base * factor).abs();
                 format!("{}.{}", base, decimals)
             };
             dyn_primitive!(array, i128, display)
         }
+        Decimal256(_, scale) => {
+            let scale = *scale as u32;
+            let factor = (ethnum::I256::ONE * 10).pow(scale);
+            let display = move |x: i256| {
+                let base = x.0 / factor;
+                let decimals = (x.0 - base * factor).abs();
+                format!("{}.{}", base, decimals)
+            };
+            dyn_primitive!(array, i256, display)
+        }
         _ => unreachable!(),
     }
 }
diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
index 33d56ed536..4e93c07a3b 100644
--- a/src/array/primitive/mod.rs
+++ b/src/array/primitive/mod.rs
@@ -7,7 +7,7 @@ use crate::{
     datatypes::*,
     error::Error,
     trusted_len::TrustedLen,
-    types::{days_ms, f16, months_days_ns, NativeType},
+    types::{days_ms, f16, i256, months_days_ns, NativeType},
 };
 
 use super::Array;
@@ -475,6 +475,8 @@ pub type Int32Array = PrimitiveArray<i32>;
 pub type Int64Array = PrimitiveArray<i64>;
 /// A type definition [`PrimitiveArray`] for `i128`
 pub type Int128Array = PrimitiveArray<i128>;
+/// A type definition [`PrimitiveArray`] for `i256`
+pub type Int256Array = PrimitiveArray<i256>;
 /// A type definition [`PrimitiveArray`] for [`days_ms`]
 pub type DaysMsArray = PrimitiveArray<days_ms>;
 /// A type definition [`PrimitiveArray`] for [`months_days_ns`]
@@ -504,6 +506,8 @@ pub type Int32Vec = MutablePrimitiveArray<i32>;
 pub type Int64Vec = MutablePrimitiveArray<i64>;
 /// A type definition [`MutablePrimitiveArray`] for `i128`
 pub type Int128Vec = MutablePrimitiveArray<i128>;
+/// A type definition [`MutablePrimitiveArray`] for `i256`
+pub type Int256Vec = MutablePrimitiveArray<i256>;
 /// A type definition [`MutablePrimitiveArray`] for [`days_ms`]
 pub type DaysMsVec = MutablePrimitiveArray<days_ms>;
 /// A type definition [`MutablePrimitiveArray`] for [`months_days_ns`]
diff --git a/src/compute/arithmetics/mod.rs b/src/compute/arithmetics/mod.rs
index b3a52cd7e1..3f09fe4d0b 100644
--- a/src/compute/arithmetics/mod.rs
+++ b/src/compute/arithmetics/mod.rs
@@ -402,13 +402,14 @@ macro_rules! with_match_negatable {(
 ) => ({
     macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
     use crate::datatypes::PrimitiveType::*;
-    use crate::types::{days_ms, months_days_ns};
+    use crate::types::{days_ms, months_days_ns, i256};
     match $key_type {
         Int8 => __with_ty__! { i8 },
         Int16 => __with_ty__! { i16 },
         Int32 => __with_ty__! { i32 },
         Int64 => __with_ty__! { i64 },
         Int128 => __with_ty__! { i128 },
+        Int256 => __with_ty__! { i256 },
         DaysMs => __with_ty__! { days_ms },
         MonthDayNano => __with_ty__! { months_days_ns },
         UInt8 | UInt16 | UInt32 | UInt64 | Float16 => todo!(),
diff --git a/src/compute/comparison/mod.rs b/src/compute/comparison/mod.rs
index ba6a84e49b..6c7714f587 100644
--- a/src/compute/comparison/mod.rs
+++ b/src/compute/comparison/mod.rs
@@ -69,12 +69,14 @@ macro_rules! match_eq_ord {(
 ) => ({
     macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
     use crate::datatypes::PrimitiveType::*;
+    use crate::types::i256;
     match $key_type {
         Int8 => __with_ty__! { i8 },
         Int16 => __with_ty__! { i16 },
         Int32 => __with_ty__! { i32 },
         Int64 => __with_ty__! { i64 },
         Int128 => __with_ty__! { i128 },
+        Int256 => __with_ty__! { i256 },
         DaysMs => todo!(),
         MonthDayNano => todo!(),
         UInt8 => __with_ty__! { u8 },
@@ -92,13 +94,14 @@ macro_rules! match_eq {(
 ) => ({
     macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
     use crate::datatypes::PrimitiveType::*;
-    use crate::types::{days_ms, months_days_ns, f16};
+    use crate::types::{days_ms, months_days_ns, f16, i256};
     match $key_type {
         Int8 => __with_ty__! { i8 },
         Int16 => __with_ty__! { i16 },
         Int32 => __with_ty__! { i32 },
         Int64 => __with_ty__! { i64 },
         Int128 => __with_ty__! { i128 },
+        Int256 => __with_ty__! { i256 },
         DaysMs => __with_ty__! { days_ms },
         MonthDayNano => __with_ty__! { months_days_ns },
         UInt8 => __with_ty__! { u8 },
@@ -612,9 +615,12 @@ fn finish_neq_validities(
                     if both_sides_invalid.values().unset_bits() != both_sides_invalid.len() {
                         // we use the `binary` kernel directly to save allocations
                         // and apply `lhs & !rhs)` in one shot.
-                        compute::boolean::binary_boolean_kernel(&lhs, &rhs, |lhs, rhs| {
-                            binary(lhs, rhs, |lhs, rhs| (lhs & !rhs))
-                        })
+
+                        compute::boolean::binary_boolean_kernel(
+                            &or,
+                            &both_sides_invalid,
+                            |lhs, rhs| binary(lhs, rhs, |lhs, rhs| (lhs & !rhs)),
+                        )
                     } else {
                         or
                     }
diff --git a/src/compute/comparison/simd/native.rs b/src/compute/comparison/simd/native.rs
index b760620a60..b8bbf9b17d 100644
--- a/src/compute/comparison/simd/native.rs
+++ b/src/compute/comparison/simd/native.rs
@@ -1,7 +1,7 @@
 use std::convert::TryInto;
 
 use super::{set, Simd8, Simd8Lanes, Simd8PartialEq, Simd8PartialOrd};
-use crate::types::{days_ms, f16, months_days_ns};
+use crate::types::{days_ms, f16, i256, months_days_ns};
 
 simd8_native_all!(u8);
 simd8_native_all!(u16);
@@ -11,6 +11,7 @@ simd8_native_all!(i8);
 simd8_native_all!(i16);
 simd8_native_all!(i32);
 simd8_native_all!(i128);
+simd8_native_all!(i256);
 simd8_native_all!(i64);
 simd8_native!(f16);
 simd8_native_partial_eq!(f16);
diff --git a/src/compute/comparison/simd/packed.rs b/src/compute/comparison/simd/packed.rs
index 898d04f5d9..937e3bb81a 100644
--- a/src/compute/comparison/simd/packed.rs
+++ b/src/compute/comparison/simd/packed.rs
@@ -2,7 +2,7 @@ use std::convert::TryInto;
 use std::simd::{SimdPartialEq, SimdPartialOrd, ToBitMask};
 
 use crate::types::simd::*;
-use crate::types::{days_ms, f16, months_days_ns};
+use crate::types::{days_ms, f16, i256, months_days_ns};
 
 use super::*;
 
@@ -71,6 +71,7 @@ simd8!(i16, i16x8);
 simd8!(i32, i32x8);
 simd8!(i64, i64x8);
 simd8_native_all!(i128);
+simd8_native_all!(i256);
 simd8_native!(f16);
 simd8_native_partial_eq!(f16);
 simd8!(f32, f32x8);
diff --git a/src/datatypes/mod.rs b/src/datatypes/mod.rs
index 4be369da0a..dafbb57848 100644
--- a/src/datatypes/mod.rs
+++ b/src/datatypes/mod.rs
@@ -154,6 +154,8 @@ pub enum DataType {
     /// scale is the number of decimal places.
     /// The number 999.99 has a precision of 5 and scale of 2.
     Decimal(usize, usize),
+    /// Decimal backed by 256 bits
+    Decimal256(usize, usize),
     /// Extension type.
     Extension(String, Box<DataType>, Option<String>),
 }
@@ -233,6 +235,7 @@ impl DataType {
                 PhysicalType::Primitive(PrimitiveType::Int64)
             }
             Decimal(_, _) => PhysicalType::Primitive(PrimitiveType::Int128),
+            Decimal256(_, _) => PhysicalType::Primitive(PrimitiveType::Int256),
             UInt8 => PhysicalType::Primitive(PrimitiveType::UInt8),
             UInt16 => PhysicalType::Primitive(PrimitiveType::UInt16),
             UInt32 => PhysicalType::Primitive(PrimitiveType::UInt32),
@@ -299,6 +302,7 @@ impl From<PrimitiveType> for DataType {
             PrimitiveType::UInt32 => DataType::UInt32,
             PrimitiveType::UInt64 => DataType::UInt64,
             PrimitiveType::Int128 => DataType::Decimal(32, 32),
+            PrimitiveType::Int256 => DataType::Decimal256(32, 32),
             PrimitiveType::Float16 => DataType::Float16,
             PrimitiveType::Float32 => DataType::Float32,
             PrimitiveType::Float64 => DataType::Float64,
diff --git a/src/ffi/schema.rs b/src/ffi/schema.rs
index 4f4fc5b8be..5e6cc4180c 100644
--- a/src/ffi/schema.rs
+++ b/src/ffi/schema.rs
@@ -331,9 +331,18 @@ unsafe fn to_data_type(schema: &ArrowSchema) -> Result<DataType> {
                                     "Decimal bit width is not a valid integer".to_string(),
                                 )
                             })?;
-                            if bit_width != 128 {
-                                return Err(Error::OutOfSpec(
-                                    "Decimal256 is not supported".to_string(),
+                            if bit_width == 256 {
+                                return Ok(DataType::Decimal256(
+                                    precision_raw.parse::<usize>().map_err(|_| {
+                                        Error::OutOfSpec(
+                                            "Decimal precision is not a valid integer".to_string(),
+                                        )
+                                    })?,
+                                    scale_raw.parse::<usize>().map_err(|_| {
+                                        Error::OutOfSpec(
+                                            "Decimal scale is not a valid integer".to_string(),
+                                        )
+                                    })?,
                                 ));
                             }
                             (precision_raw, scale_raw)
@@ -438,6 +447,7 @@ fn to_format(data_type: &DataType) -> String {
             )
         }
         DataType::Decimal(precision, scale) => format!("d:{},{}", precision, scale),
+        DataType::Decimal256(precision, scale) => format!("d:{},{},256", precision, scale),
         DataType::List(_) => "+l".to_string(),
         DataType::LargeList(_) => "+L".to_string(),
         DataType::Struct(_) => "+s".to_string(),
diff --git a/src/io/ipc/read/schema.rs b/src/io/ipc/read/schema.rs
index b1d97f0d77..5a9646ea37 100644
--- a/src/io/ipc/read/schema.rs
+++ b/src/io/ipc/read/schema.rs
@@ -327,16 +327,25 @@ fn get_data_type(
             (DataType::Duration(time_unit), IpcField::default())
         }
         Decimal(decimal) => {
-            let data_type = DataType::Decimal(
-                decimal
-                    .precision()?
-                    .try_into()
-                    .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?,
-                decimal
-                    .scale()?
-                    .try_into()
-                    .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?,
-            );
+            let bit_width: usize = decimal
+                .bit_width()?
+                .try_into()
+                .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
+            let precision: usize = decimal
+                .precision()?
+                .try_into()
+                .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
+            let scale: usize = decimal
+                .scale()?
+                .try_into()
+                .map_err(|_| Error::from(OutOfSpecKind::NegativeFooterLength))?;
+
+            let data_type = match bit_width {
+                128 => DataType::Decimal(precision, scale),
+                256 => DataType::Decimal256(precision, scale),
+                _ => return Err(Error::from(OutOfSpecKind::NegativeFooterLength)),
+            };
+
             (data_type, IpcField::default())
         }
         List(_) => deserialize_list(field)?,
diff --git a/src/io/ipc/write/schema.rs b/src/io/ipc/write/schema.rs
index 0636cb9e6c..1c4dab8e39 100644
--- a/src/io/ipc/write/schema.rs
+++ b/src/io/ipc/write/schema.rs
@@ -197,6 +197,11 @@ fn serialize_type(data_type: &DataType) -> arrow_format::ipc::Type {
             scale: *scale as i32,
             bit_width: 128,
         })),
+        Decimal256(precision, scale) => ipc::Type::Decimal(Box::new(ipc::Decimal {
+            precision: *precision as i32,
+            scale: *scale as i32,
+            bit_width: 256,
+        })),
         Binary => ipc::Type::Binary(Box::new(ipc::Binary {})),
         LargeBinary => ipc::Type::LargeBinary(Box::new(ipc::LargeBinary {})),
         Utf8 => ipc::Type::Utf8(Box::new(ipc::Utf8 {})),
@@ -281,7 +286,8 @@ fn serialize_children(data_type: &DataType, ipc_field: &IpcField) -> Vec<arrow_f
         | LargeBinary
         | Utf8
         | LargeUtf8
-        | Decimal(_, _) => vec![],
+        | Decimal(_, _)
+        | Decimal256(_, _) => vec![],
         FixedSizeList(inner, _) | LargeList(inner) | List(inner) | Map(inner, _) => {
             vec![serialize_field(inner, &ipc_field.fields[0])]
         }
diff --git a/src/io/json_integration/read/array.rs b/src/io/json_integration/read/array.rs
index 0c44cc2432..90cd709cbf 100644
--- a/src/io/json_integration/read/array.rs
+++ b/src/io/json_integration/read/array.rs
@@ -10,7 +10,7 @@ use crate::{
     datatypes::{DataType, PhysicalType, PrimitiveType, Schema},
     error::{Error, Result},
     io::ipc::IpcField,
-    types::{days_ms, months_days_ns, NativeType},
+    types::{days_ms, i256, months_days_ns, NativeType},
 };
 
 use super::super::{ArrowJsonBatch, ArrowJsonColumn, ArrowJsonDictionaryBatch};
@@ -125,6 +125,24 @@ fn to_decimal(json_col: &ArrowJsonColumn, data_type: DataType) -> PrimitiveArray
     PrimitiveArray::<i128>::new(data_type, values, validity)
 }
 
+fn to_decimal256(json_col: &ArrowJsonColumn, data_type: DataType) -> PrimitiveArray<i256> {
+    let validity = to_validity(&json_col.validity);
+    let values = json_col
+        .data
+        .as_ref()
+        .unwrap()
+        .iter()
+        .map(|value| match value {
+            Value::String(x) => i256(x.parse::<ethnum::I256>().unwrap()),
+            _ => {
+                panic!()
+            }
+        })
+        .collect();
+
+    PrimitiveArray::<i256>::new(data_type, values, validity)
+}
+
 fn to_primitive<T: NativeType + NumCast>(
     json_col: &ArrowJsonColumn,
     data_type: DataType,
@@ -280,6 +298,7 @@ pub fn to_array(
         Primitive(PrimitiveType::Int32) => Ok(Box::new(to_primitive::<i32>(json_col, data_type))),
         Primitive(PrimitiveType::Int64) => Ok(Box::new(to_primitive::<i64>(json_col, data_type))),
         Primitive(PrimitiveType::Int128) => Ok(Box::new(to_decimal(json_col, data_type))),
+        Primitive(PrimitiveType::Int256) => Ok(Box::new(to_decimal256(json_col, data_type))),
         Primitive(PrimitiveType::DaysMs) => Ok(Box::new(to_primitive_days_ms(json_col, data_type))),
         Primitive(PrimitiveType::MonthDayNano) => {
             Ok(Box::new(to_primitive_months_days_ns(json_col, data_type)))
diff --git a/src/io/json_integration/read/schema.rs b/src/io/json_integration/read/schema.rs
index aaf6912af0..06e2ecbaf4 100644
--- a/src/io/json_integration/read/schema.rs
+++ b/src/io/json_integration/read/schema.rs
@@ -170,20 +170,28 @@ fn to_data_type(item: &Value, mut children: Vec<Field>) -> Result<DataType> {
         "largeutf8" => LargeUtf8,
         "decimal" => {
             // return a list with any type as its child isn't defined in the map
-            let precision = match item.get("precision") {
-                Some(p) => Ok(p.as_u64().unwrap() as usize),
-                None => Err(Error::OutOfSpec(
-                    "Expecting a precision for decimal".to_string(),
-                )),
-            };
-            let scale = match item.get("scale") {
-                Some(s) => Ok(s.as_u64().unwrap() as usize),
-                _ => Err(Error::OutOfSpec(
-                    "Expecting a scale for decimal".to_string(),
-                )),
+            let precision = item
+                .get("precision")
+                .ok_or_else(|| Error::OutOfSpec("Expecting a precision for decimal".to_string()))?
+                .as_u64()
+                .unwrap() as usize;
+
+            let scale = item
+                .get("scale")
+                .ok_or_else(|| Error::OutOfSpec("Expecting a scale for decimal".to_string()))?
+                .as_u64()
+                .unwrap() as usize;
+
+            let bit_width = match item.get("bitWidth") {
+                Some(s) => s.as_u64().unwrap() as usize,
+                None => 128,
             };
 
-            DataType::Decimal(precision?, scale?)
+            match bit_width {
+                128 => DataType::Decimal(precision, scale),
+                256 => DataType::Decimal256(precision, scale),
+                _ => todo!(),
+            }
         }
         "floatingpoint" => match item.get("precision") {
             Some(p) if p == "HALF" => DataType::Float16,
diff --git a/src/io/json_integration/write/schema.rs b/src/io/json_integration/write/schema.rs
index f4297c3fb0..7ae67dcca5 100644
--- a/src/io/json_integration/write/schema.rs
+++ b/src/io/json_integration/write/schema.rs
@@ -89,6 +89,9 @@ fn serialize_data_type(data_type: &DataType) -> Value {
         DataType::Decimal(precision, scale) => {
             json!({"name": "decimal", "precision": precision, "scale": scale})
         }
+        DataType::Decimal256(precision, scale) => {
+            json!({"name": "decimal", "precision": precision, "scale": scale, "bit_width": 256})
+        }
         DataType::Extension(_, inner_data_type, _) => serialize_data_type(inner_data_type),
     }
 }
diff --git a/src/io/parquet/read/file.rs b/src/io/parquet/read/file.rs
index 391544685a..da4c172a12 100644
--- a/src/io/parquet/read/file.rs
+++ b/src/io/parquet/read/file.rs
@@ -57,6 +57,11 @@ impl<R: Read + Seek> FileReader<R> {
         }
         Ok(result)
     }
+
+    /// Returns the [`Schema`] associated to this file.
+    pub fn schema(&self) -> &Schema {
+        &self.row_groups.schema
+    }
 }
 
 impl<R: Read + Seek> Iterator for FileReader<R> {
diff --git a/src/types/mod.rs b/src/types/mod.rs
index 3ba86ad59e..165e4bd592 100644
--- a/src/types/mod.rs
+++ b/src/types/mod.rs
@@ -47,6 +47,8 @@ pub enum PrimitiveType {
     Int64,
     /// A signed 128-bit integer.
     Int128,
+    /// A signed 256-bit integer.
+    Int256,
     /// An unsigned 8-bit integer.
     UInt8,
     /// An unsigned 16-bit integer.
@@ -79,6 +81,7 @@ mod private {
     impl Sealed for i32 {}
     impl Sealed for i64 {}
     impl Sealed for i128 {}
+    impl Sealed for super::i256 {}
     impl Sealed for super::f16 {}
     impl Sealed for f32 {}
     impl Sealed for f64 {}
diff --git a/src/types/native.rs b/src/types/native.rs
index fc100017a7..8c82452e89 100644
--- a/src/types/native.rs
+++ b/src/types/native.rs
@@ -500,6 +500,105 @@ impl NativeType for f16 {
     }
 }
 
+/// Physical representation of a decimal
+#[derive(Clone, Copy, Default, Eq, Hash, PartialEq, PartialOrd)]
+#[allow(non_camel_case_types)]
+#[repr(C)]
+pub struct i256(pub ethnum::I256);
+
+impl i256 {
+    /// Returns a new [`i256`] from two `i128`.
+    pub fn from_words(hi: i128, lo: i128) -> Self {
+        Self(ethnum::I256::from_words(hi, lo))
+    }
+}
+
+impl Neg for i256 {
+    type Output = Self;
+
+    #[inline]
+    fn neg(self) -> Self::Output {
+        let (a, b) = self.0.into_words();
+        Self(ethnum::I256::from_words(-a, b))
+    }
+}
+
+impl std::fmt::Debug for i256 {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:?}", self.0)
+    }
+}
+
+impl std::fmt::Display for i256 {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+unsafe impl Pod for i256 {}
+unsafe impl Zeroable for i256 {}
+
+impl NativeType for i256 {
+    const PRIMITIVE: PrimitiveType = PrimitiveType::Int256;
+
+    type Bytes = [u8; 32];
+
+    #[inline]
+    fn to_le_bytes(&self) -> Self::Bytes {
+        let mut bytes = [0u8; 32];
+        let (a, b) = self.0.into_words();
+        let a = a.to_le_bytes();
+        (0..16).for_each(|i| {
+            bytes[i] = a[i];
+        });
+
+        let b = b.to_le_bytes();
+        (0..16).for_each(|i| {
+            bytes[i + 16] = b[i];
+        });
+
+        bytes
+    }
+
+    #[inline]
+    fn to_be_bytes(&self) -> Self::Bytes {
+        let mut bytes = [0u8; 32];
+        let (a, b) = self.0.into_words();
+
+        let b = b.to_be_bytes();
+        (0..16).for_each(|i| {
+            bytes[i] = b[i];
+        });
+
+        let a = a.to_be_bytes();
+        (0..16).for_each(|i| {
+            bytes[i + 16] = a[i];
+        });
+
+        bytes
+    }
+
+    #[inline]
+    fn from_be_bytes(bytes: Self::Bytes) -> Self {
+        let (b, a) = bytes.split_at(16);
+        let a: [u8; 16] = a.try_into().unwrap();
+        let b: [u8; 16] = b.try_into().unwrap();
+        let a = i128::from_be_bytes(a);
+        let b = i128::from_be_bytes(b);
+        Self(ethnum::I256::from_words(a, b))
+    }
+
+    #[inline]
+    fn from_le_bytes(bytes: Self::Bytes) -> Self {
+        let (b, a) = bytes.split_at(16);
+        let a: [u8; 16] = a.try_into().unwrap();
+        let b: [u8; 16] = b.try_into().unwrap();
+        let a = i128::from_le_bytes(a);
+        let b = i128::from_le_bytes(b);
+        Self(ethnum::I256::from_words(a, b))
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/src/types/simd/mod.rs b/src/types/simd/mod.rs
index 3e97902dcd..ba78025727 100644
--- a/src/types/simd/mod.rs
+++ b/src/types/simd/mod.rs
@@ -1,7 +1,7 @@
 //! Contains traits and implementations of multi-data used in SIMD.
 //! The actual representation is driven by the feature flag `"simd"`, which, if set,
 //! uses [`std::simd`].
-use super::{days_ms, f16, months_days_ns};
+use super::{days_ms, f16, i256, months_days_ns};
 use super::{BitChunk, BitChunkIter, NativeType};
 
 /// Describes the ability to convert itself from a [`BitChunk`].
@@ -133,6 +133,7 @@ native_simd!(f16x32, f16, 32, u32);
 native_simd!(days_msx8, days_ms, 8, u8);
 native_simd!(months_days_nsx8, months_days_ns, 8, u8);
 native_simd!(i128x8, i128, 8, u8);
+native_simd!(i256x8, i256, 8, u8);
 
 // In the native implementation, a mask is 1 bit wide, as per AVX512.
 impl<T: BitChunk> FromMaskChunk<T> for T {
@@ -162,5 +163,6 @@ native!(f16, f16x32);
 native!(f32, f32x16);
 native!(f64, f64x8);
 native!(i128, i128x8);
+native!(i256, i256x8);
 native!(days_ms, days_msx8);
 native!(months_days_ns, months_days_nsx8);


Test patch:
diff --git a/arrow-pyarrow-integration-testing/tests/test_sql.py b/arrow-pyarrow-integration-testing/tests/test_sql.py
index 3cd9c2e069..5a44bb61a1 100644
--- a/arrow-pyarrow-integration-testing/tests/test_sql.py
+++ b/arrow-pyarrow-integration-testing/tests/test_sql.py
@@ -116,6 +116,19 @@ class TestCase(unittest.TestCase):
         b = arrow_pyarrow_integration_testing.round_trip_array(a)
         self.assertEqual(a, b)
 
+    def test_decimal256_roundtrip(self):
+        """
+        Python -> Rust -> Python
+        """
+        data = [
+            round(decimal.Decimal(722.82), 2),
+            round(decimal.Decimal(-934.11), 2),
+            None,
+        ]
+        a = pyarrow.array(data, pyarrow.decimal256(5, 2))
+        b = arrow_pyarrow_integration_testing.round_trip_array(a)
+        self.assertEqual(a, b)
+
     def test_list_array(self):
         """
         Python -> Rust -> Python
diff --git a/tests/it/compute/comparison.rs b/tests/it/compute/comparison.rs
index b077e99610..eb5058d337 100644
--- a/tests/it/compute/comparison.rs
+++ b/tests/it/compute/comparison.rs
@@ -383,7 +383,7 @@ fn primitive_gt_eq() {
 }
 
 #[test]
-#[cfg(any(feature = "compute_cast", feature = "compute_boolean_kleene"))]
+#[cfg(all(feature = "compute_cast", feature = "compute_boolean_kleene"))]
 fn utf8_and_validity() {
     use arrow2::compute::cast::CastOptions;
     let a1 = Utf8Array::<i32>::from([Some("0"), Some("1"), None, Some("2")]);
@@ -401,3 +401,16 @@ fn utf8_and_validity() {
     assert_eq!(utf8::neq_and_validity(&a1, &a1), expected);
     assert_eq!(utf8::neq_and_validity(&a1, a2), expected);
 }
+
+#[test]
+#[cfg(feature = "compute_boolean_kleene")]
+fn primitive_and_validity() {
+    let a1 = Int32Array::from([Some(0), None]);
+    let a2 = Int32Array::from([Some(10), None]);
+
+    let expected = BooleanArray::from_slice([true, false]);
+    assert_eq!(primitive::neq_and_validity(&a1, &a2), expected);
+
+    let expected = BooleanArray::from_slice([false, true]);
+    assert_eq!(primitive::eq_and_validity(&a1, &a2), expected);
+}
diff --git a/tests/it/io/ipc/write/file.rs b/tests/it/io/ipc/write/file.rs
index a899aa459d..45828169ae 100644
--- a/tests/it/io/ipc/write/file.rs
+++ b/tests/it/io/ipc/write/file.rs
@@ -6,7 +6,7 @@ use arrow2::datatypes::{Field, Schema};
 use arrow2::error::Result;
 use arrow2::io::ipc::read::{read_file_metadata, FileReader};
 use arrow2::io::ipc::{write::*, IpcField};
-use arrow2::types::months_days_ns;
+use arrow2::types::{i256, months_days_ns};
 
 use crate::io::ipc::common::read_gzip_json;
 
@@ -373,3 +373,17 @@ fn write_months_days_ns() -> Result<()> {
     let columns = Chunk::try_new(vec![array])?;
     round_trip(columns, schema, None, None)
 }
+
+#[test]
+fn write_decimal256() -> Result<()> {
+    let array = Int256Array::from([
+        Some(i256::from_words(1, 0)),
+        Some(i256::from_words(-2, 0)),
+        None,
+        Some(i256::from_words(0, 1)),
+    ])
+    .boxed();
+    let schema = Schema::from(vec![Field::new("a", array.data_type().clone(), true)]);
+    let columns = Chunk::try_new(vec![array])?;
+    round_trip(columns, schema, None, None)
+}