h3ron_polars/
chunkedarray.rs

1use crate::iter::{
2    iter_indexes_nonvalidated, iter_indexes_validated, NonValidatedIndexIter, ValidatedIndexIter,
3};
4use crate::Error;
5use h3ron::{H3Cell, H3DirectedEdge, Index};
6use polars::export::arrow::bitmap::{Bitmap, MutableBitmap};
7use polars::prelude::{TakeRandom, UInt64Chunked};
8use std::marker::PhantomData;
9
10pub trait IndexValue: Index + TryFrom<u64, Error = h3ron::Error> + Clone {}
11
12impl IndexValue for H3Cell {}
13impl IndexValue for H3DirectedEdge {}
14
15#[derive(Clone)]
16pub struct IndexChunked<'a, IX: IndexValue> {
17    pub chunked_array: &'a UInt64Chunked,
18    index_phantom: PhantomData<IX>,
19}
20
21impl<'a, IX: IndexValue> IndexChunked<'a, IX> {
22    /// iterate over the `Index` values in this array.
23    ///
24    /// The contained `u64` values are validated and returned as Results
25    pub fn iter_indexes_validated(&self) -> ValidatedIndexIter<IX> {
26        iter_indexes_validated::<IX>(self.chunked_array)
27    }
28
29    /// iterate over the `Index` values in this array.
30    ///
31    /// The contained `u64` values are not validated, so there may be invalid `Index` values.
32    pub fn iter_indexes_nonvalidated(&self) -> NonValidatedIndexIter<IX> {
33        iter_indexes_nonvalidated::<IX>(self.chunked_array)
34    }
35
36    pub fn len(&self) -> usize {
37        self.chunked_array.len()
38    }
39
40    pub fn is_empty(&self) -> bool {
41        self.chunked_array.is_empty()
42    }
43
44    pub fn validity_bitmap(&self) -> Bitmap {
45        let mut mask = MutableBitmap::with_capacity(self.len());
46        for v in self.iter_indexes_nonvalidated() {
47            mask.push(match v {
48                None => false,
49                Some(index) => index.is_valid(),
50            })
51        }
52        mask.into()
53    }
54
55    pub fn to_collection<C>(&self) -> Result<C, Error>
56    where
57        C: FromIterator<IX>,
58    {
59        self.iter_indexes_validated()
60            .flatten()
61            .collect::<Result<C, _>>()
62    }
63}
64
65impl<'a, IX: IndexValue> TakeRandom for IndexChunked<'a, IX> {
66    type Item = IX;
67
68    /// get a nullable value by index. The returned `Index` is not validated.
69    fn get(&self, index: usize) -> Option<Self::Item> {
70        self.chunked_array.get(index).map(IX::new)
71    }
72
73    fn last(&self) -> Option<Self::Item> {
74        self.chunked_array.last().map(IX::new)
75    }
76}
77
78pub trait AsH3IndexChunked {
79    fn h3indexchunked<IX: IndexValue>(&self) -> IndexChunked<IX>;
80}
81
82impl AsH3IndexChunked for UInt64Chunked {
83    fn h3indexchunked<IX: IndexValue>(&self) -> IndexChunked<IX> {
84        IndexChunked {
85            chunked_array: self,
86            index_phantom: PhantomData::<IX>,
87        }
88    }
89}
90
91macro_rules! specialized_as_impl {
92    ($name:ident, $fn_name:ident, $ret_type:ty) => {
93        pub trait $name {
94            fn $fn_name(&self) -> IndexChunked<$ret_type>;
95        }
96
97        impl $name for UInt64Chunked {
98            fn $fn_name(&self) -> IndexChunked<$ret_type> {
99                self.h3indexchunked()
100            }
101        }
102    };
103}
104
105specialized_as_impl!(AsH3CellChunked, h3cell, H3Cell);
106specialized_as_impl!(AsH3DirectedEdgeChunked, h3directededge, H3DirectedEdge);
h3ron_polars/chunkedarray.rs

h3ron_polars/
chunkedarray.rs
