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::{Helper, 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                if let Ok(vector) = v.try_to_scalar_value(data_type).and_then(|scalar| {
155                    Helper::try_from_scalar_value(scalar, num_rows, Some(data_type))
156                }) {
157                    return Ok(vector);
158                }
159
160                // Some extension values, such as JSON nested in a struct, cannot safely
161                // round-trip through ScalarValue. Preserve their logical type with the
162                // type-specific vector builder instead.
163                let mut mutable_vector = data_type.create_mutable_vector(num_rows);
164                for _ in 0..num_rows {
165                    mutable_vector.try_push_value_ref(&v.as_value_ref())?;
166                }
167                Ok(mutable_vector.to_vector())
168            }
169        }
170    }
171
172    /// Create a default value for given `data_type`.
173    ///
174    /// If `is_nullable` is `true`, then this method would returns error if the created
175    /// default value is null.
176    pub fn create_default(&self, data_type: &ConcreteDataType, is_nullable: bool) -> Result<Value> {
177        match self {
178            ColumnDefaultConstraint::Function(expr) => {
179                // Functions should also ensure its return value is not null when
180                // is_nullable is true.
181                match &expr[..] {
182                    CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => {
183                        create_current_timestamp(data_type)
184                    }
185                    _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
186                }
187            }
188            ColumnDefaultConstraint::Value(v) => {
189                ensure!(is_nullable || !v.is_null(), error::NullDefaultSnafu);
190
191                Ok(v.clone())
192            }
193        }
194    }
195
196    /// Cast default value to given type
197    pub fn cast_to_datatype(&self, data_type: &ConcreteDataType) -> Result<Self> {
198        match self {
199            ColumnDefaultConstraint::Value(v) => Ok(Self::Value(cast(v.clone(), data_type)?)),
200            ColumnDefaultConstraint::Function(expr) => match &expr[..] {
201                // no need to cast, since function always require a data_type when need to create default value
202                CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => Ok(self.clone()),
203                _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
204            },
205        }
206    }
207
208    /// Only create default vector if it's impure, i.e., it's a function.
209    ///
210    /// This helps to delay creating constant default values to mito engine while also keeps impure default have consistent values
211    pub fn create_impure_default_vector(
212        &self,
213        data_type: &ConcreteDataType,
214        num_rows: usize,
215    ) -> Result<Option<VectorRef>> {
216        assert!(num_rows > 0);
217
218        match self {
219            ColumnDefaultConstraint::Function(expr) => {
220                // Functions should also ensure its return value is not null when
221                // is_nullable is true.
222                match &expr[..] {
223                    // TODO(dennis): we only supports current_timestamp right now,
224                    //   it's better to use a expression framework in future.
225                    CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => {
226                        create_current_timestamp_vector(data_type, num_rows).map(Some)
227                    }
228                    _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
229                }
230            }
231            ColumnDefaultConstraint::Value(_) => Ok(None),
232        }
233    }
234
235    /// Only create default value if it's impure, i.e., it's a function.
236    ///
237    /// This helps to delay creating constant default values to mito engine while also keeps impure default have consistent values
238    pub fn create_impure_default(&self, data_type: &ConcreteDataType) -> Result<Option<Value>> {
239        match self {
240            ColumnDefaultConstraint::Function(expr) => {
241                // Functions should also ensure its return value is not null when
242                // is_nullable is true.
243                match &expr[..] {
244                    CURRENT_TIMESTAMP | CURRENT_TIMESTAMP_FN | NOW_FN => {
245                        create_current_timestamp(data_type).map(Some)
246                    }
247                    _ => error::UnsupportedDefaultExprSnafu { expr }.fail(),
248                }
249            }
250            ColumnDefaultConstraint::Value(_) => Ok(None),
251        }
252    }
253
254    /// Returns true if this constraint might creates NULL.
255    fn maybe_null(&self) -> bool {
256        // Once we support more functions, we may return true if given function
257        // could return null.
258        matches!(self, ColumnDefaultConstraint::Value(Value::Null))
259    }
260
261    /// Returns true if this constraint is a function.
262    pub fn is_function(&self) -> bool {
263        matches!(self, ColumnDefaultConstraint::Function(_))
264    }
265}
266
267fn create_current_timestamp(data_type: &ConcreteDataType) -> Result<Value> {
268    let Some(timestamp_type) = data_type.as_timestamp() else {
269        return error::DefaultValueTypeSnafu {
270            reason: format!("Not support to assign current timestamp to {data_type:?} type"),
271        }
272        .fail();
273    };
274
275    let unit = timestamp_type.unit();
276    Ok(Value::Timestamp(Timestamp::current_time(unit)))
277}
278
279fn create_current_timestamp_vector(
280    data_type: &ConcreteDataType,
281    num_rows: usize,
282) -> Result<VectorRef> {
283    let current_timestamp_vector = TimestampMillisecondVector::from_values(std::iter::repeat_n(
284        util::current_time_millis(),
285        num_rows,
286    ));
287    if data_type.is_timestamp() {
288        current_timestamp_vector.cast(data_type)
289    } else {
290        error::DefaultValueTypeSnafu {
291            reason: format!("Not support to assign current timestamp to {data_type:?} type",),
292        }
293        .fail()
294    }
295}
296
297fn value_type_match(column_type: &ConcreteDataType, value_type: ConcreteDataType) -> bool {
298    match (column_type, value_type) {
299        (ct, vt) if ct.logical_type_id() == vt.logical_type_id() => true,
300        // Vector and Json type is encoded as binary
301        (ConcreteDataType::Vector(_) | ConcreteDataType::Json(_), ConcreteDataType::Binary(_)) => {
302            true
303        }
304        _ => false,
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use std::sync::Arc;
311
312    use super::*;
313    use crate::error::Error;
314    use crate::vectors::Int32Vector;
315
316    #[test]
317    fn test_null_default_constraint() {
318        let constraint = ColumnDefaultConstraint::null_value();
319        assert!(constraint.maybe_null());
320        let constraint = ColumnDefaultConstraint::Value(Value::Int32(10));
321        assert!(!constraint.maybe_null());
322    }
323
324    #[test]
325    fn test_validate_null_constraint() {
326        let constraint = ColumnDefaultConstraint::null_value();
327        let data_type = ConcreteDataType::int32_datatype();
328        assert!(constraint.validate(&data_type, false).is_err());
329        constraint.validate(&data_type, true).unwrap();
330    }
331
332    #[test]
333    fn test_validate_value_constraint() {
334        let constraint = ColumnDefaultConstraint::Value(Value::Int32(10));
335        let data_type = ConcreteDataType::int32_datatype();
336        constraint.validate(&data_type, false).unwrap();
337        constraint.validate(&data_type, true).unwrap();
338
339        assert!(
340            constraint
341                .validate(&ConcreteDataType::uint32_datatype(), true)
342                .is_err()
343        );
344    }
345
346    #[test]
347    fn test_validate_function_constraint() {
348        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
349        constraint
350            .validate(&ConcreteDataType::timestamp_millisecond_datatype(), false)
351            .unwrap();
352        assert!(
353            constraint
354                .validate(&ConcreteDataType::boolean_datatype(), false)
355                .is_err()
356        );
357
358        let constraint = ColumnDefaultConstraint::Function("hello()".to_string());
359        assert!(
360            constraint
361                .validate(&ConcreteDataType::timestamp_millisecond_datatype(), false)
362                .is_err()
363        );
364    }
365
366    #[test]
367    fn test_create_default_vector_by_null() {
368        let constraint = ColumnDefaultConstraint::null_value();
369        let data_type = ConcreteDataType::int32_datatype();
370        assert!(
371            constraint
372                .create_default_vector(&data_type, false, 10)
373                .is_err()
374        );
375
376        let constraint = ColumnDefaultConstraint::null_value();
377        let v = constraint
378            .create_default_vector(&data_type, true, 3)
379            .unwrap();
380        assert_eq!(3, v.len());
381        for i in 0..v.len() {
382            assert_eq!(Value::Null, v.get(i));
383        }
384    }
385
386    #[test]
387    fn test_create_default_by_value() {
388        let constraint = ColumnDefaultConstraint::Value(Value::Int32(10));
389        let data_type = ConcreteDataType::int32_datatype();
390        let v = constraint
391            .create_default_vector(&data_type, false, 4)
392            .unwrap();
393        let expect: VectorRef = Arc::new(Int32Vector::from_values(vec![10; 4]));
394        assert_eq!(expect, v);
395        let v = constraint.create_default(&data_type, false).unwrap();
396        assert_eq!(Value::Int32(10), v);
397    }
398
399    #[test]
400    fn test_struct_default_null_fields_and_json() {
401        use crate::types::{StructField, StructType};
402        use crate::value::StructValue;
403
404        let inner_type = StructType::from([StructField::new(
405            "x",
406            ConcreteDataType::int32_datatype(),
407            true,
408        )]);
409        let inner = Value::Struct(StructValue::new(vec![Value::Null], inner_type));
410        let json = crate::json::JsonSettings::default()
411            .encode(serde_json::json!({"answer": 42}))
412            .unwrap();
413        let nested_type = StructType::from([StructField::new("nested", inner.data_type(), true)]);
414        let json_type = StructType::from([StructField::new("json", json.data_type(), true)]);
415        let values = [
416            inner.clone(),
417            Value::Struct(StructValue::new(vec![inner], nested_type)),
418            Value::Struct(StructValue::new(vec![], StructType::default())),
419            Value::Struct(StructValue::new(vec![json], json_type)),
420        ];
421        for value in values {
422            let data_type = value.data_type();
423            // JSON children are read back as their underlying struct values.
424            let expected = serde_json::Value::try_from(value.clone()).unwrap();
425            for num_rows in [1, 3] {
426                let vector = ColumnDefaultConstraint::Value(value.clone())
427                    .create_default_vector(&data_type, false, num_rows)
428                    .unwrap();
429                assert_eq!(data_type, vector.data_type());
430                assert_eq!(
431                    data_type.as_arrow_type(),
432                    *vector.to_arrow_array().data_type()
433                );
434                assert_eq!(num_rows, vector.len());
435                assert_eq!(0, vector.null_count());
436                for row in 0..num_rows {
437                    assert_eq!(
438                        expected,
439                        serde_json::Value::try_from(vector.get(row)).unwrap()
440                    );
441                }
442            }
443        }
444    }
445
446    #[test]
447    fn test_string_default_preserves_batch_schema() {
448        use arrow::datatypes::{Field, Schema};
449        use arrow::record_batch::RecordBatch;
450
451        for data_type in [
452            ConcreteDataType::large_string_datatype(),
453            ConcreteDataType::utf8_view_datatype(),
454        ] {
455            let schema = Arc::new(Schema::new(vec![Field::new(
456                "tag",
457                data_type.as_arrow_type(),
458                true,
459            )]));
460            for value in [Value::from("greptime"), Value::Null] {
461                let vector = ColumnDefaultConstraint::Value(value.clone())
462                    .create_default_vector(&data_type, true, 3)
463                    .unwrap();
464                let batch =
465                    RecordBatch::try_new(schema.clone(), vec![vector.to_arrow_array()]).unwrap();
466                assert_eq!(3, batch.num_rows());
467                assert_eq!(data_type, vector.data_type());
468                for row in 0..batch.num_rows() {
469                    assert_eq!(value, vector.get(row));
470                }
471            }
472        }
473    }
474
475    #[test]
476    fn test_create_default_vector_by_func() {
477        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
478        let check_value = |v| {
479            assert!(
480                matches!(v, Value::Timestamp(_)),
481                "v {:?} is not timestamp",
482                v
483            );
484        };
485        let check_vector = |v: VectorRef| {
486            assert_eq!(4, v.len());
487            assert!(
488                matches!(v.get(0), Value::Timestamp(_)),
489                "v {:?} is not timestamp",
490                v.get(0)
491            );
492        };
493
494        // Timestamp type.
495        let data_type = ConcreteDataType::timestamp_millisecond_datatype();
496        let v = constraint
497            .create_default_vector(&data_type, false, 4)
498            .unwrap();
499        check_vector(v);
500
501        let v = constraint.create_default(&data_type, false).unwrap();
502        check_value(v);
503
504        let data_type = ConcreteDataType::timestamp_second_datatype();
505        let v = constraint
506            .create_default_vector(&data_type, false, 4)
507            .unwrap();
508        check_vector(v);
509
510        let v = constraint.create_default(&data_type, false).unwrap();
511        check_value(v);
512
513        let data_type = ConcreteDataType::timestamp_microsecond_datatype();
514        let v = constraint
515            .create_default_vector(&data_type, false, 4)
516            .unwrap();
517        check_vector(v);
518
519        let v = constraint.create_default(&data_type, false).unwrap();
520        check_value(v);
521
522        let data_type = ConcreteDataType::timestamp_nanosecond_datatype();
523        let v = constraint
524            .create_default_vector(&data_type, false, 4)
525            .unwrap();
526        check_vector(v);
527
528        let v = constraint.create_default(&data_type, false).unwrap();
529        check_value(v);
530
531        // Int64 type.
532        let data_type = ConcreteDataType::int64_datatype();
533        let v = constraint.create_default_vector(&data_type, false, 4);
534        assert!(v.is_err());
535
536        let constraint = ColumnDefaultConstraint::Function("no".to_string());
537        let data_type = ConcreteDataType::timestamp_millisecond_datatype();
538        assert!(
539            constraint
540                .create_default_vector(&data_type, false, 4)
541                .is_err()
542        );
543        assert!(constraint.create_default(&data_type, false).is_err());
544    }
545
546    #[test]
547    fn test_create_by_func_and_invalid_type() {
548        let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string());
549        let data_type = ConcreteDataType::boolean_datatype();
550        let err = constraint
551            .create_default_vector(&data_type, false, 4)
552            .unwrap_err();
553        assert!(matches!(err, Error::DefaultValueType { .. }), "{err:?}");
554    }
555}
datatypes/schema/constraint.rs

datatypes/schema/
constraint.rs
