@@ -8,31 +8,107 @@ use crate::IntoArray;
88use crate :: array:: ArrayView ;
99use crate :: arrays:: Bool ;
1010use crate :: arrays:: BoolArray ;
11- use crate :: arrays:: bool:: BoolArrayExt ;
1211use crate :: scalar_fn:: fns:: mask:: MaskReduce ;
1312use crate :: validity:: Validity ;
1413
1514impl 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