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

datatypes/
schema.rs
