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

datatypes/
schema.rs
