datatypes/
schema.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
15mod column_schema;
16pub mod constraint;
17pub mod ext;
18
19use std::collections::HashMap;
20use std::sync::Arc;
21use std::{fmt, mem};
22
23use arrow::datatypes::{Field, Schema as ArrowSchema};
24use datafusion_common::DFSchemaRef;
25use serde::{Deserialize, Deserializer, Serialize};
26use snafu::{ResultExt, ensure};
27
28use crate::error::{self, DuplicateColumnSnafu, Error, ProjectArrowSchemaSnafu, Result};
29use crate::prelude::ConcreteDataType;
30pub use crate::schema::column_schema::{
31    COLUMN_FULLTEXT_CHANGE_OPT_KEY_ENABLE, COLUMN_FULLTEXT_OPT_KEY_ANALYZER,
32    COLUMN_FULLTEXT_OPT_KEY_BACKEND, COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE,
33    COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_FULLTEXT_OPT_KEY_GRANULARITY,
34    COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE, COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY,
35    COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY, ColumnExtType, ColumnSchema, FULLTEXT_KEY,
36    FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata,
37    SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType, TIME_INDEX_KEY,
38};
39pub use crate::schema::constraint::ColumnDefaultConstraint;
40
41/// Key used to store version number of the schema in metadata.
42pub const VERSION_KEY: &str = "greptime:version";
43/// Key used to store actual column type in field metadata.
44pub const TYPE_KEY: &str = "greptime:type";
45
46/// A common schema, should be immutable.
47#[derive(Clone, PartialEq, Eq, Serialize)]
48pub struct Schema {
49    column_schemas: Vec<ColumnSchema>,
50    #[serde(skip_serializing)]
51    name_to_index: HashMap<String, usize>,
52    #[serde(skip_serializing)]
53    arrow_schema: Arc<ArrowSchema>,
54    /// Index of the timestamp key column.
55    ///
56    /// Timestamp key column is the column holds the timestamp and forms part of
57    /// the primary key. None means there is no timestamp key column.
58    #[serde(skip_serializing)]
59    timestamp_index: Option<usize>,
60    /// Version of the schema.
61    ///
62    /// Initial value is zero. The version should bump after altering schema.
63    version: u32,
64}
65
66impl<'de> Deserialize<'de> for Schema {
67    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
68    where
69        D: Deserializer<'de>,
70    {
71        use serde::de::Error;
72
73        #[derive(Deserialize)]
74        struct RawSchema {
75            column_schemas: Vec<ColumnSchema>,
76            version: u32,
77        }
78        let raw = RawSchema::deserialize(deserializer)?;
79
80        SchemaBuilder::try_from(raw.column_schemas)
81            .and_then(|x| x.version(raw.version).build())
82            .map_err(D::Error::custom)
83    }
84}
85
86impl fmt::Debug for Schema {
87    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
88        f.debug_struct("Schema")
89            .field("column_schemas", &self.column_schemas)
90            .field("name_to_index", &self.name_to_index)
91            .field("timestamp_index", &self.timestamp_index)
92            .field("version", &self.version)
93            .finish()
94    }
95}
96
97impl Schema {
98    /// Initial version of the schema.
99    pub const INITIAL_VERSION: u32 = 0;
100
101    /// Create a schema from a vector of [ColumnSchema].
102    ///
103    /// # Panics
104    /// Panics when ColumnSchema's `default_constraint` can't be serialized into json.
105    pub fn new(column_schemas: Vec<ColumnSchema>) -> Schema {
106        // Builder won't fail in this case
107        Self::new_with_version(column_schemas, Self::INITIAL_VERSION)
108    }
109
110    pub fn new_with_version(column_schemas: Vec<ColumnSchema>, version: u32) -> Schema {
111        SchemaBuilder::try_from(column_schemas)
112            .unwrap()
113            .version(version)
114            .build()
115            .unwrap()
116    }
117
118    /// Try to Create a schema from a vector of [ColumnSchema].
119    pub fn try_new(column_schemas: Vec<ColumnSchema>) -> Result<Schema> {
120        SchemaBuilder::try_from(column_schemas)?.build()
121    }
122
123    pub fn arrow_schema(&self) -> &Arc<ArrowSchema> {
124        &self.arrow_schema
125    }
126
127    pub fn column_schemas(&self) -> &[ColumnSchema] {
128        &self.column_schemas
129    }
130
131    pub fn column_schema_by_name(&self, name: &str) -> Option<&ColumnSchema> {
132        self.name_to_index
133            .get(name)
134            .map(|index| &self.column_schemas[*index])
135    }
136
137    /// Retrieve the column's name by index
138    /// # Panics
139    /// This method **may** panic if the index is out of range of column schemas.
140    pub fn column_name_by_index(&self, idx: usize) -> &str {
141        &self.column_schemas[idx].name
142    }
143
144    pub fn column_index_by_name(&self, name: &str) -> Option<usize> {
145        self.name_to_index.get(name).copied()
146    }
147
148    pub fn contains_column(&self, name: &str) -> bool {
149        self.name_to_index.contains_key(name)
150    }
151
152    pub fn num_columns(&self) -> usize {
153        self.column_schemas.len()
154    }
155
156    pub fn is_empty(&self) -> bool {
157        self.column_schemas.is_empty()
158    }
159
160    /// Returns index of the timestamp key column.
161    pub fn timestamp_index(&self) -> Option<usize> {
162        self.timestamp_index
163    }
164
165    pub fn timestamp_column(&self) -> Option<&ColumnSchema> {
166        self.timestamp_index.map(|idx| &self.column_schemas[idx])
167    }
168
169    pub fn version(&self) -> u32 {
170        self.version
171    }
172
173    pub fn metadata(&self) -> &HashMap<String, String> {
174        &self.arrow_schema.metadata
175    }
176
177    /// Returns the estimated memory footprint of this schema.
178    pub fn estimated_size(&self) -> usize {
179        mem::size_of_val(self)
180            + mem::size_of::<ColumnSchema>() * self.column_schemas.capacity()
181            + self
182                .column_schemas
183                .iter()
184                .map(|column_schema| {
185                    column_schema.estimated_size() - mem::size_of::<ColumnSchema>()
186                })
187                .sum::<usize>()
188            + mem::size_of::<(String, usize)>() * self.name_to_index.capacity()
189            + self
190                .name_to_index
191                .keys()
192                .map(|name| name.capacity())
193                .sum::<usize>()
194            + arrow_schema_size(self.arrow_schema.as_ref())
195    }
196
197    /// Generate a new projected schema
198    ///
199    /// # Panic
200    ///
201    /// If the index out ouf bound
202    pub fn try_project(&self, indices: &[usize]) -> Result<Self> {
203        let mut column_schemas = Vec::with_capacity(indices.len());
204        let mut timestamp_index = None;
205        for index in indices {
206            if let Some(ts_index) = self.timestamp_index
207                && ts_index == *index
208            {
209                timestamp_index = Some(column_schemas.len());
210            }
211            column_schemas.push(self.column_schemas[*index].clone());
212        }
213        let arrow_schema = self
214            .arrow_schema
215            .project(indices)
216            .context(ProjectArrowSchemaSnafu)?;
217        let name_to_index = column_schemas
218            .iter()
219            .enumerate()
220            .map(|(pos, column_schema)| (column_schema.name.clone(), pos))
221            .collect();
222
223        Ok(Self {
224            column_schemas,
225            name_to_index,
226            arrow_schema: Arc::new(arrow_schema),
227            timestamp_index,
228            version: self.version,
229        })
230    }
231}
232
233fn arrow_schema_size(schema: &ArrowSchema) -> usize {
234    mem::size_of_val(schema)
235        + schema.fields.size()
236        + mem::size_of::<(String, String)>() * schema.metadata.capacity()
237        + schema
238            .metadata
239            .iter()
240            .map(|(key, value)| key.capacity() + value.capacity())
241            .sum::<usize>()
242}
243
244#[derive(Default)]
245pub struct SchemaBuilder {
246    column_schemas: Vec<ColumnSchema>,
247    name_to_index: HashMap<String, usize>,
248    fields: Vec<Field>,
249    timestamp_index: Option<usize>,
250    version: u32,
251    metadata: HashMap<String, String>,
252}
253
254impl TryFrom<Vec<ColumnSchema>> for SchemaBuilder {
255    type Error = Error;
256
257    fn try_from(column_schemas: Vec<ColumnSchema>) -> Result<SchemaBuilder> {
258        SchemaBuilder::try_from_columns(column_schemas)
259    }
260}
261
262impl SchemaBuilder {
263    pub fn try_from_columns(column_schemas: Vec<ColumnSchema>) -> Result<Self> {
264        let FieldsAndIndices {
265            fields,
266            name_to_index,
267            timestamp_index,
268        } = collect_fields(&column_schemas)?;
269
270        Ok(Self {
271            column_schemas,
272            name_to_index,
273            fields,
274            timestamp_index,
275            ..Default::default()
276        })
277    }
278
279    pub fn version(mut self, version: u32) -> Self {
280        self.version = version;
281        self
282    }
283
284    /// Add key value pair to metadata.
285    ///
286    /// Old metadata with same key would be overwritten.
287    pub fn add_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
288        let _ = self.metadata.insert(key.into(), value.into());
289        self
290    }
291
292    pub fn build(mut self) -> Result<Schema> {
293        if let Some(timestamp_index) = self.timestamp_index {
294            validate_timestamp_index(&self.column_schemas, timestamp_index)?;
295        }
296
297        self.metadata
298            .insert(VERSION_KEY.to_string(), self.version.to_string());
299
300        let arrow_schema = ArrowSchema::new(self.fields).with_metadata(self.metadata);
301
302        Ok(Schema {
303            column_schemas: self.column_schemas,
304            name_to_index: self.name_to_index,
305            arrow_schema: Arc::new(arrow_schema),
306            timestamp_index: self.timestamp_index,
307            version: self.version,
308        })
309    }
310}
311
312struct FieldsAndIndices {
313    fields: Vec<Field>,
314    name_to_index: HashMap<String, usize>,
315    timestamp_index: Option<usize>,
316}
317
318fn collect_fields(column_schemas: &[ColumnSchema]) -> Result<FieldsAndIndices> {
319    let mut fields = Vec::with_capacity(column_schemas.len());
320    let mut name_to_index = HashMap::with_capacity(column_schemas.len());
321    let mut timestamp_index = None;
322    for (index, column_schema) in column_schemas.iter().enumerate() {
323        if column_schema.is_time_index() && timestamp_index.is_none() {
324            timestamp_index = Some(index);
325        }
326        let mut field = Field::try_from(column_schema)?;
327
328        // Column with type Json or Vector performs the same as binary column in Arrow, so we need to mark it
329        let extype = match column_schema.data_type {
330            ConcreteDataType::Json(_) => Some(ColumnExtType::Json),
331            ConcreteDataType::Vector(d) => Some(ColumnExtType::Vector(d.dim)),
332            _ => None,
333        };
334        if let Some(extype) = extype {
335            field
336                .metadata_mut()
337                .insert(TYPE_KEY.to_string(), extype.to_string());
338        }
339        fields.push(field);
340        ensure!(
341            name_to_index
342                .insert(column_schema.name.clone(), index)
343                .is_none(),
344            DuplicateColumnSnafu {
345                column: &column_schema.name,
346            }
347        );
348    }
349
350    Ok(FieldsAndIndices {
351        fields,
352        name_to_index,
353        timestamp_index,
354    })
355}
356
357fn validate_timestamp_index(column_schemas: &[ColumnSchema], timestamp_index: usize) -> Result<()> {
358    ensure!(
359        timestamp_index < column_schemas.len(),
360        error::InvalidTimestampIndexSnafu {
361            index: timestamp_index,
362        }
363    );
364
365    let column_schema = &column_schemas[timestamp_index];
366    ensure!(
367        column_schema.data_type.is_timestamp(),
368        error::InvalidTimestampIndexSnafu {
369            index: timestamp_index,
370        }
371    );
372    ensure!(
373        column_schema.is_time_index(),
374        error::InvalidTimestampIndexSnafu {
375            index: timestamp_index,
376        }
377    );
378
379    Ok(())
380}
381
382pub type SchemaRef = Arc<Schema>;
383
384impl TryFrom<Arc<ArrowSchema>> for Schema {
385    type Error = Error;
386
387    fn try_from(arrow_schema: Arc<ArrowSchema>) -> Result<Schema> {
388        let mut column_schemas = Vec::with_capacity(arrow_schema.fields.len());
389        let mut name_to_index = HashMap::with_capacity(arrow_schema.fields.len());
390        for field in &arrow_schema.fields {
391            let column_schema = ColumnSchema::try_from(field.as_ref())?;
392            let _ = name_to_index.insert(field.name().clone(), column_schemas.len());
393            column_schemas.push(column_schema);
394        }
395
396        let mut timestamp_index = None;
397        for (index, column_schema) in column_schemas.iter().enumerate() {
398            if column_schema.is_time_index() {
399                validate_timestamp_index(&column_schemas, index)?;
400                timestamp_index = Some(index);
401                break;
402            }
403        }
404
405        let version = try_parse_version(&arrow_schema.metadata, VERSION_KEY)?;
406
407        Ok(Self {
408            column_schemas,
409            name_to_index,
410            arrow_schema,
411            timestamp_index,
412            version,
413        })
414    }
415}
416
417impl TryFrom<ArrowSchema> for Schema {
418    type Error = Error;
419
420    fn try_from(arrow_schema: ArrowSchema) -> Result<Schema> {
421        let arrow_schema = Arc::new(arrow_schema);
422
423        Schema::try_from(arrow_schema)
424    }
425}
426
427impl TryFrom<DFSchemaRef> for Schema {
428    type Error = Error;
429
430    fn try_from(value: DFSchemaRef) -> Result<Self> {
431        value.inner().clone().try_into()
432    }
433}
434
435fn try_parse_version(metadata: &HashMap<String, String>, key: &str) -> Result<u32> {
436    if let Some(value) = metadata.get(key) {
437        let version = value
438            .parse()
439            .context(error::ParseSchemaVersionSnafu { value })?;
440
441        Ok(version)
442    } else {
443        Ok(Schema::INITIAL_VERSION)
444    }
445}
446
447#[cfg(test)]
448mod tests {
449    use super::*;
450    use crate::data_type::ConcreteDataType;
451
452    #[test]
453    fn test_build_empty_schema() {
454        let schema = SchemaBuilder::default().build().unwrap();
455        assert_eq!(0, schema.num_columns());
456        assert!(schema.is_empty());
457    }
458
459    #[test]
460    fn test_schema_no_timestamp() {
461        let column_schemas = vec![
462            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), false),
463            ColumnSchema::new("col2", ConcreteDataType::float64_datatype(), true),
464        ];
465        let schema = Schema::new(column_schemas.clone());
466
467        assert_eq!(2, schema.num_columns());
468        assert!(!schema.is_empty());
469        assert!(schema.timestamp_index().is_none());
470        assert!(schema.timestamp_column().is_none());
471        assert_eq!(Schema::INITIAL_VERSION, schema.version());
472
473        for column_schema in &column_schemas {
474            let found = schema.column_schema_by_name(&column_schema.name).unwrap();
475            assert_eq!(column_schema, found);
476        }
477        assert!(schema.column_schema_by_name("col3").is_none());
478
479        let new_schema = Schema::try_from(schema.arrow_schema().clone()).unwrap();
480
481        assert_eq!(schema, new_schema);
482        assert_eq!(column_schemas, schema.column_schemas());
483    }
484
485    #[test]
486    fn test_schema_duplicate_column() {
487        let column_schemas = vec![
488            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), false),
489            ColumnSchema::new("col1", ConcreteDataType::float64_datatype(), true),
490        ];
491        let err = Schema::try_new(column_schemas).unwrap_err();
492
493        assert!(
494            matches!(err, Error::DuplicateColumn { .. }),
495            "expect DuplicateColumn, found {}",
496            err
497        );
498    }
499
500    #[test]
501    fn test_metadata() {
502        let column_schemas = vec![ColumnSchema::new(
503            "col1",
504            ConcreteDataType::int32_datatype(),
505            false,
506        )];
507        let schema = SchemaBuilder::try_from(column_schemas)
508            .unwrap()
509            .add_metadata("k1", "v1")
510            .build()
511            .unwrap();
512
513        assert_eq!("v1", schema.metadata().get("k1").unwrap());
514    }
515
516    #[test]
517    fn test_schema_with_timestamp() {
518        let column_schemas = vec![
519            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true),
520            ColumnSchema::new(
521                "ts",
522                ConcreteDataType::timestamp_millisecond_datatype(),
523                false,
524            )
525            .with_time_index(true),
526        ];
527        let schema = SchemaBuilder::try_from(column_schemas.clone())
528            .unwrap()
529            .version(123)
530            .build()
531            .unwrap();
532
533        assert_eq!(1, schema.timestamp_index().unwrap());
534        assert_eq!(&column_schemas[1], schema.timestamp_column().unwrap());
535        assert_eq!(123, schema.version());
536
537        let new_schema = Schema::try_from(schema.arrow_schema().clone()).unwrap();
538        assert_eq!(1, schema.timestamp_index().unwrap());
539        assert_eq!(schema, new_schema);
540    }
541
542    #[test]
543    fn test_schema_wrong_timestamp() {
544        let column_schemas = vec![
545            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true)
546                .with_time_index(true),
547            ColumnSchema::new("col2", ConcreteDataType::float64_datatype(), false),
548        ];
549        assert!(
550            SchemaBuilder::try_from(column_schemas)
551                .unwrap()
552                .build()
553                .is_err()
554        );
555
556        let column_schemas = vec![
557            ColumnSchema::new("col1", ConcreteDataType::int32_datatype(), true),
558            ColumnSchema::new("col2", ConcreteDataType::float64_datatype(), false)
559                .with_time_index(true),
560        ];
561
562        assert!(
563            SchemaBuilder::try_from(column_schemas)
564                .unwrap()
565                .build()
566                .is_err()
567        );
568    }
569}
datatypes/schema.rs

datatypes/
schema.rs
