datafusion_proto_models/
to_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//! Conversions from `datafusion-common` types to the protobuf messages in this
19//! crate.
20//!
21//! See [`crate::from_proto`] for why the impls live here rather than next to
22//! the DataFusion types.
23
24use datafusion_common::DataFusionError;
25use datafusion_common::display::{PlanType, StringifiedPlan};
26use datafusion_common::{
27    JoinConstraint, JoinType, NullEquality, TableReference, UnnestOptions,
28};
29
30use crate::generated::datafusion_common::EmptyMessage;
31use crate::protobuf::{
32    self, AnalyzedLogicalPlanType, OptimizedLogicalPlanType, OptimizedPhysicalPlanType,
33    RecursionUnnestOption,
34    plan_type::PlanTypeEnum::{
35        AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan,
36        FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats,
37        InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema,
38        InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan,
39        PhysicalPlanError,
40    },
41};
42
43impl From<&UnnestOptions> for protobuf::UnnestOptions {
44    fn from(opts: &UnnestOptions) -> Self {
45        use datafusion_common::NullHandling;
46        use protobuf::unnest_options::NullHandling as ProtoNullHandling;
47        let null_handling = match opts.null_handling {
48            NullHandling::Preserve => ProtoNullHandling::Preserve,
49            NullHandling::Drop => ProtoNullHandling::Drop,
50            NullHandling::PreserveAndExpandEmpty => {
51                ProtoNullHandling::PreserveAndExpandEmpty
52            }
53        } as i32;
54        Self {
55            null_handling,
56            recursions: opts
57                .recursions
58                .iter()
59                .map(|r| RecursionUnnestOption {
60                    input_column: Some((&r.input_column).into()),
61                    output_column: Some((&r.output_column).into()),
62                    depth: r.depth as u32,
63                })
64                .collect(),
65        }
66    }
67}
68
69impl From<&StringifiedPlan> for protobuf::StringifiedPlan {
70    fn from(stringified_plan: &StringifiedPlan) -> Self {
71        Self {
72            plan_type: match stringified_plan.clone().plan_type {
73                PlanType::InitialLogicalPlan => Some(protobuf::PlanType {
74                    plan_type_enum: Some(InitialLogicalPlan(EmptyMessage {})),
75                }),
76                PlanType::AnalyzedLogicalPlan { analyzer_name } => {
77                    Some(protobuf::PlanType {
78                        plan_type_enum: Some(AnalyzedLogicalPlan(
79                            AnalyzedLogicalPlanType { analyzer_name },
80                        )),
81                    })
82                }
83                PlanType::FinalAnalyzedLogicalPlan => Some(protobuf::PlanType {
84                    plan_type_enum: Some(FinalAnalyzedLogicalPlan(EmptyMessage {})),
85                }),
86                PlanType::OptimizedLogicalPlan { optimizer_name } => {
87                    Some(protobuf::PlanType {
88                        plan_type_enum: Some(OptimizedLogicalPlan(
89                            OptimizedLogicalPlanType { optimizer_name },
90                        )),
91                    })
92                }
93                PlanType::FinalLogicalPlan => Some(protobuf::PlanType {
94                    plan_type_enum: Some(FinalLogicalPlan(EmptyMessage {})),
95                }),
96                PlanType::InitialPhysicalPlan => Some(protobuf::PlanType {
97                    plan_type_enum: Some(InitialPhysicalPlan(EmptyMessage {})),
98                }),
99                PlanType::OptimizedPhysicalPlan { optimizer_name } => {
100                    Some(protobuf::PlanType {
101                        plan_type_enum: Some(OptimizedPhysicalPlan(
102                            OptimizedPhysicalPlanType { optimizer_name },
103                        )),
104                    })
105                }
106                PlanType::FinalPhysicalPlan => Some(protobuf::PlanType {
107                    plan_type_enum: Some(FinalPhysicalPlan(EmptyMessage {})),
108                }),
109                PlanType::InitialPhysicalPlanWithStats => Some(protobuf::PlanType {
110                    plan_type_enum: Some(InitialPhysicalPlanWithStats(EmptyMessage {})),
111                }),
112                PlanType::InitialPhysicalPlanWithSchema => Some(protobuf::PlanType {
113                    plan_type_enum: Some(InitialPhysicalPlanWithSchema(EmptyMessage {})),
114                }),
115                PlanType::FinalPhysicalPlanWithStats => Some(protobuf::PlanType {
116                    plan_type_enum: Some(FinalPhysicalPlanWithStats(EmptyMessage {})),
117                }),
118                PlanType::FinalPhysicalPlanWithSchema => Some(protobuf::PlanType {
119                    plan_type_enum: Some(FinalPhysicalPlanWithSchema(EmptyMessage {})),
120                }),
121                PlanType::PhysicalPlanError => Some(protobuf::PlanType {
122                    plan_type_enum: Some(PhysicalPlanError(EmptyMessage {})),
123                }),
124            },
125            plan: stringified_plan.plan.to_string(),
126        }
127    }
128}
129
130impl From<TableReference> for protobuf::TableReference {
131    fn from(t: TableReference) -> Self {
132        use protobuf::table_reference::TableReferenceEnum;
133        let table_reference_enum = match t {
134            TableReference::Bare { table } => {
135                TableReferenceEnum::Bare(protobuf::BareTableReference {
136                    table: table.to_string(),
137                })
138            }
139            TableReference::Partial { schema, table } => {
140                TableReferenceEnum::Partial(protobuf::PartialTableReference {
141                    schema: schema.to_string(),
142                    table: table.to_string(),
143                })
144            }
145            TableReference::Full {
146                catalog,
147                schema,
148                table,
149            } => TableReferenceEnum::Full(protobuf::FullTableReference {
150                catalog: catalog.to_string(),
151                schema: schema.to_string(),
152                table: table.to_string(),
153            }),
154        };
155
156        protobuf::TableReference {
157            table_reference_enum: Some(table_reference_enum),
158        }
159    }
160}
161
162impl From<JoinType> for protobuf::JoinType {
163    fn from(t: JoinType) -> Self {
164        match t {
165            JoinType::Inner => protobuf::JoinType::Inner,
166            JoinType::Left => protobuf::JoinType::Left,
167            JoinType::Right => protobuf::JoinType::Right,
168            JoinType::Full => protobuf::JoinType::Full,
169            JoinType::LeftSemi => protobuf::JoinType::Leftsemi,
170            JoinType::RightSemi => protobuf::JoinType::Rightsemi,
171            JoinType::LeftAnti => protobuf::JoinType::Leftanti,
172            JoinType::RightAnti => protobuf::JoinType::Rightanti,
173            JoinType::LeftMark => protobuf::JoinType::Leftmark,
174            JoinType::RightMark => protobuf::JoinType::Rightmark,
175        }
176    }
177}
178
179impl From<JoinConstraint> for protobuf::JoinConstraint {
180    fn from(t: JoinConstraint) -> Self {
181        match t {
182            JoinConstraint::On => protobuf::JoinConstraint::On,
183            JoinConstraint::Using => protobuf::JoinConstraint::Using,
184        }
185    }
186}
187
188impl From<NullEquality> for protobuf::NullEquality {
189    fn from(t: NullEquality) -> Self {
190        match t {
191            NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing,
192            NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull,
193        }
194    }
195}
196
197/// Encode any slice of file-like values as a [`protobuf::FileGroup`].
198///
199/// `datafusion-datasource` cannot host this impl: `&T` is `#[fundamental]` but
200/// `[T]` is not, so `&[PartitionedFile]` counts as foreign there and the orphan
201/// rule rejects it. Here the *self* type is local, which is all the orphan rule
202/// needs — and staying generic over the element means this crate never has to
203/// name `PartitionedFile`, which lives above it in the dependency graph.
204///
205/// The element bound is satisfied by
206/// `impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile` in
207/// `datafusion-datasource`, so `protobuf::FileGroup::try_from(&files[..])`
208/// resolves for callers exactly as it did before the proto types were split out.
209impl<T> TryFrom<&[T]> for protobuf::FileGroup
210where
211    for<'a> &'a T: TryInto<protobuf::PartitionedFile, Error = DataFusionError>,
212{
213    type Error = DataFusionError;
214
215    fn try_from(files: &[T]) -> Result<Self, Self::Error> {
216        Ok(protobuf::FileGroup {
217            files: files
218                .iter()
219                .map(TryInto::try_into)
220                .collect::<Result<Vec<_>, _>>()?,
221        })
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use datafusion_common::{NullHandling, RecursionUnnestOption};
228
229    use super::*;
230
231    #[test]
232    fn table_reference_roundtrip() {
233        for reference in [
234            TableReference::bare("t"),
235            TableReference::partial("s", "t"),
236            TableReference::full("c", "s", "t"),
237        ] {
238            let encoded = protobuf::TableReference::from(reference.clone());
239            let decoded = TableReference::try_from(encoded).unwrap();
240            assert_eq!(decoded, reference);
241        }
242    }
243
244    #[test]
245    fn table_reference_from_proto_rejects_missing_oneof() {
246        let proto = protobuf::TableReference {
247            table_reference_enum: None,
248        };
249        let err = TableReference::try_from(proto).unwrap_err();
250        assert!(
251            err.to_string().contains("table_reference_enum"),
252            "unexpected error: {err}"
253        );
254    }
255
256    #[test]
257    fn join_enums_roundtrip() {
258        for join_type in [
259            JoinType::Inner,
260            JoinType::Left,
261            JoinType::Right,
262            JoinType::Full,
263            JoinType::LeftSemi,
264            JoinType::RightSemi,
265            JoinType::LeftAnti,
266            JoinType::RightAnti,
267            JoinType::LeftMark,
268            JoinType::RightMark,
269        ] {
270            assert_eq!(
271                JoinType::from(protobuf::JoinType::from(join_type)),
272                join_type
273            );
274        }
275        for constraint in [JoinConstraint::On, JoinConstraint::Using] {
276            assert_eq!(
277                JoinConstraint::from(protobuf::JoinConstraint::from(constraint)),
278                constraint
279            );
280        }
281        for null_equality in [
282            NullEquality::NullEqualsNothing,
283            NullEquality::NullEqualsNull,
284        ] {
285            assert_eq!(
286                NullEquality::from(protobuf::NullEquality::from(null_equality)),
287                null_equality
288            );
289        }
290    }
291
292    #[test]
293    fn unnest_options_roundtrip() {
294        let options = UnnestOptions {
295            null_handling: NullHandling::Drop,
296            recursions: vec![RecursionUnnestOption {
297                input_column: "a".into(),
298                output_column: "b".into(),
299                depth: 2,
300            }],
301        };
302
303        let encoded = protobuf::UnnestOptions::from(&options);
304        let decoded = UnnestOptions::from(&encoded);
305
306        assert_eq!(decoded.null_handling, options.null_handling);
307        assert_eq!(decoded.recursions, options.recursions);
308    }
309
310    #[test]
311    fn stringified_plan_roundtrip() {
312        let plan = StringifiedPlan::new(
313            PlanType::OptimizedLogicalPlan {
314                optimizer_name: "push_down_filter".to_string(),
315            },
316            "some plan",
317        );
318
319        let encoded = protobuf::StringifiedPlan::from(&plan);
320        let decoded = StringifiedPlan::from(&encoded);
321
322        assert_eq!(decoded.plan_type, plan.plan_type);
323        assert_eq!(decoded.plan, plan.plan);
324    }
325}
datafusion_proto_models/to_proto.rs

datafusion_proto_models/
to_proto.rs
