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

Code before:
                    for &[first, len] in groups {
                        original_idx.extend(first..first + len)
                    }
                },
            };

            let mut original_idx_iter = original_idx.iter().copied();

            match ac.groups().as_ref().as_ref() {
                GroupsType::Idx(groups) => {
                    for g in groups.all() {
                        idx_mapping.extend(g.iter().copied().zip(&mut original_idx_iter));
                    }
                },
                GroupsType::Slice { groups, .. } => {
                    for &[first, len] in groups {
                        idx_mapping.extend((first..first + len).zip(&mut original_idx_iter));
                    }
                },
            }
            original_idx.clear();
            take_idx = original_idx;
        }
        // SAFETY:
        // we only have unique indices ranging from 0..len
        #[cfg(all(not(target_os = "emscripten"), target_family = "wasm"))]
        {
            unsafe { perfect_sort(&idx_mapping, &mut take_idx) };
        }

        #[cfg(not(all(not(target_os = "emscripten"), target_family = "wasm")))]
        {
            POOL.with(|pool| {
                unsafe { perfect_sort(pool, &idx_mapping, &mut take_idx) };
            });
        }
        Ok(IdxCa::from_vec(PlSmallStr::EMPTY, take_idx))
    }

    #[allow(clippy::too_many_arguments)]
    fn map_by_arg_sort(
        &self,
        df: &DataFrame,
        out_column: Column,
        flattened: &Column,
        mut ac: AggregationContext,
        group_by_columns: &[Column],
        gb: GroupBy,
        cache_key: String,
        state: &ExecutionState,
    ) -> PolarsResult<Column> {

Review comment (suggestion):
This should not have `cfg`. Usage should be unaware of this.

Code after:
                    for &[first, len] in groups {
                        original_idx.extend(first..first + len)
                    }
                },
            };

            let mut original_idx_iter = original_idx.iter().copied();

            match ac.groups().as_ref().as_ref() {
                GroupsType::Idx(groups) => {
                    for g in groups.all() {
                        idx_mapping.extend(g.iter().copied().zip(&mut original_idx_iter));
                    }
                },
                GroupsType::Slice { groups, .. } => {
                    for &[first, len] in groups {
                        idx_mapping.extend((first..first + len).zip(&mut original_idx_iter));
                    }
                },
            }
            original_idx.clear();
            take_idx = original_idx;
        }
        // SAFETY:
        // we only have unique indices ranging from 0..len
        unsafe { perfect_sort(&idx_mapping, &mut take_idx) };
        Ok(IdxCa::from_vec(PlSmallStr::EMPTY, take_idx))
    }

    #[allow(clippy::too_many_arguments)]
    fn map_by_arg_sort(
        &self,
        df: &DataFrame,
        out_column: Column,
        flattened: &Column,
        mut ac: AggregationContext,
        group_by_columns: &[Column],
        gb: GroupBy,
        cache_key: String,
        state: &ExecutionState,
    ) -> PolarsResult<Column> {
        // we use an arg_sort to map the values back

        // This is a bit more complicated because the final group tuples may differ from the original
        // so we use the original indices as idx values to arg_sort the original column
        //
        // The example below shows the naive version without group tuple mapping

        // columns
        // a b a a
        //