h3ron_polars/algorithm/chunkedarray/
resolution.rs

1use crate::algorithm::chunkedarray::util::list_map_cells;
2use crate::{Error, FromIndexIterator, IndexChunked, IndexValue};
3use h3ron::error::check_valid_h3_resolution;
4use h3ron::iter::change_resolution;
5use h3ron::H3Cell;
6use polars_core::prelude::{ListChunked, UInt64Chunked, UInt8Chunked};
7use std::iter::once;
8
9/// Obtain the H3 Resolutions at the array positions where
10/// the contained `Index` values are valid.
11pub trait H3Resolution {
12    /// Obtain the H3 Resolutions at the array positions where
13    /// the contained `Index` values are valid.
14    fn h3_resolution(&self) -> UInt8Chunked;
15}
16
17impl<'a, IX: IndexValue> H3Resolution for IndexChunked<'a, IX> {
18    fn h3_resolution(&self) -> UInt8Chunked {
19        UInt8Chunked::from_iter(self.iter_indexes_validated().map(
20            |maybe_index| match maybe_index {
21                Some(Ok(index)) => Some(index.resolution()),
22                _ => None,
23            },
24        ))
25    }
26}
27
28/// Changes the resolution of the contained `H3Cell` values.
29pub trait H3ChangeResolution {
30    /// Changes the resolution of the contained `H3Cell` values
31    ///
32    /// For each cell of the input array a list of cells is produced. This list may
33    /// only contain a single element for cases where `target_resolution` is <= the array
34    /// elements resolution.
35    fn h3_change_resolution(&self, target_resolution: u8) -> Result<ListChunked, Error>;
36}
37
38impl<'a> H3ChangeResolution for IndexChunked<'a, H3Cell> {
39    fn h3_change_resolution(&self, target_resolution: u8) -> Result<ListChunked, Error> {
40        check_valid_h3_resolution(target_resolution)?;
41        list_map_cells(self, |cell| {
42            Ok(UInt64Chunked::from_index_iter(
43                change_resolution(once(cell), target_resolution)
44                    // todo: This error should not be hidden
45                    .filter_map(|cell| cell.ok()),
46            ))
47        })
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use crate::algorithm::{H3ChangeResolution, H3Resolution};
54    use crate::{AsH3CellChunked, FromIndexIterator, NamedFromIndexes};
55    use h3ron::{H3Cell, Index};
56    use polars_core::prelude::{ChunkExplode, TakeRandom, UInt64Chunked};
57
58    #[test]
59    fn cell_resolution() {
60        let expected_res = 6;
61        let cell = H3Cell::from_coordinate((4.5, 1.3).into(), expected_res).unwrap();
62        let ca = UInt64Chunked::from_index_iter([
63            Some(cell),
64            Some(H3Cell::new(55)), // invalid
65            None,
66        ]);
67        assert_eq!(ca.len(), 3);
68        let resolution_ca = ca.h3cell().h3_resolution();
69
70        assert_eq!(resolution_ca.get(0), Some(expected_res));
71        assert_eq!(resolution_ca.get(1), None);
72        assert_eq!(resolution_ca.get(2), None);
73    }
74
75    #[test]
76    fn cell_change_resolution_to_child() {
77        let cell = H3Cell::from_coordinate((4.5, 1.3).into(), 6).unwrap();
78        let ca = UInt64Chunked::new_from_indexes("", vec![cell]);
79        assert_eq!(ca.len(), 1);
80
81        let changed = ca.h3cell().h3_change_resolution(7).unwrap();
82        assert_eq!(changed.len(), 1);
83        let exploded = changed.explode().unwrap().unique().unwrap();
84        assert_eq!(exploded.len(), 7);
85    }
86}
h3ron_polars/algorithm/chunkedarray/resolution.rs

h3ron_polars/algorithm/chunkedarray/
resolution.rs
