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 smallvec::smallvec;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_panic;
12
13use crate::ArrayRef;
14use crate::array::Array;
15use crate::array::ArrayParts;
16use crate::array::TypedArrayRef;
17use crate::arrays::Slice;
18
19/// The underlying child array being sliced.
20pub(super) const CHILD_SLOT: usize = 0;
21pub(super) const NUM_SLOTS: usize = 1;
22pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["child"];
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
39pub trait SliceArrayExt: TypedArrayRef<Slice> {
40    fn child(&self) -> &ArrayRef {
41        self.as_ref().slots()[CHILD_SLOT]
42            .as_ref()
43            .vortex_expect("validated slice child slot")
44    }
45}
46impl<T: TypedArrayRef<Slice>> SliceArrayExt for T {}
47
48impl SliceData {
49    fn try_new(child_len: usize, range: Range<usize>) -> VortexResult<Self> {
50        if range.end > child_len {
51            vortex_panic!(
52                "SliceArray range out of bounds: range {:?} exceeds child array length {}",
53                range,
54                child_len
55            );
56        }
57        Ok(Self { range })
58    }
59
60    pub fn new(range: Range<usize>) -> Self {
61        Self { range }
62    }
63
64    /// Returns the length of this array.
65    pub fn len(&self) -> usize {
66        self.range.len()
67    }
68
69    /// Returns `true` if this array is empty.
70    pub fn is_empty(&self) -> bool {
71        self.len() == 0
72    }
73
74    /// The range used to slice the child array.
75    pub fn slice_range(&self) -> &Range<usize> {
76        &self.range
77    }
78
79    pub fn into_parts(self) -> SliceDataParts {
80        SliceDataParts { range: self.range }
81    }
82}
83
84impl Array<Slice> {
85    /// Constructs a new `SliceArray`.
86    pub fn try_new(child: ArrayRef, range: Range<usize>) -> VortexResult<Self> {
87        let len = range.len();
88        let dtype = child.dtype().clone();
89        let data = SliceData::try_new(child.len(), range)?;
90        Ok(unsafe {
91            Array::from_parts_unchecked(
92                ArrayParts::new(Slice, dtype, len, data).with_slots(smallvec![Some(child)]),
93            )
94        })
95    }
96
97    /// Constructs a new `SliceArray`.
98    pub fn new(child: ArrayRef, range: Range<usize>) -> Self {
99        let len = range.len();
100        let dtype = child.dtype().clone();
101        let data = SliceData::new(range);
102        unsafe {
103            Array::from_parts_unchecked(
104                ArrayParts::new(Slice, dtype, len, data).with_slots(smallvec![Some(child)]),
105            )
106        }
107    }
108}