datatypes/schema/
constraint.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::fmt::{Display, Formatter};
16
17use common_time::{util, Timestamp};
18use serde::{Deserialize, Serialize};
19use snafu::{ensure, ResultExt};
20
21use crate::data_type::{ConcreteDataType, DataType};
22use crate::error::{self, Result};
23use crate::types::cast;
24use crate::value::Value;
25use crate::vectors::operations::VectorOp;
26use crate::vectors::{TimestampMillisecondVector, VectorRef};
27
28pub const CURRENT_TIMESTAMP: &str = "current_timestamp";
29pub const CURRENT_TIMESTAMP_FN: &str = "current_timestamp()";
30pub const NOW_FN: &str = "now()";
31
32/// Column's default constraint.
33#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
34pub enum ColumnDefaultConstraint {
35    // A function invocation
36    // TODO(dennis): we save the function expression here, maybe use a struct in future.
37    Function(String),
38    // A value
39    Value(Value),
40}
41
42impl TryFrom<&[u8]> for ColumnDefaultConstraint {
43    type Error = error::Error;
44
45    fn try_from(bytes: &[u8]) -> Result<Self> {
46        let json = String::from_utf8_lossy(bytes);
47        serde_json::from_str(&json).context(error::DeserializeSnafu { json })
48    }
49}
50
51impl TryFrom<ColumnDefaultConstraint> for Vec<u8> {
52    type Error = error::Error;
53
54    fn try_from(value: ColumnDefaultConstraint) -> std::result::Result<Self, Self::Error> {
55        let s = serde_json::to_string(&value).context(error::SerializeSnafu)?;
56        Ok(s.into_bytes())
57    }
58}
59
60impl Display for ColumnDefaultConstraint {
61    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
62        match self {
63            ColumnDefaultConstraint::Function(expr) => write!(f, "{expr}"),
64            ColumnDefaultConstraint::Value(v) => write!(f, "{v}"),
65        }
66    }
67}
68
69impl ColumnDefaultConstraint {
70    /// Returns a default null constraint.
71    pub fn null_value() -> ColumnDefaultConstraint {
72        ColumnDefaultConstraint::Value(Value::Null)
73    }
74
75    /// Check whether the constraint is valid for columns with given `data_type`
76    /// and `is_nullable` attributes.
77    pub fn validate(&self, data_type: &ConcreteDataType, is_nullable: bool) -> Result<()> {
78        ensure!(is_nullable || !self.maybe_null(), error::NullDefaultSnafu);
79
80        match self {
81            ColumnDefaultConstraint::Function(expr) => {
82                ensure!(
83                    expr == CURRENT_TIMESTAMP || expr == CURRENT_TIMESTAMP_FN || expr == NOW_FN,
84                    error::UnsupportedDefaultExprSnafu { expr }
85                );
86                ensure!(
87                    data_type.is_timestamp(),
88                    error::DefaultValueTypeSnafu {
89                        reason: "return value of the function must has timestamp type",
90                    }
91                );
92            }
93            ColumnDefaultConstraint::Value(v) => {
94                if !v.is_null() {
95                    // Whether the value could be nullable has been checked before, only need
96                    // to check the type compatibility here.
97                    ensure!(
98                        value_type_match(data_type, v.data_type()),
99                        error::DefaultValueTypeSnafu {
100                            reason: format!(
101                                "column has type {:?} but default value has type {:?}",
102                                data_type.logical_type_id(),
103                                v.logical_type_id()
104                            ),
105                        }
106                    );
107                }
108            }
109        }
110
111        Ok(())
112    }
113
114    /// Create a vector that contains `num_rows` default values for given `data_type`.
115    ///
116    /// If `is_nullable` is `true`, then this method would returns error if the created
117    /// default value is null.
118    ///
119    /// # Panics
120    /// Panics if `num_rows == 0`.
121    pub fn create_default_vector(
122        &self,
123        data_type: &ConcreteDataType,
124        is_nullable: bool,
125        num_rows: usize,
126    ) -> Result<VectorRef> {
127        assert!(num_rows > 0);
128
129        match self {
130            ColumnDefaultConstraint::Function(expr) => {
131                // Functions should also ensure its return value is not null when
132                // is_nullable is true.
133                match &expr[..] {
134                    // TODO(dennis): we only supports current_timestamp right now,
135                    //   it's better to use a expression framework in future.
136                    CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => {
137                        create_current_timestamp_vector(data_type, num_rows)
138                    }
139                    _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
140                }
141            }
142            ColumnDefaultConstraint::Value(v) => {
143                ensure!(is_nullable || !v.is_null(), error::NullDefaultSnafu);
144
145                // TODO(yingwen):
146                // 1. For null value, we could use NullVector once it supports custom logical type.
147                // 2. For non null value, we could use ConstantVector, but it would cause all codes
148                //  attempt to downcast the vector fail if they don't check whether the vector is const
149                //  first.
150                let mut mutable_vector = data_type.create_mutable_vector(1);
151                mutable_vector.try_push_value_ref(v.as_value_ref())?;
152                let base_vector = mutable_vector.to_vector();
153                Ok(base_vector.replicate(&[num_rows]))
154            }
155        }
156    }
157
158    /// Create a default value for given `data_type`.
159    ///
160    /// If `is_nullable` is `true`, then this method would returns error if the created
161    /// default value is null.
162    pub fn create_default(&self, data_type: &ConcreteDataType, is_nullable: bool) -> Result<Value> {
163        match self {
164            ColumnDefaultConstraint::Function(expr) => {
165                // Functions should also ensure its return value is not null when
166                // is_nullable is true.
167                match &expr[..] {
168                    CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => {
169                        create_current_timestamp(data_type)
170                    }
171                    _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
172                }
173            }
174            ColumnDefaultConstraint::Value(v) => {
175                ensure!(is_nullable || !v.is_null(), error::NullDefaultSnafu);
176
177                Ok(v.clone())
178            }
179        }
180    }
181
182    /// Cast default value to given type
183    pub fn cast_to_datatype(&self, data_type: &ConcreteDataType) -> Result<Self> {
184        match self {
185            ColumnDefaultConstraint::Value(v) => Ok(Self::Value(cast(v.clone(), data_type)?)),
186            ColumnDefaultConstraint::Function(expr) => match &expr[..] {
187                // no need to cast, since function always require a data_type when need to create default value
188                CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => Ok(self.clone()),
189                _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
190            },
191        }
192    }
193
194    /// Only create default vector if it's impure, i.e., it's a function.
195    ///
196    /// This helps to delay creating constant default values to mito engine while also keeps impure default have consistent values
197    pub fn create_impure_default_vector(
198        &self,
199        data_type: &ConcreteDataType,
200        num_rows: usize,
201    ) -> Result<Option<VectorRef>> {
202        assert!(num_rows > 0);
203
204        match self {
205            ColumnDefaultConstraint::Function(expr) => {
206                // Functions should also ensure its return value is not null when
207                // is_nullable is true.
208                match &expr[..] {
209                    // TODO(dennis): we only supports current_timestamp right now,
210                    //   it's better to use a expression framework in future.
211                    CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => {
212                        create_current_timestamp_vector(data_type, num_rows).map(Some)
213                    }
214                    _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
215                }
216            }
217            ColumnDefaultConstraint::Value(_) => Ok(None),
218        }
219    }
220
221    /// Only create default value if it's impure, i.e., it's a function.
222    ///
223    /// This helps to delay creating constant default values to mito engine while also keeps impure default have consistent values
224    pub fn create_impure_default(&self, data_type: &ConcreteDataType) -> Result<Option<Value>> {
225        match self {
226            ColumnDefaultConstraint::Function(expr) => {
227                // Functions should also ensure its return value is not null when
228                // is_nullable is true.
229                match &expr[..] {
230                    CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => {
231                        create_current_timestamp(data_type).map(Some)
232                    }
233                    _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
234                }
235            }
236            ColumnDefaultConstraint::Value(_) => Ok(None),
237        }
238    }
239
240    /// Returns true if this constraint might creates NULL.
241    fn maybe_null(&self) -> bool {
242        // Once we support more functions, we may return true if given function
243        // could return null.
244        matches!(self, ColumnDefaultConstraint::Value(Value::Null))
245    }
246
247    /// Returns true if this constraint is a function.
248    pub fn is_function(&self) -> bool {
249        matches!(self, ColumnDefaultConstraint::Function(_))
250    }
251}
252
253fn create_current_timestamp(data_type: &ConcreteDataType) -> Result<Value> {
254    let Some(timestamp_type) = data_type.as_timestamp() else {
255        return error::DefaultValueTypeSnafu {
256            reason: format!("Not support to assign current timestamp to {data_type:?} type"),
257        }
258        .fail();
259    };
260
261    let unit = timestamp_type.unit();
262    Ok(Value::Timestamp(Timestamp::current_time(unit)))
263}
264
265fn create_current_timestamp_vector(
266    data_type: &ConcreteDataType,
267    num_rows: usize,
268) -> Result<VectorRef> {
269    let current_timestamp_vector = TimestampMillisecondVector::from_values(std::iter::repeat_n(
270        util::current_time_millis(),
271        num_rows,
272    ));
273    if data_type.is_timestamp() {
274        current_timestamp_vector.cast(data_type)
275    } else {
276        error::DefaultValueTypeSnafu {
277            reason: format!("Not support to assign current timestamp to {data_type:?} type",),
278        }
279        .fail()
280    }
281}
282
283fn value_type_match(column_type: &ConcreteDataType, value_type: ConcreteDataType) -> bool {
284    match (column_type, value_type) {
285        (ct, vt) if ct.logical_type_id() == vt.logical_type_id() => true,
286        // Vector and Json type is encoded as binary
287        (ConcreteDataType::Vector(_) | ConcreteDataType::Json(_), ConcreteDataType::Binary(_)) => {
288            true
289        }
290        _ => false,
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use std::sync::Arc;
297
298    use super::*;
299    use crate::error::Error;
300    use crate::vectors::Int32Vector;
301
302    #[test]
303    fn test_null_default_constraint() {
304        let constraint = ColumnDefaultConstraint::null_value();
305        assert!(constraint.maybe_null());
306        let constraint = ColumnDefaultConstraint::Value(Value::Int32(10));
307        assert!(!constraint.maybe_null());
308    }
309
310    #[test]
311    fn test_validate_null_constraint() {
312        let constraint = ColumnDefaultConstraint::null_value();
313        let data_type = ConcreteDataType::int32_datatype();
314        assert!(constraint.validate(&data_type, false).is_err());
315        constraint.validate(&data_type, true).unwrap();
316    }
317
318    #[test]
319    fn test_validate_value_constraint() {
320        let constraint = ColumnDefaultConstraint::Value(Value::Int32(10));
321        let data_type = ConcreteDataType::int32_datatype();
322        constraint.validate(&data_type, false).unwrap();
323        constraint.validate(&data_type, true).unwrap();
324
325        assert!(constraint
326            .validate(&ConcreteDataType::uint32_datatype(), true)
327            .is_err());
328    }
329
330    #[test]
331    fn test_validate_function_constraint() {
332        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
333        constraint
334            .validate(&ConcreteDataType::timestamp_millisecond_datatype(), false)
335            .unwrap();
336        assert!(constraint
337            .validate(&ConcreteDataType::boolean_datatype(), false)
338            .is_err());
339
340        let constraint = ColumnDefaultConstraint::Function("hello()".to_string());
341        assert!(constraint
342            .validate(&ConcreteDataType::timestamp_millisecond_datatype(), false)
343            .is_err());
344    }
345
346    #[test]
347    fn test_create_default_vector_by_null() {
348        let constraint = ColumnDefaultConstraint::null_value();
349        let data_type = ConcreteDataType::int32_datatype();
350        assert!(constraint
351            .create_default_vector(&data_type, false, 10)
352            .is_err());
353
354        let constraint = ColumnDefaultConstraint::null_value();
355        let v = constraint
356            .create_default_vector(&data_type, true, 3)
357            .unwrap();
358        assert_eq!(3, v.len());
359        for i in 0..v.len() {
360            assert_eq!(Value::Null, v.get(i));
361        }
362    }
363
364    #[test]
365    fn test_create_default_by_value() {
366        let constraint = ColumnDefaultConstraint::Value(Value::Int32(10));
367        let data_type = ConcreteDataType::int32_datatype();
368        let v = constraint
369            .create_default_vector(&data_type, false, 4)
370            .unwrap();
371        let expect: VectorRef = Arc::new(Int32Vector::from_values(vec![10; 4]));
372        assert_eq!(expect, v);
373        let v = constraint.create_default(&data_type, false).unwrap();
374        assert_eq!(Value::Int32(10), v);
375    }
376
377    #[test]
378    fn test_create_default_vector_by_func() {
379        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
380        let check_value = |v| {
381            assert!(
382                matches!(v, Value::Timestamp(_)),
383                "v {:?} is not timestamp",
384                v
385            );
386        };
387        let check_vector = |v: VectorRef| {
388            assert_eq!(4, v.len());
389            assert!(
390                matches!(v.get(0), Value::Timestamp(_)),
391                "v {:?} is not timestamp",
392                v.get(0)
393            );
394        };
395
396        // Timestamp type.
397        let data_type = ConcreteDataType::timestamp_millisecond_datatype();
398        let v = constraint
399            .create_default_vector(&data_type, false, 4)
400            .unwrap();
401        check_vector(v);
402
403        let v = constraint.create_default(&data_type, false).unwrap();
404        check_value(v);
405
406        let data_type = ConcreteDataType::timestamp_second_datatype();
407        let v = constraint
408            .create_default_vector(&data_type, false, 4)
409            .unwrap();
410        check_vector(v);
411
412        let v = constraint.create_default(&data_type, false).unwrap();
413        check_value(v);
414
415        let data_type = ConcreteDataType::timestamp_microsecond_datatype();
416        let v = constraint
417            .create_default_vector(&data_type, false, 4)
418            .unwrap();
419        check_vector(v);
420
421        let v = constraint.create_default(&data_type, false).unwrap();
422        check_value(v);
423
424        let data_type = ConcreteDataType::timestamp_nanosecond_datatype();
425        let v = constraint
426            .create_default_vector(&data_type, false, 4)
427            .unwrap();
428        check_vector(v);
429
430        let v = constraint.create_default(&data_type, false).unwrap();
431        check_value(v);
432
433        // Int64 type.
434        let data_type = ConcreteDataType::int64_datatype();
435        let v = constraint.create_default_vector(&data_type, false, 4);
436        assert!(v.is_err());
437
438        let constraint = ColumnDefaultConstraint::Function("no".to_string());
439        let data_type = ConcreteDataType::timestamp_millisecond_datatype();
440        assert!(constraint
441            .create_default_vector(&data_type, false, 4)
442            .is_err());
443        assert!(constraint.create_default(&data_type, false).is_err());
444    }
445
446    #[test]
447    fn test_create_by_func_and_invalid_type() {
448        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
449        let data_type = ConcreteDataType::boolean_datatype();
450        let err = constraint
451            .create_default_vector(&data_type, false, 4)
452            .unwrap_err();
453        assert!(matches!(err, Error::DefaultValueType { .. }), "{err:?}");
454    }
455}
datatypes/schema/constraint.rs

datatypes/schema/
constraint.rs
