vortex_array/arrays/constant/vtable/
pipeline.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::hash::{Hash, Hasher};
6use std::sync::Arc;
7
8use vortex_dtype::{DType, NativePType, match_each_native_ptype};
9use vortex_error::{VortexExpect, VortexResult};
10
11use crate::arrays::{ConstantArray, ConstantVTable};
12use crate::operator::{LengthBounds, Operator, OperatorEq, OperatorHash, OperatorId, OperatorRef};
13use crate::pipeline::bits::BitView;
14use crate::pipeline::vec::Selection;
15use crate::pipeline::view::ViewMut;
16use crate::pipeline::{
17    BindContext, Element, Kernel, KernelContext, N, PipelinedOperator, RowSelection,
18};
19use crate::vtable::PipelineVTable;
20
21impl PipelineVTable<ConstantVTable> for ConstantVTable {
22    fn to_operator(array: &ConstantArray) -> VortexResult<Option<OperatorRef>> {
23        Ok(Some(Arc::new(array.clone())))
24    }
25}
26
27impl OperatorHash for ConstantArray {
28    fn operator_hash<H: Hasher>(&self, state: &mut H) {
29        self.scalar.hash(state);
30        self.len.hash(state);
31    }
32}
33
34impl OperatorEq for ConstantArray {
35    fn operator_eq(&self, other: &Self) -> bool {
36        self.scalar == other.scalar && self.len == other.len
37    }
38}
39
40impl Operator for ConstantArray {
41    fn id(&self) -> OperatorId {
42        self.encoding_id()
43    }
44
45    fn as_any(&self) -> &dyn Any {
46        self
47    }
48
49    fn dtype(&self) -> &DType {
50        self.scalar.dtype()
51    }
52
53    fn bounds(&self) -> LengthBounds {
54        self.len.into()
55    }
56
57    fn children(&self) -> &[OperatorRef] {
58        &[]
59    }
60
61    fn with_children(self: Arc<Self>, _children: Vec<OperatorRef>) -> VortexResult<OperatorRef> {
62        Ok(self)
63    }
64}
65
66impl PipelinedOperator for ConstantArray {
67    fn row_selection(&self) -> RowSelection {
68        RowSelection::Domain(self.len)
69    }
70
71    fn bind(&self, _ctx: &dyn BindContext) -> VortexResult<Box<dyn Kernel>> {
72        debug_assert!(matches!(
73            self.dtype(),
74            DType::Bool(_) | DType::Primitive(..)
75        ));
76        match self.scalar.dtype() {
77            DType::Bool(_) => Ok(Box::new(BoolConstantKernel {
78                value: self
79                    .scalar
80                    .as_bool()
81                    .value()
82                    .vortex_expect("scalar value not bool"),
83            })),
84            DType::Primitive(..) => Ok(match_each_native_ptype!(
85                self.scalar.as_primitive().ptype(),
86                |T| {
87                    Box::new(ConstantKernel::<T> {
88                        value: self
89                            .scalar
90                            .as_primitive()
91                            .typed_value::<T>()
92                            .vortex_expect("scalar value not of type T"),
93                    })
94                }
95            )),
96            _ => todo!(
97                "Unsupported scalar type for constant: {:?}",
98                self.scalar.dtype()
99            ),
100        }
101    }
102
103    fn vector_children(&self) -> Vec<usize> {
104        vec![]
105    }
106
107    fn batch_children(&self) -> Vec<usize> {
108        vec![]
109    }
110}
111
112/// Kernel that produces constant primitive values.
113pub struct ConstantKernel<T: NativePType> {
114    value: T,
115}
116
117/// Kernel that produces constant boolean values.
118pub struct BoolConstantKernel {
119    value: bool,
120}
121
122impl<T: Element + NativePType> Kernel for ConstantKernel<T> {
123    fn step(
124        &self,
125        _ctx: &KernelContext,
126        _chunk_idx: usize,
127        _selection: &BitView,
128        out: &mut ViewMut,
129    ) -> VortexResult<()> {
130        // TODO(ngates): benchmark whether to populate the true indices, or the entire vector.
131        out.as_array_mut::<T>()[..N].fill(self.value);
132        out.set_selection(Selection::Prefix);
133        Ok(())
134    }
135}
136
137impl Kernel for BoolConstantKernel {
138    fn step(
139        &self,
140        _ctx: &KernelContext,
141        _chunk_idx: usize,
142        _selection: &BitView,
143        out: &mut ViewMut,
144    ) -> VortexResult<()> {
145        // TODO(ngates): benchmark whether to populate the true indices, or the entire vector.
146        out.as_array_mut::<bool>()[..N].fill(self.value);
147        out.set_selection(Selection::Prefix);
148        Ok(())
149    }
150}
vortex_array/arrays/constant/vtable/pipeline.rs

vortex_array/arrays/constant/vtable/
pipeline.rs
