Repository: pola-rs/polars
File: crates/polars-expr/src/expressions/mod.rs (Rust)

Code before:
    /// Do the group indices reference all values in the aggregation state.
    fn groups_cover_all_values(&mut self) -> bool {
        if self.original_len || self.is_literal() {
            return true;
        }

        let num_values = self.flat_naive_length();
        match self.groups().as_ref().as_ref() {
            GroupsType::Idx(groups) => {
                let mut seen = MutableBitmap::from_len_zeroed(num_values);
                for (_, g) in groups {
                    for i in g.iter() {
                        unsafe { seen.set_unchecked(*i as usize, true) };
                    }
                }
                seen.unset_bits() == 0
            },
            GroupsType::Slice {
                groups,
                overlapping: true,
            } => {
                let mut seen = MutableBitmap::from_len_zeroed(num_values);
                for [start, length] in groups {
                    for i in 0..*length {
                        unsafe { seen.set_unchecked((*start + i) as usize, true) };
                    }
                }
                seen.unset_bits() == 0
            },

            // If we don't have overlapping data, we can just do a count.
            GroupsType::Slice {
                groups,
                overlapping: false,
            } => groups.iter().map(|[_, l]| *l as usize).sum::<usize>() == num_values,
        }
    }
}

/// Take a DataFrame and evaluate the expressions.
/// Implement this for Column, lt, eq, etc
pub trait PhysicalExpr: Send + Sync {
    fn as_expression(&self) -> Option<&Expr> {
        None
    }

    /// Take a DataFrame and evaluate the expression.
    fn evaluate(&self, df: &DataFrame, _state: &ExecutionState) -> PolarsResult<Column>;

    /// Some expression that are not aggregations can be done per group
    /// Think of sort, slice, filter, shift, etc.

Review comment (performance):
This may be expensive if there is much overlap. However, GroupsType::Slice is sorted, at least per the documentation. It is therefore possible to implement this much more efficiently, including a cheap assert if the slices are not sorted.

Code after:
    /// Do the group indices reference all values in the aggregation state.
    fn groups_cover_all_values(&mut self) -> bool {
        if self.original_len
            || matches!(
                self.state,
                AggState::LiteralScalar(_) | AggState::AggregatedScalar(_)
            )
        {
            return true;
        }

        let num_values = self.flat_naive_length();
        match self.groups().as_ref().as_ref() {
            GroupsType::Idx(groups) => {
                let mut seen = MutableBitmap::from_len_zeroed(num_values);
                for (_, g) in groups {
                    for i in g.iter() {
                        unsafe { seen.set_unchecked(*i as usize, true) };
                    }
                }
                seen.unset_bits() == 0
            },
            GroupsType::Slice {
                groups,
                overlapping: true,
            } => {
                // @NOTE: Slice groups are sorted by their `start` value.
                let mut offset = 0;
                let mut covers_all = true;
                for [start, length] in groups {
                    covers_all &= *start <= offset;
                    offset = start + length;
                }
                covers_all && offset == num_values as IdxSize
            },

            // If we don't have overlapping data, we can just do a count.
            GroupsType::Slice {
                groups,
                overlapping: false,
            } => groups.iter().map(|[_, l]| *l as usize).sum::<usize>() == num_values,
        }
    }
}

/// Take a DataFrame and evaluate the expressions.
/// Implement this for Column, lt, eq, etc
pub trait PhysicalExpr: Send + Sync {
    fn as_expression(&self) -> Option<&Expr> {
        None
    }