Skip to content

Commit 26231bb

Browse files
authored
Preserve Boolean buffer handles during mask reduction (#10018)
## Summary Preserves the buffer handle and bit offset when masking a Boolean array. The previous reducer called `to_bit_buffer()`, which requires host memory and panics for device-backed values. ## Changes Rebuilds through `try_new_from_handle` and adds one device-buffer regression that checks buffer identity, bit offset, length, and attached validity. Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent d55f2ab commit 26231bb

1 file changed

Lines changed: 81 additions & 5 deletions

File tree

  • vortex-array/src/arrays/bool/compute

‎vortex-array/src/arrays/bool/compute/mask.rs‎

Lines changed: 81 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,31 +8,107 @@ use crate::IntoArray;
88
use crate::array::ArrayView;
99
use crate::arrays::Bool;
1010
use crate::arrays::BoolArray;
11-
use crate::arrays::bool::BoolArrayExt;
1211
use crate::scalar_fn::fns::mask::MaskReduce;
1312
use crate::validity::Validity;
1413

1514
impl MaskReduce for Bool {
1615
fn mask(array: ArrayView<'_, Bool>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
1716
Ok(Some(
18-
BoolArray::new(
19-
array.to_bit_buffer(),
17+
BoolArray::try_new_from_handle(
18+
array.bits.clone(),
19+
array.meta.offset(),
20+
array.len(),
2021
array.validity()?.and(Validity::Array(mask.clone()))?,
21-
)
22+
)?
2223
.into_array(),
2324
))
2425
}
2526
}
2627

2728
#[cfg(test)]
28-
mod test {
29+
mod tests {
30+
use std::any::Any;
31+
use std::ops::Range;
32+
use std::sync::Arc;
33+
34+
use futures::future::BoxFuture;
2935
use rstest::rstest;
36+
use vortex_buffer::Alignment;
37+
use vortex_buffer::ByteBuffer;
38+
use vortex_error::VortexResult;
39+
use vortex_error::vortex_bail;
3040

41+
use crate::ArrayEq;
42+
use crate::ArrayRef;
43+
use crate::EqMode;
3144
use crate::IntoArray;
3245
use crate::VortexSessionExecute;
3346
use crate::array_session;
47+
use crate::arrays::Bool;
3448
use crate::arrays::BoolArray;
49+
use crate::buffer::BufferHandle;
50+
use crate::buffer::DeviceBuffer;
3551
use crate::compute::conformance::mask::test_mask_conformance;
52+
use crate::scalar_fn::fns::mask::MaskReduce;
53+
use crate::validity::Validity;
54+
55+
// Any access beyond buffer metadata would invalidate the metadata-only reduction.
56+
#[derive(Debug, PartialEq, Eq, Hash)]
57+
struct MetadataOnlyDeviceBuffer(usize);
58+
59+
impl DeviceBuffer for MetadataOnlyDeviceBuffer {
60+
fn as_any(&self) -> &dyn Any {
61+
self
62+
}
63+
64+
fn len(&self) -> usize {
65+
self.0
66+
}
67+
68+
fn alignment(&self) -> Alignment {
69+
Alignment::of::<u8>()
70+
}
71+
72+
fn copy_to_host_sync(&self, _alignment: Alignment) -> VortexResult<ByteBuffer> {
73+
vortex_bail!("mask reduction must not copy device values")
74+
}
75+
76+
fn copy_to_host(
77+
&self,
78+
_alignment: Alignment,
79+
) -> VortexResult<BoxFuture<'static, VortexResult<ByteBuffer>>> {
80+
vortex_bail!("mask reduction must not copy device values")
81+
}
82+
83+
fn slice(&self, _range: Range<usize>) -> Arc<dyn DeviceBuffer> {
84+
panic!("mask reduction must preserve the original buffer");
85+
}
86+
87+
fn aligned(self: Arc<Self>, alignment: Alignment) -> VortexResult<Arc<dyn DeviceBuffer>> {
88+
assert_eq!(self.alignment(), alignment);
89+
Ok(self)
90+
}
91+
}
92+
93+
#[test]
94+
fn mask_preserves_device_buffer_and_bit_offset() -> VortexResult<()> {
95+
let handle = BufferHandle::new_device(Arc::new(MetadataOnlyDeviceBuffer(1)));
96+
let input = BoolArray::try_new_from_handle(handle, 3, 5, Validity::AllValid)?;
97+
let mask = BoolArray::from_iter([true, false, true, false, true]).into_array();
98+
let masked = <Bool as MaskReduce>::mask(input.as_ref().as_::<Bool>(), &mask)?
99+
.expect("Boolean masking must reduce");
100+
let output = masked.as_::<Bool>();
101+
102+
assert!(output.bits.array_eq(&input.bits, EqMode::Ptr));
103+
assert_eq!(output.meta.offset(), input.meta.offset());
104+
assert_eq!(output.len(), input.len());
105+
let Validity::Array(validity) = masked.validity()? else {
106+
panic!("mask must remain array-backed");
107+
};
108+
assert!(ArrayRef::ptr_eq(&validity, &mask));
109+
110+
Ok(())
111+
}
36112

37113
#[rstest]
38114
#[case(BoolArray::from_iter([true, false, true, true, false]))]

0 commit comments

Comments
 (0)