Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simplify Validation/Alignment APIs of ArrayDataBuilder: validate and align #6966

Merged
merged 5 commits into from
Jan 13, 2025
Merged

Simplify Validation/Alignment APIs of ArrayDataBuilder: validate and align #6966

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
da8c07c
Simplify ArrayDataBuilder API for validation and alignment
alamb Jan 10, 2025
8c3153f
Update arrow-data/src/data.rs
alamb Jan 11, 2025
85f2381
Rename functions for consistency
alamb Jan 11, 2025
d0740bd
Improve documentation
alamb Jan 11, 2025
53aff81
Merge remote-tracking branch 'apache/main' into alamb/simplify_builde…
alamb Jan 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 88 additions & 34 deletions arrow-data/src/data.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,8 +264,12 @@ impl ArrayData {
offset,
buffers,
child_data,
align_buffers: false,
// SAFETY: caller responsible for ensuring data is valid
skip_validation: true,
}
.build_unchecked()
.build()
.unwrap()
}

/// Create a new ArrayData, validating that the provided buffers form a valid
Expand Down Expand Up @@ -1775,7 +1779,7 @@ impl PartialEq for ArrayData {
}
}

/// Builder for `ArrayData` type
/// Builder for [`ArrayData`] type
#[derive(Debug)]
pub struct ArrayDataBuilder {
data_type: DataType,
Expand All @@ -1786,6 +1790,10 @@ pub struct ArrayDataBuilder {
offset: usize,
buffers: Vec<Buffer>,
child_data: Vec<ArrayData>,
/// Should buffers be realigned (copying if necessary)? Defaults to false.
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the point of this PR -- to add these flags to the builder directly and then check them during calls to build()

align_buffers: bool,
/// Should data validation be skipped for this [`ArrayData`]? Defaults to false.
skip_validation: bool,
}

impl ArrayDataBuilder {
Expand All @@ -1801,6 +1809,8 @@ impl ArrayDataBuilder {
offset: 0,
buffers: vec![],
child_data: vec![],
align_buffers: false,
skip_validation: false,
}
}

Expand Down Expand Up @@ -1877,51 +1887,78 @@ impl ArrayDataBuilder {

/// Creates an array data, without any validation
///
/// This is shorthand for `self.with_skip_validation(true).build()`
///
/// # Safety
///
/// The same caveats as [`ArrayData::new_unchecked`]
/// apply.
#[allow(clippy::let_and_return)]
pub unsafe fn build_unchecked(self) -> ArrayData {
let data = self.build_impl();
// Provide a force_validate mode
#[cfg(feature = "force_validate")]
data.validate_data().unwrap();
data
self.with_skip_validation(true).build().unwrap()
}

/// Same as [`Self::build_unchecked`] but ignoring `force_validate` feature flag
unsafe fn build_impl(self) -> ArrayData {
let nulls = self
.nulls
/// Creates an `ArrayData`, consuming `self`
pub fn build(self) -> Result<ArrayData, ArrowError> {
let Self {
data_type,
len,
null_count,
null_bit_buffer,
nulls,
offset,
buffers,
child_data,
align_buffers,
skip_validation,
} = self;

let nulls = nulls
.or_else(|| {
let buffer = self.null_bit_buffer?;
let buffer = BooleanBuffer::new(buffer, self.offset, self.len);
Some(match self.null_count {
Some(n) => NullBuffer::new_unchecked(buffer, n),
let buffer = null_bit_buffer?;
let buffer = BooleanBuffer::new(buffer, offset, len);
Some(match null_count {
Some(n) => {
// SAFETY: if validation is on, call to `data.validate_data()` will validate the null buffer
unsafe { NullBuffer::new_unchecked(buffer, n) }
}
None => NullBuffer::new(buffer),
})
})
.filter(|b| b.null_count() != 0);

ArrayData {
data_type: self.data_type,
len: self.len,
offset: self.offset,
buffers: self.buffers,
child_data: self.child_data,
let mut data = ArrayData {
data_type,
len,
offset,
buffers,
child_data,
nulls,
};

if align_buffers {
data.align_buffers();
}
}

/// Creates an array data, validating all inputs
pub fn build(self) -> Result<ArrayData, ArrowError> {
let data = unsafe { self.build_impl() };
data.validate_data()?;
#[cfg(feature = "force_validate")]
let force_validation = true;
#[cfg(not(feature = "force_validate"))]
let force_validation = false;
// force validation in testing mode
let skip_validation = skip_validation && !force_validation;

if !skip_validation {
data.validate_data()?;
}
Ok(data)
}

/// Creates an array data, validating all inputs, and aligning any buffers
#[deprecated(since = "54.1.0", note = "Use ArrayData::with_align_buffers instead")]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

deprecated one API in favor of setting a flag

pub fn build_aligned(self) -> Result<ArrayData, ArrowError> {
self.with_align_buffers(true).build()
}

/// Ensure that all buffers are aligned, copying data if necessary
///
/// Rust requires that arrays are aligned to their corresponding primitive,
/// see [`Layout::array`](std::alloc::Layout::array) and [`std::mem::align_of`].
Expand All @@ -1930,17 +1967,32 @@ impl ArrayDataBuilder {
/// to allow for [slice](std::slice) based APIs. See [`BufferSpec::FixedWidth`].
///
/// As this alignment is architecture specific, and not guaranteed by all arrow implementations,
/// this method is provided to automatically copy buffers to a new correctly aligned allocation
/// this flag is provided to automatically copy buffers to a new correctly aligned allocation
/// when necessary, making it useful when interacting with buffers produced by other systems,
/// e.g. IPC or FFI.
///
/// This is unlike `[Self::build`] which will instead return an error on encountering
/// If this flag is not enabled, `[Self::build`] return an error on encountering
/// insufficiently aligned buffers.
pub fn build_aligned(self) -> Result<ArrayData, ArrowError> {
let mut data = unsafe { self.build_impl() };
data.align_buffers();
data.validate_data()?;
Ok(data)
pub fn with_align_buffers(mut self, align_buffers: bool) -> Self {
self.align_buffers = align_buffers;
self
}

/// Skips validation of the data.
///
/// If this flag is enabled, `[Self::build`] will skip validation of the
/// data
///
/// If this flag is not enabled, `[Self::build`] will validate that all
/// buffers are valid and will return an error if any data is invalid.
/// Validation can be expensive.
///
/// # Safety
/// If validation is skipped, the buffers must form a valid Array array,
/// otherwise undefined behavior will result
pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
self.skip_validation = skip_validation;
self
}
}

Expand All @@ -1955,6 +2007,8 @@ impl From<ArrayData> for ArrayDataBuilder {
nulls: d.nulls,
null_bit_buffer: None,
null_count: None,
align_buffers: false,
skip_validation: false,
}
}
}
Expand Down
48 changes: 14 additions & 34 deletions arrow-ipc/src/reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,17 +177,13 @@ fn create_array(
let values = create_array(reader, values_field, variadic_counts, require_alignment)?;

let run_array_length = run_node.length() as usize;
let builder = ArrayData::builder(data_type.clone())
let array_data = ArrayData::builder(data_type.clone())
.len(run_array_length)
.offset(0)
.add_child_data(run_ends.into_data())
.add_child_data(values.into_data());

let array_data = if require_alignment {
builder.build()?
} else {
builder.build_aligned()?
};
.add_child_data(values.into_data())
.with_align_buffers(!require_alignment)
.build()?;

Ok(make_array(array_data))
}
Expand Down Expand Up @@ -257,15 +253,11 @@ fn create_array(
)));
}

let builder = ArrayData::builder(data_type.clone())
let array_data = ArrayData::builder(data_type.clone())
.len(length as usize)
.offset(0);

let array_data = if require_alignment {
builder.build()?
} else {
builder.build_aligned()?
};
.offset(0)
.with_align_buffers(!require_alignment)
.build()?;

// no buffer increases
Ok(Arc::new(NullArray::from(array_data)))
Expand Down Expand Up @@ -311,11 +303,7 @@ fn create_primitive_array(
t => unreachable!("Data type {:?} either unsupported or not primitive", t),
};

let array_data = if require_alignment {
builder.build()?
} else {
builder.build_aligned()?
};
let array_data = builder.with_align_buffers(!require_alignment).build()?;

Ok(make_array(array_data))
}
Expand Down Expand Up @@ -347,11 +335,7 @@ fn create_list_array(
_ => unreachable!("Cannot create list or map array from {:?}", data_type),
};

let array_data = if require_alignment {
builder.build()?
} else {
builder.build_aligned()?
};
let array_data = builder.with_align_buffers(!require_alignment).build()?;

Ok(make_array(array_data))
}
Expand All @@ -367,17 +351,13 @@ fn create_dictionary_array(
) -> Result<ArrayRef, ArrowError> {
if let Dictionary(_, _) = *data_type {
let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
let builder = ArrayData::builder(data_type.clone())
let array_data = ArrayData::builder(data_type.clone())
.len(field_node.length() as usize)
.add_buffer(buffers[1].clone())
.add_child_data(value_array.into_data())
.null_bit_buffer(null_buffer);

let array_data = if require_alignment {
builder.build()?
} else {
builder.build_aligned()?
};
.null_bit_buffer(null_buffer)
.with_align_buffers(!require_alignment)
.build()?;

Ok(make_array(array_data))
} else {
Expand Down
Loading