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

datatypes/schema/
constraint.rs
