parquet/arrow/schema/
virtual_type.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! RowNumber
19//!
20
21use arrow_schema::{ArrowError, DataType, Field, extension::ExtensionType};
22
23/// Prefix for virtual column extension type names.
24macro_rules! VIRTUAL_PREFIX {
25    () => {
26        "parquet.virtual."
27    };
28}
29
30/// The extension type for row group indices
31///
32/// Extension name: `parquet.virtual.row_group_index`
33///
34/// This virtual column has storage type `Int64` and uses empty string metadata
35#[derive(Debug, Default, Clone, Copy, PartialEq)]
36pub struct RowGroupIndex;
37
38impl ExtensionType for RowGroupIndex {
39    const NAME: &'static str = concat!(VIRTUAL_PREFIX!(), "row_group_index");
40    type Metadata = &'static str;
41
42    fn metadata(&self) -> &Self::Metadata {
43        &""
44    }
45
46    fn serialize_metadata(&self) -> Option<String> {
47        Some(String::default())
48    }
49
50    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
51        if metadata.is_some_and(str::is_empty) {
52            Ok("")
53        } else {
54            Err(ArrowError::InvalidArgumentError(
55                "Virtual column extension type expects an empty string as metadata".to_owned(),
56            ))
57        }
58    }
59
60    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
61        match data_type {
62            DataType::Int64 => Ok(()),
63            data_type => Err(ArrowError::InvalidArgumentError(format!(
64                "Virtual column data type mismatch, expected Int64, found {data_type}"
65            ))),
66        }
67    }
68
69    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self, ArrowError> {
70        Self.supports_data_type(data_type).map(|()| Self)
71    }
72
73    fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<(), ArrowError> {
74        Self.supports_data_type(data_type)
75    }
76}
77
78/// The extension type for row numbers.
79///
80/// Extension name: `parquet.virtual.row_number`.
81///
82/// This virtual column has storage type `Int64` and uses empty string metadata.
83#[derive(Debug, Default, Clone, Copy, PartialEq)]
84pub struct RowNumber;
85
86impl ExtensionType for RowNumber {
87    const NAME: &'static str = concat!(VIRTUAL_PREFIX!(), "row_number");
88    type Metadata = &'static str;
89
90    fn metadata(&self) -> &Self::Metadata {
91        &""
92    }
93
94    fn serialize_metadata(&self) -> Option<String> {
95        Some(String::default())
96    }
97
98    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
99        if metadata.is_some_and(str::is_empty) {
100            Ok("")
101        } else {
102            Err(ArrowError::InvalidArgumentError(
103                "Virtual column extension type expects an empty string as metadata".to_owned(),
104            ))
105        }
106    }
107
108    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
109        match data_type {
110            DataType::Int64 => Ok(()),
111            data_type => Err(ArrowError::InvalidArgumentError(format!(
112                "Virtual column data type mismatch, expected Int64, found {data_type}"
113            ))),
114        }
115    }
116
117    fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self, ArrowError> {
118        Self.supports_data_type(data_type).map(|()| Self)
119    }
120
121    fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<(), ArrowError> {
122        Self.supports_data_type(data_type)
123    }
124}
125
126/// Returns `true` if the field is a virtual column.
127///
128/// Virtual columns have extension type names starting with `parquet.virtual.`.
129pub fn is_virtual_column(field: &Field) -> bool {
130    field
131        .extension_type_name()
132        .is_some_and(|name| name.starts_with(VIRTUAL_PREFIX!()))
133}
134
135#[cfg(test)]
136mod tests {
137    use arrow_schema::{
138        ArrowError, DataType, Field,
139        extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY},
140    };
141
142    use super::*;
143
144    #[test]
145    fn row_number_valid() -> Result<(), ArrowError> {
146        let mut field = Field::new("", DataType::Int64, false);
147        field.try_with_extension_type(RowNumber)?;
148        field.try_extension_type::<RowNumber>()?;
149
150        Ok(())
151    }
152
153    #[test]
154    #[should_panic(expected = "Extension type name missing")]
155    fn row_number_missing_name() {
156        let field = Field::new("", DataType::Int64, false)
157            .with_metadata([(EXTENSION_TYPE_METADATA_KEY, "")]);
158        field.extension_type::<RowNumber>();
159    }
160
161    #[test]
162    #[should_panic(expected = "expected Int64, found Int32")]
163    fn row_number_invalid_type() {
164        Field::new("", DataType::Int32, false).with_extension_type(RowNumber);
165    }
166
167    #[test]
168    #[should_panic(expected = "Virtual column extension type expects an empty string as metadata")]
169    fn row_number_missing_metadata() {
170        let field = Field::new("", DataType::Int64, false)
171            .with_metadata([(EXTENSION_TYPE_NAME_KEY, RowNumber::NAME)]);
172        field.extension_type::<RowNumber>();
173    }
174
175    #[test]
176    #[should_panic(expected = "Virtual column extension type expects an empty string as metadata")]
177    fn row_number_invalid_metadata() {
178        let field = Field::new("", DataType::Int64, false).with_metadata([
179            (EXTENSION_TYPE_NAME_KEY, RowNumber::NAME),
180            (EXTENSION_TYPE_METADATA_KEY, "non-empty"),
181        ]);
182        field.extension_type::<RowNumber>();
183    }
184
185    #[test]
186    fn row_group_index_valid() -> Result<(), ArrowError> {
187        let mut field = Field::new("", DataType::Int64, false);
188        field.try_with_extension_type(RowGroupIndex)?;
189        field.try_extension_type::<RowGroupIndex>()?;
190
191        Ok(())
192    }
193
194    #[test]
195    #[should_panic(expected = "Extension type name missing")]
196    fn row_group_index_missing_name() {
197        let field = Field::new("", DataType::Int64, false)
198            .with_metadata([(EXTENSION_TYPE_METADATA_KEY, "")]);
199        field.extension_type::<RowGroupIndex>();
200    }
201
202    #[test]
203    #[should_panic(expected = "expected Int64, found Int32")]
204    fn row_group_index_invalid_type() {
205        Field::new("", DataType::Int32, false).with_extension_type(RowGroupIndex);
206    }
207
208    #[test]
209    #[should_panic(expected = "Virtual column extension type expects an empty string as metadata")]
210    fn row_group_index_missing_metadata() {
211        let field = Field::new("", DataType::Int64, false)
212            .with_metadata([(EXTENSION_TYPE_NAME_KEY, RowGroupIndex::NAME)]);
213        field.extension_type::<RowGroupIndex>();
214    }
215
216    #[test]
217    #[should_panic(expected = "Virtual column extension type expects an empty string as metadata")]
218    fn row_group_index_invalid_metadata() {
219        let field = Field::new("", DataType::Int64, false).with_metadata([
220            (EXTENSION_TYPE_NAME_KEY, RowGroupIndex::NAME),
221            (EXTENSION_TYPE_METADATA_KEY, "non-empty"),
222        ]);
223        field.extension_type::<RowGroupIndex>();
224    }
225}
parquet/arrow/schema/virtual_type.rs

parquet/arrow/schema/
virtual_type.rs
