Skip to main content

vortex_array/
columnar.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6
7use crate::AnyCanonical;
8use crate::ArrayRef;
9use crate::Canonical;
10use crate::CanonicalView;
11use crate::DynArray;
12use crate::Executable;
13use crate::ExecutionCtx;
14use crate::IntoArray;
15use crate::arrays::ConstantArray;
16use crate::arrays::ConstantVTable;
17use crate::dtype::DType;
18use crate::matcher::Matcher;
19use crate::scalar::Scalar;
20
21/// Represents a columnnar array of data, either in canonical form or as a constant array.
22///
23/// Since the [`Canonical`] enum has one variant per logical data type, it is inefficient for
24/// representing constant arrays. The [`Columnar`] enum allows holding an array of data in either
25/// canonical or constant form enabling efficient handling of constants during execution.
26pub enum Columnar {
27    /// A columnar array in canonical form.
28    Canonical(Canonical),
29    /// A columnar array in constant form.
30    Constant(ConstantArray),
31}
32
33impl Columnar {
34    /// Creates a new columnar array from a scalar.
35    pub fn constant<S: Into<Scalar>>(scalar: S, len: usize) -> Self {
36        Columnar::Constant(ConstantArray::new(scalar.into(), len))
37    }
38
39    /// Returns the length of this columnar array.
40    pub fn len(&self) -> usize {
41        match self {
42            Columnar::Canonical(canonical) => canonical.len(),
43            Columnar::Constant(constant) => constant.len(),
44        }
45    }
46
47    /// Returns true if this columnar array has length zero.
48    pub fn is_empty(&self) -> bool {
49        self.len() == 0
50    }
51
52    /// Returns the data type of this columnar array.
53    pub fn dtype(&self) -> &DType {
54        match self {
55            Columnar::Canonical(canonical) => canonical.dtype(),
56            Columnar::Constant(constant) => constant.dtype(),
57        }
58    }
59}
60
61impl IntoArray for Columnar {
62    fn into_array(self) -> ArrayRef {
63        match self {
64            Columnar::Canonical(canonical) => canonical.into_array(),
65            Columnar::Constant(constant) => constant.into_array(),
66        }
67    }
68}
69
70/// Execute into [`Columnar`] by running `execute_until` with the [`AnyColumnar`] matcher.
71impl Executable for Columnar {
72    fn execute(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
73        let result = array.execute_until::<AnyColumnar>(ctx)?;
74        if let Some(constant) = result.as_opt::<ConstantVTable>() {
75            Ok(Columnar::Constant(constant.clone()))
76        } else {
77            Ok(Columnar::Canonical(
78                result
79                    .as_opt::<AnyCanonical>()
80                    .map(Canonical::from)
81                    .vortex_expect("execute_until::<AnyColumnar> must return a columnar array"),
82            ))
83        }
84    }
85}
86
87pub enum ColumnarView<'a> {
88    Canonical(CanonicalView<'a>),
89    Constant(&'a ConstantArray),
90}
91
92impl<'a> AsRef<dyn DynArray> for ColumnarView<'a> {
93    fn as_ref(&self) -> &dyn DynArray {
94        match self {
95            ColumnarView::Canonical(canonical) => canonical.as_ref(),
96            ColumnarView::Constant(constant) => constant.as_ref(),
97        }
98    }
99}
100
101pub struct AnyColumnar;
102impl Matcher for AnyColumnar {
103    type Match<'a> = ColumnarView<'a>;
104
105    fn try_match<'a>(array: &'a dyn DynArray) -> Option<Self::Match<'a>> {
106        if let Some(constant) = array.as_opt::<ConstantVTable>() {
107            Some(ColumnarView::Constant(constant))
108        } else {
109            array.as_opt::<AnyCanonical>().map(ColumnarView::Canonical)
110        }
111    }
112}