nd_arrow_array/compute/
shape.rs

1//! Layout kernels: broadcast, transpose, squeeze, expand, select, take, concat.
2//!
3//! Axes are addressed by dimension name. Layout kernels keep the attributes.
4
5use std::sync::Arc;
6
7use arrow::array::{Array, MutableArrayData, make_array};
8
9use crate::NdArrowArray;
10use crate::compute::layout::{Gather, GatherAxis};
11use crate::dimensions::{Dimension, Dimensions};
12use crate::error::{Result, nd_err};
13
14/// A Python-style slice: `start:stop:step`. Negative indices count from the end.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct Slice {
17    /// First index. `None` is the start of the axis (the end for a negative step).
18    pub start: Option<isize>,
19    /// Index after the last one. `None` is past the end of the axis.
20    pub stop: Option<isize>,
21    /// Distance between indices. Negative steps go backward. Zero is an error.
22    pub step: isize,
23}
24
25impl Slice {
26    /// A slice `start:stop:step`.
27    pub fn new(start: Option<isize>, stop: Option<isize>, step: isize) -> Self {
28        Self { start, stop, step }
29    }
30
31    /// `start..stop` with step 1.
32    pub fn range(start: isize, stop: isize) -> Self {
33        Self::new(Some(start), Some(stop), 1)
34    }
35
36    /// `::step`.
37    pub fn step(step: isize) -> Self {
38        Self::new(None, None, step)
39    }
40
41    /// The whole axis, `:`.
42    pub fn full() -> Self {
43        Self::step(1)
44    }
45
46    /// Resolve against an axis of `len`: (first index, count, step), like `slice.indices`.
47    pub fn resolve(&self, len: usize) -> Result<(usize, usize, isize)> {
48        let step = self.step;
49        if step == 0 {
50            return Err(nd_err!(InvalidArgument, "slice step cannot be zero"));
51        }
52        let len = len as isize;
53        let clamp = |value: Option<isize>, default: isize, low: isize, high: isize| match value {
54            None => default,
55            Some(v) if v < 0 => (v + len).max(low),
56            Some(v) => v.min(high),
57        };
58        let (start, stop) = if step > 0 {
59            (clamp(self.start, 0, 0, len), clamp(self.stop, len, 0, len))
60        } else {
61            (
62                clamp(self.start, len - 1, -1, len - 1),
63                clamp(self.stop, -1, -1, len - 1),
64            )
65        };
66        let count = if step > 0 {
67            if stop > start {
68                (stop - start + step - 1) / step
69            } else {
70                0
71            }
72        } else if start > stop {
73            (start - stop - step - 1) / -step
74        } else {
75            0
76        };
77        Ok((start.max(0) as usize, count as usize, step))
78    }
79}
80
81impl From<std::ops::Range<isize>> for Slice {
82    fn from(range: std::ops::Range<isize>) -> Self {
83        Self::range(range.start, range.end)
84    }
85}
86
87impl From<std::ops::RangeFrom<isize>> for Slice {
88    fn from(range: std::ops::RangeFrom<isize>) -> Self {
89        Self::new(Some(range.start), None, 1)
90    }
91}
92
93impl From<std::ops::RangeTo<isize>> for Slice {
94    fn from(range: std::ops::RangeTo<isize>) -> Self {
95        Self::new(None, Some(range.end), 1)
96    }
97}
98
99impl From<std::ops::RangeFull> for Slice {
100    fn from(_: std::ops::RangeFull) -> Self {
101        Self::full()
102    }
103}
104
105/// A selection on one axis for [`isel`].
106#[derive(Debug, Clone, PartialEq, Eq)]
107pub enum Selector {
108    /// Keep a slice of the axis.
109    Slice(Slice),
110    /// Keep one index and drop the axis. Negative indices count from the end.
111    Index(isize),
112    /// Keep the given indices, in order. Negative indices count from the end.
113    Take(Vec<isize>),
114}
115
116impl From<Slice> for Selector {
117    fn from(slice: Slice) -> Self {
118        Self::Slice(slice)
119    }
120}
121
122impl From<isize> for Selector {
123    fn from(index: isize) -> Self {
124        Self::Index(index)
125    }
126}
127
128impl From<std::ops::Range<isize>> for Selector {
129    fn from(range: std::ops::Range<isize>) -> Self {
130        Self::Slice(range.into())
131    }
132}
133
134impl From<Vec<isize>> for Selector {
135    fn from(indices: Vec<isize>) -> Self {
136        Self::Take(indices)
137    }
138}
139
140pub(crate) fn resolve_index(index: isize, len: usize, dim: &str) -> Result<usize> {
141    let resolved = if index < 0 {
142        index + len as isize
143    } else {
144        index
145    };
146    if resolved < 0 || resolved as usize >= len {
147        return Err(nd_err!(
148            InvalidArgument,
149            "index {index} is out of bounds for dimension '{dim}' of size {len}"
150        ));
151    }
152    Ok(resolved as usize)
153}
154
155/// Broadcast `a` onto `target` without a copy where possible.
156///
157/// Axes align by name. Each axis of `a` must be in `target` with the same size
158/// or size 1. The result has the axis order of `target`.
159pub fn broadcast_to(a: &NdArrowArray, target: &Dimensions) -> Result<NdArrowArray> {
160    if a.dims() == target {
161        return Ok(a.clone());
162    }
163    a.dims().check_broadcast_to(target)?;
164    let strides = a.dims().c_strides();
165    let axes = target
166        .iter()
167        .map(|dim| {
168            let stride = match a.dims().position(dim.name()) {
169                Some(axis) if a.dims().as_slice()[axis].size() == dim.size() => strides[axis],
170                _ => 0,
171            };
172            GatherAxis {
173                dim: dim.clone(),
174                offsets: (0..dim.size()).map(|i| i * stride).collect(),
175            }
176        })
177        .collect();
178    Gather { base: 0, axes }.apply(a)
179}
180
181/// Broadcast all arrays onto their shared dimensions, see [`Dimensions::broadcast`].
182pub fn broadcast_arrays(arrays: &[&NdArrowArray]) -> Result<Vec<NdArrowArray>> {
183    let dims: Vec<&Dimensions> = arrays.iter().map(|a| a.dims()).collect();
184    let target = Dimensions::broadcast(&dims)?;
185    arrays.iter().map(|a| broadcast_to(a, &target)).collect()
186}
187
188/// Reorder the axes. `order` must name every dimension once.
189pub fn transpose(a: &NdArrowArray, order: &[&str]) -> Result<NdArrowArray> {
190    if order.len() != a.rank() {
191        return Err(nd_err!(
192            InvalidArgument,
193            "transpose order {order:?} does not match dimensions {}",
194            a.dims()
195        ));
196    }
197    let identity = Gather::identity(a.dims());
198    let mut axes = Vec::with_capacity(order.len());
199    for name in order {
200        let axis = a.dims().axis(name)?;
201        axes.push(identity.axes[axis].clone());
202    }
203    let plan = Gather { base: 0, axes };
204    plan.dims()?;
205    plan.apply(a)
206}
207
208/// Reverse the axis order, like `ndarray.T`.
209pub fn transpose_reversed(a: &NdArrowArray) -> Result<NdArrowArray> {
210    let names: Vec<&str> = a.dims().names().into_iter().rev().collect();
211    transpose(a, &names)
212}
213
214/// Move dimension `name` to position `to`.
215pub fn move_dim(a: &NdArrowArray, name: &str, to: usize) -> Result<NdArrowArray> {
216    let mut names = a.dims().names();
217    let from = a.dims().axis(name)?;
218    if to >= names.len() {
219        return Err(nd_err!(InvalidArgument, "position {to} is out of bounds"));
220    }
221    let moved = names.remove(from);
222    names.insert(to, moved);
223    transpose(a, &names)
224}
225
226/// Remove size-1 axes without a copy. `None` removes every size-1 axis.
227pub fn squeeze(a: &NdArrowArray, dims: Option<&[&str]>) -> Result<NdArrowArray> {
228    if let Some(names) = dims {
229        for name in names {
230            let axis = a.dims().axis(name)?;
231            let size = a.dims().as_slice()[axis].size();
232            if size != 1 {
233                return Err(nd_err!(
234                    InvalidArgument,
235                    "cannot squeeze dimension '{name}' of size {size}"
236                ));
237            }
238        }
239    }
240    let kept = a.dims().iter().filter(|d| match dims {
241        None => d.size() != 1,
242        Some(names) => !names.contains(&d.name()),
243    });
244    replace_dims(a, Dimensions::try_new(kept.cloned())?)
245}
246
247/// Insert a new axis of size 1 at position `axis`, without a copy.
248pub fn expand_dims(a: &NdArrowArray, name: &str, axis: usize) -> Result<NdArrowArray> {
249    if axis > a.rank() {
250        return Err(nd_err!(InvalidArgument, "position {axis} is out of bounds"));
251    }
252    let mut dims: Vec<Dimension> = a.dims().iter().cloned().collect();
253    dims.insert(axis, Dimension::new(name, 1));
254    replace_dims(a, Dimensions::try_new(dims)?)
255}
256
257/// Rename dimensions without a copy.
258pub fn rename_dims(a: &NdArrowArray, renames: &[(&str, &str)]) -> Result<NdArrowArray> {
259    for (old, _) in renames {
260        a.dims().axis(old)?;
261    }
262    let dims = a
263        .dims()
264        .iter()
265        .map(|d| match renames.iter().find(|(old, _)| *old == d.name()) {
266            Some((_, new)) => Dimension::new(*new, d.size()),
267            None => d.clone(),
268        });
269    replace_dims(a, Dimensions::try_new(dims)?)
270}
271
272fn replace_dims(a: &NdArrowArray, dims: Dimensions) -> Result<NdArrowArray> {
273    Ok(NdArrowArray::try_new(a.values().clone(), dims)?.with_attributes(a.attributes().clone()))
274}
275
276/// Select by position on one or more axes, like xarray `isel`, in one copy.
277///
278/// A slice keeps the axis, an index drops it, a take keeps the listed indices.
279pub fn isel(a: &NdArrowArray, selections: &[(&str, Selector)]) -> Result<NdArrowArray> {
280    let mut plan = Gather::identity(a.dims());
281    let mut dropped = vec![false; a.rank()];
282    for (name, selector) in selections {
283        let axis_index = a.dims().axis(name)?;
284        let axis = &mut plan.axes[axis_index];
285        let len = axis.offsets.len();
286        let stride = if len > 1 {
287            axis.offsets[1] - axis.offsets[0]
288        } else {
289            0
290        };
291        let source: Vec<usize> = match selector {
292            Selector::Slice(slice) => {
293                let (start, count, step) = slice.resolve(len)?;
294                (0..count)
295                    .map(|i| (start as isize + i as isize * step) as usize)
296                    .collect()
297            }
298            Selector::Index(index) => {
299                dropped[axis_index] = true;
300                vec![resolve_index(*index, len, name)?]
301            }
302            Selector::Take(indices) => indices
303                .iter()
304                .map(|&i| resolve_index(i, len, name))
305                .collect::<Result<_>>()?,
306        };
307        axis.dim = Dimension::new(axis.dim.name(), source.len());
308        axis.offsets = source.into_iter().map(|i| i * stride).collect();
309    }
310    let mut kept = Vec::with_capacity(plan.axes.len());
311    for (axis, dropped) in plan.axes.into_iter().zip(dropped) {
312        if dropped {
313            plan.base += axis.offsets[0];
314        } else {
315            kept.push(axis);
316        }
317    }
318    plan.axes = kept;
319    plan.apply(a)
320}
321
322/// Take `indices` along `dim`, like `numpy.take(a, indices, axis)`.
323pub fn take(a: &NdArrowArray, dim: &str, indices: &[isize]) -> Result<NdArrowArray> {
324    isel(a, &[(dim, Selector::Take(indices.to_vec()))])
325}
326
327/// Reverse the order along `dim`, like `numpy.flip`.
328pub fn flip(a: &NdArrowArray, dim: &str) -> Result<NdArrowArray> {
329    isel(a, &[(dim, Selector::Slice(Slice::step(-1)))])
330}
331
332/// Roll the elements along `dim` by `shift`, like `numpy.roll`.
333pub fn roll(a: &NdArrowArray, dim: &str, shift: isize) -> Result<NdArrowArray> {
334    let len = a
335        .dims()
336        .size_of(dim)
337        .ok_or_else(|| nd_err!(InvalidArgument, "dimension '{dim}' is not in {}", a.dims()))?
338        as isize;
339    if len == 0 {
340        return Ok(a.clone());
341    }
342    let indices: Vec<isize> = (0..len).map(|i| (i - shift).rem_euclid(len)).collect();
343    take(a, dim, &indices)
344}
345
346/// How [`pad`] fills the new positions, like the `mode` of `numpy.pad`.
347#[derive(Debug, Clone)]
348pub enum PadMode {
349    /// A constant: a length-1 array, or `None` for null.
350    Constant(Option<arrow::array::ArrayRef>),
351    /// Repeat the edge value.
352    Edge,
353    /// Mirror without the edge value: `[3 2 | 1 2 3 | 2 1]`.
354    Reflect,
355    /// Mirror with the edge value: `[2 1 | 1 2 3 | 3 2]`.
356    Symmetric,
357    /// Wrap around: `[2 3 | 1 2 3 | 1 2]`.
358    Wrap,
359}
360
361fn dim_size(a: &NdArrowArray, dim: &str) -> Result<usize> {
362    a.dims()
363        .size_of(dim)
364        .ok_or_else(|| nd_err!(InvalidArgument, "dimension '{dim}' is not in {}", a.dims()))
365}
366
367/// Add `before` and `after` positions along `dim`, like `numpy.pad`.
368pub fn pad(
369    a: &NdArrowArray,
370    dim: &str,
371    before: usize,
372    after: usize,
373    mode: PadMode,
374) -> Result<NdArrowArray> {
375    let n = dim_size(a, dim)? as isize;
376    let index = |i: isize| -> Result<isize> {
377        if (0..n).contains(&i) {
378            return Ok(i);
379        }
380        if n == 0 {
381            return Err(nd_err!(
382                InvalidArgument,
383                "cannot pad an empty dimension '{dim}' with {mode:?}"
384            ));
385        }
386        Ok(match &mode {
387            PadMode::Edge => i.clamp(0, n - 1),
388            PadMode::Wrap => i.rem_euclid(n),
389            PadMode::Symmetric => {
390                let period = 2 * n;
391                let k = i.rem_euclid(period);
392                if k < n { k } else { period - 1 - k }
393            }
394            PadMode::Reflect => {
395                if n == 1 {
396                    0
397                } else {
398                    let period = 2 * (n - 1);
399                    let k = i.rem_euclid(period);
400                    if k < n { k } else { period - k }
401                }
402            }
403            PadMode::Constant(_) => unreachable!(),
404        })
405    };
406    let (before, after) = (before as isize, after as isize);
407    if let PadMode::Constant(value) = &mode {
408        let block = |size: usize| -> Result<NdArrowArray> {
409            let dims = a.dims().iter().map(|d| {
410                if d.name() == dim {
411                    Dimension::new(dim, size)
412                } else {
413                    d.clone()
414                }
415            });
416            let dims = Dimensions::try_new(dims)?;
417            let fill = match value {
418                None => arrow::array::new_null_array(a.data_type(), 1),
419                Some(v) => arrow::compute::cast(v.slice(0, 1).as_ref(), a.data_type())?,
420            };
421            broadcast_to(&NdArrowArray::try_new_scalar(fill)?, &dims)
422        };
423        let head = block(before as usize)?;
424        let tail = block(after as usize)?;
425        return concat(&[&head, a, &tail], dim);
426    }
427    let indices = (-before..n + after)
428        .map(index)
429        .collect::<Result<Vec<_>>>()?;
430    take(a, dim, &indices)
431}
432
433/// Repeat the whole array `reps` times along `dim`, like `numpy.tile`.
434pub fn tile(a: &NdArrowArray, dim: &str, reps: usize) -> Result<NdArrowArray> {
435    let n = dim_size(a, dim)? as isize;
436    let indices: Vec<isize> = (0..reps).flat_map(|_| 0..n).collect();
437    take(a, dim, &indices)
438}
439
440/// Repeat each element `repeats` times along `dim`, like `numpy.repeat`.
441pub fn repeat(a: &NdArrowArray, dim: &str, repeats: usize) -> Result<NdArrowArray> {
442    let n = dim_size(a, dim)? as isize;
443    let indices: Vec<isize> = (0..n)
444        .flat_map(|i| std::iter::repeat_n(i, repeats))
445        .collect();
446    take(a, dim, &indices)
447}
448
449/// Join arrays along an existing dimension, like `numpy.concatenate`.
450///
451/// All arrays need the same dimensions, except for the size of `dim`. Arrays
452/// with another axis order are transposed to the order of the first array.
453/// The result keeps the attributes of the first array.
454pub fn concat(arrays: &[&NdArrowArray], dim: &str) -> Result<NdArrowArray> {
455    let Some(first) = arrays.first() else {
456        return Err(nd_err!(InvalidArgument, "concat needs at least one array"));
457    };
458    let axis = first.dims().axis(dim)?;
459    let order = first.dims().names();
460    let mut aligned: Vec<NdArrowArray> = Vec::with_capacity(arrays.len());
461    let mut total = 0usize;
462    for array in arrays {
463        if array.data_type() != first.data_type() {
464            return Err(nd_err!(
465                SchemaMismatch,
466                "concat value types differ: {} and {}",
467                first.data_type(),
468                array.data_type()
469            ));
470        }
471        let array = if array.dims().names() == order {
472            (*array).clone()
473        } else {
474            transpose(array, &order)?
475        };
476        for (d, f) in array.dims().iter().zip(first.dims().iter()) {
477            if d.name() != dim && d.size() != f.size() {
478                return Err(nd_err!(
479                    ShapeMismatch,
480                    "concat along '{dim}': dimension '{}' has sizes {} and {}",
481                    d.name(),
482                    f.size(),
483                    d.size()
484                ));
485            }
486        }
487        total += array.dims().as_slice()[axis].size();
488        aligned.push(array);
489    }
490
491    let outer: usize = first.dims().as_slice()[..axis]
492        .iter()
493        .map(Dimension::size)
494        .product();
495    let inner: usize = first.dims().as_slice()[axis + 1..]
496        .iter()
497        .map(Dimension::size)
498        .product();
499    let data: Vec<_> = aligned.iter().map(|a| a.values().to_data()).collect();
500    let mut out = MutableArrayData::new(data.iter().collect(), false, outer * total * inner);
501    for o in 0..outer {
502        for (i, array) in aligned.iter().enumerate() {
503            let block = array.dims().as_slice()[axis].size() * inner;
504            out.extend(i, o * block, (o + 1) * block);
505        }
506    }
507    let values = make_array(out.freeze());
508    let dims = first.dims().iter().map(|d| {
509        if d.name() == dim {
510            Dimension::new(dim, total)
511        } else {
512            d.clone()
513        }
514    });
515    Ok(NdArrowArray::try_new(
516        Arc::new(values) as Arc<dyn Array>,
517        Dimensions::try_new(dims)?,
518    )?
519    .with_attributes(first.attributes().clone()))
520}
521
522/// Join arrays along a new dimension at position `axis`, like `numpy.stack`.
523pub fn stack(arrays: &[&NdArrowArray], dim: &str, axis: usize) -> Result<NdArrowArray> {
524    let expanded = arrays
525        .iter()
526        .map(|a| expand_dims(a, dim, axis))
527        .collect::<Result<Vec<_>>>()?;
528    concat(&expanded.iter().collect::<Vec<_>>(), dim)
529}
530
531#[cfg(test)]
532mod tests {
533    use arrow::array::{AsArray, Int32Array, StringArray};
534    use arrow::datatypes::Int32Type;
535
536    use super::*;
537
538    fn arange(dims: &[(&str, usize)]) -> NdArrowArray {
539        let dims = Dimensions::try_new(dims.iter().copied()).unwrap();
540        let n = dims.num_elements() as i32;
541        NdArrowArray::try_new(Arc::new(Int32Array::from_iter_values(0..n)), dims).unwrap()
542    }
543
544    fn values(a: &NdArrowArray) -> Vec<i32> {
545        a.values().as_primitive::<Int32Type>().values().to_vec()
546    }
547
548    fn dims(spec: &[(&str, usize)]) -> Dimensions {
549        Dimensions::try_new(spec.iter().copied()).unwrap()
550    }
551
552    #[test]
553    fn slice_resolve_matches_python() {
554        let cases = [
555            (Slice::full(), 5, (0, 5, 1)),
556            (Slice::range(1, 4), 5, (1, 3, 1)),
557            (Slice::range(-2, 100), 5, (3, 2, 1)),
558            (Slice::step(2), 5, (0, 3, 2)),
559            (Slice::step(-1), 5, (4, 5, -1)),
560            (Slice::new(Some(3), Some(0), -2), 5, (3, 2, -2)),
561            (Slice::range(4, 1), 5, (4, 0, 1)),
562        ];
563        for (slice, len, expected) in cases {
564            assert_eq!(slice.resolve(len).unwrap(), expected, "{slice:?}");
565        }
566        assert!(Slice::step(0).resolve(3).is_err());
567    }
568
569    #[test]
570    fn broadcast_inserts_and_stretches_axes() {
571        let a = arange(&[("lat", 3)]);
572        let b = broadcast_to(&a, &dims(&[("time", 2), ("lat", 3)])).unwrap();
573        assert_eq!(values(&b), vec![0, 1, 2, 0, 1, 2]);
574
575        let col = arange(&[("time", 2), ("lat", 1)]);
576        let c = broadcast_to(&col, &dims(&[("lat", 3), ("time", 2)])).unwrap();
577        assert_eq!(values(&c), vec![0, 1, 0, 1, 0, 1]);
578
579        assert!(broadcast_to(&a, &dims(&[("lat", 4)])).is_err());
580    }
581
582    #[test]
583    fn broadcast_arrays_unions_dims() {
584        let a = arange(&[("x", 2)]);
585        let b = arange(&[("y", 3)]);
586        let out = broadcast_arrays(&[&a, &b]).unwrap();
587        assert_eq!(out[0].dims(), &dims(&[("x", 2), ("y", 3)]));
588        assert_eq!(values(&out[0]), vec![0, 0, 0, 1, 1, 1]);
589        assert_eq!(values(&out[1]), vec![0, 1, 2, 0, 1, 2]);
590    }
591
592    #[test]
593    fn transpose_and_move_dim() {
594        let a = arange(&[("x", 2), ("y", 3), ("z", 2)]);
595        let t = transpose(&a, &["z", "x", "y"]).unwrap();
596        assert_eq!(t.shape(), vec![2, 2, 3]);
597        assert_eq!(values(&t), vec![0, 2, 4, 6, 8, 10, 1, 3, 5, 7, 9, 11]);
598        assert_eq!(values(&move_dim(&a, "z", 0).unwrap()), values(&t));
599        assert!(transpose(&a, &["x", "y"]).is_err());
600        assert!(transpose(&a, &["x", "x", "y"]).is_err());
601        let r = transpose_reversed(&arange(&[("x", 2), ("y", 3)])).unwrap();
602        assert_eq!(values(&r), vec![0, 3, 1, 4, 2, 5]);
603    }
604
605    #[test]
606    fn squeeze_expand_rename() {
607        let a = arange(&[("x", 1), ("y", 3), ("z", 1)]);
608        assert_eq!(squeeze(&a, None).unwrap().dims(), &dims(&[("y", 3)]));
609        assert_eq!(
610            squeeze(&a, Some(&["z"])).unwrap().dims(),
611            &dims(&[("x", 1), ("y", 3)])
612        );
613        assert!(squeeze(&a, Some(&["y"])).is_err());
614        let e = expand_dims(&a, "w", 3).unwrap();
615        assert_eq!(e.dims().names(), vec!["x", "y", "z", "w"]);
616        let r = rename_dims(&a, &[("y", "lat")]).unwrap();
617        assert_eq!(r.dims().names(), vec!["x", "lat", "z"]);
618        assert!(rename_dims(&a, &[("q", "r")]).is_err());
619    }
620
621    #[test]
622    fn isel_slices_indexes_and_takes() {
623        let a = arange(&[("x", 3), ("y", 4)]);
624        let s = isel(&a, &[("y", Slice::range(1, 3).into())]).unwrap();
625        assert_eq!(values(&s), vec![1, 2, 5, 6, 9, 10]);
626
627        let i = isel(&a, &[("x", Selector::Index(-1))]).unwrap();
628        assert_eq!(i.dims(), &dims(&[("y", 4)]));
629        assert_eq!(values(&i), vec![8, 9, 10, 11]);
630
631        let both = isel(
632            &a,
633            &[
634                ("x", Slice::step(-2).into()),
635                ("y", Selector::Take(vec![3, 0])),
636            ],
637        )
638        .unwrap();
639        assert_eq!(values(&both), vec![11, 8, 3, 0]);
640        assert!(isel(&a, &[("x", Selector::Index(3))]).is_err());
641    }
642
643    #[test]
644    fn flip_and_roll() {
645        let a = arange(&[("x", 5)]);
646        assert_eq!(values(&flip(&a, "x").unwrap()), vec![4, 3, 2, 1, 0]);
647        assert_eq!(values(&roll(&a, "x", 2).unwrap()), vec![3, 4, 0, 1, 2]);
648        assert_eq!(values(&roll(&a, "x", -1).unwrap()), vec![1, 2, 3, 4, 0]);
649    }
650
651    #[test]
652    fn pad_modes_match_numpy() {
653        let a = arange(&[("x", 3)]);
654        let padded = |mode| values(&pad(&a, "x", 2, 2, mode).unwrap());
655        assert_eq!(padded(PadMode::Edge), vec![0, 0, 0, 1, 2, 2, 2]);
656        assert_eq!(padded(PadMode::Reflect), vec![2, 1, 0, 1, 2, 1, 0]);
657        assert_eq!(padded(PadMode::Symmetric), vec![1, 0, 0, 1, 2, 2, 1]);
658        assert_eq!(padded(PadMode::Wrap), vec![1, 2, 0, 1, 2, 0, 1]);
659        let constant = pad(&a, "x", 1, 0, PadMode::Constant(None)).unwrap();
660        assert_eq!(constant.values().null_count(), 1);
661        assert_eq!(constant.shape(), vec![4]);
662    }
663
664    #[test]
665    fn tile_and_repeat() {
666        let a = arange(&[("x", 2), ("y", 2)]);
667        assert_eq!(
668            values(&tile(&a, "y", 2).unwrap()),
669            vec![0, 1, 0, 1, 2, 3, 2, 3]
670        );
671        assert_eq!(
672            values(&repeat(&a, "x", 2).unwrap()),
673            vec![0, 1, 0, 1, 2, 3, 2, 3]
674        );
675        assert_eq!(
676            values(&repeat(&a, "y", 2).unwrap()),
677            vec![0, 0, 1, 1, 2, 2, 3, 3]
678        );
679    }
680
681    #[test]
682    fn concat_and_stack() {
683        let a = arange(&[("x", 2), ("y", 2)]);
684        let b = arange(&[("x", 1), ("y", 2)]);
685        let c = concat(&[&a, &b], "x").unwrap();
686        assert_eq!(c.dims(), &dims(&[("x", 3), ("y", 2)]));
687        assert_eq!(values(&c), vec![0, 1, 2, 3, 0, 1]);
688
689        let d = concat(&[&a, &a], "y").unwrap();
690        assert_eq!(values(&d), vec![0, 1, 0, 1, 2, 3, 2, 3]);
691
692        let t = transpose(&a, &["y", "x"]).unwrap();
693        assert_eq!(
694            values(&concat(&[&a, &t], "x").unwrap()),
695            vec![0, 1, 2, 3, 0, 1, 2, 3]
696        );
697
698        let s = stack(&[&a, &a], "member", 0).unwrap();
699        assert_eq!(s.dims(), &dims(&[("member", 2), ("x", 2), ("y", 2)]));
700        assert!(concat(&[&a, &arange(&[("x", 2), ("y", 3)])], "x").is_err());
701    }
702
703    #[test]
704    fn layout_kernels_work_for_strings() {
705        let a = NdArrowArray::try_new(
706            Arc::new(StringArray::from(vec!["a", "b", "c", "d"])),
707            dims(&[("x", 2), ("y", 2)]),
708        )
709        .unwrap();
710        let t = transpose(&a, &["y", "x"]).unwrap();
711        let s: Vec<&str> = t
712            .values()
713            .as_string::<i32>()
714            .iter()
715            .map(Option::unwrap)
716            .collect();
717        assert_eq!(s, vec!["a", "c", "b", "d"]);
718    }
719}
nd_arrow_array/compute/shape.rs

nd_arrow_array/compute/
shape.rs
