1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

use polars_arrow::trusted_len::PushUnchecked;

use super::*;

pub(super) fn arg_where(s: &mut [Series]) -> PolarsResult<Series> {
    let predicate = s[0].bool()?;

    if predicate.is_empty() {
        Ok(Series::full_null(predicate.name(), 0, &IDX_DTYPE))
    } else {
        let capacity = predicate.sum().unwrap();
        let mut out = Vec::with_capacity(capacity as usize);
        let mut cnt = 0 as IdxSize;

        predicate.downcast_iter().for_each(|arr| {
            let values = match arr.validity() {
                Some(validity) => validity & arr.values(),
                None => arr.values().clone(),
            };

            // todo! could use chunkiter from arrow here
            for bit in values.iter() {
                if bit {
                    // safety:
                    // we allocated enough slots
                    unsafe { out.push_unchecked(cnt) }
                }
                cnt += 1;
            }
        });
        let arr = Box::new(IdxArr::from_vec(out)) as ArrayRef;
        Ok(IdxCa::from_chunks(predicate.name(), vec![arr]).into_series())
    }
}