datafusion_datasource_parquet/
virtual_column.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//! Typed wrapper for parquet virtual columns.
19//!
20//! arrow-rs identifies virtual columns via arrow extension types carried on
21//! the `FieldRef`. [`ParquetVirtualColumn`] lifts that contract into the type
22//! system so callers validate at the boundary (via `TryFrom<&FieldRef>`)
23//! rather than string-comparing extension-type names deep inside the reader.
24
25use arrow::datatypes::FieldRef;
26use arrow_schema::extension::ExtensionType;
27use datafusion_common::{DataFusionError, Result, not_impl_err};
28use parquet::arrow::RowNumber;
29use std::sync::Arc;
30
31/// A parquet virtual column validated to have a supported arrow extension
32/// type.
33///
34/// Construct via [`TryFrom<&FieldRef>`]; add a new variant (and update the
35/// `TryFrom` impl) when DataFusion gains support for another arrow-rs virtual
36/// extension type.
37#[derive(Debug, Clone)]
38pub enum ParquetVirtualColumn {
39    /// Absolute row number within the parquet file. Backed by arrow-rs's
40    /// [`RowNumber`] extension type.
41    RowNumber(FieldRef),
42}
43
44impl ParquetVirtualColumn {
45    pub fn field(&self) -> &FieldRef {
46        match self {
47            Self::RowNumber(field) => field,
48        }
49    }
50}
51
52impl From<ParquetVirtualColumn> for FieldRef {
53    fn from(col: ParquetVirtualColumn) -> Self {
54        match col {
55            ParquetVirtualColumn::RowNumber(field) => field,
56        }
57    }
58}
59
60impl TryFrom<&FieldRef> for ParquetVirtualColumn {
61    type Error = DataFusionError;
62
63    fn try_from(field: &FieldRef) -> Result<Self> {
64        let Some(name) = field.extension_type_name() else {
65            return not_impl_err!(
66                "Virtual column '{}' is missing an Arrow extension type; \
67                 supported extension types: [{}]",
68                field.name(),
69                RowNumber::NAME
70            );
71        };
72        match name {
73            n if n == RowNumber::NAME => Ok(Self::RowNumber(Arc::clone(field))),
74            other => not_impl_err!(
75                "Virtual column '{}' uses unsupported Arrow extension type '{}'; \
76                 supported types: [{}]. Add a ParquetVirtualColumn variant and \
77                 a test for this type before wiring it through.",
78                field.name(),
79                other,
80                RowNumber::NAME
81            ),
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use arrow::datatypes::{DataType, Field};
90
91    #[test]
92    fn row_number_field_converts() {
93        let field: FieldRef = Arc::new(
94            Field::new("row_number", DataType::Int64, false)
95                .with_extension_type(RowNumber),
96        );
97        let col = ParquetVirtualColumn::try_from(&field).expect("valid row_number");
98        assert!(matches!(col, ParquetVirtualColumn::RowNumber(_)));
99        assert_eq!(col.field().name(), "row_number");
100    }
101
102    #[test]
103    fn missing_extension_type_rejected() {
104        let field: FieldRef = Arc::new(Field::new("plain", DataType::Int64, false));
105        let err = ParquetVirtualColumn::try_from(&field).unwrap_err();
106        assert!(
107            err.to_string().contains("missing an Arrow extension type"),
108            "got: {err}"
109        );
110    }
111
112    #[test]
113    fn unsupported_extension_type_rejected() {
114        // RowGroupIndex is a real arrow-rs virtual type not yet in our enum.
115        let field: FieldRef = Arc::new(
116            Field::new("row_group_index", DataType::Int64, false)
117                .with_extension_type(parquet::arrow::RowGroupIndex),
118        );
119        let err = ParquetVirtualColumn::try_from(&field).unwrap_err();
120        assert!(
121            err.to_string().contains("parquet.virtual.row_group_index"),
122            "error should name the offending extension type, got: {err}"
123        );
124    }
125}
datafusion_datasource_parquet/virtual_column.rs

datafusion_datasource_parquet/
virtual_column.rs
