datafusion_datasource/file_sink_config/
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 conversion for the format-independent [`FileSinkConfig`].
19
20use std::sync::Arc;
21
22use datafusion_common::{DataFusionError, Result, internal_datafusion_err};
23use datafusion_execution::object_store::ObjectStoreUrl;
24use datafusion_expr::dml::InsertOp;
25use datafusion_proto_models::protobuf;
26
27use crate::ListingTableUrl;
28use crate::file_groups::FileGroup;
29use crate::file_sink_config::{FileOutputMode, FileSinkConfig};
30
31impl TryFrom<&FileSinkConfig> for protobuf::FileSinkConfig {
32    type Error = DataFusionError;
33
34    /// Serialize this shared file-sink configuration without format-specific
35    /// writer options.
36    fn try_from(config: &FileSinkConfig) -> Result<Self> {
37        let file_groups = config
38            .file_group
39            .iter()
40            .map(TryInto::try_into)
41            .collect::<Result<Vec<_>>>()?;
42        let table_paths = config
43            .table_paths
44            .iter()
45            .map(ToString::to_string)
46            .collect::<Vec<_>>();
47        let table_partition_cols = config
48            .table_partition_cols
49            .iter()
50            .map(|(name, data_type)| {
51                Ok(protobuf::PartitionColumn {
52                    name: name.to_owned(),
53                    arrow_type: Some(data_type.try_into()?),
54                })
55            })
56            .collect::<Result<Vec<_>>>()?;
57        let insert_op = match config.insert_op {
58            InsertOp::Append => protobuf::InsertOp::Append,
59            InsertOp::Overwrite => protobuf::InsertOp::Overwrite,
60            InsertOp::Replace => protobuf::InsertOp::Replace,
61        };
62        let file_output_mode = match config.file_output_mode {
63            FileOutputMode::Automatic => protobuf::FileOutputMode::Automatic,
64            FileOutputMode::SingleFile => protobuf::FileOutputMode::SingleFile,
65            FileOutputMode::Directory => protobuf::FileOutputMode::Directory,
66        };
67
68        Ok(protobuf::FileSinkConfig {
69            object_store_url: config.object_store_url.to_string(),
70            file_groups,
71            table_paths,
72            output_schema: Some(config.output_schema.as_ref().try_into()?),
73            table_partition_cols,
74            keep_partition_by_columns: config.keep_partition_by_columns,
75            insert_op: insert_op.into(),
76            file_extension: config.file_extension.clone(),
77            file_output_mode: file_output_mode.into(),
78        })
79    }
80}
81
82impl TryFrom<&protobuf::FileSinkConfig> for FileSinkConfig {
83    type Error = DataFusionError;
84
85    /// Reconstruct a shared file-sink configuration from protobuf.
86    fn try_from(conf: &protobuf::FileSinkConfig) -> Result<Self> {
87        let file_group = FileGroup::new(
88            conf.file_groups
89                .iter()
90                .map(TryInto::try_into)
91                .collect::<Result<Vec<_>>>()?,
92        );
93        let table_paths = conf
94            .table_paths
95            .iter()
96            .map(ListingTableUrl::parse)
97            .collect::<Result<Vec<_>>>()?;
98        let table_partition_cols = conf
99            .table_partition_cols
100            .iter()
101            .map(|protobuf::PartitionColumn { name, arrow_type }| {
102                let data_type = arrow_type
103                    .as_ref()
104                    .ok_or_else(|| {
105                        internal_datafusion_err!(
106                            "PartitionColumn is missing required field 'arrow_type'"
107                        )
108                    })?
109                    .try_into()?;
110                Ok((name.clone(), data_type))
111            })
112            .collect::<Result<Vec<_>>>()?;
113        let insert_op = protobuf::InsertOp::try_from(conf.insert_op).map_err(|_| {
114            internal_datafusion_err!(
115                "Received a FileSinkConfig message with unknown InsertOp {}",
116                conf.insert_op
117            )
118        })?;
119        let insert_op = match insert_op {
120            protobuf::InsertOp::Append => InsertOp::Append,
121            protobuf::InsertOp::Overwrite => InsertOp::Overwrite,
122            protobuf::InsertOp::Replace => InsertOp::Replace,
123        };
124        let file_output_mode = protobuf::FileOutputMode::try_from(conf.file_output_mode)
125            .map_err(|_| {
126                internal_datafusion_err!(
127                    "Received a FileSinkConfig message with unknown FileOutputMode {}",
128                    conf.file_output_mode
129                )
130            })?;
131        let file_output_mode = match file_output_mode {
132            protobuf::FileOutputMode::Automatic => FileOutputMode::Automatic,
133            protobuf::FileOutputMode::SingleFile => FileOutputMode::SingleFile,
134            protobuf::FileOutputMode::Directory => FileOutputMode::Directory,
135        };
136        let output_schema = conf.output_schema.as_ref().ok_or_else(|| {
137            internal_datafusion_err!(
138                "FileSinkConfig is missing required field 'output_schema'"
139            )
140        })?;
141
142        Ok(Self {
143            original_url: String::default(),
144            object_store_url: ObjectStoreUrl::parse(&conf.object_store_url)?,
145            file_group,
146            table_paths,
147            output_schema: Arc::new(output_schema.try_into()?),
148            table_partition_cols,
149            insert_op,
150            keep_partition_by_columns: conf.keep_partition_by_columns,
151            file_extension: conf.file_extension.clone(),
152            file_output_mode,
153        })
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use arrow::datatypes::Schema;
160
161    use super::*;
162
163    fn valid_file_sink_config() -> protobuf::FileSinkConfig {
164        protobuf::FileSinkConfig {
165            object_store_url: ObjectStoreUrl::local_filesystem().to_string(),
166            output_schema: Some(
167                (&Schema::empty())
168                    .try_into()
169                    .expect("empty schema should serialize"),
170            ),
171            insert_op: protobuf::InsertOp::Append.into(),
172            file_output_mode: protobuf::FileOutputMode::Automatic.into(),
173            ..Default::default()
174        }
175    }
176
177    fn assert_decode_error(
178        mutate: impl FnOnce(&mut protobuf::FileSinkConfig),
179        expected: impl AsRef<str>,
180    ) {
181        let mut conf = valid_file_sink_config();
182        mutate(&mut conf);
183
184        let error =
185            FileSinkConfig::try_from(&conf).expect_err("invalid config should fail");
186        match error {
187            DataFusionError::Internal(message) => {
188                let message = message
189                    .split_once(DataFusionError::BACK_TRACE_SEP)
190                    .map_or(message.as_str(), |(message, _)| message);
191                assert_eq!(message, expected.as_ref());
192            }
193            error => panic!("expected internal error, got {error}"),
194        }
195    }
196
197    #[test]
198    fn rejects_unknown_insert_op() {
199        assert_decode_error(
200            |conf| conf.insert_op = i32::MAX,
201            format!(
202                "Received a FileSinkConfig message with unknown InsertOp {}",
203                i32::MAX
204            ),
205        );
206    }
207
208    #[test]
209    fn rejects_unknown_file_output_mode() {
210        assert_decode_error(
211            |conf| conf.file_output_mode = i32::MAX,
212            format!(
213                "Received a FileSinkConfig message with unknown FileOutputMode {}",
214                i32::MAX
215            ),
216        );
217    }
218
219    #[test]
220    fn rejects_missing_output_schema() {
221        assert_decode_error(
222            |conf| conf.output_schema = None,
223            "FileSinkConfig is missing required field 'output_schema'",
224        );
225    }
226
227    #[test]
228    fn rejects_partition_column_without_arrow_type() {
229        assert_decode_error(
230            |conf| {
231                conf.table_partition_cols.push(protobuf::PartitionColumn {
232                    name: "partition".to_string(),
233                    arrow_type: None,
234                });
235            },
236            "PartitionColumn is missing required field 'arrow_type'",
237        );
238    }
239}
datafusion_datasource/file_sink_config/proto.rs

datafusion_datasource/file_sink_config/
proto.rs
