h3ron_polars/algorithm/frame/
valid.rs

1use crate::frame::H3DataFrame;
2use crate::{AsH3IndexChunked, Error, IndexValue};
3use polars::export::arrow::array::BooleanArray;
4use polars::prelude::{BooleanChunked, DataFrame, FromData};
5
6pub trait FilterH3IsValid {
7    fn filter_h3_is_valid<IX, S>(&self, index_column_name: S) -> Result<Self, Error>
8    where
9        Self: Sized,
10        IX: IndexValue,
11        S: AsRef<str>;
12}
13
14fn filter_h3_is_valid<IX>(df: &DataFrame, index_column_name: &str) -> Result<DataFrame, Error>
15where
16    IX: IndexValue,
17{
18    let indexchunked = df.column(index_column_name)?.u64()?.h3indexchunked::<IX>();
19    let ba = BooleanArray::from_data_default(indexchunked.validity_bitmap(), None);
20
21    Ok(df.filter(&BooleanChunked::from(ba))?)
22}
23
24impl FilterH3IsValid for DataFrame {
25    /// Remove all rows where the contained `Index` in column `index_column_name` is invalid.
26    ///
27    /// # Example
28    ///
29    /// ```
30    /// use polars::prelude::{DataFrame, Series, NamedFrom, TakeRandom};
31    /// use h3ron::{H3Cell, Index};
32    /// use h3ron_polars::algorithm::frame::FilterH3IsValid;
33    ///
34    /// let h3index = H3Cell::from_coordinate((56.4, 23.2).into(), 5)
35    ///     .unwrap()
36    ///     .h3index();
37    ///
38    /// let df = DataFrame::new(vec![Series::new("x", vec![Some(56), Some(h3index), None])]).unwrap();
39    /// let df = df.filter_h3_is_valid::<H3Cell, _>("x").unwrap();
40    ///
41    /// assert_eq!(df.shape().0, 1);
42    /// assert_eq!(df.column("x").unwrap().u64().unwrap().get(0), Some(h3index));
43    /// ```
44    fn filter_h3_is_valid<IX, S>(&self, index_column_name: S) -> Result<Self, Error>
45    where
46        Self: Sized,
47        IX: IndexValue,
48        S: AsRef<str>,
49    {
50        filter_h3_is_valid::<IX>(self, index_column_name.as_ref())
51    }
52}
53
54impl<IX: IndexValue> H3DataFrame<IX> {
55    pub fn filter_h3_is_valid(&self) -> Result<Self, Error> {
56        filter_h3_is_valid::<IX>(self.dataframe(), self.h3index_column_name())
57            .map(|df| H3DataFrame::from_dataframe_nonvalidated(df, self.h3index_column_name()))
58    }
59}
h3ron_polars/algorithm/frame/valid.rs

h3ron_polars/algorithm/frame/
valid.rs
