This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 36229593fe perf: Optimize bool <-> numeric casts (#11242)
36229593fe is described below
commit 36229593fe00b03045ef5a8169333dbc8f943637
Author: Neil Conway <[email protected]>
AuthorDate: Tue Sep 29 04:49:57 2026 -0400
perf: Optimize bool <-> numeric casts (#11242)
# Which issue does this PR close?
- N/A
# Rationale for this change
Casting an array of numeric to boolean iterated over each array element,
checking for validity and then appending the new value and validity bit.
Using `BooleanArray::from_unary` is much faster, for two reasons:
1. It reuses the input's validity array instead of reconstructing it
bit-by-bit
2. Avoiding per-element null checks makes it easier for the compiler to
vectorize the loop.
Casting an array of boolean to numeric also used a similar
value-by-value loop. As before, we can do better by reusing the input's
validity bitmap. We also play a further trick: we chunk the input bits
into 64-bit words and process one chunk at a time; the fixed-size loop
helps LLVM vectorize more effectively.
Benchmarks: (M4 Max)
- Int32 → Boolean, ~10% nulls: 17.310 → 0.899 µs — 19.26× faster
- Boolean → Int32, ~10% nulls: 12.578 → 0.584 µs — 21.54× faster
# What changes are included in this PR?
* Optimize `cast_numeric_to_bool` and `cast_bool_to_numeric` as
described above
* Add unit tests
* Add benchmarks
# Are these changes tested?
Yes; new tests added.
# Are there any user-facing changes?
No.
# AI usage
Developed with Codex (Astra 6), reviewed and revised with Claude Code
(Opus 5.5). I reviewed and understand the resulting code.
---
arrow-cast/src/cast/mod.rs | 104 +++++++++++++++++++++++++++++-------------
arrow/benches/cast_kernels.rs | 8 ++++
2 files changed, 80 insertions(+), 32 deletions(-)
diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs
index 6eed69976e..96b9e4764c 100644
--- a/arrow-cast/src/cast/mod.rs
+++ b/arrow-cast/src/cast/mod.rs
@@ -2629,24 +2629,10 @@ fn cast_numeric_to_bool<FROM>(from: &dyn Array) ->
Result<ArrayRef, ArrowError>
where
FROM: ArrowPrimitiveType,
{
- numeric_to_bool_cast::<FROM>(from.as_primitive::<FROM>()).map(|to|
Arc::new(to) as ArrayRef)
-}
-
-fn numeric_to_bool_cast<T>(from: &PrimitiveArray<T>) -> Result<BooleanArray,
ArrowError>
-where
- T: ArrowPrimitiveType,
-{
- let mut b = BooleanBuilder::with_capacity(from.len());
-
- for i in 0..from.len() {
- if from.is_null(i) {
- b.append_null();
- } else {
- b.append_value(cast_num_to_bool::<T::Native>(from.value(i)));
- }
- }
-
- Ok(b.finish())
+ Ok(Arc::new(BooleanArray::from_unary(
+ from.as_primitive::<FROM>(),
+ cast_num_to_bool,
+ )))
}
/// Cast numeric types to boolean
@@ -2667,7 +2653,6 @@ fn cast_bool_to_numeric<TO>(
) -> Result<ArrayRef, ArrowError>
where
TO: ArrowPrimitiveType,
- TO::Native: num_traits::cast::NumCast,
{
Ok(Arc::new(bool_to_numeric_cast::<TO>(
from.as_any().downcast_ref::<BooleanArray>().unwrap(),
@@ -2678,20 +2663,24 @@ where
fn bool_to_numeric_cast<T>(from: &BooleanArray, _cast_options: &CastOptions)
-> PrimitiveArray<T>
where
T: ArrowPrimitiveType,
- T::Native: num_traits::NumCast,
{
- let iter = (0..from.len()).map(|i| {
- if from.is_null(i) {
- None
+ let to_numeric = |bits: u64, i: usize| {
+ if bits & (1 << i) != 0 {
+ T::Native::ONE
} else {
- single_bool_to_numeric::<T::Native>(from.value(i))
- }
- });
- // Benefit:
- // 20% performance improvement
- // Soundness:
- // The iterator is trustedLen because it comes from a Range
- unsafe { PrimitiveArray::<T>::from_trusted_len_iter(iter) }
+ T::Native::ZERO
+ }
+ };
+ // Unpack a 64-bit word at a time: the fixed-size inner loop helps the
+ // compiler vectorize the conversion.
+ let chunks = from.values().bit_chunks();
+ let mut values = Vec::with_capacity(from.len());
+ for bits in &chunks {
+ values.extend((0..64).map(|i| to_numeric(bits, i)));
+ }
+ let bits = chunks.remainder_bits();
+ values.extend((0..chunks.remainder_len()).map(|i| to_numeric(bits, i)));
+ PrimitiveArray::new(values.into(), from.nulls().cloned())
}
/// Cast single bool value to numeric value.
@@ -2872,7 +2861,7 @@ mod tests {
use crate::parse::parse_decimal;
use DataType::*;
use arrow_array::{Int64Array, RunArray, StringArray};
- use arrow_buffer::{Buffer, IntervalDayTime, NullBuffer};
+ use arrow_buffer::{BooleanBuffer, Buffer, IntervalDayTime, NullBuffer};
use arrow_buffer::{ScalarBuffer, i256};
use arrow_schema::{DataType, Field};
use chrono::NaiveDate;
@@ -5013,6 +5002,57 @@ mod tests {
assert!(!c.is_valid(2));
}
+ #[test]
+ fn test_cast_bool_numeric_sliced() {
+ for nulls in [
+ None,
+ Some(NullBuffer::from(
+ [false, true, true, false, true, true].repeat(12),
+ )),
+ Some(NullBuffer::new_null(72)),
+ ] {
+ // Include true values under nulls and slices crossing a bitmap
word boundary.
+ let booleans = BooleanArray::new(
+ BooleanBuffer::from([true, false, true, false, true,
false].repeat(12)),
+ nulls.clone(),
+ );
+ let numbers = Int32Array::new([1, 0, 1, 0, 1,
0].repeat(12).into(), nulls);
+ for data_type in [
+ Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64,
Float16, Float32, Float64,
+ ] {
+ let expected = cast(&numbers, &data_type).unwrap();
+ for offset in [0, 3, 63, 72] {
+ let input = booleans.slice(offset, 72 - offset);
+ let expected = expected.slice(offset, 72 - offset);
+ let actual = cast(&input, &data_type).unwrap();
+ assert_eq!(actual.as_ref(), expected.as_ref());
+
+ let actual = cast(expected.as_ref(), &Boolean).unwrap();
+ assert_eq!(actual.as_ref(), &input);
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn test_cast_float_to_bool_special_values() {
+ let numbers = Float64Array::from(vec![
+ 0.0,
+ -0.0,
+ 1.5,
+ -1.5,
+ f64::NAN,
+ f64::INFINITY,
+ f64::NEG_INFINITY,
+ ]);
+ let expected = BooleanArray::from(vec![false, false, true, true, true,
true, true]);
+ for data_type in [Float16, Float32, Float64] {
+ let input = cast(&numbers, &data_type).unwrap();
+ let actual = cast(input.as_ref(), &Boolean).unwrap();
+ assert_eq!(actual.as_ref(), &expected);
+ }
+ }
+
#[test]
fn test_cast_bool_to_utf8view() {
let array = BooleanArray::from(vec![Some(true), Some(false), None]);
diff --git a/arrow/benches/cast_kernels.rs b/arrow/benches/cast_kernels.rs
index 5956bdb89e..2f895069fc 100644
--- a/arrow/benches/cast_kernels.rs
+++ b/arrow/benches/cast_kernels.rs
@@ -259,6 +259,8 @@ fn cast_array(array: &ArrayRef, to_type: DataType) {
fn add_benchmark(c: &mut Criterion) {
let i32_array = build_array::<Int32Type>(512);
+ let i32_array_8192 = build_array::<Int32Type>(8192);
+ let bool_array: ArrayRef = Arc::new(create_boolean_array(8192, 0.1, 0.5));
let i64_array = build_array::<Int64Type>(512);
let f32_array = build_array::<Float32Type>(512);
let f32_utf8_array = cast(&build_array::<Float32Type>(512),
&DataType::Utf8).unwrap();
@@ -319,6 +321,12 @@ fn add_benchmark(c: &mut Criterion) {
c.bench_function("cast int32 to int64 512", |b| {
b.iter(|| cast_array(&i32_array, DataType::Int64))
});
+ c.bench_function("cast int32 to bool 8192", |b| {
+ b.iter(|| cast_array(&i32_array_8192, DataType::Boolean))
+ });
+ c.bench_function("cast bool to int32 8192", |b| {
+ b.iter(|| cast_array(&bool_array, DataType::Int32))
+ });
c.bench_function("cast float32 to int32 512", |b| {
b.iter(|| cast_array(&f32_array, DataType::Int32))
});
