nd_arrow_array/compute/
search.rs

1//! Sorting and searching: `searchsorted`, `sort`, `argsort`, `unique`.
2
3use std::cmp::Ordering;
4use std::sync::Arc;
5
6use arrow::array::make_comparator;
7use arrow::array::{Array, ArrayRef, UInt32Array, UInt64Array};
8use arrow::compute::{SortOptions, cast as arrow_cast, sort_to_indices, take as arrow_take};
9
10use crate::NdArrowArray;
11use crate::compute::shape::move_dim;
12use crate::dimensions::Dimensions;
13use crate::error::{Result, nd_err};
14
15/// Which insertion point [`searchsorted`] returns for equal values.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
17pub enum SearchSide {
18    /// The first suitable position, like `side="left"`.
19    #[default]
20    Left,
21    /// The last suitable position, like `side="right"`.
22    Right,
23}
24
25const ASCENDING: SortOptions = SortOptions {
26    descending: false,
27    nulls_first: false,
28};
29
30/// Insertion points of `values` into the ascending 1-D array `sorted`, like
31/// `numpy.searchsorted`. The result has the dimensions of `values` and type
32/// `UInt64`. A null value gives a null.
33pub fn searchsorted(
34    sorted: &NdArrowArray,
35    values: &NdArrowArray,
36    side: SearchSide,
37) -> Result<NdArrowArray> {
38    if sorted.rank() != 1 {
39        return Err(nd_err!(
40            InvalidArgument,
41            "searchsorted needs a 1-D sorted array"
42        ));
43    }
44    let haystack = sorted.values();
45    let needles = arrow_cast(values.values().as_ref(), haystack.data_type())?;
46    let compare = make_comparator(haystack.as_ref(), needles.as_ref(), ASCENDING)?;
47    let out: UInt64Array = (0..needles.len())
48        .map(|j| {
49            needles.is_valid(j).then(|| {
50                partition_point(haystack.len(), |i| match side {
51                    SearchSide::Left => compare(i, j) == Ordering::Less,
52                    SearchSide::Right => compare(i, j) != Ordering::Greater,
53                }) as u64
54            })
55        })
56        .collect();
57    NdArrowArray::try_new(Arc::new(out), values.dims().clone())
58}
59
60/// First index in `0..len` where `pred` is false. `pred` must be true then false.
61pub(crate) fn partition_point(len: usize, pred: impl Fn(usize) -> bool) -> usize {
62    let (mut low, mut high) = (0usize, len);
63    while low < high {
64        let mid = low + (high - low) / 2;
65        if pred(mid) {
66            low = mid + 1;
67        } else {
68            high = mid;
69        }
70    }
71    low
72}
73
74/// Per-row sort indices with `dim` moved last. Returns the moved array and the
75/// global take indices.
76fn sort_rows(
77    a: &NdArrowArray,
78    dim: &str,
79    descending: bool,
80) -> Result<(NdArrowArray, Vec<u32>, usize)> {
81    let last = a
82        .rank()
83        .checked_sub(1)
84        .ok_or_else(|| nd_err!(InvalidArgument, "cannot sort a scalar"))?;
85    let moved = move_dim(a, dim, last)?;
86    let row = moved.dims().as_slice()[last].size();
87    let options = SortOptions {
88        descending,
89        nulls_first: false,
90    };
91    let mut indices = Vec::with_capacity(moved.len());
92    if row > 0 {
93        for start in (0..moved.len()).step_by(row) {
94            let slice = moved.values().slice(start, row);
95            let local = sort_to_indices(slice.as_ref(), Some(options), None)?;
96            indices.extend(local.values().iter().map(|&i| i + start as u32));
97        }
98    }
99    Ok((moved, indices, row))
100}
101
102/// Sort along `dim`, like `numpy.sort`. Nulls go last. Keeps the attributes.
103pub fn sort(a: &NdArrowArray, dim: &str, descending: bool) -> Result<NdArrowArray> {
104    let (moved, indices, _) = sort_rows(a, dim, descending)?;
105    let taken = arrow_take(moved.values().as_ref(), &UInt32Array::from(indices), None)?;
106    let sorted =
107        NdArrowArray::try_new(taken, moved.dims().clone())?.with_attributes(a.attributes().clone());
108    move_dim(&sorted, dim, a.dims().axis(dim)?)
109}
110
111/// Indices that sort along `dim`, like `numpy.argsort`. Type `UInt64`.
112pub fn argsort(a: &NdArrowArray, dim: &str, descending: bool) -> Result<NdArrowArray> {
113    let (moved, indices, row) = sort_rows(a, dim, descending)?;
114    let local: UInt64Array = indices
115        .iter()
116        .map(|&i| (i as usize % row.max(1)) as u64)
117        .collect();
118    let out = NdArrowArray::try_new(Arc::new(local), moved.dims().clone())?;
119    move_dim(&out, dim, a.dims().axis(dim)?)
120}
121
122/// Sorted unique non-null values of all elements, like `numpy.unique`, as a
123/// 1-D array along a new dimension `dim`.
124pub fn unique(a: &NdArrowArray, dim: &str) -> Result<NdArrowArray> {
125    let values = a.values();
126    let order = sort_to_indices(values.as_ref(), Some(ASCENDING), None)?;
127    let sorted: ArrayRef = arrow_take(values.as_ref(), &order, None)?;
128    let compare = make_comparator(sorted.as_ref(), sorted.as_ref(), ASCENDING)?;
129    let keep: UInt32Array = (0..sorted.len())
130        .filter(|&i| sorted.is_valid(i) && (i == 0 || compare(i - 1, i) != Ordering::Equal))
131        .map(|i| i as u32)
132        .collect();
133    let out = arrow_take(sorted.as_ref(), &keep, None)?;
134    let len = out.len();
135    NdArrowArray::try_new(out, Dimensions::try_new([(dim, len)])?)
136}
137
138/// Indices of the `true` (or non-zero) elements, like `numpy.argwhere`.
139///
140/// The result is `Int64` with dimensions `(point, axis)`: one row per match,
141/// one column per axis of `a`, in C order. Nulls do not match.
142pub fn argwhere(a: &NdArrowArray) -> Result<NdArrowArray> {
143    let mask = truthy(a)?;
144    let strides = a.dims().c_strides();
145    let shape = a.shape();
146    let mut out: Vec<i64> = Vec::new();
147    for (position, value) in mask.iter().enumerate() {
148        if value == Some(true) {
149            out.extend(
150                strides
151                    .iter()
152                    .zip(&shape)
153                    .map(|(&stride, &size)| ((position / stride.max(1)) % size.max(1)) as i64),
154            );
155        }
156    }
157    let rank = a.rank();
158    let points = out.len().checked_div(rank).unwrap_or(mask.true_count());
159    NdArrowArray::try_new(
160        Arc::new(arrow::array::Int64Array::from(out)),
161        Dimensions::try_new([("point", points), ("axis", rank)])?,
162    )
163}
164
165fn truthy(a: &NdArrowArray) -> Result<arrow::array::BooleanArray> {
166    if a.data_type() == &arrow::datatypes::DataType::Boolean {
167        return Ok(a
168            .values()
169            .as_any()
170            .downcast_ref::<arrow::array::BooleanArray>()
171            .unwrap()
172            .clone());
173    }
174    let zero = arrow_cast(&arrow::array::Int8Array::from(vec![0]), a.data_type())?;
175    Ok(arrow::compute::kernels::cmp::neq(
176        a.values(),
177        &arrow::array::Scalar::new(zero),
178    )?)
179}
180
181/// `true` where the element is in `test_values`, like `numpy.isin`.
182/// A null element gives `false`.
183pub fn isin(a: &NdArrowArray, test_values: &ArrayRef) -> Result<NdArrowArray> {
184    let tests = arrow_cast(test_values.as_ref(), a.data_type())?;
185    let order = sort_to_indices(tests.as_ref(), Some(ASCENDING), None)?;
186    let sorted = arrow_take(tests.as_ref(), &order, None)?;
187    let valid = sorted.len() - sorted.null_count();
188    let compare = make_comparator(sorted.as_ref(), a.values().as_ref(), ASCENDING)?;
189    let out: arrow::array::BooleanArray = (0..a.len())
190        .map(|j| {
191            Some(
192                a.values().is_valid(j) && {
193                    let i = partition_point(valid, |i| compare(i, j) == Ordering::Less);
194                    i < valid && compare(i, j) == Ordering::Equal
195                },
196            )
197        })
198        .collect();
199    NdArrowArray::try_new(Arc::new(out), a.dims().clone())
200}
201
202/// The bin of each element for increasing `bins` edges, like `numpy.digitize`.
203///
204/// With `right == false`, element `x` gets `i` where `bins[i-1] <= x < bins[i]`.
205/// The result is `UInt64` with the dimensions of `a`.
206pub fn digitize(a: &NdArrowArray, bins: &NdArrowArray, right: bool) -> Result<NdArrowArray> {
207    let side = if right {
208        SearchSide::Left
209    } else {
210        SearchSide::Right
211    };
212    searchsorted(bins, a, side)
213}
214
215#[cfg(test)]
216mod tests {
217    use arrow::array::{AsArray, Float64Array, Int32Array, StringArray};
218    use arrow::datatypes::{Int32Type, UInt64Type};
219
220    use super::*;
221
222    fn nd(values: ArrayRef, dims: &[(&str, usize)]) -> NdArrowArray {
223        NdArrowArray::try_new(values, Dimensions::try_new(dims.iter().copied()).unwrap()).unwrap()
224    }
225
226    fn u64s(a: &NdArrowArray) -> Vec<Option<u64>> {
227        a.values().as_primitive::<UInt64Type>().iter().collect()
228    }
229
230    #[test]
231    fn searchsorted_sides() {
232        let sorted = nd(Arc::new(Int32Array::from(vec![1, 2, 2, 3])), &[("x", 4)]);
233        let needles = nd(
234            Arc::new(Float64Array::from(vec![
235                Some(2.0),
236                Some(0.0),
237                Some(9.0),
238                None,
239            ])),
240            &[("q", 4)],
241        );
242        let left = searchsorted(&sorted, &needles, SearchSide::Left).unwrap();
243        assert_eq!(u64s(&left), vec![Some(1), Some(0), Some(4), None]);
244        let right = searchsorted(&sorted, &needles, SearchSide::Right).unwrap();
245        assert_eq!(u64s(&right)[0], Some(3));
246    }
247
248    #[test]
249    fn sort_and_argsort_along_axis() {
250        let a = nd(
251            Arc::new(Int32Array::from(vec![3, 1, 2, 9, 7, 8])),
252            &[("x", 2), ("y", 3)],
253        );
254        let s = sort(&a, "y", false).unwrap();
255        assert_eq!(
256            s.values().as_primitive::<Int32Type>().values(),
257            &[1, 2, 3, 7, 8, 9]
258        );
259        let sx = sort(&a, "x", true).unwrap();
260        assert_eq!(sx.dims().names(), vec!["x", "y"]);
261        assert_eq!(
262            sx.values().as_primitive::<Int32Type>().values(),
263            &[9, 7, 8, 3, 1, 2]
264        );
265        let i = argsort(&a, "y", false).unwrap();
266        assert_eq!(u64s(&i), [1, 2, 0, 1, 2, 0].map(Some).to_vec());
267    }
268
269    #[test]
270    fn argwhere_isin_digitize() {
271        let a = nd(
272            Arc::new(Int32Array::from(vec![0, 3, 0, 5, 6, 0])),
273            &[("y", 2), ("x", 3)],
274        );
275        let w = argwhere(&a).unwrap();
276        assert_eq!(w.shape(), vec![3, 2]);
277        let v: Vec<i64> = w
278            .values()
279            .as_primitive::<arrow::datatypes::Int64Type>()
280            .values()
281            .to_vec();
282        assert_eq!(v, vec![0, 1, 1, 0, 1, 1]);
283
284        let tests: ArrayRef = Arc::new(Int32Array::from(vec![Some(6), None, Some(3)]));
285        let m = isin(&a, &tests).unwrap();
286        let m: Vec<bool> = m.values().as_boolean().iter().map(Option::unwrap).collect();
287        assert_eq!(m, vec![false, true, false, false, true, false]);
288
289        let bins = nd(
290            Arc::new(Float64Array::from(vec![0.0, 1.0, 2.5])),
291            &[("edge", 3)],
292        );
293        let x = nd(
294            Arc::new(Float64Array::from(vec![-1.0, 0.0, 1.0, 2.0, 3.0])),
295            &[("x", 5)],
296        );
297        assert_eq!(
298            u64s(&digitize(&x, &bins, false).unwrap()),
299            [0, 1, 2, 2, 3].map(Some).to_vec()
300        );
301        assert_eq!(
302            u64s(&digitize(&x, &bins, true).unwrap()),
303            [0, 0, 1, 2, 3].map(Some).to_vec()
304        );
305    }
306
307    #[test]
308    fn unique_strings() {
309        let a = nd(
310            Arc::new(StringArray::from(vec![
311                Some("b"),
312                Some("a"),
313                None,
314                Some("b"),
315            ])),
316            &[("x", 4)],
317        );
318        let u = unique(&a, "label").unwrap();
319        let v: Vec<&str> = u
320            .values()
321            .as_string::<i32>()
322            .iter()
323            .map(Option::unwrap)
324            .collect();
325        assert_eq!(v, vec!["a", "b"]);
326        assert_eq!(u.dims().names(), vec!["label"]);
327    }
328}
nd_arrow_array/compute/search.rs

nd_arrow_array/compute/
search.rs
