@@ -180,13 +180,41 @@ where
180
180
181
181
// partition indices into valid and null indices
182
182
fn 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)
189
215
}
216
+ // faster path
217
+ _ => ( ( 0 ..( len as u32 ) ) . collect ( ) , vec ! [ ] ) ,
190
218
}
191
219
}
192
220
0 commit comments