polars_core/frame/group_by/aggregations/

boolean.rs

use arrow::bitmap::bitmask::BitMask;

use super::*;
use crate::chunked_array::cast::CastOptions;
use crate::chunked_array::{arg_max_bool, arg_min_bool};

pub fn _agg_helper_idx_bool<F>(groups: &GroupsIdx, f: F) -> Series
where
    F: Fn((IdxSize, &IdxVec)) -> Option<bool> + Send + Sync,
{
    let ca: BooleanChunked = RAYON.install(|| groups.into_par_iter().map(f).collect());
    ca.into_series()
}

pub fn _agg_helper_slice_bool<F>(groups: &[[IdxSize; 2]], f: F) -> Series
where
    F: Fn([IdxSize; 2]) -> Option<bool> + Send + Sync,
{
    let ca: BooleanChunked = RAYON.install(|| groups.par_iter().copied().map(f).collect());
    ca.into_series()
}

#[cfg(feature = "bitwise")]
impl BooleanChunked {
    pub(crate) unsafe fn agg_and(&self, groups: &GroupsType) -> BooleanChunked {
        self.agg_all(groups, true)
    }

    pub(crate) unsafe fn agg_or(&self, groups: &GroupsType) -> BooleanChunked {
        self.agg_any(groups, true)
    }

    pub(crate) unsafe fn agg_xor(&self, groups: &GroupsType) -> BooleanChunked {
        self.bool_agg(
            groups,
            true,
            |values, idxs| {
                idxs.iter()
                    .map(|i| {
                        <IdxSize as From<bool>>::from(unsafe {
                            values.get_bit_unchecked(*i as usize)
                        })
                    })
                    .sum::<IdxSize>()
                    % 2
                    == 1
            },
            |values, validity, idxs| {
                idxs.iter()
                    .map(|i| {
                        <IdxSize as From<bool>>::from(unsafe {
                            validity.get_bit_unchecked(*i as usize)
                                & values.get_bit_unchecked(*i as usize)
                        })
                    })
                    .sum::<IdxSize>()
                    % 2
                    == 1
            },
            |_, _, _| unreachable!(),
            |values, start, length| {
                unsafe { values.sliced_unchecked(start as usize, length as usize) }.set_bits() % 2
                    == 1
            },
            |values, validity, start, length| {
                let values = unsafe { values.sliced_unchecked(start as usize, length as usize) };
                let validity =
                    unsafe { validity.sliced_unchecked(start as usize, length as usize) };
                values.num_intersections_with(validity) % 2 == 1
            },
            |_, _, _, _| unreachable!(),
        )
    }
}

impl BooleanChunked {
    pub(crate) unsafe fn agg_min(&self, groups: &GroupsType) -> Series {
        // faster paths
        if !self.has_nulls() || matches!(groups, GroupsType::Slice { .. }) {
            match self.is_sorted_flag() {
                IsSorted::Ascending => {
                    return self.clone().into_series().agg_first_non_null(groups);
                },
                IsSorted::Descending => {
                    return self.clone().into_series().agg_last_non_null(groups);
                },
                _ => {},
            }
        }
        let ca_self = self.rechunk();
        let arr = ca_self.downcast_iter().next().unwrap();
        let no_nulls = arr.null_count() == 0;
        match groups {
            GroupsType::Idx(groups) => _agg_helper_idx_bool(groups, |(first, idx)| {
                debug_assert!(idx.len() <= self.len());
                if idx.is_empty() {
                    None
                } else if idx.len() == 1 {
                    arr.get(first as usize)
                } else if no_nulls {
                    take_arg_min_bool_iter_unchecked_no_nulls(arr, idx2usize(idx))
                        .map(|p| arr.value_unchecked(idx[p] as usize))
                } else {
                    take_arg_min_bool_iter_unchecked_nulls(arr, idx2usize(idx))
                        .map(|p| arr.value_unchecked(idx[p] as usize))
                }
            }),
            GroupsType::Slice {
                groups: groups_slice,
                ..
            } => _agg_helper_slice_bool(groups_slice, |[first, len]| {
                debug_assert!(len <= self.len() as IdxSize);
                match len {
                    0 => None,
                    1 => self.get(first as usize),
                    _ => {
                        let arr_group = _slice_from_offsets(self, first, len);
                        arr_group.min()
                    },
                }
            }),
        }
    }
    pub(crate) unsafe fn agg_max(&self, groups: &GroupsType) -> Series {
        // faster paths
        if !self.has_nulls() || matches!(groups, GroupsType::Slice { .. }) {
            match self.is_sorted_flag() {
                IsSorted::Ascending => return self.clone().into_series().agg_last_non_null(groups),
                IsSorted::Descending => {
                    return self.clone().into_series().agg_first_non_null(groups);
                },
                _ => {},
            }
        }

        let ca_self = self.rechunk();
        let arr = ca_self.downcast_iter().next().unwrap();
        let no_nulls = arr.null_count() == 0;
        match groups {
            GroupsType::Idx(groups) => _agg_helper_idx_bool(groups, |(first, idx)| {
                debug_assert!(idx.len() <= self.len());
                if idx.is_empty() {
                    None
                } else if idx.len() == 1 {
                    self.get(first as usize)
                } else if no_nulls {
                    take_arg_max_bool_iter_unchecked_no_nulls(arr, idx2usize(idx))
                        .map(|p| arr.value_unchecked(idx[p] as usize))
                } else {
                    take_arg_max_bool_iter_unchecked_nulls(arr, idx2usize(idx))
                        .map(|p| arr.value_unchecked(idx[p] as usize))
                }
            }),
            GroupsType::Slice {
                groups: groups_slice,
                ..
            } => _agg_helper_slice_bool(groups_slice, |[first, len]| {
                debug_assert!(len <= self.len() as IdxSize);
                match len {
                    0 => None,
                    1 => self.get(first as usize),
                    _ => {
                        let arr_group = _slice_from_offsets(self, first, len);
                        arr_group.max()
                    },
                }
            }),
        }
    }

