Skip to content

Commit 7302888

Browse files
authored
Introduce UnsafeFlag to manage disabling ArrayData validation (apache#7027)
* Introduce UnsafeFlag to manage disabling validation * fix docs
1 parent 7a15e4b commit 7302888

File tree

1 file changed

+41
-7
lines changed

1 file changed

+41
-7
lines changed

arrow-data/src/data.rs

+41-7
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ use std::mem;
2828
use std::ops::Range;
2929
use std::sync::Arc;
3030

31+
use crate::data::private::UnsafeFlag;
3132
use crate::{equal, validate_binary_view, validate_string_view};
3233

3334
#[inline]
@@ -255,6 +256,10 @@ impl ArrayData {
255256
buffers: Vec<Buffer>,
256257
child_data: Vec<ArrayData>,
257258
) -> Self {
259+
let mut skip_validation = UnsafeFlag::new();
260+
// SAFETY: caller responsible for ensuring data is valid
261+
skip_validation.set(true);
262+
258263
ArrayDataBuilder {
259264
data_type,
260265
len,
@@ -265,8 +270,7 @@ impl ArrayData {
265270
buffers,
266271
child_data,
267272
align_buffers: false,
268-
// SAFETY: caller responsible for ensuring data is valid
269-
skip_validation: true,
273+
skip_validation,
270274
}
271275
.build()
272276
.unwrap()
@@ -1779,6 +1783,36 @@ impl PartialEq for ArrayData {
17791783
}
17801784
}
17811785

1786+
mod private {
1787+
/// A boolean flag that cannot be mutated outside of unsafe code.
1788+
///
1789+
/// Defaults to a value of false.
1790+
///
1791+
/// This structure is used to enforce safety in the [`ArrayDataBuilder`]
1792+
///
1793+
/// [`ArrayDataBuilder`]: super::ArrayDataBuilder
1794+
#[derive(Debug)]
1795+
pub struct UnsafeFlag(bool);
1796+
1797+
impl UnsafeFlag {
1798+
/// Creates a new `UnsafeFlag` with the value set to `false`
1799+
#[inline]
1800+
pub const fn new() -> Self {
1801+
Self(false)
1802+
}
1803+
1804+
#[inline]
1805+
pub unsafe fn set(&mut self, val: bool) {
1806+
self.0 = val;
1807+
}
1808+
1809+
#[inline]
1810+
pub fn get(&self) -> bool {
1811+
self.0
1812+
}
1813+
}
1814+
}
1815+
17821816
/// Builder for [`ArrayData`] type
17831817
#[derive(Debug)]
17841818
pub struct ArrayDataBuilder {
@@ -1803,7 +1837,7 @@ pub struct ArrayDataBuilder {
18031837
/// This flag can only be set to true using `unsafe` APIs. However, once true
18041838
/// subsequent calls to `build()` may result in undefined behavior if the data
18051839
/// is not valid.
1806-
skip_validation: bool,
1840+
skip_validation: UnsafeFlag,
18071841
}
18081842

18091843
impl ArrayDataBuilder {
@@ -1820,7 +1854,7 @@ impl ArrayDataBuilder {
18201854
buffers: vec![],
18211855
child_data: vec![],
18221856
align_buffers: false,
1823-
skip_validation: false,
1857+
skip_validation: UnsafeFlag::new(),
18241858
}
18251859
}
18261860

@@ -1957,7 +1991,7 @@ impl ArrayDataBuilder {
19571991
}
19581992

19591993
// SAFETY: `skip_validation` is only set to true using `unsafe` APIs
1960-
if !skip_validation || cfg!(feature = "force_validate") {
1994+
if !skip_validation.get() || cfg!(feature = "force_validate") {
19611995
data.validate_data()?;
19621996
}
19631997
Ok(data)
@@ -2003,7 +2037,7 @@ impl ArrayDataBuilder {
20032037
/// If validation is skipped, the buffers must form a valid Arrow array,
20042038
/// otherwise undefined behavior will result
20052039
pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2006-
self.skip_validation = skip_validation;
2040+
self.skip_validation.set(skip_validation);
20072041
self
20082042
}
20092043
}
@@ -2020,7 +2054,7 @@ impl From<ArrayData> for ArrayDataBuilder {
20202054
null_bit_buffer: None,
20212055
null_count: None,
20222056
align_buffers: false,
2023-
skip_validation: false,
2057+
skip_validation: UnsafeFlag::new(),
20242058
}
20252059
}
20262060
}

0 commit comments

Comments
 (0)