datafusion_datasource/
proto.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//! Protobuf conversions for the file-scan leaf types owned by this crate:
19//! [`FileRange`], [`PartitionedFile`] and [`FileGroup`].
20//!
21//! These are the single copy of that wire logic, used both by the central
22//! serializer in `datafusion-proto` and by the per-source `try_to_proto` hooks,
23//! so the format cannot drift between them.
24//!
25//! None of these conversions need a codec or an encode/decode context: every
26//! field is plain data or goes through `datafusion-proto-common`. That is why
27//! they are plain [`TryFrom`] impls rather than the `try_to_proto(ctx)` /
28//! `try_from_proto(node, ctx)` hooks used for plans, expressions and scan
29//! configs: the standard trait can express a conversion that takes nothing but
30//! the value, and the orphan rule allows it here because one side of each
31//! conversion is a type this crate owns.
32
33use std::sync::Arc;
34
35use chrono::{TimeZone, Utc};
36use datafusion_common::{DataFusionError, Result, internal_datafusion_err};
37use datafusion_proto_models::protobuf;
38use object_store::ObjectMeta;
39use object_store::path::Path;
40
41use crate::file_groups::FileGroup;
42use crate::{FileRange, PartitionedFile};
43
44impl TryFrom<&FileRange> for protobuf::FileRange {
45    type Error = DataFusionError;
46
47    fn try_from(range: &FileRange) -> Result<Self> {
48        Ok(protobuf::FileRange {
49            start: range.start,
50            end: range.end,
51        })
52    }
53}
54
55impl TryFrom<&protobuf::FileRange> for FileRange {
56    type Error = DataFusionError;
57
58    fn try_from(range: &protobuf::FileRange) -> Result<Self> {
59        Ok(FileRange {
60            start: range.start,
61            end: range.end,
62        })
63    }
64}
65
66impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile {
67    type Error = DataFusionError;
68
69    fn try_from(file: &PartitionedFile) -> Result<Self> {
70        let last_modified = file.object_meta.last_modified;
71        let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| {
72            DataFusionError::Plan(format!(
73                "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}"
74            ))
75        })? as u64;
76        Ok(protobuf::PartitionedFile {
77            arrow_schema: file
78                .arrow_schema
79                .as_ref()
80                .map(|s| s.as_ref().try_into())
81                .transpose()?,
82            path: file.object_meta.location.as_ref().to_owned(),
83            size: file.object_meta.size,
84            last_modified_ns,
85            partition_values: file
86                .partition_values
87                .iter()
88                .map(|v| v.try_into())
89                .collect::<Result<Vec<_>, _>>()?,
90            range: file.range.as_ref().map(TryInto::try_into).transpose()?,
91            statistics: file.statistics.as_ref().map(|s| s.as_ref().into()),
92        })
93    }
94}
95
96impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile {
97    type Error = DataFusionError;
98
99    fn try_from(file: &protobuf::PartitionedFile) -> Result<Self> {
100        let mut pf = PartitionedFile::new_from_meta(ObjectMeta {
101            location: Path::parse(file.path.as_str()).map_err(|e| {
102                internal_datafusion_err!("Invalid object_store path: {e}")
103            })?,
104            last_modified: Utc.timestamp_nanos(file.last_modified_ns as i64),
105            size: file.size,
106            e_tag: None,
107            version: None,
108        })
109        .with_partition_values(
110            file.partition_values
111                .iter()
112                .map(|v| v.try_into())
113                .collect::<Result<Vec<_>, _>>()?,
114        );
115        if let Some(proto_schema) = file.arrow_schema.as_ref() {
116            pf = pf.with_arrow_schema(Arc::new(
117                proto_schema.try_into().map_err(DataFusionError::from)?,
118            ));
119        }
120        if let Some(range) = file.range.as_ref() {
121            let range = FileRange::try_from(range)?;
122            pf = pf.with_range(range.start, range.end);
123        }
124        if let Some(proto_stats) = file.statistics.as_ref() {
125            // The wire format carries statistics for the full table schema (file + partition
126            // columns), so assign directly — `with_statistics` would append the partition
127            // column stats a second time.
128            pf.statistics = Some(Arc::new(proto_stats.try_into()?));
129        }
130        Ok(pf)
131    }
132}
133
134impl TryFrom<&FileGroup> for protobuf::FileGroup {
135    type Error = DataFusionError;
136
137    fn try_from(group: &FileGroup) -> Result<Self> {
138        Ok(protobuf::FileGroup {
139            files: group
140                .files()
141                .iter()
142                .map(TryInto::try_into)
143                .collect::<Result<Vec<_>>>()?,
144        })
145    }
146}
147
148impl TryFrom<&protobuf::FileGroup> for FileGroup {
149    type Error = DataFusionError;
150
151    fn try_from(group: &protobuf::FileGroup) -> Result<Self> {
152        Ok(FileGroup::new(
153            group
154                .files
155                .iter()
156                .map(TryInto::try_into)
157                .collect::<Result<Vec<_>>>()?,
158        ))
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use arrow::datatypes::{DataType, Field, Schema};
165    use datafusion_common::{ScalarValue, Statistics};
166
167    use super::*;
168
169    #[test]
170    fn partitioned_file_roundtrip_preserves_all_fields() -> Result<()> {
171        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
172        let pf = PartitionedFile::new_from_meta(ObjectMeta {
173            location: Path::parse("foo/bar.parquet")?,
174            last_modified: Utc.timestamp_nanos(1_000_000_000),
175            size: 1234,
176            e_tag: None,
177            version: None,
178        })
179        .with_partition_values(vec![ScalarValue::from("2024-01-01")])
180        .with_range(10, 20)
181        .with_arrow_schema(Arc::clone(&schema))
182        .with_statistics(Arc::new(Statistics::new_unknown(&schema)));
183
184        let encoded = protobuf::PartitionedFile::try_from(&pf)?;
185        let decoded = PartitionedFile::try_from(&encoded)?;
186
187        assert_eq!(decoded.object_meta.location, pf.object_meta.location);
188        assert_eq!(decoded.object_meta.size, pf.object_meta.size);
189        assert_eq!(
190            decoded.object_meta.last_modified,
191            pf.object_meta.last_modified
192        );
193        assert_eq!(decoded.partition_values, pf.partition_values);
194        assert_eq!(decoded.range, pf.range);
195        assert_eq!(decoded.arrow_schema.as_deref(), Some(schema.as_ref()));
196        // Statistics span the full table schema (file columns followed by one
197        // entry per partition column), and survive the round trip intact.
198        assert_eq!(
199            pf.statistics.as_ref().unwrap().column_statistics.len(),
200            schema.fields().len() + pf.partition_values.len()
201        );
202        assert_eq!(decoded.statistics, pf.statistics);
203        Ok(())
204    }
205
206    #[test]
207    fn partitioned_file_path_roundtrip_percent_encoded() -> Result<()> {
208        // The wire format carries the *encoded* path, so a location that already
209        // contains percent escapes must survive without a second round of
210        // encoding or decoding.
211        let path_str = "foo/foo%2Fbar/baz%252Fqux";
212        let pf = PartitionedFile::new_from_meta(ObjectMeta {
213            location: Path::parse(path_str)?,
214            last_modified: Utc.timestamp_nanos(1_000),
215            size: 42,
216            e_tag: None,
217            version: None,
218        });
219
220        let encoded = protobuf::PartitionedFile::try_from(&pf)?;
221        assert_eq!(encoded.path, path_str);
222
223        let decoded = PartitionedFile::try_from(&encoded)?;
224        assert_eq!(decoded.object_meta.location.as_ref(), path_str);
225        assert_eq!(decoded.object_meta.location, pf.object_meta.location);
226        Ok(())
227    }
228
229    #[test]
230    fn partitioned_file_arrow_schema_roundtrip_preserves_metadata() -> Result<()> {
231        use std::collections::HashMap;
232
233        let arrow_schema = Arc::new(Schema::new_with_metadata(
234            vec![
235                Field::new("id", DataType::Int64, false),
236                Field::new("value", DataType::Utf8, true).with_metadata(HashMap::from([
237                    ("field_meta".to_string(), "field_value".to_string()),
238                ])),
239            ],
240            HashMap::from([("schema_meta".to_string(), "schema_value".to_string())]),
241        ));
242        let pf = PartitionedFile::new("foo/bar.parquet", 10)
243            .with_arrow_schema(Arc::clone(&arrow_schema));
244
245        let encoded = protobuf::PartitionedFile::try_from(&pf)?;
246        assert!(encoded.arrow_schema.is_some());
247
248        let decoded = PartitionedFile::try_from(&encoded)?;
249        assert_eq!(decoded.arrow_schema.as_deref(), Some(arrow_schema.as_ref()));
250        Ok(())
251    }
252
253    #[test]
254    fn partitioned_file_from_proto_rejects_invalid_path() {
255        let proto = protobuf::PartitionedFile {
256            path: "foo//bar.parquet".to_string(),
257            ..Default::default()
258        };
259
260        let err = PartitionedFile::try_from(&proto).unwrap_err();
261        assert!(
262            err.to_string().contains("Invalid object_store path"),
263            "unexpected error: {err}"
264        );
265    }
266
267    #[test]
268    fn file_group_from_slice_matches_file_group() -> Result<()> {
269        // `protobuf::FileGroup: TryFrom<&[T]>` lives in `datafusion-proto-models`,
270        // generic over the element so that crate never names `PartitionedFile`.
271        // This is the caller-visible half: the bound resolves via
272        // `TryFrom<&PartitionedFile> for protobuf::PartitionedFile` above.
273        let files = vec![
274            PartitionedFile::new("a.parquet", 1),
275            PartitionedFile::new("b.parquet", 2),
276        ];
277
278        let from_slice = protobuf::FileGroup::try_from(&files[..])?;
279        let from_group = protobuf::FileGroup::try_from(&FileGroup::new(files))?;
280
281        assert_eq!(from_slice, from_group);
282        assert_eq!(from_slice.files.len(), 2);
283        Ok(())
284    }
285
286    #[test]
287    fn file_group_roundtrip() -> Result<()> {
288        let group = FileGroup::new(vec![
289            PartitionedFile::new("a.parquet", 1),
290            PartitionedFile::new("b.parquet", 2),
291        ]);
292
293        let encoded = protobuf::FileGroup::try_from(&group)?;
294        let decoded = FileGroup::try_from(&encoded)?;
295
296        assert_eq!(decoded.len(), 2);
297        assert_eq!(
298            decoded.files()[1].object_meta.location,
299            group.files()[1].object_meta.location
300        );
301        Ok(())
302    }
303}
datafusion_datasource/proto.rs

datafusion_datasource/
proto.rs