    pub(crate) unsafe fn agg_arg_min(&self, groups: &GroupsType) -> Series {
        // faster paths
        if !self.has_nulls() || matches!(groups, GroupsType::Slice { .. }) {
            match self.is_sorted_flag() {
                IsSorted::Ascending => {
                    return self.clone().into_series().agg_arg_first_non_null(groups);
                },
                IsSorted::Descending => {
                    return self.clone().into_series().agg_arg_last_non_null(groups);
                },
                _ => {},
            }
        }

        let ca_self = self.rechunk();
        let arr = ca_self.downcast_iter().next().unwrap();
        let no_nulls = arr.null_count() == 0;
        match groups {
            GroupsType::Idx(groups) => agg_helper_idx_on_all::<IdxType, _>(groups, |idx| {
                debug_assert!(idx.len() <= ca_self.len());
                if idx.is_empty() {
                    None
                } else if idx.len() == 1 {
                    arr.get(idx[0] as usize).map(|_| 0)
                } else if no_nulls {
                    take_arg_min_bool_iter_unchecked_no_nulls(arr, idx2usize(idx))
                        .map(|p| p as IdxSize)
                } else {
                    take_arg_min_bool_iter_unchecked_nulls(arr, idx2usize(idx))
                        .map(|p| p as IdxSize)
                }
            }),
            GroupsType::Slice {
                groups: groups_slice,
                ..
            } => _agg_helper_slice::<IdxType, _>(groups_slice, |[first, len]| {
                debug_assert!(len <= self.len() as IdxSize);
                match len {
                    0 => None,
                    1 => self.get(first as usize).map(|_| 0),
                    _ => {
                        let group_ca = _slice_from_offsets(self, first, len);
                        arg_min_bool(&group_ca).map(|p| p as IdxSize)
                    },
                }
            }),
        }
    }

    pub(crate) unsafe fn agg_arg_max(&self, groups: &GroupsType) -> Series {
        // faster paths
        if !self.has_nulls() || matches!(groups, GroupsType::Slice { .. }) {
            match self.is_sorted_flag() {
                IsSorted::Ascending => {
                    return self.clone().into_series().agg_arg_last_non_null(groups);
                },
                IsSorted::Descending => {
                    return self.clone().into_series().agg_arg_first_non_null(groups);
                },
                _ => {},
            }
        }

        let ca_self = self.rechunk();
        let arr = ca_self.downcast_iter().next().unwrap();
        let no_nulls = arr.null_count() == 0;
        match groups {
            GroupsType::Idx(groups) => agg_helper_idx_on_all::<IdxType, _>(groups, |idx| {
                debug_assert!(idx.len() <= ca_self.len());
                if idx.is_empty() {
                    None
                } else if idx.len() == 1 {
                    arr.get(idx[0] as usize).map(|_| 0)
                } else if no_nulls {
                    take_arg_max_bool_iter_unchecked_no_nulls(arr, idx2usize(idx))
                        .map(|p| p as IdxSize)
                } else {
                    take_arg_max_bool_iter_unchecked_nulls(arr, idx2usize(idx))
                        .map(|p| p as IdxSize)
                }
            }),
            GroupsType::Slice {
                groups: groups_slice,
                ..
            } => _agg_helper_slice::<IdxType, _>(groups_slice, |[first, len]| {
                debug_assert!(len <= self.len() as IdxSize);
                match len {
                    0 => None,
                    1 => self.get(first as usize).map(|_| 0),
                    _ => {
                        let group_ca = _slice_from_offsets(self, first, len);
                        arg_max_bool(&group_ca).map(|p| p as IdxSize)
                    },
                }
            }),
        }
    }

