@@ -180,13 +180,41 @@ where
180180
181181// partition indices into valid and null indices
182182fn partition_validity ( array : & dyn Array ) -> ( Vec < u32 > , Vec < u32 > ) {
183- match array. null_count ( ) {
184- // faster path
185- 0 => ( ( 0 ..( array. len ( ) as u32 ) ) . collect ( ) , vec ! [ ] ) ,
186- _ => {
187- let indices = 0 ..( array. len ( ) as u32 ) ;
188- indices. partition ( |index| array. is_valid ( * index as usize ) )
183+ let len = array. len ( ) ;
184+ let null_count = array. null_count ( ) ;
185+ match array. nulls ( ) {
186+ Some ( nulls) if null_count > 0 => {
187+ let mut valid_indices = Vec :: with_capacity ( len - null_count) ;
188+ let mut null_indices = Vec :: with_capacity ( null_count) ;
189+
190+ let valid_slice = valid_indices. spare_capacity_mut ( ) ;
191+ let null_slice = null_indices. spare_capacity_mut ( ) ;
192+ let mut valid_idx = 0 ;
193+ let mut null_idx = 0 ;
194+
195+ nulls. into_iter ( ) . enumerate ( ) . for_each ( |( i, v) | {
196+ if v {
197+ valid_slice[ valid_idx] . write ( i as u32 ) ;
198+ valid_idx += 1 ;
199+ } else {
200+ null_slice[ null_idx] . write ( i as u32 ) ;
201+ null_idx += 1 ;
202+ }
203+ } ) ;
204+
205+ assert_eq ! ( null_idx, null_count) ;
206+ assert_eq ! ( valid_idx, len - null_count) ;
207+ // Safety: The new lengths match the initial capacity as asserted above,
208+ // the bounds checks while writing also ensure they less than or equal to the capacity.
209+ unsafe {
210+ valid_indices. set_len ( valid_idx) ;
211+ null_indices. set_len ( null_idx) ;
212+ }
213+
214+ ( valid_indices, null_indices)
189215 }
216+ // faster path
217+ _ => ( ( 0 ..( len as u32 ) ) . collect ( ) , vec ! [ ] ) ,
190218 }
191219}
192220
0 commit comments