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, ColumnDefaultConstraint, ColumnSchema, FULLTEXT_KEY, FulltextAnalyzer,
20    FulltextBackend, FulltextOptions, INVERTED_INDEX_KEY, Metadata, SKIPPING_INDEX_KEY,
21    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
39const COLUMN_OPTION_MAPPINGS: [(&str, &str); 5] = [
40    (FULLTEXT_GRPC_KEY, FULLTEXT_KEY),
41    (INVERTED_INDEX_GRPC_KEY, INVERTED_INDEX_KEY),
42    (SKIPPING_INDEX_GRPC_KEY, SKIPPING_INDEX_KEY),
43    (EXTENSION_TYPE_NAME_KEY, EXTENSION_TYPE_NAME_KEY),
44    (EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_METADATA_KEY),
45];
46
47/// Tries to construct a `ColumnSchema` from the given  `ColumnDef`.
48pub fn try_as_column_schema(column_def: &ColumnDef) -> Result<ColumnSchema> {
49    let data_type = ColumnDataTypeWrapper::try_new(
50        column_def.data_type,
51        column_def.datatype_extension.clone(),
52    )?;
53
54    let constraint = if column_def.default_constraint.is_empty() {
55        None
56    } else {
57        Some(
58            ColumnDefaultConstraint::try_from(column_def.default_constraint.as_slice()).context(
59                error::ConvertColumnDefaultConstraintSnafu {
60                    column: &column_def.name,
61                },
62            )?,
63        )
64    };
65
66    let mut metadata = HashMap::new();
67    if !column_def.comment.is_empty() {
68        metadata.insert(COMMENT_KEY.to_string(), column_def.comment.clone());
69    }
70    if let Some(options) = column_def.options.as_ref() {
71        if let Some(fulltext) = options.options.get(FULLTEXT_GRPC_KEY) {
72            metadata.insert(FULLTEXT_KEY.to_string(), fulltext.to_owned());
73        }
74        if let Some(inverted_index) = options.options.get(INVERTED_INDEX_GRPC_KEY) {
75            metadata.insert(INVERTED_INDEX_KEY.to_string(), inverted_index.to_owned());
76        }
77        if let Some(skipping_index) = options.options.get(SKIPPING_INDEX_GRPC_KEY) {
78            metadata.insert(SKIPPING_INDEX_KEY.to_string(), skipping_index.to_owned());
79        }
80        if let Some(extension_name) = options.options.get(EXTENSION_TYPE_NAME_KEY) {
81            metadata.insert(EXTENSION_TYPE_NAME_KEY.to_string(), extension_name.clone());
82        }
83        if let Some(extension_metadata) = options.options.get(EXTENSION_TYPE_METADATA_KEY) {
84            metadata.insert(
85                EXTENSION_TYPE_METADATA_KEY.to_string(),
86                extension_metadata.clone(),
87            );
88        }
89    }
90
91    ColumnSchema::new(&column_def.name, data_type.into(), column_def.is_nullable)
92        .with_metadata(metadata)
93        .with_time_index(column_def.semantic_type() == SemanticType::Timestamp)
94        .with_default_constraint(constraint)
95        .context(error::InvalidColumnDefaultConstraintSnafu {
96            column: &column_def.name,
97        })
98}
99
100/// Tries to construct a `ColumnDef` from the given `ColumnSchema`.
101///
102/// TODO(weny): Add tests for this function.
103pub fn try_as_column_def(column_schema: &ColumnSchema, is_primary_key: bool) -> Result<ColumnDef> {
104    let column_datatype =
105        ColumnDataTypeWrapper::try_from(column_schema.data_type.clone()).map(|w| w.to_parts())?;
106
107    let semantic_type = if column_schema.is_time_index() {
108        SemanticType::Timestamp
109    } else if is_primary_key {
110        SemanticType::Tag
111    } else {
112        SemanticType::Field
113    } as i32;
114    let comment = column_schema
115        .metadata()
116        .get(COMMENT_KEY)
117        .cloned()
118        .unwrap_or_default();
119
120    let default_constraint = match column_schema.default_constraint() {
121        None => vec![],
122        Some(v) => v
123            .clone()
124            .try_into()
125            .context(ConvertColumnDefaultConstraintSnafu {
126                column: &column_schema.name,
127            })?,
128    };
129    let options = options_from_column_schema(column_schema);
130    Ok(ColumnDef {
131        name: column_schema.name.clone(),
132        data_type: column_datatype.0 as i32,
133        is_nullable: column_schema.is_nullable(),
134        default_constraint,
135        semantic_type,
136        comment,
137        datatype_extension: column_datatype.1,
138        options,
139    })
140}
141
142/// Collect the [ColumnOptions] into the [Metadata] that can be used in, for example, [ColumnSchema].
143pub fn collect_column_options(column_options: Option<&ColumnOptions>) -> Metadata {
144    let Some(ColumnOptions { options }) = column_options else {
145        return Metadata::default();
146    };
147
148    let mut metadata = Metadata::with_capacity(options.len());
149    for (x, y) in COLUMN_OPTION_MAPPINGS {
150        if let Some(v) = options.get(x) {
151            metadata.insert(y.to_string(), v.clone());
152        }
153    }
154    metadata
155}
156
157/// Constructs a `ColumnOptions` from the given `ColumnSchema`.
158pub fn options_from_column_schema(column_schema: &ColumnSchema) -> Option<ColumnOptions> {
159    let mut options = ColumnOptions::default();
160    if let Some(fulltext) = column_schema.metadata().get(FULLTEXT_KEY) {
161        options
162            .options
163            .insert(FULLTEXT_GRPC_KEY.to_string(), fulltext.to_owned());
164    }
165    if let Some(inverted_index) = column_schema.metadata().get(INVERTED_INDEX_KEY) {
166        options
167            .options
168            .insert(INVERTED_INDEX_GRPC_KEY.to_string(), inverted_index.clone());
169    }
170    if let Some(skipping_index) = column_schema.metadata().get(SKIPPING_INDEX_KEY) {
171        options
172            .options
173            .insert(SKIPPING_INDEX_GRPC_KEY.to_string(), skipping_index.clone());
174    }
175    if let Some(extension_name) = column_schema.metadata().get(EXTENSION_TYPE_NAME_KEY) {
176        options
177            .options
178            .insert(EXTENSION_TYPE_NAME_KEY.to_string(), extension_name.clone());
179    }
180    if let Some(extension_metadata) = column_schema.metadata().get(EXTENSION_TYPE_METADATA_KEY) {
181        options.options.insert(
182            EXTENSION_TYPE_METADATA_KEY.to_string(),
183            extension_metadata.clone(),
184        );
185    }
186
187    (!options.options.is_empty()).then_some(options)
188}
189
190/// Checks if the `ColumnOptions` contains fulltext options.
191pub fn contains_fulltext(options: &Option<ColumnOptions>) -> bool {
192    options
193        .as_ref()
194        .is_some_and(|o| o.options.contains_key(FULLTEXT_GRPC_KEY))
195}
196
197/// Checks if the `ColumnOptions` contains skipping index options.
198pub fn contains_skipping(options: &Option<ColumnOptions>) -> bool {
199    options
200        .as_ref()
201        .is_some_and(|o| o.options.contains_key(SKIPPING_INDEX_GRPC_KEY))
202}
203
204/// Tries to construct a `ColumnOptions` from the given `FulltextOptions`.
205pub fn options_from_fulltext(fulltext: &FulltextOptions) -> Result<Option<ColumnOptions>> {
206    let mut options = ColumnOptions::default();
207
208    let v = serde_json::to_string(fulltext).context(error::SerializeJsonSnafu)?;
209    options.options.insert(FULLTEXT_GRPC_KEY.to_string(), v);
210
211    Ok((!options.options.is_empty()).then_some(options))
212}
213
214/// Tries to construct a `ColumnOptions` from the given `SkippingIndexOptions`.
215pub fn options_from_skipping(skipping: &SkippingIndexOptions) -> Result<Option<ColumnOptions>> {
216    let mut options = ColumnOptions::default();
217
218    let v = serde_json::to_string(skipping).context(error::SerializeJsonSnafu)?;
219    options
220        .options
221        .insert(SKIPPING_INDEX_GRPC_KEY.to_string(), v);
222
223    Ok((!options.options.is_empty()).then_some(options))
224}
225
226/// Tries to construct a `ColumnOptions` for inverted index.
227pub fn options_from_inverted() -> ColumnOptions {
228    let mut options = ColumnOptions::default();
229    options
230        .options
231        .insert(INVERTED_INDEX_GRPC_KEY.to_string(), "true".to_string());
232    options
233}
234
235/// Tries to construct a `FulltextAnalyzer` from the given analyzer.
236pub fn as_fulltext_option_analyzer(analyzer: Analyzer) -> FulltextAnalyzer {
237    match analyzer {
238        Analyzer::English => FulltextAnalyzer::English,
239        Analyzer::Chinese => FulltextAnalyzer::Chinese,
240    }
241}
242
243/// Tries to construct a `FulltextBackend` from the given backend.
244pub fn as_fulltext_option_backend(backend: PbFulltextBackend) -> FulltextBackend {
245    match backend {
246        PbFulltextBackend::Bloom => FulltextBackend::Bloom,
247        PbFulltextBackend::Tantivy => FulltextBackend::Tantivy,
248    }
249}
250
251/// Tries to construct a `SkippingIndexType` from the given skipping index type.
252pub fn as_skipping_index_type(skipping_index_type: PbSkippingIndexType) -> SkippingIndexType {
253    match skipping_index_type {
254        PbSkippingIndexType::BloomFilter => SkippingIndexType::BloomFilter,
255    }
256}
257
258#[cfg(test)]
259mod tests {
260
261    use datatypes::data_type::ConcreteDataType;
262    use datatypes::schema::{FulltextAnalyzer, FulltextBackend};
263
264    use super::*;
265    use crate::v1::ColumnDataType;
266
267    #[test]
268    fn test_try_as_column_schema() {
269        let column_def = ColumnDef {
270            name: "test".to_string(),
271            data_type: ColumnDataType::String as i32,
272            is_nullable: true,
273            default_constraint: ColumnDefaultConstraint::Value("test_default".into())
274                .try_into()
275                .unwrap(),
276            semantic_type: SemanticType::Field as i32,
277            comment: "test_comment".to_string(),
278            datatype_extension: None,
279            options: Some(ColumnOptions {
280                options: HashMap::from([
281                    (
282                        FULLTEXT_GRPC_KEY.to_string(),
283                        "{\"enable\":true}".to_string(),
284                    ),
285                    (INVERTED_INDEX_GRPC_KEY.to_string(), "true".to_string()),
286                ]),
287            }),
288        };
289
290        let schema = try_as_column_schema(&column_def).unwrap();
291        assert_eq!(schema.name, "test");
292        assert_eq!(schema.data_type, ConcreteDataType::string_datatype());
293        assert!(!schema.is_time_index());
294        assert!(schema.is_nullable());
295        assert_eq!(
296            schema.default_constraint().unwrap(),
297            &ColumnDefaultConstraint::Value("test_default".into())
298        );
299        assert_eq!(schema.metadata().get(COMMENT_KEY).unwrap(), "test_comment");
300        assert_eq!(
301            schema.fulltext_options().unwrap().unwrap(),
302            FulltextOptions {
303                enable: true,
304                ..Default::default()
305            }
306        );
307        assert!(schema.is_inverted_indexed());
308    }
309
310    #[test]
311    fn test_options_from_column_schema() {
312        let schema = ColumnSchema::new("test", ConcreteDataType::string_datatype(), true);
313        let options = options_from_column_schema(&schema);
314        assert!(options.is_none());
315
316        let mut schema = ColumnSchema::new("test", ConcreteDataType::string_datatype(), true)
317            .with_fulltext_options(FulltextOptions::new_unchecked(
318                true,
319                FulltextAnalyzer::English,
320                false,
321                FulltextBackend::Bloom,
322                10240,
323                0.01,
324            ))
325            .unwrap();
326        schema.set_inverted_index(true);
327        let options = options_from_column_schema(&schema).unwrap();
328        assert_eq!(
329            options.options.get(FULLTEXT_GRPC_KEY).unwrap(),
330            "{\"enable\":true,\"analyzer\":\"English\",\"case-sensitive\":false,\"backend\":\"bloom\",\"granularity\":10240,\"false-positive-rate-in-10000\":100}"
331        );
332        assert_eq!(
333            options.options.get(INVERTED_INDEX_GRPC_KEY).unwrap(),
334            "true"
335        );
336    }
337
338    #[test]
339    fn test_options_with_fulltext() {
340        let fulltext = FulltextOptions::new_unchecked(
341            true,
342            FulltextAnalyzer::English,
343            false,
344            FulltextBackend::Bloom,
345            10240,
346            0.01,
347        );
348        let options = options_from_fulltext(&fulltext).unwrap().unwrap();
349        assert_eq!(
350            options.options.get(FULLTEXT_GRPC_KEY).unwrap(),
351            "{\"enable\":true,\"analyzer\":\"English\",\"case-sensitive\":false,\"backend\":\"bloom\",\"granularity\":10240,\"false-positive-rate-in-10000\":100}"
352        );
353    }
354
355    #[test]
356    fn test_contains_fulltext() {
357        let options = ColumnOptions {
358            options: HashMap::from([(
359                FULLTEXT_GRPC_KEY.to_string(),
360                "{\"enable\":true}".to_string(),
361            )]),
362        };
363        assert!(contains_fulltext(&Some(options)));
364
365        let options = ColumnOptions {
366            options: HashMap::new(),
367        };
368        assert!(!contains_fulltext(&Some(options)));
369
370        assert!(!contains_fulltext(&None));
371    }
372}