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::{Timestamp, util};
18use serde::{Deserialize, Serialize};
19use snafu::{ResultExt, ensure};
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!(
335            constraint
336                .validate(&ConcreteDataType::uint32_datatype(), true)
337                .is_err()
338        );
339    }
340
341    #[test]
342    fn test_validate_function_constraint() {
343        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
344        constraint
345            .validate(&ConcreteDataType::timestamp_millisecond_datatype(), false)
346            .unwrap();
347        assert!(
348            constraint
349                .validate(&ConcreteDataType::boolean_datatype(), false)
350                .is_err()
351        );
352
353        let constraint = ColumnDefaultConstraint::Function("hello()".to_string());
354        assert!(
355            constraint
356                .validate(&ConcreteDataType::timestamp_millisecond_datatype(), false)
357                .is_err()
358        );
359    }
360
361    #[test]
362    fn test_create_default_vector_by_null() {
363        let constraint = ColumnDefaultConstraint::null_value();
364        let data_type = ConcreteDataType::int32_datatype();
365        assert!(
366            constraint
367                .create_default_vector(&data_type, false, 10)
368                .is_err()
369        );
370
371        let constraint = ColumnDefaultConstraint::null_value();
372        let v = constraint
373            .create_default_vector(&data_type, true, 3)
374            .unwrap();
375        assert_eq!(3, v.len());
376        for i in 0..v.len() {
377            assert_eq!(Value::Null, v.get(i));
378        }
379    }
380
381    #[test]
382    fn test_create_default_by_value() {
383        let constraint = ColumnDefaultConstraint::Value(Value::Int32(10));
384        let data_type = ConcreteDataType::int32_datatype();
385        let v = constraint
386            .create_default_vector(&data_type, false, 4)
387            .unwrap();
388        let expect: VectorRef = Arc::new(Int32Vector::from_values(vec![10; 4]));
389        assert_eq!(expect, v);
390        let v = constraint.create_default(&data_type, false).unwrap();
391        assert_eq!(Value::Int32(10), v);
392    }
393
394    #[test]
395    fn test_create_default_vector_by_func() {
396        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
397        let check_value = |v| {
398            assert!(
399                matches!(v, Value::Timestamp(_)),
400                "v {:?} is not timestamp",
401                v
402            );
403        };
404        let check_vector = |v: VectorRef| {
405            assert_eq!(4, v.len());
406            assert!(
407                matches!(v.get(0), Value::Timestamp(_)),
408                "v {:?} is not timestamp",
409                v.get(0)
410            );
411        };
412
413        // Timestamp type.
414        let data_type = ConcreteDataType::timestamp_millisecond_datatype();
415        let v = constraint
416            .create_default_vector(&data_type, false, 4)
417            .unwrap();
418        check_vector(v);
419
420        let v = constraint.create_default(&data_type, false).unwrap();
421        check_value(v);
422
423        let data_type = ConcreteDataType::timestamp_second_datatype();
424        let v = constraint
425            .create_default_vector(&data_type, false, 4)
426            .unwrap();
427        check_vector(v);
428
429        let v = constraint.create_default(&data_type, false).unwrap();
430        check_value(v);
431
432        let data_type = ConcreteDataType::timestamp_microsecond_datatype();
433        let v = constraint
434            .create_default_vector(&data_type, false, 4)
435            .unwrap();
436        check_vector(v);
437
438        let v = constraint.create_default(&data_type, false).unwrap();
439        check_value(v);
440
441        let data_type = ConcreteDataType::timestamp_nanosecond_datatype();
442        let v = constraint
443            .create_default_vector(&data_type, false, 4)
444            .unwrap();
445        check_vector(v);
446
447        let v = constraint.create_default(&data_type, false).unwrap();
448        check_value(v);
449
450        // Int64 type.
451        let data_type = ConcreteDataType::int64_datatype();
452        let v = constraint.create_default_vector(&data_type, false, 4);
453        assert!(v.is_err());
454
455        let constraint = ColumnDefaultConstraint::Function("no".to_string());
456        let data_type = ConcreteDataType::timestamp_millisecond_datatype();
457        assert!(
458            constraint
459                .create_default_vector(&data_type, false, 4)
460                .is_err()
461        );
462        assert!(constraint.create_default(&data_type, false).is_err());
463    }
464
465    #[test]
466    fn test_create_by_func_and_invalid_type() {
467        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
468        let data_type = ConcreteDataType::boolean_datatype();
469        let err = constraint
470            .create_default_vector(&data_type, false, 4)
471            .unwrap_err();
472        assert!(matches!(err, Error::DefaultValueType { .. }), "{err:?}");
473    }
474}
datatypes/schema/constraint.rs

datatypes/schema/
constraint.rs
