Skip to content

Commit

Permalink
iter str values from categorical
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Nov 8, 2021
1 parent ad92f1e commit 82f2ae5
Showing 1 changed file with 30 additions and 0 deletions.
30 changes: 30 additions & 0 deletions polars/polars-core/src/chunked_array/categorical/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,38 @@ impl CategoricalChunked {
self.categorical_map = other.categorical_map.clone();
self
}

/// Create an `[Iterator]` that iterates over the `&str` values of the `[CategoricalChunked]`.
pub fn iter_str(&self) -> CatIter<'_> {
let iter = self.deref().into_iter();
CatIter {
rev: self.categorical_map.as_ref().unwrap(),
iter,
}
}
}

pub struct CatIter<'a> {
rev: &'a RevMapping,
iter: Box<dyn PolarsIterator<Item = Option<u32>> + 'a>,
}

impl<'a> Iterator for CatIter<'a> {
type Item = Option<&'a str>;

fn next(&mut self) -> Option<Self::Item> {
self.iter
.next()
.map(|item| item.map(|idx| self.rev.get(idx)))
}

fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
}

impl<'a> ExactSizeIterator for CatIter<'a> {}

impl Deref for CategoricalChunked {
type Target = UInt32Chunked;

Expand Down

0 comments on commit 82f2ae5

Please sign in to comment.