datatypes/vectors/
constant.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::Any;
16use std::fmt;
17use std::sync::Arc;
18
19use arrow::array::{Array, ArrayRef, UInt32Array};
20use snafu::{ensure, ResultExt};
21
22use crate::data_type::ConcreteDataType;
23use crate::error::{self, Result, SerializeSnafu};
24use crate::serialize::Serializable;
25use crate::value::{Value, ValueRef};
26use crate::vectors::{BooleanVector, Helper, UInt32Vector, Validity, Vector, VectorRef};
27
28#[derive(Clone)]
29pub struct ConstantVector {
30    length: usize,
31    vector: VectorRef,
32}
33
34impl ConstantVector {
35    /// Create a new [ConstantVector].
36    ///
37    /// # Panics
38    /// Panics if `vector.len() != 1`.
39    pub fn new(vector: VectorRef, length: usize) -> Self {
40        assert_eq!(1, vector.len());
41
42        // Avoid const recursion.
43        if vector.is_const() {
44            let vec: &ConstantVector = unsafe { Helper::static_cast(&vector) };
45            return Self::new(vec.inner().clone(), length);
46        }
47        Self { vector, length }
48    }
49
50    pub fn inner(&self) -> &VectorRef {
51        &self.vector
52    }
53
54    /// Returns the constant value.
55    pub fn get_constant_ref(&self) -> ValueRef {
56        self.vector.get_ref(0)
57    }
58
59    pub(crate) fn replicate_vector(&self, offsets: &[usize]) -> VectorRef {
60        assert_eq!(offsets.len(), self.len());
61
62        if offsets.is_empty() {
63            return self.slice(0, 0);
64        }
65
66        Arc::new(ConstantVector::new(
67            self.vector.clone(),
68            *offsets.last().unwrap(),
69        ))
70    }
71
72    pub(crate) fn filter_vector(&self, filter: &BooleanVector) -> Result<VectorRef> {
73        let length = self.len() - filter.false_count();
74        if length == self.len() {
75            return Ok(Arc::new(self.clone()));
76        }
77        Ok(Arc::new(ConstantVector::new(self.inner().clone(), length)))
78    }
79
80    pub(crate) fn cast_vector(&self, to_type: &ConcreteDataType) -> Result<VectorRef> {
81        Ok(Arc::new(ConstantVector::new(
82            self.inner().cast(to_type)?,
83            self.length,
84        )))
85    }
86
87    pub(crate) fn take_vector(&self, indices: &UInt32Vector) -> Result<VectorRef> {
88        if indices.is_empty() {
89            return Ok(self.slice(0, 0));
90        }
91        ensure!(
92            indices.null_count() == 0,
93            error::UnsupportedOperationSnafu {
94                op: "taking a null index",
95                vector_type: self.vector_type_name(),
96            }
97        );
98
99        let len = self.len();
100        let arr = indices.to_arrow_array();
101        let indices_arr = arr.as_any().downcast_ref::<UInt32Array>().unwrap();
102        if !arrow::compute::min_boolean(
103            &arrow::compute::kernels::cmp::lt(indices_arr, &UInt32Array::new_scalar(len as u32))
104                .unwrap(),
105        )
106        .unwrap()
107        {
108            panic!("Array index out of bounds, cannot take index out of the length of the array: {len}");
109        }
110
111        Ok(Arc::new(ConstantVector::new(
112            self.inner().clone(),
113            indices.len(),
114        )))
115    }
116}
117
118impl Vector for ConstantVector {
119    fn data_type(&self) -> ConcreteDataType {
120        self.vector.data_type()
121    }
122
123    fn vector_type_name(&self) -> String {
124        "ConstantVector".to_string()
125    }
126
127    fn as_any(&self) -> &dyn Any {
128        self
129    }
130
131    fn len(&self) -> usize {
132        self.length
133    }
134
135    fn to_arrow_array(&self) -> ArrayRef {
136        let v = self.vector.replicate(&[self.length]);
137        v.to_arrow_array()
138    }
139
140    fn to_boxed_arrow_array(&self) -> Box<dyn Array> {
141        let v = self.vector.replicate(&[self.length]);
142        v.to_boxed_arrow_array()
143    }
144
145    fn is_const(&self) -> bool {
146        true
147    }
148
149    fn validity(&self) -> Validity {
150        if self.vector.is_null(0) {
151            Validity::all_null(self.length)
152        } else {
153            Validity::all_valid(self.length)
154        }
155    }
156
157    fn memory_size(&self) -> usize {
158        self.vector.memory_size()
159    }
160
161    fn is_null(&self, _row: usize) -> bool {
162        self.vector.is_null(0)
163    }
164
165    fn only_null(&self) -> bool {
166        self.vector.is_null(0)
167    }
168
169    fn slice(&self, _offset: usize, length: usize) -> VectorRef {
170        Arc::new(Self {
171            vector: self.vector.clone(),
172            length,
173        })
174    }
175
176    fn get(&self, _index: usize) -> Value {
177        self.vector.get(0)
178    }
179
180    fn get_ref(&self, _index: usize) -> ValueRef {
181        self.vector.get_ref(0)
182    }
183
184    fn null_count(&self) -> usize {
185        if self.only_null() {
186            self.len()
187        } else {
188            0
189        }
190    }
191}
192
193impl fmt::Debug for ConstantVector {
194    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
195        write!(f, "ConstantVector([{:?}; {}])", self.get(0), self.len())
196    }
197}
198
199impl Serializable for ConstantVector {
200    fn serialize_to_json(&self) -> Result<Vec<serde_json::Value>> {
201        std::iter::repeat_n(self.get(0), self.len())
202            .map(serde_json::Value::try_from)
203            .collect::<serde_json::Result<_>>()
204            .context(SerializeSnafu)
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use arrow::datatypes::DataType as ArrowDataType;
211
212    use super::*;
213    use crate::vectors::Int32Vector;
214
215    #[test]
216    fn test_constant_vector_misc() {
217        let a = Int32Vector::from_slice(vec![1]);
218        let c = ConstantVector::new(Arc::new(a), 10);
219
220        assert_eq!("ConstantVector", c.vector_type_name());
221        assert!(c.is_const());
222        assert_eq!(10, c.len());
223        assert!(c.validity().is_all_valid());
224        assert!(!c.only_null());
225        assert_eq!(4, c.memory_size());
226
227        for i in 0..10 {
228            assert!(!c.is_null(i));
229            assert_eq!(Value::Int32(1), c.get(i));
230        }
231
232        let arrow_arr = c.to_arrow_array();
233        assert_eq!(10, arrow_arr.len());
234        assert_eq!(&ArrowDataType::Int32, arrow_arr.data_type());
235    }
236
237    #[test]
238    fn test_debug_null_array() {
239        let a = Int32Vector::from_slice(vec![1]);
240        let c = ConstantVector::new(Arc::new(a), 10);
241
242        let s = format!("{c:?}");
243        assert_eq!(s, "ConstantVector([Int32(1); 10])");
244    }
245
246    #[test]
247    fn test_serialize_json() {
248        let a = Int32Vector::from_slice(vec![1]);
249        let c = ConstantVector::new(Arc::new(a), 10);
250
251        let s = serde_json::to_string(&c.serialize_to_json().unwrap()).unwrap();
252        assert_eq!(s, "[1,1,1,1,1,1,1,1,1,1]");
253    }
254}
datatypes/vectors/constant.rs

datatypes/vectors/
constant.rs
