Skip to main content

vortex_array/arrays/slice/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6use std::ops::Range;
7
8use vortex_error::VortexResult;
9use vortex_error::vortex_panic;
10
11use crate::ArrayRef;
12use crate::array::Array;
13use crate::array::ArrayParts;
14use crate::array_slots;
15use crate::arrays::Slice;
16
17#[array_slots(Slice)]
18pub struct SliceSlots {
19    /// The underlying child array being sliced.
20    #[slot(0)]
21    pub child: ArrayRef,
22}
23
24#[derive(Clone, Debug)]
25pub struct SliceData {
26    pub(super) range: Range<usize>,
27}
28
29impl Display for SliceData {
30    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
31        write!(f, "range: {}..{}", self.range.start, self.range.end)
32    }
33}
34
35pub struct SliceDataParts {
36    pub range: Range<usize>,
37}
38
39impl SliceData {
40    fn try_new(child_len: usize, range: Range<usize>) -> VortexResult<Self> {
41        if range.end > child_len {
42            vortex_panic!(
43                "SliceArray range out of bounds: range {:?} exceeds child array length {}",
44                range,
45                child_len
46            );
47        }
48        Ok(Self { range })
49    }
50
51    pub fn new(range: Range<usize>) -> Self {
52        Self { range }
53    }
54
55    /// Returns the length of this array.
56    pub fn len(&self) -> usize {
57        self.range.len()
58    }
59
60    /// Returns `true` if this array is empty.
61    pub fn is_empty(&self) -> bool {
62        self.len() == 0
63    }
64
65    /// The range used to slice the child array.
66    pub fn slice_range(&self) -> &Range<usize> {
67        &self.range
68    }
69
70    pub fn into_parts(self) -> SliceDataParts {
71        SliceDataParts { range: self.range }
72    }
73}
74
75impl Array<Slice> {
76    /// Constructs a new `SliceArray`.
77    pub fn try_new(child: ArrayRef, range: Range<usize>) -> VortexResult<Self> {
78        let len = range.len();
79        let dtype = child.dtype().clone();
80        let data = SliceData::try_new(child.len(), range)?;
81        Ok(unsafe {
82            Array::from_parts_unchecked(ArrayParts::new(
83                Slice,
84                dtype,
85                len,
86                data,
87                SliceSlots { child }.into_slots(),
88            ))
89        })
90    }
91
92    /// Constructs a new `SliceArray`.
93    pub fn new(child: ArrayRef, range: Range<usize>) -> Self {
94        let len = range.len();
95        let dtype = child.dtype().clone();
96        let data = SliceData::new(range);
97        unsafe {
98            Array::from_parts_unchecked(ArrayParts::new(
99                Slice,
100                dtype,
101                len,
102                data,
103                SliceSlots { child }.into_slots(),
104            ))
105        }
106    }
107}