    pub(crate) unsafe fn agg_sum(&self, groups: &GroupsType) -> Series {
        self.cast_with_options(&IDX_DTYPE, CastOptions::Overflowing)
            .unwrap()
            .agg_sum(groups)
    }

    /// # Safety
    ///
    /// Groups should be in correct.
    #[expect(clippy::too_many_arguments)]
    unsafe fn bool_agg(
        &self,
        groups: &GroupsType,
        ignore_nulls: bool,

        idx_no_valid: impl Fn(BitMask, &[IdxSize]) -> bool + Send + Sync,
        idx_validity: impl Fn(BitMask, BitMask, &[IdxSize]) -> bool + Send + Sync,
        idx_kleene: impl Fn(BitMask, BitMask, &[IdxSize]) -> Option<bool> + Send + Sync,

        slice_no_valid: impl Fn(BitMask, IdxSize, IdxSize) -> bool + Send + Sync,
        slice_validity: impl Fn(BitMask, BitMask, IdxSize, IdxSize) -> bool + Send + Sync,
        slice_kleene: impl Fn(BitMask, BitMask, IdxSize, IdxSize) -> Option<bool> + Send + Sync,
    ) -> BooleanChunked {
        let name = self.name().clone();
        let values = self.rechunk();
        let values = values.downcast_as_array();

        let groups_len = groups.len();

        RAYON.install(|| {
            let validity = values
                .validity()
                .filter(|v| v.unset_bits() > 0)
                .map(BitMask::from_bitmap);
            let values = BitMask::from_bitmap(values.values());

            if !ignore_nulls && let Some(validity) = validity {
                match groups {
                    GroupsType::Idx(idx) => {
                        let all = idx.all();
                        collect_bool_opt_par(name, groups_len, |g| {
                            idx_kleene(values, validity, &all[g])
                        })
                    },
                    GroupsType::Slice { groups, .. } => {
                        collect_bool_opt_par(name, groups_len, |g| {
                            let [s, l] = groups[g];
                            slice_kleene(values, validity, s, l)
                        })
                    },
                }
            } else {
                match groups {
                    GroupsType::Idx(idx) => {
                        let all = idx.all();
                        match validity {
                            None => collect_bool_par(name, groups_len, |g| {
                                idx_no_valid(values, &all[g])
                            }),
                            Some(validity) => collect_bool_par(name, groups_len, |g| {
                                idx_validity(values, validity, &all[g])
                            }),
                        }
                    },
                    GroupsType::Slice { groups, .. } => match validity {
                        None => collect_bool_par(name, groups_len, |g| {
                            let [s, l] = groups[g];
                            slice_no_valid(values, s, l)
                        }),
                        Some(validity) => collect_bool_par(name, groups_len, |g| {
                            let [s, l] = groups[g];
                            slice_validity(values, validity, s, l)
                        }),
                    },
                }
            }
        })
    }

    /// # Safety
    ///
    /// Groups should be in correct.
    pub unsafe fn agg_any(&self, groups: &GroupsType, ignore_nulls: bool) -> BooleanChunked {
        self.bool_agg(
            groups,
            ignore_nulls,
            |values, idxs| {
                idxs.iter()
                    .any(|i| unsafe { values.get_bit_unchecked(*i as usize) })
            },
            |values, validity, idxs| {
                idxs.iter().any(|i| unsafe {
                    validity.get_bit_unchecked(*i as usize) & values.get_bit_unchecked(*i as usize)
                })
            },
            |values, validity, idxs| {
                let mut saw_null = false;
                for i in idxs.iter() {
                    let is_valid = unsafe { validity.get_bit_unchecked(*i as usize) };
                    let is_true = unsafe { values.get_bit_unchecked(*i as usize) };

                    if is_valid & is_true {
                        return Some(true);
                    }
                    saw_null |= !is_valid;
                }
                (!saw_null).then_some(false)
            },
            |values, start, length| {
                unsafe { values.sliced_unchecked(start as usize, length as usize) }.leading_zeros()
                    < length as usize
            },
            |values, validity, start, length| {
                let values = unsafe { values.sliced_unchecked(start as usize, length as usize) };
                let validity =
                    unsafe { validity.sliced_unchecked(start as usize, length as usize) };
                values.intersects_with(validity)
            },
            |values, validity, start, length| {
                let values = unsafe { values.sliced_unchecked(start as usize, length as usize) };
                let validity =
                    unsafe { validity.sliced_unchecked(start as usize, length as usize) };

                if values.intersects_with(validity) {
                    Some(true)
                } else if validity.unset_bits() == 0 {
                    Some(false)
                } else {
                    None
                }
            },
        )
    }

