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, GeneralDataFusionSnafu, 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
47                    .to_array_of_size(num_rows)
48                    .context(GeneralDataFusionSnafu)?;
49                let data_type = v.data_type().clone();
50                Helper::try_into_vector(v).context(IntoVectorSnafu { data_type })?
51            }
52        })
53    }
54}
55
56impl TryFrom<&DfColumnarValue> for ColumnarValue {
57    type Error = error::Error;
58    fn try_from(value: &DfColumnarValue) -> Result<Self> {
59        Ok(match value {
60            DfColumnarValue::Scalar(v) => ColumnarValue::Scalar(v.clone()),
61            DfColumnarValue::Array(v) => {
62                ColumnarValue::Vector(Helper::try_into_vector(v.clone()).with_context(|_| {
63                    IntoVectorSnafu {
64                        data_type: v.data_type().clone(),
65                    }
66                })?)
67            }
68        })
69    }
70}
71
72impl From<ColumnarValue> for DfColumnarValue {
73    fn from(columnar_value: ColumnarValue) -> Self {
74        match columnar_value {
75            ColumnarValue::Scalar(v) => DfColumnarValue::Scalar(v),
76            ColumnarValue::Vector(v) => DfColumnarValue::Array(v.to_arrow_array()),
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use std::sync::Arc;
84
85    use datatypes::arrow::datatypes::DataType as ArrowDataType;
86    use datatypes::vectors::BooleanVector;
87
88    use super::*;
89
90    #[test]
91    fn test_scalar_columnar_value() {
92        let value = ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)));
93
94        assert_eq!(ConcreteDataType::boolean_datatype(), value.data_type());
95        let v = value.clone().try_into_vector(1).unwrap();
96        assert_eq!(1, v.len());
97
98        let df_value = DfColumnarValue::from(value);
99        assert_eq!(ArrowDataType::Boolean, df_value.data_type());
100    }
101
102    #[test]
103    fn test_vector_columnar_value() {
104        let value = ColumnarValue::Vector(Arc::new(BooleanVector::from(vec![true, false, true])));
105
106        assert_eq!(ConcreteDataType::boolean_datatype(), value.data_type());
107        let v = value.clone().try_into_vector(1).unwrap();
108        assert_eq!(3, v.len());
109
110        let df_value = DfColumnarValue::from(value);
111        assert_eq!(ArrowDataType::Boolean, df_value.data_type());
112    }
113}
common_query/columnar_value.rs

common_query/
columnar_value.rs
