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    pub child: ArrayRef,
21}
22
23#[derive(Clone, Debug)]
24pub struct SliceData {
25    pub(super) range: Range<usize>,
26}
27
28impl Display for SliceData {
29    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
30        write!(f, "range: {}..{}", self.range.start, self.range.end)
31    }
32}
33
34pub struct SliceDataParts {
35    pub range: Range<usize>,
36}
37
38impl SliceData {
39    fn try_new(child_len: usize, range: Range<usize>) -> VortexResult<Self> {
40        if range.end > child_len {
41            vortex_panic!(
42                "SliceArray range out of bounds: range {:?} exceeds child array length {}",
43                range,
44                child_len
45            );
46        }
47        Ok(Self { range })
48    }
49
50    pub fn new(range: Range<usize>) -> Self {
51        Self { range }
52    }
53
54    /// Returns the length of this array.
55    pub fn len(&self) -> usize {
56        self.range.len()
57    }
58
59    /// Returns `true` if this array is empty.
60    pub fn is_empty(&self) -> bool {
61        self.len() == 0
62    }
63
64    /// The range used to slice the child array.
65    pub fn slice_range(&self) -> &Range<usize> {
66        &self.range
67    }
68
69    pub fn into_parts(self) -> SliceDataParts {
70        SliceDataParts { range: self.range }
71    }
72}
73
74impl Array<Slice> {
75    /// Constructs a new `SliceArray`.
76    pub fn try_new(child: ArrayRef, range: Range<usize>) -> VortexResult<Self> {
77        let len = range.len();
78        let dtype = child.dtype().clone();
79        let data = SliceData::try_new(child.len(), range)?;
80        Ok(unsafe {
81            Array::from_parts_unchecked(
82                ArrayParts::new(Slice, dtype, len, data)
83                    .with_slots(SliceSlots { child }.into_slots()),
84            )
85        })
86    }
87
88    /// Constructs a new `SliceArray`.
89    pub fn new(child: ArrayRef, range: Range<usize>) -> Self {
90        let len = range.len();
91        let dtype = child.dtype().clone();
92        let data = SliceData::new(range);
93        unsafe {
94            Array::from_parts_unchecked(
95                ArrayParts::new(Slice, dtype, len, data)
96                    .with_slots(SliceSlots { child }.into_slots()),
97            )
98        }
99    }
100}