    /// # Safety
    ///
    /// Groups should be in correct.
    pub unsafe fn agg_all(&self, groups: &GroupsType, ignore_nulls: bool) -> BooleanChunked {
        self.bool_agg(
            groups,
            ignore_nulls,
            |values, idxs| {
                idxs.iter()
                    .all(|i| unsafe { values.get_bit_unchecked(*i as usize) })
            },
            |values, validity, idxs| {
                idxs.iter().all(|i| unsafe {
                    !validity.get_bit_unchecked(*i as usize) | values.get_bit_unchecked(*i as usize)
                })
            },
            |values, validity, idxs| {
                let mut saw_null = false;
                for i in idxs.iter() {
                    let is_valid = unsafe { validity.get_bit_unchecked(*i as usize) };
                    let is_true = unsafe { values.get_bit_unchecked(*i as usize) };

                    if is_valid & !is_true {
                        return Some(false);
                    }
                    saw_null |= !is_valid;
                }
                (!saw_null).then_some(true)
            },
            |values, start, length| {
                let values = unsafe { values.sliced_unchecked(start as usize, length as usize) };
                values.unset_bits() == 0
            },
            |values, validity, start, length| {
                let values = unsafe { values.sliced_unchecked(start as usize, length as usize) };
                let validity =
                    unsafe { validity.sliced_unchecked(start as usize, length as usize) };
                values.num_intersections_with(validity) == validity.set_bits()
            },
            |values, validity, start, length| {
                let values = unsafe { values.sliced_unchecked(start as usize, length as usize) };
                let validity =
                    unsafe { validity.sliced_unchecked(start as usize, length as usize) };

                let num_non_nulls = validity.set_bits();

                if values.num_intersections_with(validity) < num_non_nulls {
                    Some(false)
                } else if num_non_nulls < values.len() {
                    None
                } else {
                    Some(true)
                }
            },
        )
    }
}

pub fn collect_bool_par<F>(name: PlSmallStr, len: usize, f: F) -> BooleanChunked
where
    F: Fn(usize) -> bool + Send + Sync,
{
    let n_bytes = len.div_ceil(8);
    let mut values: Vec<u8> = Vec::with_capacity(n_bytes);

    (0..n_bytes)
        .into_par_iter()
        .map(|b| {
            let lo = b * 8;
            let hi = (lo + 8).min(len);
            let mut v = 0u8;
            for (bit, g) in (lo..hi).enumerate() {
                v |= (f(g) as u8) << bit;
            }
            v
        })
        .collect_into_vec(&mut values);

    BooleanChunked::from_bitmap(name, Bitmap::from_u8_vec(values, len))
}

pub fn collect_bool_opt_par<F>(name: PlSmallStr, len: usize, f: F) -> BooleanChunked
where
    F: Fn(usize) -> Option<bool> + Send + Sync,
{
    let n_bytes = len.div_ceil(8);
    let mut values: Vec<u8> = Vec::with_capacity(n_bytes);
    let mut validity: Vec<u8> = Vec::with_capacity(n_bytes);

    (0..n_bytes)
        .into_par_iter()
        .map(|b| {
            let lo = b * 8;
            let hi = (lo + 8).min(len);
            let (mut v, mut m) = (0u8, 0u8);
            for (bit, g) in (lo..hi).enumerate() {
                if let Some(x) = f(g) {
                    m |= 1 << bit;
                    v |= (x as u8) << bit;
                }
            }
            (v, m)
        })
        .unzip_into_vecs(&mut values, &mut validity);

    let values = Bitmap::from_u8_vec(values, len);
    let validity = Bitmap::from_u8_vec(validity, len);
    let validity = (validity.unset_bits() > 0).then_some(validity);
    BooleanChunked::with_chunk(
        name,
        BooleanArray::new(ArrowDataType::Boolean, values, validity),
    )
}