datatypes/extension/
histogram.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
15//! Arrow extension type for native-histogram struct columns.
16
17use arrow_schema::extension::ExtensionType;
18use arrow_schema::{ArrowError, DataType, FieldRef};
19
20/// Arrow extension type identifying a native-histogram struct column.
21///
22/// Applied to the struct field at parquet-write time so that readers can
23/// identify native-histogram columns by extension (`greptime.histogram`) rather
24/// than relying on the field name.
25#[derive(Debug, Clone, Default)]
26pub struct HistogramExtensionType;
27
28impl ExtensionType for HistogramExtensionType {
29    const NAME: &'static str = "greptime.histogram";
30    type Metadata = ();
31
32    fn metadata(&self) -> &Self::Metadata {
33        &()
34    }
35
36    fn serialize_metadata(&self) -> Option<String> {
37        None
38    }
39
40    fn deserialize_metadata(_metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
41        Ok(())
42    }
43
44    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
45        match data_type {
46            DataType::Struct(_) => Ok(()),
47            dt => Err(ArrowError::SchemaError(format!(
48                "Unexpected data type {dt}"
49            ))),
50        }
51    }
52
53    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self, ArrowError> {
54        let ext = Self;
55        ext.supports_data_type(data_type)?;
56        Ok(ext)
57    }
58}
59
60/// Check if this field is a native-histogram extension type.
61pub fn is_histogram_extension_type(field: &FieldRef) -> bool {
62    field.extension_type_name() == Some(HistogramExtensionType::NAME)
63}
64
65#[cfg(test)]
66mod tests {
67    use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY;
68    use arrow_schema::{DataType, Field, Fields};
69
70    use super::*;
71
72    #[test]
73    fn test_extension_name_and_detection() {
74        assert_eq!(HistogramExtensionType::NAME, "greptime.histogram");
75
76        // A plain struct field is not a histogram extension type.
77        let empty: Fields = Vec::<Field>::new().into();
78        let plain = std::sync::Arc::new(Field::new("s", DataType::Struct(empty), true));
79        assert!(!is_histogram_extension_type(&plain));
80
81        // Tagging the field with the extension makes it detectable.
82        let mut tagged = (*plain).clone();
83        tagged.metadata_mut().insert(
84            EXTENSION_TYPE_NAME_KEY.to_string(),
85            HistogramExtensionType::NAME.to_string(),
86        );
87        let tagged = std::sync::Arc::new(tagged);
88        assert!(is_histogram_extension_type(&tagged));
89    }
90
91    #[test]
92    fn test_supports_struct_only() {
93        let ext = HistogramExtensionType;
94        let empty: Fields = Vec::<Field>::new().into();
95        assert!(ext.supports_data_type(&DataType::Struct(empty)).is_ok());
96        assert!(ext.supports_data_type(&DataType::Int32).is_err());
97    }
98}
datatypes/extension/histogram.rs

datatypes/extension/
histogram.rs
