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::{ResultExt, ensure};
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!(
109                "Array index out of bounds, cannot take index out of the length of the array: {len}"
110            );
111        }
112
113        Ok(Arc::new(ConstantVector::new(
114            self.inner().clone(),
115            indices.len(),
116        )))
117    }
118}
119
120impl Vector for ConstantVector {
121    fn data_type(&self) -> ConcreteDataType {
122        self.vector.data_type()
123    }
124
125    fn vector_type_name(&self) -> String {
126        "ConstantVector".to_string()
127    }
128
129    fn as_any(&self) -> &dyn Any {
130        self
131    }
132
133    fn len(&self) -> usize {
134        self.length
135    }
136
137    fn to_arrow_array(&self) -> ArrayRef {
138        let v = self.vector.replicate(&[self.length]);
139        v.to_arrow_array()
140    }
141
142    fn to_boxed_arrow_array(&self) -> Box<dyn Array> {
143        let v = self.vector.replicate(&[self.length]);
144        v.to_boxed_arrow_array()
145    }
146
147    fn is_const(&self) -> bool {
148        true
149    }
150
151    fn validity(&self) -> Validity {
152        if self.vector.is_null(0) {
153            Validity::all_null(self.length)
154        } else {
155            Validity::all_valid(self.length)
156        }
157    }
158
159    fn memory_size(&self) -> usize {
160        self.vector.memory_size()
161    }
162
163    fn is_null(&self, _row: usize) -> bool {
164        self.vector.is_null(0)
165    }
166
167    fn only_null(&self) -> bool {
168        self.vector.is_null(0)
169    }
170
171    fn slice(&self, _offset: usize, length: usize) -> VectorRef {
172        Arc::new(Self {
173            vector: self.vector.clone(),
174            length,
175        })
176    }
177
178    fn get(&self, _index: usize) -> Value {
179        self.vector.get(0)
180    }
181
182    fn get_ref(&self, _index: usize) -> ValueRef {
183        self.vector.get_ref(0)
184    }
185
186    fn null_count(&self) -> usize {
187        if self.only_null() { self.len() } else { 0 }
188    }
189}
190
191impl fmt::Debug for ConstantVector {
192    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
193        write!(f, "ConstantVector([{:?}; {}])", self.get(0), self.len())
194    }
195}
196
197impl Serializable for ConstantVector {
198    fn serialize_to_json(&self) -> Result<Vec<serde_json::Value>> {
199        std::iter::repeat_n(self.get(0), self.len())
200            .map(serde_json::Value::try_from)
201            .collect::<serde_json::Result<_>>()
202            .context(SerializeSnafu)
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use arrow::datatypes::DataType as ArrowDataType;
209
210    use super::*;
211    use crate::vectors::Int32Vector;
212
213    #[test]
214    fn test_constant_vector_misc() {
215        let a = Int32Vector::from_slice(vec![1]);
216        let c = ConstantVector::new(Arc::new(a), 10);
217
218        assert_eq!("ConstantVector", c.vector_type_name());
219        assert!(c.is_const());
220        assert_eq!(10, c.len());
221        assert!(c.validity().is_all_valid());
222        assert!(!c.only_null());
223        assert_eq!(4, c.memory_size());
224
225        for i in 0..10 {
226            assert!(!c.is_null(i));
227            assert_eq!(Value::Int32(1), c.get(i));
228        }
229
230        let arrow_arr = c.to_arrow_array();
231        assert_eq!(10, arrow_arr.len());
232        assert_eq!(&ArrowDataType::Int32, arrow_arr.data_type());
233    }
234
235    #[test]
236    fn test_debug_null_array() {
237        let a = Int32Vector::from_slice(vec![1]);
238        let c = ConstantVector::new(Arc::new(a), 10);
239
240        let s = format!("{c:?}");
241        assert_eq!(s, "ConstantVector([Int32(1); 10])");
242    }
243
244    #[test]
245    fn test_serialize_json() {
246        let a = Int32Vector::from_slice(vec![1]);
247        let c = ConstantVector::new(Arc::new(a), 10);
248
249        let s = serde_json::to_string(&c.serialize_to_json().unwrap()).unwrap();
250        assert_eq!(s, "[1,1,1,1,1,1,1,1,1,1]");
251    }
252}
datatypes/vectors/constant.rs

datatypes/vectors/
constant.rs
