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(
83                ArrayParts::new(Slice, dtype, len, data)
84                    .with_slots(SliceSlots { child }.into_slots()),
85            )
86        })
87    }
88
89    /// Constructs a new `SliceArray`.
90    pub fn new(child: ArrayRef, range: Range<usize>) -> Self {
91        let len = range.len();
92        let dtype = child.dtype().clone();
93        let data = SliceData::new(range);
94        unsafe {
95            Array::from_parts_unchecked(
96                ArrayParts::new(Slice, dtype, len, data)
97                    .with_slots(SliceSlots { child }.into_slots()),
98            )
99        }
100    }
101}