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

Code before:
    pub(crate) fn flat_naive(&self) -> Cow<'_, Column> {
        match &self.state {
            AggState::NotAggregated(c) => Cow::Borrowed(c),
            AggState::AggregatedList(c) => {
                if cfg!(debug_assertions) {
                    // Warning, so we find cases where we accidentally explode overlapping groups
                    // We don't want this as this can create a lot of data
                    if let GroupsType::Slice {
                        overlapping: true, ..
                    } = self.groups.as_ref().as_ref()
                    {
                        polars_warn!(
                            "performance - an aggregated list with overlapping groups may consume excessive memory"
                        )
                    }
                }

                // We should not insert nulls, otherwise the offsets in the groups will not be correct.
                Cow::Owned(c.explode(true).unwrap())
            },
            AggState::AggregatedScalar(c) => Cow::Borrowed(c),
            AggState::LiteralScalar(c) => Cow::Borrowed(c),
        }
    }

    fn flat_naive_length(&self) -> usize {
        match &self.state {
            AggState::NotAggregated(c) => c.len(),
            AggState::AggregatedList(c) => c.list().unwrap().inner_length(),
            AggState::AggregatedScalar(c) => c.len(),
            AggState::LiteralScalar(_) => 1,
        }
    }

    /// Take the series.
    pub(crate) fn take(&mut self) -> Column {
        let c = match &mut self.state {
            AggState::NotAggregated(c)
            | AggState::AggregatedScalar(c)
            | AggState::AggregatedList(c) => c,
            AggState::LiteralScalar(c) => c,
        };
        std::mem::take(c)
    }

    /// 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;
        }


Review comment (suggestion):
Doing a `flat_naive` just to get the length is quite costly. I have ran into this problem several times, so I think having a function for it is better.

Code after:
    pub(crate) fn flat_naive(&self) -> Cow<'_, Column> {
        match &self.state {
            AggState::NotAggregated(c) => Cow::Borrowed(c),
            AggState::AggregatedList(c) => {
                if cfg!(debug_assertions) {
                    // Warning, so we find cases where we accidentally explode overlapping groups
                    // We don't want this as this can create a lot of data
                    if let GroupsType::Slice {
                        overlapping: true, ..
                    } = self.groups.as_ref().as_ref()
                    {
                        polars_warn!(
                            "performance - an aggregated list with overlapping groups may consume excessive memory"
                        )
                    }
                }

                // We should not insert nulls, otherwise the offsets in the groups will not be correct.
                Cow::Owned(c.explode(true).unwrap())
            },
            AggState::AggregatedScalar(c) => Cow::Borrowed(c),
            AggState::LiteralScalar(c) => Cow::Borrowed(c),
        }
    }

    fn flat_naive_length(&self) -> usize {
        match &self.state {
            AggState::NotAggregated(c) => c.len(),
            AggState::AggregatedList(c) => c.list().unwrap().inner_length(),
            AggState::AggregatedScalar(c) => c.len(),
            AggState::LiteralScalar(_) => 1,
        }
    }

    /// Take the series.
    pub(crate) fn take(&mut self) -> Column {
        let c = match &mut self.state {
            AggState::NotAggregated(c)
            | AggState::AggregatedScalar(c)
            | AggState::AggregatedList(c) => c,
            AggState::LiteralScalar(c) => c,
        };
        std::mem::take(c)
    }

    /// 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(_)