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