Skip to main content

api/v1/
column_def.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::collections::HashMap;
16
17use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY};
18use datatypes::schema::{
19    COMMENT_KEY, COMPRESSION_KEY, ColumnDefaultConstraint, ColumnSchema, CompressionOptions,
20    FULLTEXT_KEY, FulltextAnalyzer, FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata,
21    SKIPPING_INDEX_KEY, SkippingIndexOptions, SkippingIndexType,
22};
23use greptime_proto::v1::{
24    Analyzer, FulltextBackend as PbFulltextBackend, SkippingIndexType as PbSkippingIndexType,
25};
26use snafu::ResultExt;
27
28use crate::error::{self, ConvertColumnDefaultConstraintSnafu, Result};
29use crate::helper::ColumnDataTypeWrapper;
30use crate::v1::{ColumnDef, ColumnOptions, SemanticType};
31
32/// Key used to store fulltext options in gRPC column options.
33const FULLTEXT_GRPC_KEY: &str = "fulltext";
34/// Key used to store inverted index options in gRPC column options.
35const INVERTED_INDEX_GRPC_KEY: &str = "inverted_index";
36/// Key used to store skip index options in gRPC column options.
37const SKIPPING_INDEX_GRPC_KEY: &str = "skipping_index";
38/// Key used to store SST compression options in gRPC column options.
39const COMPRESSION_GRPC_KEY: &str = "compression";
40
41const COLUMN_OPTION_MAPPINGS: [(&str, &str); 6] = [
42    (FULLTEXT_GRPC_KEY, FULLTEXT_KEY),
43    (INVERTED_INDEX_GRPC_KEY, INVERTED_INDEX_KEY),
44    (SKIPPING_INDEX_GRPC_KEY, SKIPPING_INDEX_KEY),
45    (COMPRESSION_GRPC_KEY, COMPRESSION_KEY),
46    (EXTENSION_TYPE_NAME_KEY, EXTENSION_TYPE_NAME_KEY),
47    (EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_METADATA_KEY),
48];
49
50/// Tries to construct a `ColumnSchema` from the given  `ColumnDef`.
51pub fn try_as_column_schema(column_def: &ColumnDef) -> Result<ColumnSchema> {
52    let data_type = ColumnDataTypeWrapper::try_new(
53        column_def.data_type,
54        column_def.datatype_extension.clone(),
55    )?;
56
57    let constraint = if column_def.default_constraint.is_empty() {
58        None
59    } else {
60        Some(
61            ColumnDefaultConstraint::try_from(column_def.default_constraint.as_slice()).context(
62                error::ConvertColumnDefaultConstraintSnafu {
63                    column: &column_def.name,
64                },
65            )?,
66        )
67    };
68
69    let mut metadata = HashMap::new();
70    if !column_def.comment.is_empty() {
71        metadata.insert(COMMENT_KEY.to_string(), column_def.comment.clone());
72    }
73    if let Some(options) = column_def.options.as_ref() {
74        if let Some(fulltext) = options.options.get(FULLTEXT_GRPC_KEY) {
75            metadata.insert(FULLTEXT_KEY.to_string(), fulltext.to_owned());
76        }
77        if let Some(inverted_index) = options.options.get(INVERTED_INDEX_GRPC_KEY) {
78            metadata.insert(INVERTED_INDEX_KEY.to_string(), inverted_index.to_owned());
79        }
80        if let Some(skipping_index) = options.options.get(SKIPPING_INDEX_GRPC_KEY) {
81            metadata.insert(SKIPPING_INDEX_KEY.to_string(), skipping_index.to_owned());
82        }
83        if let Some(compression) = options.options.get(COMPRESSION_GRPC_KEY) {
84            metadata.insert(COMPRESSION_KEY.to_string(), compression.to_owned());
85        }
86        if let Some(extension_name) = options.options.get(EXTENSION_TYPE_NAME_KEY) {
87            metadata.insert(EXTENSION_TYPE_NAME_KEY.to_string(), extension_name.clone());
88        }
89        if let Some(extension_metadata) = options.options.get(EXTENSION_TYPE_METADATA_KEY) {
90            metadata.insert(
91                EXTENSION_TYPE_METADATA_KEY.to_string(),
92                extension_metadata.clone(),
93            );
94        }
95    }
96
97    let schema = ColumnSchema::new(&column_def.name, data_type.into(), column_def.is_nullable)
98        .with_metadata(metadata);
99    schema
100        .compression_options()
101        .context(error::InvalidColumnCompressionSnafu {
102            column: &column_def.name,
103        })?;
104    schema
105        .with_time_index(column_def.semantic_type() == SemanticType::Timestamp)
106        .with_default_constraint(constraint)
107        .context(error::InvalidColumnDefaultConstraintSnafu {
108            column: &column_def.name,
109        })
110}
111
112/// Tries to construct a `ColumnDef` from the given `ColumnSchema`.
113///
114/// TODO(weny): Add tests for this function.
115pub fn try_as_column_def(column_schema: &ColumnSchema, is_primary_key: bool) -> Result<ColumnDef> {
116    let column_datatype =
117        ColumnDataTypeWrapper::try_from(column_schema.data_type.clone()).map(|w| w.to_parts())?;
118
119    let semantic_type = if column_schema.is_time_index() {
120        SemanticType::Timestamp
121    } else if is_primary_key {
122        SemanticType::Tag
123    } else {
124        SemanticType::Field
125    } as i32;
126    let comment = column_schema
127        .metadata()
128        .get(COMMENT_KEY)
129        .cloned()
130        .unwrap_or_default();
131
132    let default_constraint = match column_schema.default_constraint() {
133        None => vec![],
134        Some(v) => v
135            .clone()
136            .try_into()
137            .context(ConvertColumnDefaultConstraintSnafu {
138                column: &column_schema.name,
139            })?,
140    };
141    let options = options_from_column_schema(column_schema);
142    Ok(ColumnDef {
143        name: column_schema.name.clone(),
144        data_type: column_datatype.0 as i32,
145        is_nullable: column_schema.is_nullable(),
146        default_constraint,
147        semantic_type,
148        comment,
149        datatype_extension: column_datatype.1,
150        options,
151    })
152}
153
154/// Collect the [ColumnOptions] into the [Metadata] that can be used in, for example, [ColumnSchema].
155pub fn collect_column_options(column_options: Option<&ColumnOptions>) -> Metadata {
156    let Some(ColumnOptions { options }) = column_options else {
157        return Metadata::default();
158    };
159
160    let mut metadata = Metadata::with_capacity(options.len());
161    for (x, y) in COLUMN_OPTION_MAPPINGS {
162        if let Some(v) = options.get(x) {
163            metadata.insert(y.to_string(), v.clone());
164        }
165    }
166    metadata
167}
168
169/// Constructs a `ColumnOptions` from the given `ColumnSchema`.
170pub fn options_from_column_schema(column_schema: &ColumnSchema) -> Option<ColumnOptions> {
171    let mut options = ColumnOptions::default();
172    if let Some(fulltext) = column_schema.metadata().get(FULLTEXT_KEY) {
173        options
174            .options
175            .insert(FULLTEXT_GRPC_KEY.to_string(), fulltext.to_owned());
176    }
177    if let Some(inverted_index) = column_schema.metadata().get(INVERTED_INDEX_KEY) {
178        options
179            .options
180            .insert(INVERTED_INDEX_GRPC_KEY.to_string(), inverted_index.clone());
181    }
182    if let Some(skipping_index) = column_schema.metadata().get(SKIPPING_INDEX_KEY) {
183        options
184            .options
185            .insert(SKIPPING_INDEX_GRPC_KEY.to_string(), skipping_index.clone());
186    }
187    if let Some(compression) = column_schema.metadata().get(COMPRESSION_KEY) {
188        options
189            .options
190            .insert(COMPRESSION_GRPC_KEY.to_string(), compression.clone());
191    }
192    if let Some(extension_name) = column_schema.metadata().get(EXTENSION_TYPE_NAME_KEY) {
193        options
194            .options
195            .insert(EXTENSION_TYPE_NAME_KEY.to_string(), extension_name.clone());
196    }
197    if let Some(extension_metadata) = column_schema.metadata().get(EXTENSION_TYPE_METADATA_KEY) {
198        options.options.insert(
199            EXTENSION_TYPE_METADATA_KEY.to_string(),
200            extension_metadata.clone(),
201        );
202    }
203
204    (!options.options.is_empty()).then_some(options)
205}
206
207/// Checks if the `ColumnOptions` contains fulltext options.
208pub fn contains_fulltext(options: &Option<ColumnOptions>) -> bool {
209    options
210        .as_ref()
211        .is_some_and(|o| o.options.contains_key(FULLTEXT_GRPC_KEY))
212}
213
214/// Checks if the `ColumnOptions` contains skipping index options.
215pub fn contains_skipping(options: &Option<ColumnOptions>) -> bool {
216    options
217        .as_ref()
218        .is_some_and(|o| o.options.contains_key(SKIPPING_INDEX_GRPC_KEY))
219}
220
221/// Tries to construct a `ColumnOptions` from the given `FulltextOptions`.
222pub fn options_from_fulltext(fulltext: &FulltextOptions) -> Result<Option<ColumnOptions>> {
223    let mut options = ColumnOptions::default();
224
225    let v = serde_json::to_string(fulltext).context(error::SerializeJsonSnafu)?;
226    options.options.insert(FULLTEXT_GRPC_KEY.to_string(), v);
227
228    Ok((!options.options.is_empty()).then_some(options))
229}
230
231/// Tries to construct a `ColumnOptions` from the given `SkippingIndexOptions`.
232pub fn options_from_skipping(skipping: &SkippingIndexOptions) -> Result<Option<ColumnOptions>> {
233    let mut options = ColumnOptions::default();
234
235    let v = serde_json::to_string(skipping).context(error::SerializeJsonSnafu)?;
236    options
237        .options
238        .insert(SKIPPING_INDEX_GRPC_KEY.to_string(), v);
239
240    Ok((!options.options.is_empty()).then_some(options))
241}
242
243/// Tries to construct a `ColumnOptions` from the given `CompressionOptions`.
244pub fn options_from_compression(compression: &CompressionOptions) -> Result<ColumnOptions> {
245    let mut options = ColumnOptions::default();
246    let v = serde_json::to_string(compression).context(error::SerializeJsonSnafu)?;
247    options.options.insert(COMPRESSION_GRPC_KEY.to_string(), v);
248    Ok(options)
249}
250
251/// Tries to construct a `ColumnOptions` for inverted index.
252pub fn options_from_inverted() -> ColumnOptions {
253    let mut options = ColumnOptions::default();
254    options
255        .options
256        .insert(INVERTED_INDEX_GRPC_KEY.to_string(), "true".to_string());
257    options
258}
259
260/// Tries to construct a `FulltextAnalyzer` from the given analyzer.
261pub fn as_fulltext_option_analyzer(analyzer: Analyzer) -> FulltextAnalyzer {
262    match analyzer {
263        Analyzer::English => FulltextAnalyzer::English,
264        Analyzer::Chinese => FulltextAnalyzer::Chinese,
265    }
266}
267
268/// Tries to construct a `FulltextBackend` from the given backend.
269pub fn as_fulltext_option_backend(backend: PbFulltextBackend) -> FulltextBackend {
270    match backend {
271        PbFulltextBackend::Bloom => FulltextBackend::Bloom,
272        PbFulltextBackend::Tantivy => FulltextBackend::Tantivy,
273    }
274}
275
276/// Tries to construct a `SkippingIndexType` from the given skipping index type.
277pub fn as_skipping_index_type(skipping_index_type: PbSkippingIndexType) -> SkippingIndexType {
278    match skipping_index_type {
279        PbSkippingIndexType::BloomFilter => SkippingIndexType::BloomFilter,
280    }
281}
282
283/// Converts fulltext options for a JSON2 type hint to the wire representation.
284pub fn json_fulltext_options_to_proto(
285    options: FulltextOptions,
286) -> crate::v1::JsonFulltextIndexOptions {
287    crate::v1::JsonFulltextIndexOptions {
288        analyzer: match options.analyzer {
289            FulltextAnalyzer::English => Analyzer::English.into(),
290            FulltextAnalyzer::Chinese => Analyzer::Chinese.into(),
291        },
292        case_sensitive: options.case_sensitive,
293        backend: match options.backend {
294            FulltextBackend::Bloom => PbFulltextBackend::Bloom.into(),
295            FulltextBackend::Tantivy => PbFulltextBackend::Tantivy.into(),
296        },
297        granularity: options.granularity,
298        false_positive_rate: options.false_positive_rate(),
299    }
300}
301
302/// Validates fulltext options received for a JSON2 type hint.
303/// The presence of these options enables the fulltext index.
304pub fn json_fulltext_options_from_proto(
305    options: crate::v1::JsonFulltextIndexOptions,
306) -> datatypes::error::Result<FulltextOptions> {
307    let analyzer = Analyzer::try_from(options.analyzer).map_err(|_| {
308        datatypes::error::InvalidFulltextOptionSnafu {
309            msg: format!("Unknown fulltext analyzer: {}", options.analyzer),
310        }
311        .build()
312    })?;
313    let backend = PbFulltextBackend::try_from(options.backend).map_err(|_| {
314        datatypes::error::InvalidFulltextOptionSnafu {
315            msg: format!("Unknown fulltext backend: {}", options.backend),
316        }
317        .build()
318    })?;
319    FulltextOptions::new(
320        true,
321        as_fulltext_option_analyzer(analyzer),
322        options.case_sensitive,
323        as_fulltext_option_backend(backend),
324        options.granularity,
325        options.false_positive_rate,
326    )
327}
328
329/// Validates skipping options received for a JSON2 type hint.
330pub fn json_skipping_options_from_proto(
331    options: crate::v1::JsonSkippingIndexOptions,
332) -> datatypes::error::Result<SkippingIndexOptions> {
333    let index_type = PbSkippingIndexType::try_from(options.index_type).map_err(|_| {
334        datatypes::error::InvalidSkippingIndexOptionSnafu {
335            msg: format!("Unknown skipping index type: {}", options.index_type),
336        }
337        .build()
338    })?;
339    SkippingIndexOptions::new(
340        options.granularity,
341        options.false_positive_rate,
342        as_skipping_index_type(index_type),
343    )
344}
345
346#[cfg(test)]
347mod tests {
348    use common_base::readable_size::ReadableSize;
349    use datatypes::data_type::ConcreteDataType;
350    use datatypes::schema::{CompressionOptions, FulltextAnalyzer, FulltextBackend};
351
352    use super::*;
353    use crate::v1::ColumnDataType;
354
355    #[test]
356    fn json_fulltext_options_roundtrip() {
357        for backend in [FulltextBackend::Bloom, FulltextBackend::Tantivy] {
358            let options =
359                FulltextOptions::new(true, FulltextAnalyzer::Chinese, true, backend, 8, 0.05)
360                    .unwrap();
361            let wire = json_fulltext_options_to_proto(options.clone());
362            assert_eq!(options, json_fulltext_options_from_proto(wire).unwrap());
363            let mut invalid = wire;
364            invalid.backend = 99;
365            assert!(json_fulltext_options_from_proto(invalid).is_err());
366            let mut invalid = wire;
367            invalid.analyzer = 99;
368            assert!(json_fulltext_options_from_proto(invalid).is_err());
369            let mut invalid = wire;
370            invalid.granularity = 0;
371            assert!(json_fulltext_options_from_proto(invalid).is_err());
372        }
373    }
374
375    #[test]
376    fn test_try_as_column_schema() {
377        let column_def = ColumnDef {
378            name: "test".to_string(),
379            data_type: ColumnDataType::String as i32,
380            is_nullable: true,
381            default_constraint: ColumnDefaultConstraint::Value("test_default".into())
382                .try_into()
383                .unwrap(),
384            semantic_type: SemanticType::Field as i32,
385            comment: "test_comment".to_string(),
386            datatype_extension: None,
387            options: Some(ColumnOptions {
388                options: HashMap::from([
389                    (
390                        FULLTEXT_GRPC_KEY.to_string(),
391                        "{\"enable\":true}".to_string(),
392                    ),
393                    (INVERTED_INDEX_GRPC_KEY.to_string(), "true".to_string()),
394                ]),
395            }),
396        };
397
398        let schema = try_as_column_schema(&column_def).unwrap();
399        assert_eq!(schema.name, "test");
400        assert_eq!(schema.data_type, ConcreteDataType::string_datatype());
401        assert!(!schema.is_time_index());
402        assert!(schema.is_nullable());
403        assert_eq!(
404            schema.default_constraint().unwrap(),
405            &ColumnDefaultConstraint::Value("test_default".into())
406        );
407        assert_eq!(schema.metadata().get(COMMENT_KEY).unwrap(), "test_comment");
408        assert_eq!(
409            schema.fulltext_options().unwrap().unwrap(),
410            FulltextOptions {
411                enable: true,
412                ..Default::default()
413            }
414        );
415        assert!(schema.is_inverted_indexed());
416    }
417
418    #[test]
419    fn test_compression_options_round_trip() {
420        let options = CompressionOptions {
421            level: Some(3),
422            page_size: Some(ReadableSize::mb(8)),
423            ..Default::default()
424        };
425        let schema = ColumnSchema::new("test", ConcreteDataType::string_datatype(), true)
426            .with_compression_options(options)
427            .unwrap();
428
429        let column_def = try_as_column_def(&schema, false).unwrap();
430        let schema = try_as_column_schema(&column_def).unwrap();
431        assert_eq!(Some(options), schema.compression_options().unwrap());
432
433        let metadata = collect_column_options(column_def.options.as_ref());
434        assert!(metadata.contains_key(COMPRESSION_KEY));
435    }
436
437    #[test]
438    fn test_reject_invalid_compression_metadata() {
439        use common_error::ext::ErrorExt;
440        use common_error::status_code::StatusCode;
441
442        for (data_type, json) in [
443            (ConcreteDataType::string_datatype(), r#"{"level":23}"#),
444            (ConcreteDataType::string_datatype(), r#"{"level":0}"#),
445            (ConcreteDataType::string_datatype(), r#"{"page_size":"0B"}"#),
446            (ConcreteDataType::string_datatype(), r#"{"type":"lz4"}"#),
447            (ConcreteDataType::string_datatype(), "not json"),
448            (
449                ConcreteDataType::list_datatype(std::sync::Arc::new(
450                    ConcreteDataType::string_datatype(),
451                )),
452                r#"{"level":3}"#,
453            ),
454        ] {
455            let schema =
456                ColumnSchema::new("test", data_type, true).with_metadata(HashMap::from([(
457                    COMPRESSION_KEY.to_string(),
458                    json.to_string(),
459                )]));
460            let column_def = try_as_column_def(&schema, false).unwrap();
461            let err = try_as_column_schema(&column_def).unwrap_err();
462            assert_eq!(StatusCode::InvalidArguments, err.status_code(), "{json}");
463        }
464    }
465
466    #[test]
467    fn test_options_from_column_schema() {
468        let schema = ColumnSchema::new("test", ConcreteDataType::string_datatype(), true);
469        let options = options_from_column_schema(&schema);
470        assert!(options.is_none());
471
472        let mut schema = ColumnSchema::new("test", ConcreteDataType::string_datatype(), true)
473            .with_fulltext_options(FulltextOptions::new_unchecked(
474                true,
475                FulltextAnalyzer::English,
476                false,
477                FulltextBackend::Bloom,
478                10240,
479                0.01,
480            ))
481            .unwrap();
482        schema.set_inverted_index(true);
483        let options = options_from_column_schema(&schema).unwrap();
484        assert_eq!(
485            options.options.get(FULLTEXT_GRPC_KEY).unwrap(),
486            "{\"enable\":true,\"analyzer\":\"English\",\"case-sensitive\":false,\"backend\":\"bloom\",\"granularity\":10240,\"false-positive-rate-in-10000\":100}"
487        );
488        assert_eq!(
489            options.options.get(INVERTED_INDEX_GRPC_KEY).unwrap(),
490            "true"
491        );
492    }
493
494    #[test]
495    fn test_options_with_fulltext() {
496        let fulltext = FulltextOptions::new_unchecked(
497            true,
498            FulltextAnalyzer::English,
499            false,
500            FulltextBackend::Bloom,
501            10240,
502            0.01,
503        );
504        let options = options_from_fulltext(&fulltext).unwrap().unwrap();
505        assert_eq!(
506            options.options.get(FULLTEXT_GRPC_KEY).unwrap(),
507            "{\"enable\":true,\"analyzer\":\"English\",\"case-sensitive\":false,\"backend\":\"bloom\",\"granularity\":10240,\"false-positive-rate-in-10000\":100}"
508        );
509    }
510
511    #[test]
512    fn test_contains_fulltext() {
513        let options = ColumnOptions {
514            options: HashMap::from([(
515                FULLTEXT_GRPC_KEY.to_string(),
516                "{\"enable\":true}".to_string(),
517            )]),
518        };
519        assert!(contains_fulltext(&Some(options)));
520
521        let options = ColumnOptions {
522            options: HashMap::new(),
523        };
524        assert!(!contains_fulltext(&Some(options)));
525
526        assert!(!contains_fulltext(&None));
527    }
528}