parquet/arrow/schema/
virtual_type.rs

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

parquet/arrow/schema/
virtual_type.rs
