common_query/
columnar_value.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 datafusion_expr::ColumnarValue as DfColumnarValue;
16use datatypes::prelude::ConcreteDataType;
17use datatypes::vectors::{Helper, VectorRef};
18use snafu::ResultExt;
19
20use crate::error::{self, IntoVectorSnafu, Result};
21use crate::prelude::ScalarValue;
22
23/// Represents the result from an expression
24#[derive(Clone)]
25pub enum ColumnarValue {
26    Vector(VectorRef),
27    /// A single value
28    Scalar(ScalarValue),
29}
30
31impl ColumnarValue {
32    pub fn data_type(&self) -> ConcreteDataType {
33        match self {
34            ColumnarValue::Vector(vector) => vector.data_type(),
35            ColumnarValue::Scalar(scalar_value) => {
36                ConcreteDataType::from_arrow_type(&scalar_value.data_type())
37            }
38        }
39    }
40
41    /// Convert a columnar value into an VectorRef
42    pub fn try_into_vector(self, num_rows: usize) -> Result<VectorRef> {
43        Ok(match self {
44            ColumnarValue::Vector(v) => v,
45            ColumnarValue::Scalar(s) => {
46                let v = s.to_array_of_size(num_rows)?;
47                let data_type = v.data_type().clone();
48                Helper::try_into_vector(v).context(IntoVectorSnafu { data_type })?
49            }
50        })
51    }
52}
53
54impl TryFrom<&DfColumnarValue> for ColumnarValue {
55    type Error = error::Error;
56    fn try_from(value: &DfColumnarValue) -> Result<Self> {
57        Ok(match value {
58            DfColumnarValue::Scalar(v) => ColumnarValue::Scalar(v.clone()),
59            DfColumnarValue::Array(v) => {
60                ColumnarValue::Vector(Helper::try_into_vector(v.clone()).with_context(|_| {
61                    IntoVectorSnafu {
62                        data_type: v.data_type().clone(),
63                    }
64                })?)
65            }
66        })
67    }
68}
69
70impl From<ColumnarValue> for DfColumnarValue {
71    fn from(columnar_value: ColumnarValue) -> Self {
72        match columnar_value {
73            ColumnarValue::Scalar(v) => DfColumnarValue::Scalar(v),
74            ColumnarValue::Vector(v) => DfColumnarValue::Array(v.to_arrow_array()),
75        }
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use std::sync::Arc;
82
83    use datatypes::arrow::datatypes::DataType as ArrowDataType;
84    use datatypes::vectors::BooleanVector;
85
86    use super::*;
87
88    #[test]
89    fn test_scalar_columnar_value() {
90        let value = ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)));
91
92        assert_eq!(ConcreteDataType::boolean_datatype(), value.data_type());
93        let v = value.clone().try_into_vector(1).unwrap();
94        assert_eq!(1, v.len());
95
96        let df_value = DfColumnarValue::from(value);
97        assert_eq!(ArrowDataType::Boolean, df_value.data_type());
98    }
99
100    #[test]
101    fn test_vector_columnar_value() {
102        let value = ColumnarValue::Vector(Arc::new(BooleanVector::from(vec![true, false, true])));
103
104        assert_eq!(ConcreteDataType::boolean_datatype(), value.data_type());
105        let v = value.clone().try_into_vector(1).unwrap();
106        assert_eq!(3, v.len());
107
108        let df_value = DfColumnarValue::from(value);
109        assert_eq!(ArrowDataType::Boolean, df_value.data_type());
110    }
111}
common_query/columnar_value.rs

common_query/
columnar_value.rs
