datafusion_proto_models/
from_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 the protobuf messages in this crate to their
19//! `datafusion-common` counterparts.
20//!
21//! The DataFusion side of these conversions lives *below* this crate in the
22//! dependency graph, so it cannot host the impls itself. They live here
23//! instead, on the local proto type — the same arrangement
24//! `datafusion-proto-common` uses for `ScalarValue` and `Statistics`.
25
26use std::sync::Arc;
27
28use datafusion_common::config::{
29    CsvOptions, JsonOptions, MaxRowGroupBytes, ParquetCdcOptions, ParquetColumnOptions,
30    ParquetOptions, TableParquetOptions,
31};
32use datafusion_common::display::{PlanType, StringifiedPlan};
33use datafusion_common::parsers::{CompressionTypeVariant, CsvQuoteStyle};
34use datafusion_common::{
35    JoinConstraint, JoinType, NullEquality, RecursionUnnestOption, TableReference,
36    UnnestOptions,
37};
38use datafusion_proto_common::FromProtoError as Error;
39
40use crate::protobuf::{
41    self, AnalyzedLogicalPlanType, CsvOptions as CsvOptionsProto,
42    CsvQuoteStyle as CsvQuoteStyleProto, JsonOptions as JsonOptionsProto,
43    OptimizedLogicalPlanType, OptimizedPhysicalPlanType,
44    ParquetCdcOptions as ParquetCdcOptionsProto,
45    ParquetColumnOptions as ParquetColumnOptionsProto,
46    ParquetOptions as ParquetOptionsProto,
47    TableParquetOptions as TableParquetOptionsProto, parquet_column_options,
48    parquet_options,
49    plan_type::PlanTypeEnum::{
50        AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan,
51        FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats,
52        InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema,
53        InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan,
54        PhysicalPlanError,
55    },
56};
57
58impl From<&protobuf::UnnestOptions> for UnnestOptions {
59    fn from(opts: &protobuf::UnnestOptions) -> Self {
60        use datafusion_common::NullHandling;
61        use protobuf::unnest_options::NullHandling as ProtoNullHandling;
62        let null_handling = match ProtoNullHandling::try_from(opts.null_handling) {
63            Ok(ProtoNullHandling::Preserve) => NullHandling::Preserve,
64            Ok(ProtoNullHandling::Drop) => NullHandling::Drop,
65            Ok(ProtoNullHandling::PreserveAndExpandEmpty) => {
66                NullHandling::PreserveAndExpandEmpty
67            }
68            // Unknown enum values fall back to the default (Preserve), which
69            // matches DataFusion's historical behavior.
70            Err(_) => NullHandling::Preserve,
71        };
72        Self {
73            null_handling,
74            recursions: opts
75                .recursions
76                .iter()
77                .map(|r| RecursionUnnestOption {
78                    input_column: r.input_column.as_ref().unwrap().into(),
79                    output_column: r.output_column.as_ref().unwrap().into(),
80                    depth: r.depth as usize,
81                })
82                .collect::<Vec<_>>(),
83        }
84    }
85}
86
87impl TryFrom<protobuf::TableReference> for TableReference {
88    type Error = Error;
89
90    fn try_from(value: protobuf::TableReference) -> Result<Self, Self::Error> {
91        use protobuf::table_reference::TableReferenceEnum;
92        let table_reference_enum = value
93            .table_reference_enum
94            .ok_or_else(|| Error::required("table_reference_enum"))?;
95
96        match table_reference_enum {
97            TableReferenceEnum::Bare(protobuf::BareTableReference { table }) => {
98                Ok(TableReference::bare(table))
99            }
100            TableReferenceEnum::Partial(protobuf::PartialTableReference {
101                schema,
102                table,
103            }) => Ok(TableReference::partial(schema, table)),
104            TableReferenceEnum::Full(protobuf::FullTableReference {
105                catalog,
106                schema,
107                table,
108            }) => Ok(TableReference::full(catalog, schema, table)),
109        }
110    }
111}
112
113impl From<&protobuf::StringifiedPlan> for StringifiedPlan {
114    fn from(stringified_plan: &protobuf::StringifiedPlan) -> Self {
115        Self {
116            plan_type: match stringified_plan
117                .plan_type
118                .as_ref()
119                .and_then(|pt| pt.plan_type_enum.as_ref())
120                .unwrap_or_else(|| {
121                    panic!(
122                        "Cannot create protobuf::StringifiedPlan from {stringified_plan:?}"
123                    )
124                }) {
125                InitialLogicalPlan(_) => PlanType::InitialLogicalPlan,
126                AnalyzedLogicalPlan(AnalyzedLogicalPlanType { analyzer_name }) => {
127                    PlanType::AnalyzedLogicalPlan {
128                        analyzer_name:analyzer_name.clone()
129                    }
130                }
131                FinalAnalyzedLogicalPlan(_) => PlanType::FinalAnalyzedLogicalPlan,
132                OptimizedLogicalPlan(OptimizedLogicalPlanType { optimizer_name }) => {
133                    PlanType::OptimizedLogicalPlan {
134                        optimizer_name: optimizer_name.clone(),
135                    }
136                }
137                FinalLogicalPlan(_) => PlanType::FinalLogicalPlan,
138                InitialPhysicalPlan(_) => PlanType::InitialPhysicalPlan,
139                InitialPhysicalPlanWithStats(_) => PlanType::InitialPhysicalPlanWithStats,
140                InitialPhysicalPlanWithSchema(_) => PlanType::InitialPhysicalPlanWithSchema,
141                OptimizedPhysicalPlan(OptimizedPhysicalPlanType { optimizer_name }) => {
142                    PlanType::OptimizedPhysicalPlan {
143                        optimizer_name: optimizer_name.clone(),
144                    }
145                }
146                FinalPhysicalPlan(_) => PlanType::FinalPhysicalPlan,
147                FinalPhysicalPlanWithStats(_) => PlanType::FinalPhysicalPlanWithStats,
148                FinalPhysicalPlanWithSchema(_) => PlanType::FinalPhysicalPlanWithSchema,
149                PhysicalPlanError(_) => PlanType::PhysicalPlanError,
150            },
151            plan: Arc::new(stringified_plan.plan.clone()),
152        }
153    }
154}
155
156impl From<protobuf::JoinType> for JoinType {
157    fn from(t: protobuf::JoinType) -> Self {
158        match t {
159            protobuf::JoinType::Inner => JoinType::Inner,
160            protobuf::JoinType::Left => JoinType::Left,
161            protobuf::JoinType::Right => JoinType::Right,
162            protobuf::JoinType::Full => JoinType::Full,
163            protobuf::JoinType::Leftsemi => JoinType::LeftSemi,
164            protobuf::JoinType::Rightsemi => JoinType::RightSemi,
165            protobuf::JoinType::Leftanti => JoinType::LeftAnti,
166            protobuf::JoinType::Rightanti => JoinType::RightAnti,
167            protobuf::JoinType::Leftmark => JoinType::LeftMark,
168            protobuf::JoinType::Rightmark => JoinType::RightMark,
169        }
170    }
171}
172
173impl From<protobuf::JoinConstraint> for JoinConstraint {
174    fn from(t: protobuf::JoinConstraint) -> Self {
175        match t {
176            protobuf::JoinConstraint::On => JoinConstraint::On,
177            protobuf::JoinConstraint::Using => JoinConstraint::Using,
178        }
179    }
180}
181
182impl From<protobuf::NullEquality> for NullEquality {
183    fn from(t: protobuf::NullEquality) -> Self {
184        match t {
185            protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing,
186            protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull,
187        }
188    }
189}
190
191impl From<&CsvOptionsProto> for CsvOptions {
192    fn from(proto: &CsvOptionsProto) -> Self {
193        CsvOptions {
194            has_header: if !proto.has_header.is_empty() {
195                Some(proto.has_header[0] != 0)
196            } else {
197                None
198            },
199            delimiter: proto.delimiter.first().copied().unwrap_or(b','),
200            quote: proto.quote.first().copied().unwrap_or(b'"'),
201            terminator: if !proto.terminator.is_empty() {
202                Some(proto.terminator[0])
203            } else {
204                None
205            },
206            escape: if !proto.escape.is_empty() {
207                Some(proto.escape[0])
208            } else {
209                None
210            },
211            double_quote: if !proto.double_quote.is_empty() {
212                Some(proto.double_quote[0] != 0)
213            } else {
214                None
215            },
216            compression: match proto.compression {
217                0 => CompressionTypeVariant::GZIP,
218                1 => CompressionTypeVariant::BZIP2,
219                2 => CompressionTypeVariant::XZ,
220                3 => CompressionTypeVariant::ZSTD,
221                _ => CompressionTypeVariant::UNCOMPRESSED,
222            },
223            schema_infer_max_rec: proto.schema_infer_max_rec.map(|v| v as usize),
224            date_format: if proto.date_format.is_empty() {
225                None
226            } else {
227                Some(proto.date_format.clone())
228            },
229            datetime_format: if proto.datetime_format.is_empty() {
230                None
231            } else {
232                Some(proto.datetime_format.clone())
233            },
234            timestamp_format: if proto.timestamp_format.is_empty() {
235                None
236            } else {
237                Some(proto.timestamp_format.clone())
238            },
239            timestamp_tz_format: if proto.timestamp_tz_format.is_empty() {
240                None
241            } else {
242                Some(proto.timestamp_tz_format.clone())
243            },
244            time_format: if proto.time_format.is_empty() {
245                None
246            } else {
247                Some(proto.time_format.clone())
248            },
249            null_value: if proto.null_value.is_empty() {
250                None
251            } else {
252                Some(proto.null_value.clone())
253            },
254            null_regex: if proto.null_regex.is_empty() {
255                None
256            } else {
257                Some(proto.null_regex.clone())
258            },
259            comment: if !proto.comment.is_empty() {
260                Some(proto.comment[0])
261            } else {
262                None
263            },
264            newlines_in_values: if proto.newlines_in_values.is_empty() {
265                None
266            } else {
267                Some(proto.newlines_in_values[0] != 0)
268            },
269            truncated_rows: if proto.truncated_rows.is_empty() {
270                None
271            } else {
272                Some(proto.truncated_rows[0] != 0)
273            },
274            compression_level: proto.compression_level,
275            quote_style: match CsvQuoteStyleProto::try_from(proto.quote_style) {
276                Ok(CsvQuoteStyleProto::Always) => CsvQuoteStyle::Always,
277                Ok(CsvQuoteStyleProto::NonNumeric) => CsvQuoteStyle::NonNumeric,
278                Ok(CsvQuoteStyleProto::Never) => CsvQuoteStyle::Never,
279                Ok(CsvQuoteStyleProto::Necessary) => CsvQuoteStyle::Necessary,
280                _ => CsvQuoteStyle::Necessary,
281            },
282            ignore_leading_whitespace: if proto.ignore_leading_whitespace.is_empty() {
283                None
284            } else {
285                Some(proto.ignore_leading_whitespace[0] != 0)
286            },
287            ignore_trailing_whitespace: if proto.ignore_trailing_whitespace.is_empty() {
288                None
289            } else {
290                Some(proto.ignore_trailing_whitespace[0] != 0)
291            },
292        }
293    }
294}
295
296impl From<&JsonOptionsProto> for JsonOptions {
297    fn from(proto: &JsonOptionsProto) -> Self {
298        JsonOptions {
299            compression: match proto.compression {
300                0 => CompressionTypeVariant::GZIP,
301                1 => CompressionTypeVariant::BZIP2,
302                2 => CompressionTypeVariant::XZ,
303                3 => CompressionTypeVariant::ZSTD,
304                _ => CompressionTypeVariant::UNCOMPRESSED,
305            },
306            schema_infer_max_rec: proto.schema_infer_max_rec.map(|v| v as usize),
307            compression_level: proto.compression_level,
308            newline_delimited: proto.newline_delimited.unwrap_or(true),
309        }
310    }
311}
312
313impl From<ParquetCdcOptionsProto> for ParquetCdcOptions {
314    fn from(value: ParquetCdcOptionsProto) -> Self {
315        ParquetCdcOptions {
316            enabled: value.enabled,
317            min_chunk_size: value.min_chunk_size as usize,
318            max_chunk_size: value.max_chunk_size as usize,
319            norm_level: value.norm_level,
320        }
321    }
322}
323
324impl TryFrom<&ParquetOptionsProto> for ParquetOptions {
325    type Error = datafusion_common::DataFusionError;
326
327    fn try_from(
328        proto: &ParquetOptionsProto,
329    ) -> datafusion_common::Result<Self, Self::Error> {
330        let writer_version = match proto.writer_version.as_str() {
331            // Proto3 decodes an omitted string field as the empty string. The
332            // schema documents writer_version's logical default as "1.0", so
333            // preserve that default when the field is absent on the wire.
334            "" => ParquetOptions::default().writer_version,
335            version => version.parse()?,
336        };
337
338        Ok(ParquetOptions {
339            enable_page_index: proto.enable_page_index,
340            pruning: proto.pruning,
341            skip_metadata: proto.skip_metadata,
342            metadata_size_hint: proto
343                .metadata_size_hint_opt
344                .as_ref()
345                .map(|opt| match opt {
346                    parquet_options::MetadataSizeHintOpt::MetadataSizeHint(size) => {
347                        *size as usize
348                    }
349                }),
350            pushdown_filters: proto.pushdown_filters,
351            reorder_filters: proto.reorder_filters,
352            force_filter_selections: proto.force_filter_selections,
353            data_pagesize_limit: proto.data_pagesize_limit as usize,
354            write_batch_size: proto.write_batch_size as usize,
355            writer_version,
356            compression: proto.compression_opt.as_ref().map(|opt| match opt {
357                parquet_options::CompressionOpt::Compression(compression) => {
358                    compression.clone()
359                }
360            }),
361            dictionary_enabled: proto.dictionary_enabled_opt.as_ref().map(|opt| {
362                match opt {
363                    parquet_options::DictionaryEnabledOpt::DictionaryEnabled(
364                        enabled,
365                    ) => *enabled,
366                }
367            }),
368            dictionary_page_size_limit: proto.dictionary_page_size_limit as usize,
369            statistics_enabled: proto.statistics_enabled_opt.as_ref().map(
370                |opt| match opt {
371                    parquet_options::StatisticsEnabledOpt::StatisticsEnabled(
372                        statistics,
373                    ) => statistics.clone(),
374                },
375            ),
376            max_row_group_size: proto.max_row_group_size as usize,
377            max_in_list_size: proto.max_in_list_size as usize,
378            created_by: proto.created_by.clone(),
379            column_index_truncate_length: proto
380                .column_index_truncate_length_opt
381                .as_ref()
382                .map(|opt| match opt {
383                    parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length) => *length as usize,
384                }),
385            statistics_truncate_length: proto
386                .statistics_truncate_length_opt
387                .as_ref()
388                .map(|opt| match opt {
389                    parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(length) => *length as usize,
390                }),
391            data_page_row_count_limit: proto.data_page_row_count_limit as usize,
392            encoding: proto.encoding_opt.as_ref().map(|opt| match opt {
393                parquet_options::EncodingOpt::Encoding(encoding) => {
394                    encoding.clone()
395                }
396            }),
397            bloom_filter_on_read: proto.bloom_filter_on_read,
398            bloom_filter_on_write: proto.bloom_filter_on_write,
399            bloom_filter_fpp: proto
400                .bloom_filter_fpp_opt
401                .as_ref()
402                .map(|opt| match opt {
403                    parquet_options::BloomFilterFppOpt::BloomFilterFpp(fpp) => *fpp,
404                }),
405            bloom_filter_ndv: proto
406                .bloom_filter_ndv_opt
407                .as_ref()
408                .map(|opt| match opt {
409                    parquet_options::BloomFilterNdvOpt::BloomFilterNdv(ndv) => *ndv,
410                }),
411            allow_single_file_parallelism: proto.allow_single_file_parallelism,
412            maximum_parallel_row_group_writers: proto
413                .maximum_parallel_row_group_writers
414                as usize,
415            maximum_buffered_record_batches_per_stream: proto
416                .maximum_buffered_record_batches_per_stream
417                as usize,
418            schema_force_view_types: proto.schema_force_view_types,
419            binary_as_string: proto.binary_as_string,
420            skip_arrow_metadata: proto.skip_arrow_metadata,
421            coerce_int96: proto.coerce_int96_opt.as_ref().map(|opt| match opt {
422                parquet_options::CoerceInt96Opt::CoerceInt96(coerce_int96) => {
423                    coerce_int96.clone()
424                }
425            }),
426            coerce_int96_tz: proto
427                .coerce_int96_tz_opt
428                .as_ref()
429                .map(|opt| match opt {
430                    parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(tz) => {
431                        tz.clone()
432                    }
433                }),
434            max_predicate_cache_size: proto
435                .max_predicate_cache_size_opt
436                .as_ref()
437                .map(|opt| match opt {
438                    parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(
439                        size,
440                    ) => *size as usize,
441                }),
442            max_row_group_bytes: proto
443                .max_row_group_bytes_opt
444                .as_ref()
445                .and_then(|opt| match opt {
446                    parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(size) => {
447                        MaxRowGroupBytes::try_new(*size as usize).ok()
448                    }
449                }),
450            content_defined_chunking: proto
451                .content_defined_chunking
452                .map(ParquetCdcOptions::from)
453                .unwrap_or_default(),
454        })
455    }
456}
457
458impl From<ParquetColumnOptionsProto> for ParquetColumnOptions {
459    fn from(proto: ParquetColumnOptionsProto) -> Self {
460        ParquetColumnOptions {
461            bloom_filter_enabled: proto.bloom_filter_enabled_opt.map(
462                |parquet_column_options::BloomFilterEnabledOpt::BloomFilterEnabled(v)| v,
463            ),
464            encoding: proto
465                .encoding_opt
466                .map(|parquet_column_options::EncodingOpt::Encoding(v)| v),
467            dictionary_enabled: proto.dictionary_enabled_opt.map(
468                |parquet_column_options::DictionaryEnabledOpt::DictionaryEnabled(v)| v,
469            ),
470            compression: proto
471                .compression_opt
472                .map(|parquet_column_options::CompressionOpt::Compression(v)| v),
473            statistics_enabled: proto.statistics_enabled_opt.map(
474                |parquet_column_options::StatisticsEnabledOpt::StatisticsEnabled(v)| v,
475            ),
476            bloom_filter_fpp: proto
477                .bloom_filter_fpp_opt
478                .map(|parquet_column_options::BloomFilterFppOpt::BloomFilterFpp(v)| v),
479            bloom_filter_ndv: proto
480                .bloom_filter_ndv_opt
481                .map(|parquet_column_options::BloomFilterNdvOpt::BloomFilterNdv(v)| v),
482        }
483    }
484}
485
486impl TryFrom<&TableParquetOptionsProto> for TableParquetOptions {
487    type Error = datafusion_common::DataFusionError;
488
489    fn try_from(
490        proto: &TableParquetOptionsProto,
491    ) -> datafusion_common::Result<Self, Self::Error> {
492        Ok(TableParquetOptions {
493            global: proto
494                .global
495                .as_ref()
496                .map(ParquetOptions::try_from)
497                .transpose()?
498                .unwrap_or_default(),
499            column_specific_options: proto
500                .column_specific_options
501                .iter()
502                .map(|parquet_column_options| {
503                    (
504                        parquet_column_options.column_name.clone(),
505                        ParquetColumnOptions::from(
506                            parquet_column_options.options.clone().unwrap_or_default(),
507                        ),
508                    )
509                })
510                .collect(),
511            key_value_metadata: proto
512                .key_value_metadata
513                .iter()
514                .map(|(k, v)| (k.clone(), Some(v.clone())))
515                .collect(),
516            ..Default::default()
517        })
518    }
519}
datafusion_proto_models/from_proto.rs

datafusion_proto_models/
from_proto.rs
