1use std::collections::HashMap;
19use std::fmt::Debug;
20use std::sync::Arc;
21
22use crate::protobuf::logical_plan_node::LogicalPlanType::CustomScan;
23use crate::protobuf::{
24 ColumnUnnestListItem, ColumnUnnestListRecursion, CteWorkTableScanNode,
25 CustomTableScanNode, DmlNode, SortExprNodeCollection, dml_node,
26};
27use crate::{
28 convert_required,
29 protobuf::{
30 self, LogicalExtensionNode, LogicalPlanNode,
31 listing_table_scan_node::FileFormatType, logical_plan_node::LogicalPlanType,
32 },
33};
34
35use crate::protobuf::{ToProtoError, proto_error};
36use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, SchemaRef};
37use datafusion_catalog::cte_worktable::CteWorkTable;
38use datafusion_catalog::empty::EmptyTable;
39use datafusion_common::file_options::file_type::FileType;
40use datafusion_common::format::{
41 ExplainAnalyzeCategories, ExplainFormat, MetricCategory, MetricType,
42};
43use datafusion_common::{
44 NullEquality, Result, TableReference, assert_or_internal_err, context,
45 internal_datafusion_err, internal_err, not_impl_err, plan_err,
46};
47use datafusion_datasource::file_format::FileFormat;
48use datafusion_datasource::file_format::{
49 FileFormatFactory, file_type_to_format, format_as_file_type,
50};
51use datafusion_datasource_arrow::file_format::{ArrowFormat, ArrowFormatFactory};
52#[cfg(feature = "avro")]
53use datafusion_datasource_avro::file_format::AvroFormat;
54use datafusion_datasource_csv::file_format::{CsvFormat, CsvFormatFactory};
55use datafusion_datasource_json::file_format::{
56 JsonFormat as OtherNdJsonFormat, JsonFormatFactory,
57};
58#[cfg(feature = "parquet")]
59use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory};
60use datafusion_expr::dml::InsertOp;
61use datafusion_expr::{
62 AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning,
63 RecursiveQuery, SkipType, TableSource, Unnest, WriteOp,
64};
65use datafusion_expr::{
66 DistinctOn, DropView, Expr, JoinConstraint, LogicalPlan, LogicalPlanBuilder,
67 ScalarUDF, SortExpr, Statement, WindowUDF, dml,
68 logical_plan::{
69 Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView,
70 DdlStatement, Distinct, EmptyRelation, Extension, Join, Prepare, Projection,
71 Repartition, Sort, SubqueryAlias, TableScan, TableScanBuilder, Values, Window,
72 builder::project,
73 },
74};
75use datafusion_proto_common::protobuf_common;
76
77use self::to_proto::{serialize_expr, serialize_exprs};
78use crate::logical_plan::to_proto::serialize_range_split_point;
79use crate::logical_plan::to_proto::serialize_sorts;
80use datafusion_catalog::TableProvider;
81use datafusion_catalog::default_table_source::{provider_as_source, source_as_provider};
82use datafusion_catalog::view::ViewTable;
83use datafusion_catalog_listing::{ListingOptions, ListingTable, ListingTableConfig};
84use datafusion_datasource::ListingTableUrl;
85use datafusion_execution::TaskContext;
86use prost::Message;
87use prost::bytes::BufMut;
88
89pub mod file_formats;
90pub mod from_proto;
91pub mod to_proto;
92
93pub trait AsLogicalPlan: Debug + Send + Sync + Clone {
94 fn try_decode(buf: &[u8]) -> Result<Self>
95 where
96 Self: Sized;
97
98 fn try_encode<B>(&self, buf: &mut B) -> Result<()>
99 where
100 B: BufMut,
101 Self: Sized;
102
103 fn try_into_logical_plan(
104 &self,
105 ctx: &TaskContext,
106 extension_codec: &dyn LogicalExtensionCodec,
107 ) -> Result<LogicalPlan>;
108
109 fn try_from_logical_plan(
110 plan: &LogicalPlan,
111 extension_codec: &dyn LogicalExtensionCodec,
112 ) -> Result<Self>
113 where
114 Self: Sized;
115}
116
117#[cfg_attr(debug_assertions, inline(never))]
121fn serialize_logical_plan_arm<F>(serializer: F) -> Result<LogicalPlanNode>
122where
123 F: FnOnce() -> Result<LogicalPlanNode>,
124{
125 serializer()
126}
127
128macro_rules! dispatch_logical_plan {
129 ($plan:expr, { $($pattern:pat => $body:expr $(,)?)+ }) => {
130 match $plan {
131 $(
132 $pattern => serialize_logical_plan_arm(|| -> Result<LogicalPlanNode> {
133 $body
134 }),
135 )+
136 }
137 };
138}
139
140pub trait LogicalExtensionCodec: Debug + Send + Sync + std::any::Any {
141 fn try_decode(
142 &self,
143 buf: &[u8],
144 inputs: &[LogicalPlan],
145 ctx: &TaskContext,
146 ) -> Result<Extension>;
147
148 fn try_encode(&self, node: &Extension, buf: &mut Vec<u8>) -> Result<()>;
149
150 fn try_decode_table_provider(
151 &self,
152 buf: &[u8],
153 table_ref: &TableReference,
154 schema: SchemaRef,
155 ctx: &TaskContext,
156 ) -> Result<Arc<dyn TableProvider>>;
157
158 fn try_encode_table_provider(
159 &self,
160 table_ref: &TableReference,
161 node: Arc<dyn TableProvider>,
162 buf: &mut Vec<u8>,
163 ) -> Result<()>;
164
165 fn try_decode_file_format(
166 &self,
167 _buf: &[u8],
168 _ctx: &TaskContext,
169 ) -> Result<Arc<dyn FileFormatFactory>> {
170 not_impl_err!("LogicalExtensionCodec is not provided for file format")
171 }
172
173 fn try_encode_file_format(
174 &self,
175 _buf: &mut Vec<u8>,
176 _node: Arc<dyn FileFormatFactory>,
177 ) -> Result<()> {
178 Ok(())
179 }
180
181 fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result<Arc<ScalarUDF>> {
182 not_impl_err!("LogicalExtensionCodec is not provided for scalar function {name}")
183 }
184
185 fn try_encode_udf(&self, _node: &ScalarUDF, _buf: &mut Vec<u8>) -> Result<()> {
186 Ok(())
187 }
188
189 fn try_decode_higher_order_function(
190 &self,
191 name: &str,
192 _buf: &[u8],
193 ) -> Result<Arc<HigherOrderUDF>> {
194 not_impl_err!(
195 "LogicalExtensionCodec is not provided for higher order function {name}"
196 )
197 }
198
199 fn try_encode_higher_order_function(
200 &self,
201 _node: &HigherOrderUDF,
202 _buf: &mut Vec<u8>,
203 ) -> Result<()> {
204 Ok(())
205 }
206
207 fn try_decode_udaf(&self, name: &str, _buf: &[u8]) -> Result<Arc<AggregateUDF>> {
208 not_impl_err!(
209 "LogicalExtensionCodec is not provided for aggregate function {name}"
210 )
211 }
212
213 fn try_encode_udaf(&self, _node: &AggregateUDF, _buf: &mut Vec<u8>) -> Result<()> {
214 Ok(())
215 }
216
217 fn try_decode_udwf(&self, name: &str, _buf: &[u8]) -> Result<Arc<WindowUDF>> {
218 not_impl_err!("LogicalExtensionCodec is not provided for window function {name}")
219 }
220
221 fn try_encode_udwf(&self, _node: &WindowUDF, _buf: &mut Vec<u8>) -> Result<()> {
222 Ok(())
223 }
224}
225
226#[derive(Debug, Clone)]
227pub struct DefaultLogicalExtensionCodec {}
228
229impl LogicalExtensionCodec for DefaultLogicalExtensionCodec {
230 fn try_decode(
231 &self,
232 _buf: &[u8],
233 _inputs: &[LogicalPlan],
234 _ctx: &TaskContext,
235 ) -> Result<Extension> {
236 not_impl_err!("LogicalExtensionCodec is not provided")
237 }
238
239 fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
240 not_impl_err!("LogicalExtensionCodec is not provided")
241 }
242
243 fn try_decode_table_provider(
244 &self,
245 _buf: &[u8],
246 _table_ref: &TableReference,
247 _schema: SchemaRef,
248 _ctx: &TaskContext,
249 ) -> Result<Arc<dyn TableProvider>> {
250 not_impl_err!("LogicalExtensionCodec is not provided")
251 }
252
253 fn try_encode_table_provider(
254 &self,
255 _table_ref: &TableReference,
256 _node: Arc<dyn TableProvider>,
257 _buf: &mut Vec<u8>,
258 ) -> Result<()> {
259 not_impl_err!("LogicalExtensionCodec is not provided")
260 }
261
262 fn try_decode_file_format(
263 &self,
264 buf: &[u8],
265 ctx: &TaskContext,
266 ) -> Result<Arc<dyn FileFormatFactory>> {
267 let proto = protobuf::FileFormatProto::decode(buf).map_err(|e| {
268 internal_datafusion_err!("Failed to decode FileFormatProto: {e}")
269 })?;
270
271 let kind = protobuf::FileFormatKind::try_from(proto.kind).map_err(|_| {
272 internal_datafusion_err!("Unknown FileFormatKind: {}", proto.kind)
273 })?;
274
275 match kind {
276 protobuf::FileFormatKind::Csv => file_formats::CsvLogicalExtensionCodec
277 .try_decode_file_format(&proto.encoded_file_format, ctx),
278 protobuf::FileFormatKind::Json => file_formats::JsonLogicalExtensionCodec
279 .try_decode_file_format(&proto.encoded_file_format, ctx),
280 #[cfg(feature = "parquet")]
281 protobuf::FileFormatKind::Parquet => {
282 file_formats::ParquetLogicalExtensionCodec
283 .try_decode_file_format(&proto.encoded_file_format, ctx)
284 }
285 protobuf::FileFormatKind::Arrow => file_formats::ArrowLogicalExtensionCodec
286 .try_decode_file_format(&proto.encoded_file_format, ctx),
287 protobuf::FileFormatKind::Avro => file_formats::AvroLogicalExtensionCodec
288 .try_decode_file_format(&proto.encoded_file_format, ctx),
289 #[cfg(not(feature = "parquet"))]
290 protobuf::FileFormatKind::Parquet => {
291 not_impl_err!("Parquet support requires the 'parquet' feature")
292 }
293 protobuf::FileFormatKind::Unspecified => {
294 not_impl_err!("Unspecified file format kind")
295 }
296 }
297 }
298
299 fn try_encode_file_format(
300 &self,
301 buf: &mut Vec<u8>,
302 node: Arc<dyn FileFormatFactory>,
303 ) -> Result<()> {
304 let mut encoded_file_format = Vec::new();
305
306 let kind = if node.downcast_ref::<CsvFormatFactory>().is_some() {
307 file_formats::CsvLogicalExtensionCodec
308 .try_encode_file_format(&mut encoded_file_format, Arc::clone(&node))?;
309 protobuf::FileFormatKind::Csv
310 } else if node.downcast_ref::<JsonFormatFactory>().is_some() {
311 file_formats::JsonLogicalExtensionCodec
312 .try_encode_file_format(&mut encoded_file_format, Arc::clone(&node))?;
313 protobuf::FileFormatKind::Json
314 } else if node.downcast_ref::<ArrowFormatFactory>().is_some() {
315 file_formats::ArrowLogicalExtensionCodec
316 .try_encode_file_format(&mut encoded_file_format, Arc::clone(&node))?;
317 protobuf::FileFormatKind::Arrow
318 } else {
319 #[cfg(feature = "parquet")]
320 {
321 if node.downcast_ref::<ParquetFormatFactory>().is_some() {
322 file_formats::ParquetLogicalExtensionCodec.try_encode_file_format(
323 &mut encoded_file_format,
324 Arc::clone(&node),
325 )?;
326 protobuf::FileFormatKind::Parquet
327 } else {
328 return not_impl_err!(
329 "Unsupported FileFormatFactory type for DefaultLogicalExtensionCodec"
330 );
331 }
332 }
333 #[cfg(not(feature = "parquet"))]
334 {
335 return not_impl_err!(
336 "Unsupported FileFormatFactory type for DefaultLogicalExtensionCodec"
337 );
338 }
339 };
340
341 let proto = protobuf::FileFormatProto {
342 kind: kind as i32,
343 encoded_file_format,
344 };
345 proto.encode(buf).map_err(|e| {
346 internal_datafusion_err!("Failed to encode FileFormatProto: {e}")
347 })?;
348 Ok(())
349 }
350}
351
352#[macro_export]
353macro_rules! into_logical_plan {
354 ($PB:expr, $CTX:expr, $CODEC:expr) => {{
355 if let Some(field) = $PB.as_ref() {
356 field.as_ref().try_into_logical_plan($CTX, $CODEC)
357 } else {
358 Err(proto_error("Missing required field in protobuf"))
359 }
360 }};
361}
362
363fn from_table_reference(
364 table_ref: Option<&protobuf::TableReference>,
365 error_context: &str,
366) -> Result<TableReference> {
367 let table_ref = table_ref.ok_or_else(|| {
368 internal_datafusion_err!(
369 "Protobuf deserialization error, {error_context} was missing required field name."
370 )
371 })?;
372
373 Ok(TableReference::try_from(table_ref.clone())?)
374}
375
376fn to_table_source(
380 node: &Option<Box<LogicalPlanNode>>,
381 ctx: &TaskContext,
382 extension_codec: &dyn LogicalExtensionCodec,
383) -> Result<Arc<dyn TableSource>> {
384 if let Some(node) = node {
385 match node.try_into_logical_plan(ctx, extension_codec)? {
386 LogicalPlan::TableScan(TableScan { source, .. }) => Ok(source),
387 _ => plan_err!("expected TableScan node"),
388 }
389 } else {
390 plan_err!("LogicalPlanNode should be provided")
391 }
392}
393
394fn from_table_source(
398 table_name: TableReference,
399 target: Arc<dyn TableSource>,
400 extension_codec: &dyn LogicalExtensionCodec,
401) -> Result<LogicalPlanNode> {
402 let r = LogicalPlan::TableScan(TableScanBuilder::new(table_name, target).build()?);
403
404 LogicalPlanNode::try_from_logical_plan(&r, extension_codec)
405}
406
407fn metric_type_from_proto(value: i32) -> Result<MetricType> {
408 let pb = protobuf_common::MetricType::try_from(value)
409 .map_err(|_| proto_error(format!("Unknown MetricType discriminant: {value}")))?;
410 Ok(match pb {
411 protobuf_common::MetricType::Summary => MetricType::Summary,
412 protobuf_common::MetricType::Dev => MetricType::Dev,
413 })
414}
415
416fn metric_type_to_proto(value: MetricType) -> protobuf_common::MetricType {
417 match value {
418 MetricType::Summary => protobuf_common::MetricType::Summary,
419 MetricType::Dev => protobuf_common::MetricType::Dev,
420 }
421}
422
423fn metric_category_from_proto(value: i32) -> Result<MetricCategory> {
424 let pb = protobuf_common::MetricCategory::try_from(value).map_err(|_| {
425 proto_error(format!("Unknown MetricCategory discriminant: {value}"))
426 })?;
427 Ok(match pb {
428 protobuf_common::MetricCategory::Rows => MetricCategory::Rows,
429 protobuf_common::MetricCategory::Bytes => MetricCategory::Bytes,
430 protobuf_common::MetricCategory::Timing => MetricCategory::Timing,
431 protobuf_common::MetricCategory::Uncategorized => MetricCategory::Uncategorized,
432 })
433}
434
435fn metric_category_to_proto(value: MetricCategory) -> protobuf_common::MetricCategory {
436 match value {
437 MetricCategory::Rows => protobuf_common::MetricCategory::Rows,
438 MetricCategory::Bytes => protobuf_common::MetricCategory::Bytes,
439 MetricCategory::Timing => protobuf_common::MetricCategory::Timing,
440 MetricCategory::Uncategorized => protobuf_common::MetricCategory::Uncategorized,
441 }
442}
443
444fn explain_analyze_categories_from_proto(
445 node: &protobuf_common::ExplainAnalyzeCategoriesNode,
446) -> Result<ExplainAnalyzeCategories> {
447 if node.all {
448 Ok(ExplainAnalyzeCategories::All)
449 } else {
450 let cats = node
451 .only
452 .iter()
453 .copied()
454 .map(metric_category_from_proto)
455 .collect::<Result<Vec<_>>>()?;
456 Ok(ExplainAnalyzeCategories::Only(cats))
457 }
458}
459
460fn explain_analyze_categories_to_proto(
461 value: &ExplainAnalyzeCategories,
462) -> protobuf_common::ExplainAnalyzeCategoriesNode {
463 match value {
464 ExplainAnalyzeCategories::All => protobuf_common::ExplainAnalyzeCategoriesNode {
465 all: true,
466 only: vec![],
467 },
468 ExplainAnalyzeCategories::Only(cats) => {
469 protobuf_common::ExplainAnalyzeCategoriesNode {
470 all: false,
471 only: cats
472 .iter()
473 .copied()
474 .map(|c| metric_category_to_proto(c) as i32)
475 .collect(),
476 }
477 }
478 }
479}
480
481impl AsLogicalPlan for LogicalPlanNode {
482 fn try_decode(buf: &[u8]) -> Result<Self>
483 where
484 Self: Sized,
485 {
486 LogicalPlanNode::decode(buf)
487 .map_err(|e| internal_datafusion_err!("failed to decode logical plan: {e:?}"))
488 }
489
490 fn try_encode<B>(&self, buf: &mut B) -> Result<()>
491 where
492 B: BufMut,
493 Self: Sized,
494 {
495 self.encode(buf)
496 .map_err(|e| internal_datafusion_err!("failed to encode logical plan: {e:?}"))
497 }
498
499 fn try_into_logical_plan(
500 &self,
501 ctx: &TaskContext,
502 extension_codec: &dyn LogicalExtensionCodec,
503 ) -> Result<LogicalPlan> {
504 let plan = self.logical_plan_type.as_ref().ok_or_else(|| {
505 proto_error(format!(
506 "logical_plan::from_proto() Unsupported logical plan '{self:?}'"
507 ))
508 })?;
509 match plan {
510 LogicalPlanType::Values(values) => {
511 let n_cols = values.n_cols as usize;
512 let values: Vec<Vec<Expr>> = if values.values_list.is_empty() {
513 Ok(Vec::new())
514 } else if values.values_list.len() % n_cols != 0 {
515 internal_err!(
516 "Invalid values list length, expect {} to be divisible by {}",
517 values.values_list.len(),
518 n_cols
519 )
520 } else {
521 values
522 .values_list
523 .chunks_exact(n_cols)
524 .map(|r| from_proto::parse_exprs(r, ctx, extension_codec))
525 .collect::<Result<Vec<_>, _>>()
526 .map_err(|e| e.into())
527 }?;
528
529 LogicalPlanBuilder::values(values)?.build()
530 }
531 LogicalPlanType::Projection(projection) => {
532 let input: LogicalPlan =
533 into_logical_plan!(projection.input, ctx, extension_codec)?;
534 let expr: Vec<Expr> =
535 from_proto::parse_exprs(&projection.expr, ctx, extension_codec)?;
536
537 let new_proj = project(input, expr)?;
538 match projection.optional_alias.as_ref() {
539 Some(a) => match a {
540 protobuf::projection_node::OptionalAlias::Alias(alias) => {
541 Ok(LogicalPlan::SubqueryAlias(SubqueryAlias::try_new(
542 Arc::new(new_proj),
543 alias.clone(),
544 )?))
545 }
546 },
547 _ => Ok(new_proj),
548 }
549 }
550 LogicalPlanType::Selection(selection) => {
551 let input: LogicalPlan =
552 into_logical_plan!(selection.input, ctx, extension_codec)?;
553 let expr: Expr = selection
554 .expr
555 .as_ref()
556 .map(|expr| from_proto::parse_expr(expr, ctx, extension_codec))
557 .transpose()?
558 .ok_or_else(|| proto_error("expression required"))?;
559 LogicalPlanBuilder::from(input).filter(expr)?.build()
560 }
561 LogicalPlanType::Window(window) => {
562 let input: LogicalPlan =
563 into_logical_plan!(window.input, ctx, extension_codec)?;
564 let window_expr =
565 from_proto::parse_exprs(&window.window_expr, ctx, extension_codec)?;
566 LogicalPlanBuilder::from(input).window(window_expr)?.build()
567 }
568 LogicalPlanType::Aggregate(aggregate) => {
569 let input: LogicalPlan =
570 into_logical_plan!(aggregate.input, ctx, extension_codec)?;
571 let group_expr =
572 from_proto::parse_exprs(&aggregate.group_expr, ctx, extension_codec)?;
573 let aggr_expr =
574 from_proto::parse_exprs(&aggregate.aggr_expr, ctx, extension_codec)?;
575 LogicalPlanBuilder::from(input)
576 .aggregate(group_expr, aggr_expr)?
577 .build()
578 }
579 LogicalPlanType::ListingScan(scan) => {
580 let schema: Schema = convert_required!(scan.schema)?;
581
582 let filters =
583 from_proto::parse_exprs(&scan.filters, ctx, extension_codec)?;
584
585 let mut all_sort_orders = vec![];
586 for order in &scan.file_sort_order {
587 all_sort_orders.push(from_proto::parse_sorts(
588 &order.sort_expr_nodes,
589 ctx,
590 extension_codec,
591 )?)
592 }
593
594 let file_format: Arc<dyn FileFormat> =
595 match scan.file_format_type.as_ref().ok_or_else(|| {
596 proto_error(format!(
597 "logical_plan::from_proto() Unsupported file format '{self:?}'"
598 ))
599 })? {
600 #[cfg_attr(not(feature = "parquet"), allow(unused_variables))]
601 FileFormatType::Parquet(protobuf::ParquetFormat {options}) => {
602 #[cfg(feature = "parquet")]
603 {
604 let mut parquet = ParquetFormat::default();
605 if let Some(options) = options {
606 parquet = parquet.with_options(options.try_into()?)
607 }
608 Arc::new(parquet)
609 }
610 #[cfg(not(feature = "parquet"))]
611 panic!("Unable to process parquet file since `parquet` feature is not enabled");
612 }
613 FileFormatType::Csv(protobuf::CsvFormat {
614 options
615 }) => {
616 let mut csv = CsvFormat::default();
617 if let Some(options) = options {
618 csv = csv.with_options(options.try_into()?)
619 }
620 Arc::new(csv)
621 },
622 FileFormatType::Json(protobuf::NdJsonFormat {
623 options
624 }) => {
625 let mut json = OtherNdJsonFormat::default();
626 if let Some(options) = options {
627 json = json.with_options(options.try_into()?)
628 }
629 Arc::new(json)
630 }
631 FileFormatType::Avro(..) => {
632 #[cfg(feature = "avro")]
633 {
634 Arc::new(AvroFormat)
635 }
636 #[cfg(not(feature = "avro"))]
637 {
638 panic!(
639 "Unable to process avro file since `avro` feature is not enabled"
640 );
641 }
642 }
643 FileFormatType::Arrow(..) => {
644 Arc::new(ArrowFormat)
645 }
646 };
647
648 let table_paths = &scan
649 .paths
650 .iter()
651 .map(ListingTableUrl::parse)
652 .collect::<Result<Vec<_>, _>>()?;
653
654 let partition_columns = scan
655 .table_partition_cols
656 .iter()
657 .map(|col| {
658 let Some(arrow_type) = col.arrow_type.as_ref() else {
659 return Err(proto_error(
660 "Missing Arrow type in partition columns",
661 ));
662 };
663 let arrow_type = DataType::try_from(arrow_type).map_err(|e| {
664 proto_error(format!("Received an unknown ArrowType: {e}"))
665 })?;
666 Ok((col.name.clone(), arrow_type))
667 })
668 .collect::<Result<Vec<_>>>()?;
669
670 let options = ListingOptions::new(file_format)
671 .with_file_extension(&scan.file_extension)
672 .with_table_partition_cols(partition_columns)
673 .with_file_sort_order(all_sort_orders);
674
675 let config =
676 ListingTableConfig::new_with_multi_paths(table_paths.clone())
677 .with_listing_options(options)
678 .with_schema(Arc::new(schema));
679
680 let provider = ListingTable::try_new(config)?.with_cache(
681 ctx.runtime_env().cache_manager.get_file_statistic_cache(),
682 );
683
684 let table_name =
685 from_table_reference(scan.table_name.as_ref(), "ListingTableScan")?;
686
687 let mut projection = None;
688 if let Some(columns) = &scan.projection {
689 let column_indices = columns
690 .columns
691 .iter()
692 .map(|name| provider.schema().index_of(name))
693 .collect::<Result<Vec<usize>, _>>()?;
694 projection = Some(column_indices);
695 }
696
697 LogicalPlanBuilder::scan_with_filters(
698 table_name,
699 provider_as_source(Arc::new(provider)),
700 projection,
701 filters,
702 )?
703 .build()
704 }
705 CustomScan(scan) => {
706 let schema: Schema = convert_required!(scan.schema)?;
707 let schema = Arc::new(schema);
708 let mut projection = None;
709 if let Some(columns) = &scan.projection {
710 let column_indices = columns
711 .columns
712 .iter()
713 .map(|name| schema.index_of(name))
714 .collect::<Result<Vec<usize>, _>>()?;
715 projection = Some(column_indices);
716 }
717
718 let filters =
719 from_proto::parse_exprs(&scan.filters, ctx, extension_codec)?;
720
721 let table_name =
722 from_table_reference(scan.table_name.as_ref(), "CustomScan")?;
723
724 let provider = extension_codec.try_decode_table_provider(
725 &scan.custom_table_data,
726 &table_name,
727 schema,
728 ctx,
729 )?;
730
731 LogicalPlanBuilder::scan_with_filters(
732 table_name,
733 provider_as_source(provider),
734 projection,
735 filters,
736 )?
737 .build()
738 }
739 LogicalPlanType::Sort(sort) => {
740 let input: LogicalPlan =
741 into_logical_plan!(sort.input, ctx, extension_codec)?;
742 let sort_expr: Vec<SortExpr> =
743 from_proto::parse_sorts(&sort.expr, ctx, extension_codec)?;
744 let fetch: Option<usize> = sort.fetch.try_into().ok();
745 LogicalPlanBuilder::from(input)
746 .sort_with_limit(sort_expr, fetch)?
747 .build()
748 }
749 LogicalPlanType::Repartition(repartition) => {
750 use datafusion_expr::Partitioning;
751 let input: LogicalPlan =
752 into_logical_plan!(repartition.input, ctx, extension_codec)?;
753 use protobuf::repartition_node::PartitionMethod;
754 let pb_partition_method = repartition.partition_method.as_ref().ok_or_else(|| {
755 internal_datafusion_err!(
756 "Protobuf deserialization error, RepartitionNode was missing required field 'partition_method'"
757 )
758 })?;
759
760 let partitioning_scheme = match pb_partition_method {
761 PartitionMethod::Hash(protobuf::HashRepartition {
762 hash_expr: pb_hash_expr,
763 partition_count,
764 }) => Partitioning::Hash(
765 from_proto::parse_exprs(pb_hash_expr, ctx, extension_codec)?,
766 *partition_count as usize,
767 ),
768 PartitionMethod::RoundRobin(partition_count) => {
769 Partitioning::RoundRobinBatch(*partition_count as usize)
770 }
771 PartitionMethod::Range(protobuf::RangeRepartition {
772 sort_expr: pb_sort_expr,
773 split_point,
774 }) => Partitioning::Range(RangePartitioning::try_new(
775 from_proto::parse_sorts(pb_sort_expr, ctx, extension_codec)?,
776 split_point
777 .iter()
778 .map(from_proto::parse_protobuf_range_split_point)
779 .collect::<Result<Vec<_>, _>>()?,
780 )?),
781 };
782
783 LogicalPlanBuilder::from(input)
784 .repartition(partitioning_scheme)?
785 .build()
786 }
787 LogicalPlanType::EmptyRelation(empty_relation) => {
788 LogicalPlanBuilder::empty(empty_relation.produce_one_row).build()
789 }
790 LogicalPlanType::CreateExternalTable(create_extern_table) => {
791 let pb_schema = (create_extern_table.schema.clone()).ok_or_else(|| {
792 internal_datafusion_err!(
793 "Protobuf deserialization error, CreateExternalTableNode was missing required field schema."
794 )
795 })?;
796
797 let constraints = (create_extern_table.constraints.clone()).ok_or_else(|| {
798 internal_datafusion_err!(
799 "Protobuf deserialization error, CreateExternalTableNode was missing required table constraints."
800 )
801 })?;
802 let definition = if !create_extern_table.definition.is_empty() {
803 Some(create_extern_table.definition.clone())
804 } else {
805 None
806 };
807
808 let mut order_exprs = vec![];
809 for expr in &create_extern_table.order_exprs {
810 order_exprs.push(from_proto::parse_sorts(
811 &expr.sort_expr_nodes,
812 ctx,
813 extension_codec,
814 )?);
815 }
816
817 let mut column_defaults =
818 HashMap::with_capacity(create_extern_table.column_defaults.len());
819 for (col_name, expr) in &create_extern_table.column_defaults {
820 let expr = from_proto::parse_expr(expr, ctx, extension_codec)?;
821 column_defaults.insert(col_name.clone(), expr);
822 }
823
824 let locations = if !create_extern_table.locations.is_empty() {
825 create_extern_table.locations.clone()
826 } else if !create_extern_table.location.is_empty() {
827 vec![create_extern_table.location.clone()]
828 } else {
829 return Err(proto_error(
830 "CreateExternalTableNode requires at least one location",
831 ));
832 };
833 let location = locations[0].clone();
834
835 Ok(LogicalPlan::Ddl(DdlStatement::CreateExternalTable(
836 Box::new(
837 CreateExternalTable::builder(
838 from_table_reference(
839 create_extern_table.name.as_ref(),
840 "CreateExternalTable",
841 )?,
842 location,
843 create_extern_table.file_type.clone(),
844 pb_schema.try_into()?,
845 )
846 .with_locations(locations)
847 .with_partition_cols(
848 create_extern_table.table_partition_cols.clone(),
849 )
850 .with_order_exprs(order_exprs)
851 .with_if_not_exists(create_extern_table.if_not_exists)
852 .with_or_replace(create_extern_table.or_replace)
853 .with_temporary(create_extern_table.temporary)
854 .with_definition(definition)
855 .with_unbounded(create_extern_table.unbounded)
856 .with_options(create_extern_table.options.clone())
857 .with_constraints(constraints.into())
858 .with_column_defaults(column_defaults)
859 .build(),
860 ),
861 )))
862 }
863 LogicalPlanType::CreateView(create_view) => {
864 let plan = create_view
865 .input.clone().ok_or_else(|| internal_datafusion_err!(
866 "Protobuf deserialization error, CreateViewNode has invalid LogicalPlan input."
867 ))?
868 .try_into_logical_plan(ctx, extension_codec)?;
869 let definition = if !create_view.definition.is_empty() {
870 Some(create_view.definition.clone())
871 } else {
872 None
873 };
874
875 Ok(LogicalPlan::Ddl(DdlStatement::CreateView(CreateView {
876 name: from_table_reference(create_view.name.as_ref(), "CreateView")?,
877 temporary: create_view.temporary,
878 input: Arc::new(plan),
879 or_replace: create_view.or_replace,
880 definition,
881 })))
882 }
883 LogicalPlanType::CreateCatalogSchema(create_catalog_schema) => {
884 let pb_schema = (create_catalog_schema.schema.clone()).ok_or_else(|| {
885 internal_datafusion_err!(
886 "Protobuf deserialization error, CreateCatalogSchemaNode was missing required field schema."
887 )
888 })?;
889
890 Ok(LogicalPlan::Ddl(DdlStatement::CreateCatalogSchema(
891 CreateCatalogSchema {
892 schema_name: create_catalog_schema.schema_name.clone(),
893 if_not_exists: create_catalog_schema.if_not_exists,
894 schema: pb_schema.try_into()?,
895 },
896 )))
897 }
898 LogicalPlanType::CreateCatalog(create_catalog) => {
899 let pb_schema = (create_catalog.schema.clone()).ok_or_else(|| {
900 internal_datafusion_err!(
901 "Protobuf deserialization error, CreateCatalogNode was missing required field schema."
902 )
903 })?;
904
905 Ok(LogicalPlan::Ddl(DdlStatement::CreateCatalog(
906 CreateCatalog {
907 catalog_name: create_catalog.catalog_name.clone(),
908 if_not_exists: create_catalog.if_not_exists,
909 schema: pb_schema.try_into()?,
910 },
911 )))
912 }
913 LogicalPlanType::Analyze(analyze) => {
914 let input: LogicalPlan =
915 into_logical_plan!(analyze.input, ctx, extension_codec)?;
916 let analyze_level = analyze
917 .analyze_level
918 .map(metric_type_from_proto)
919 .transpose()?;
920 let analyze_categories = analyze
921 .analyze_categories
922 .as_ref()
923 .map(explain_analyze_categories_from_proto)
924 .transpose()?;
925 let pb_format = protobuf::ExplainFormat::try_from(analyze.format)
926 .map_err(|_| {
927 proto_error(format!(
928 "Received an AnalyzeNode message with unknown ExplainFormat {}",
929 analyze.format
930 ))
931 })?;
932 let analyze_format = match pb_format {
933 protobuf::ExplainFormat::Indent => ExplainFormat::Indent,
934 protobuf::ExplainFormat::Tree => ExplainFormat::Tree,
935 protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON,
936 protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz,
937 };
938 let explain_option =
939 datafusion_expr::logical_plan::ExplainOption::default()
940 .with_verbose(analyze.verbose)
941 .with_analyze(true)
942 .with_analyze_level(analyze_level)
943 .with_analyze_categories(analyze_categories)
944 .with_format(analyze_format);
945 LogicalPlanBuilder::from(input)
946 .explain_option_format(explain_option)?
947 .build()
948 }
949 LogicalPlanType::Explain(explain) => {
950 let input: LogicalPlan =
951 into_logical_plan!(explain.input, ctx, extension_codec)?;
952 let pb_format = protobuf::ExplainFormat::try_from(explain.format)
953 .map_err(|_| {
954 proto_error(format!(
955 "Received an ExplainNode message with unknown ExplainFormat {}",
956 explain.format
957 ))
958 })?;
959 let explain_format = match pb_format {
960 protobuf::ExplainFormat::Indent => ExplainFormat::Indent,
961 protobuf::ExplainFormat::Tree => ExplainFormat::Tree,
962 protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON,
963 protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz,
964 };
965 let explain_option =
966 datafusion_expr::logical_plan::ExplainOption::default()
967 .with_verbose(explain.verbose)
968 .with_format(explain_format)
969 .with_show_statistics(explain.show_statistics);
970 LogicalPlanBuilder::from(input)
971 .explain_option_format(explain_option)?
972 .build()
973 }
974 LogicalPlanType::SubqueryAlias(aliased_relation) => {
975 let input: LogicalPlan =
976 into_logical_plan!(aliased_relation.input, ctx, extension_codec)?;
977 let alias = from_table_reference(
978 aliased_relation.alias.as_ref(),
979 "SubqueryAlias",
980 )?;
981 LogicalPlanBuilder::from(input).alias(alias)?.build()
982 }
983 LogicalPlanType::Limit(limit) => {
984 let input: LogicalPlan =
985 into_logical_plan!(limit.input, ctx, extension_codec)?;
986 let skip = limit.skip.max(0) as usize;
987
988 let fetch = if limit.fetch < 0 {
989 None
990 } else {
991 Some(limit.fetch as usize)
992 };
993
994 LogicalPlanBuilder::from(input).limit(skip, fetch)?.build()
995 }
996 LogicalPlanType::Join(join) => {
997 let left_keys: Vec<Expr> =
998 from_proto::parse_exprs(&join.left_join_key, ctx, extension_codec)?;
999 let right_keys: Vec<Expr> =
1000 from_proto::parse_exprs(&join.right_join_key, ctx, extension_codec)?;
1001 if left_keys.len() != right_keys.len() {
1002 return Err(proto_error(format!(
1003 "Received a JoinNode message with left_join_key and right_join_key of different lengths: {} and {}",
1004 left_keys.len(),
1005 right_keys.len()
1006 )));
1007 }
1008 let join_type =
1009 protobuf::JoinType::try_from(join.join_type).map_err(|_| {
1010 proto_error(format!(
1011 "Received a JoinNode message with unknown JoinType {}",
1012 join.join_type
1013 ))
1014 })?;
1015 let join_constraint = protobuf::JoinConstraint::try_from(
1016 join.join_constraint,
1017 )
1018 .map_err(|_| {
1019 proto_error(format!(
1020 "Received a JoinNode message with unknown JoinConstraint {}",
1021 join.join_constraint
1022 ))
1023 })?;
1024 let null_equality = protobuf::NullEquality::try_from(join.null_equality)
1025 .map_err(|_| {
1026 proto_error(format!(
1027 "Received a JoinNode message with unknown NullEquality {}",
1028 join.null_equality
1029 ))
1030 })?;
1031 let filter: Option<Expr> = join
1032 .filter
1033 .as_ref()
1034 .map(|expr| from_proto::parse_expr(expr, ctx, extension_codec))
1035 .map_or(Ok(None), |v| v.map(Some))?;
1036 let left = into_logical_plan!(join.left, ctx, extension_codec)?;
1037 let right = into_logical_plan!(join.right, ctx, extension_codec)?;
1038 let on: Vec<(Expr, Expr)> =
1039 left_keys.into_iter().zip(right_keys).collect();
1040
1041 Ok(LogicalPlan::Join(Join::try_new(
1048 Arc::new(left),
1049 Arc::new(right),
1050 on,
1051 filter,
1052 datafusion_expr::JoinType::from(join_type),
1053 JoinConstraint::from(join_constraint),
1054 NullEquality::from(null_equality),
1055 join.null_aware,
1056 )?))
1057 }
1058 LogicalPlanType::Union(union) => {
1059 assert_or_internal_err!(
1060 union.inputs.len() >= 2,
1061 "Protobuf deserialization error, Union requires at least two inputs."
1062 );
1063 let (first, rest) = union.inputs.split_first().unwrap();
1064 let mut builder = LogicalPlanBuilder::from(
1065 first.try_into_logical_plan(ctx, extension_codec)?,
1066 );
1067
1068 for i in rest {
1069 let plan = i.try_into_logical_plan(ctx, extension_codec)?;
1070 builder = builder.union(plan)?;
1071 }
1072 builder.build()
1073 }
1074 LogicalPlanType::CrossJoin(crossjoin) => {
1075 let left = into_logical_plan!(crossjoin.left, ctx, extension_codec)?;
1076 let right = into_logical_plan!(crossjoin.right, ctx, extension_codec)?;
1077
1078 LogicalPlanBuilder::from(left).cross_join(right)?.build()
1079 }
1080 LogicalPlanType::Extension(LogicalExtensionNode { node, inputs }) => {
1081 let input_plans: Vec<LogicalPlan> = inputs
1082 .iter()
1083 .map(|i| i.try_into_logical_plan(ctx, extension_codec))
1084 .collect::<Result<_>>()?;
1085
1086 let extension_node =
1087 extension_codec.try_decode(node, &input_plans, ctx)?;
1088 Ok(LogicalPlan::Extension(extension_node))
1089 }
1090 LogicalPlanType::Distinct(distinct) => {
1091 let input: LogicalPlan =
1092 into_logical_plan!(distinct.input, ctx, extension_codec)?;
1093 LogicalPlanBuilder::from(input).distinct()?.build()
1094 }
1095 LogicalPlanType::DistinctOn(distinct_on) => {
1096 let input: LogicalPlan =
1097 into_logical_plan!(distinct_on.input, ctx, extension_codec)?;
1098 let on_expr =
1099 from_proto::parse_exprs(&distinct_on.on_expr, ctx, extension_codec)?;
1100 let select_expr = from_proto::parse_exprs(
1101 &distinct_on.select_expr,
1102 ctx,
1103 extension_codec,
1104 )?;
1105 let sort_expr = match distinct_on.sort_expr.len() {
1106 0 => None,
1107 _ => Some(from_proto::parse_sorts(
1108 &distinct_on.sort_expr,
1109 ctx,
1110 extension_codec,
1111 )?),
1112 };
1113 LogicalPlanBuilder::from(input)
1114 .distinct_on(on_expr, select_expr, sort_expr)?
1115 .build()
1116 }
1117 LogicalPlanType::ViewScan(scan) => {
1118 let schema: Schema = convert_required!(scan.schema)?;
1119
1120 let mut projection = None;
1121 if let Some(columns) = &scan.projection {
1122 let column_indices = columns
1123 .columns
1124 .iter()
1125 .map(|name| schema.index_of(name))
1126 .collect::<Result<Vec<usize>, _>>()?;
1127 projection = Some(column_indices);
1128 }
1129
1130 let input: LogicalPlan =
1131 into_logical_plan!(scan.input, ctx, extension_codec)?;
1132
1133 let definition = if !scan.definition.is_empty() {
1134 Some(scan.definition.clone())
1135 } else {
1136 None
1137 };
1138
1139 let provider = ViewTable::new(input, definition);
1140
1141 let table_name =
1142 from_table_reference(scan.table_name.as_ref(), "ViewScan")?;
1143
1144 LogicalPlanBuilder::scan(
1145 table_name,
1146 provider_as_source(Arc::new(provider)),
1147 projection,
1148 )?
1149 .build()
1150 }
1151 LogicalPlanType::Prepare(prepare) => {
1152 let input: LogicalPlan =
1153 into_logical_plan!(prepare.input, ctx, extension_codec)?;
1154 let data_types: Vec<DataType> = prepare
1155 .data_types
1156 .iter()
1157 .map(DataType::try_from)
1158 .collect::<Result<_, _>>()?;
1159 let fields: Vec<Field> = prepare
1160 .fields
1161 .iter()
1162 .map(Field::try_from)
1163 .collect::<Result<_, _>>()?;
1164
1165 if fields.is_empty() {
1169 LogicalPlanBuilder::from(input)
1170 .prepare(
1171 prepare.name.clone(),
1172 data_types
1173 .into_iter()
1174 .map(|dt| Field::new("", dt, true).into())
1175 .collect(),
1176 )?
1177 .build()
1178 } else {
1179 LogicalPlanBuilder::from(input)
1180 .prepare(
1181 prepare.name.clone(),
1182 fields.into_iter().map(|f| f.into()).collect(),
1183 )?
1184 .build()
1185 }
1186 }
1187 LogicalPlanType::DropView(dropview) => {
1188 Ok(LogicalPlan::Ddl(DdlStatement::DropView(DropView {
1189 name: from_table_reference(dropview.name.as_ref(), "DropView")?,
1190 if_exists: dropview.if_exists,
1191 schema: Arc::new(convert_required!(dropview.schema)?),
1192 })))
1193 }
1194 LogicalPlanType::CopyTo(copy) => {
1195 let input: LogicalPlan =
1196 into_logical_plan!(copy.input, ctx, extension_codec)?;
1197
1198 let file_type: Arc<dyn FileType> = format_as_file_type(
1199 extension_codec.try_decode_file_format(©.file_type, ctx)?,
1200 );
1201
1202 Ok(LogicalPlan::Copy(dml::CopyTo::new(
1203 Arc::new(input),
1204 copy.output_url.clone(),
1205 copy.partition_by.clone(),
1206 file_type,
1207 Default::default(),
1208 )))
1209 }
1210 LogicalPlanType::Unnest(unnest) => {
1211 let input: LogicalPlan =
1212 into_logical_plan!(unnest.input, ctx, extension_codec)?;
1213
1214 LogicalPlanBuilder::from(input)
1215 .unnest_columns_with_options(
1216 unnest.exec_columns.iter().map(|c| c.into()).collect(),
1217 unnest
1218 .options
1219 .as_ref()
1220 .map(datafusion_common::UnnestOptions::from)
1221 .ok_or_else(|| {
1222 proto_error("Missing required field in protobuf")
1223 })?,
1224 )?
1225 .build()
1226 }
1227 LogicalPlanType::RecursiveQuery(recursive_query_node) => {
1228 let static_term = recursive_query_node
1229 .static_term
1230 .as_ref()
1231 .ok_or_else(|| internal_datafusion_err!(
1232 "Protobuf deserialization error, RecursiveQueryNode was missing required field static_term."
1233 ))?
1234 .try_into_logical_plan(ctx, extension_codec)?;
1235
1236 let recursive_term = recursive_query_node
1237 .recursive_term
1238 .as_ref()
1239 .ok_or_else(|| internal_datafusion_err!(
1240 "Protobuf deserialization error, RecursiveQueryNode was missing required field recursive_term."
1241 ))?
1242 .try_into_logical_plan(ctx, extension_codec)?;
1243
1244 RecursiveQuery::try_new(
1247 recursive_query_node.name.clone(),
1248 Arc::new(static_term),
1249 Arc::new(recursive_term),
1250 recursive_query_node.is_distinct,
1251 )
1252 .map(LogicalPlan::RecursiveQuery)
1253 }
1254 LogicalPlanType::CteWorkTableScan(cte_work_table_scan_node) => {
1255 let CteWorkTableScanNode { name, schema } = cte_work_table_scan_node;
1256 let schema = convert_required!(*schema)?;
1257 let cte_work_table = CteWorkTable::new(name.as_str(), Arc::new(schema));
1258 LogicalPlanBuilder::scan(
1259 name.as_str(),
1260 provider_as_source(Arc::new(cte_work_table)),
1261 None,
1262 )?
1263 .build()
1264 }
1265 LogicalPlanType::EmptyTableScan(scan) => {
1266 let schema: Schema = convert_required!(scan.schema)?;
1267 let schema = Arc::new(schema);
1268 let mut projection = None;
1269 if let Some(columns) = &scan.projection {
1270 let column_indices = columns
1271 .columns
1272 .iter()
1273 .map(|name| schema.index_of(name))
1274 .collect::<Result<Vec<usize>, _>>()?;
1275 projection = Some(column_indices);
1276 }
1277
1278 let filters =
1279 from_proto::parse_exprs(&scan.filters, ctx, extension_codec)?;
1280
1281 let table_name =
1282 from_table_reference(scan.table_name.as_ref(), "EmptyTableScan")?;
1283
1284 let provider = Arc::new(EmptyTable::new(Arc::clone(&schema)));
1285
1286 LogicalPlanBuilder::scan_with_filters(
1287 table_name,
1288 provider_as_source(provider),
1289 projection,
1290 filters,
1291 )?
1292 .build()
1293 }
1294 LogicalPlanType::Dml(dml_node) => {
1295 let table_name =
1296 from_table_reference(dml_node.table_name.as_ref(), "DML ")?;
1297 let target = to_table_source(&dml_node.target, ctx, extension_codec)?;
1298 let write_op =
1299 from_proto::parse_write_op(dml_node, ctx, extension_codec)?;
1300 Ok(LogicalPlan::Dml(DmlStatement::new(
1301 table_name,
1302 target,
1303 write_op,
1304 Arc::new(into_logical_plan!(dml_node.input, ctx, extension_codec)?),
1305 )))
1306 }
1307 }
1308 }
1309
1310 #[cfg_attr(feature = "recursive_protection", recursive::recursive)]
1311 fn try_from_logical_plan(
1312 plan: &LogicalPlan,
1313 extension_codec: &dyn LogicalExtensionCodec,
1314 ) -> Result<Self>
1315 where
1316 Self: Sized,
1317 {
1318 dispatch_logical_plan!(plan, {
1319 LogicalPlan::Values(Values { values, .. }) => {
1320 let n_cols = if values.is_empty() {
1321 0
1322 } else {
1323 values[0].len()
1324 } as u64;
1325 let values_list =
1326 serialize_exprs(values.iter().flatten(), extension_codec)?;
1327 Ok(LogicalPlanNode {
1328 logical_plan_type: Some(LogicalPlanType::Values(
1329 protobuf::ValuesNode {
1330 n_cols,
1331 values_list,
1332 },
1333 )),
1334 })
1335 }
1336 LogicalPlan::TableScan(TableScan {
1337 table_name,
1338 source,
1339 filters,
1340 projection,
1341 ..
1342 }) => {
1343 let provider = source_as_provider(source)?;
1344 let schema = provider.schema();
1345
1346 let projection = match projection {
1347 None => None,
1348 Some(columns) => {
1349 let column_names = columns
1350 .iter()
1351 .map(|i| schema.field(*i).name().to_owned())
1352 .collect();
1353 Some(protobuf::ProjectionColumns {
1354 columns: column_names,
1355 })
1356 }
1357 };
1358
1359 let filters: Vec<protobuf::LogicalExprNode> =
1360 serialize_exprs(filters, extension_codec)?;
1361
1362 if let Some(listing_table) = provider.downcast_ref::<ListingTable>() {
1363 let format = listing_table.options().format.as_ref();
1364 let file_format_type = {
1365 let mut maybe_some_type = None;
1366
1367 #[cfg(feature = "parquet")]
1368 if let Some(parquet) = format.downcast_ref::<ParquetFormat>() {
1369 let options = parquet.options();
1370 maybe_some_type =
1371 Some(FileFormatType::Parquet(protobuf::ParquetFormat {
1372 options: Some(options.try_into()?),
1373 }));
1374 };
1375
1376 if let Some(csv) = format.downcast_ref::<CsvFormat>() {
1377 let options = csv.options();
1378 maybe_some_type =
1379 Some(FileFormatType::Csv(protobuf::CsvFormat {
1380 options: Some(options.try_into()?),
1381 }));
1382 }
1383
1384 if let Some(json) = format.downcast_ref::<OtherNdJsonFormat>() {
1385 let options = json.options();
1386 maybe_some_type =
1387 Some(FileFormatType::Json(protobuf::NdJsonFormat {
1388 options: Some(options.try_into()?),
1389 }))
1390 }
1391
1392 #[cfg(feature = "avro")]
1393 if format.is::<AvroFormat>() {
1394 maybe_some_type =
1395 Some(FileFormatType::Avro(protobuf::AvroFormat {}))
1396 }
1397
1398 if format.is::<ArrowFormat>() {
1399 maybe_some_type =
1400 Some(FileFormatType::Arrow(protobuf::ArrowFormat {}))
1401 }
1402
1403 if let Some(file_format_type) = maybe_some_type {
1404 file_format_type
1405 } else {
1406 return Err(proto_error(format!(
1407 "Error deserializing unknown file format: {:?}",
1408 listing_table.options().format
1409 )));
1410 }
1411 };
1412
1413 let options = listing_table.options();
1414
1415 let mut builder = SchemaBuilder::from(schema.as_ref());
1416 for (idx, field) in schema.fields().iter().enumerate().rev() {
1417 if options
1418 .table_partition_cols
1419 .iter()
1420 .any(|(name, _)| name == field.name())
1421 {
1422 builder.remove(idx);
1423 }
1424 }
1425
1426 let schema = builder.finish();
1427
1428 let schema: protobuf::Schema = (&schema).try_into()?;
1429
1430 let mut exprs_vec: Vec<SortExprNodeCollection> = vec![];
1431 for order in &options.file_sort_order {
1432 let expr_vec = SortExprNodeCollection {
1433 sort_expr_nodes: serialize_sorts(order, extension_codec)?,
1434 };
1435 exprs_vec.push(expr_vec);
1436 }
1437
1438 let partition_columns = options
1439 .table_partition_cols
1440 .iter()
1441 .map(|(name, arrow_type)| {
1442 let arrow_type = protobuf::ArrowType::try_from(arrow_type)
1443 .map_err(|e| {
1444 proto_error(format!(
1445 "Received an unknown ArrowType: {e}"
1446 ))
1447 })?;
1448 Ok(protobuf::PartitionColumn {
1449 name: name.clone(),
1450 arrow_type: Some(arrow_type),
1451 })
1452 })
1453 .collect::<Result<Vec<_>>>()?;
1454
1455 Ok(LogicalPlanNode {
1456 logical_plan_type: Some(LogicalPlanType::ListingScan(
1457 protobuf::ListingTableScanNode {
1458 file_format_type: Some(file_format_type),
1459 table_name: Some(protobuf::TableReference::from(
1460 table_name.clone(),
1461 )),
1462 file_extension: options.file_extension.clone(),
1463 table_partition_cols: partition_columns,
1464 paths: listing_table
1465 .table_paths()
1466 .iter()
1467 .map(|x| x.to_string())
1468 .collect(),
1469 schema: Some(schema),
1470 projection,
1471 filters,
1472 file_sort_order: exprs_vec,
1473 },
1474 )),
1475 })
1476 } else if let Some(view_table) = provider.downcast_ref::<ViewTable>() {
1477 let schema: protobuf::Schema = schema.as_ref().try_into()?;
1478 Ok(LogicalPlanNode {
1479 logical_plan_type: Some(LogicalPlanType::ViewScan(Box::new(
1480 protobuf::ViewTableScanNode {
1481 table_name: Some(protobuf::TableReference::from(
1482 table_name.clone(),
1483 )),
1484 input: Some(Box::new(
1485 LogicalPlanNode::try_from_logical_plan(
1486 view_table.logical_plan(),
1487 extension_codec,
1488 )?,
1489 )),
1490 schema: Some(schema),
1491 projection,
1492 definition: view_table
1493 .definition()
1494 .map(|s| s.to_string())
1495 .unwrap_or_default(),
1496 },
1497 ))),
1498 })
1499 } else if let Some(cte_work_table) =
1500 provider.downcast_ref::<CteWorkTable>()
1501 {
1502 let name = cte_work_table.name().to_string();
1503 let schema = cte_work_table.schema();
1504 let schema: protobuf::Schema = schema.as_ref().try_into()?;
1505
1506 Ok(LogicalPlanNode {
1507 logical_plan_type: Some(LogicalPlanType::CteWorkTableScan(
1508 CteWorkTableScanNode {
1509 name,
1510 schema: Some(schema),
1511 },
1512 )),
1513 })
1514 } else if provider.downcast_ref::<EmptyTable>().is_some() {
1515 let schema: protobuf::Schema = schema.as_ref().try_into()?;
1516
1517 Ok(LogicalPlanNode {
1518 logical_plan_type: Some(LogicalPlanType::EmptyTableScan(
1519 protobuf::EmptyTableScanNode {
1520 table_name: Some(protobuf::TableReference::from(
1521 table_name.clone(),
1522 )),
1523 schema: Some(schema),
1524 projection,
1525 filters,
1526 },
1527 )),
1528 })
1529 } else {
1530 let schema: protobuf::Schema = schema.as_ref().try_into()?;
1531 let mut bytes = vec![];
1532 extension_codec
1533 .try_encode_table_provider(table_name, provider, &mut bytes)
1534 .map_err(|e| context!("Error serializing custom table", e))?;
1535 let scan = CustomScan(CustomTableScanNode {
1536 table_name: Some(protobuf::TableReference::from(
1537 table_name.clone(),
1538 )),
1539 projection,
1540 schema: Some(schema),
1541 filters,
1542 custom_table_data: bytes,
1543 });
1544 let node = LogicalPlanNode {
1545 logical_plan_type: Some(scan),
1546 };
1547 Ok(node)
1548 }
1549 }
1550 LogicalPlan::Projection(Projection { expr, input, .. }) => {
1551 Ok(LogicalPlanNode {
1552 logical_plan_type: Some(LogicalPlanType::Projection(Box::new(
1553 protobuf::ProjectionNode {
1554 input: Some(Box::new(
1555 LogicalPlanNode::try_from_logical_plan(
1556 input.as_ref(),
1557 extension_codec,
1558 )?,
1559 )),
1560 expr: serialize_exprs(expr, extension_codec)?,
1561 optional_alias: None,
1562 },
1563 ))),
1564 })
1565 }
1566 LogicalPlan::Filter(filter) => {
1567 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1568 filter.input.as_ref(),
1569 extension_codec,
1570 )?;
1571 Ok(LogicalPlanNode {
1572 logical_plan_type: Some(LogicalPlanType::Selection(Box::new(
1573 protobuf::SelectionNode {
1574 input: Some(Box::new(input)),
1575 expr: Some(Box::new(serialize_expr(
1576 &filter.predicate,
1577 extension_codec,
1578 )?)),
1579 },
1580 ))),
1581 })
1582 }
1583 LogicalPlan::Distinct(Distinct::All(input)) => {
1584 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1585 input.as_ref(),
1586 extension_codec,
1587 )?;
1588 Ok(LogicalPlanNode {
1589 logical_plan_type: Some(LogicalPlanType::Distinct(Box::new(
1590 protobuf::DistinctNode {
1591 input: Some(Box::new(input)),
1592 },
1593 ))),
1594 })
1595 }
1596 LogicalPlan::Distinct(Distinct::On(DistinctOn {
1597 on_expr,
1598 select_expr,
1599 sort_expr,
1600 input,
1601 ..
1602 })) => {
1603 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1604 input.as_ref(),
1605 extension_codec,
1606 )?;
1607 let sort_expr = match sort_expr {
1608 None => vec![],
1609 Some(sort_expr) => serialize_sorts(sort_expr, extension_codec)?,
1610 };
1611 Ok(LogicalPlanNode {
1612 logical_plan_type: Some(LogicalPlanType::DistinctOn(Box::new(
1613 protobuf::DistinctOnNode {
1614 on_expr: serialize_exprs(on_expr, extension_codec)?,
1615 select_expr: serialize_exprs(select_expr, extension_codec)?,
1616 sort_expr,
1617 input: Some(Box::new(input)),
1618 },
1619 ))),
1620 })
1621 }
1622 LogicalPlan::Window(Window {
1623 input, window_expr, ..
1624 }) => {
1625 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1626 input.as_ref(),
1627 extension_codec,
1628 )?;
1629 Ok(LogicalPlanNode {
1630 logical_plan_type: Some(LogicalPlanType::Window(Box::new(
1631 protobuf::WindowNode {
1632 input: Some(Box::new(input)),
1633 window_expr: serialize_exprs(window_expr, extension_codec)?,
1634 },
1635 ))),
1636 })
1637 }
1638 LogicalPlan::Aggregate(Aggregate {
1639 group_expr,
1640 aggr_expr,
1641 input,
1642 ..
1643 }) => {
1644 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1645 input.as_ref(),
1646 extension_codec,
1647 )?;
1648 Ok(LogicalPlanNode {
1649 logical_plan_type: Some(LogicalPlanType::Aggregate(Box::new(
1650 protobuf::AggregateNode {
1651 input: Some(Box::new(input)),
1652 group_expr: serialize_exprs(group_expr, extension_codec)?,
1653 aggr_expr: serialize_exprs(aggr_expr, extension_codec)?,
1654 },
1655 ))),
1656 })
1657 }
1658 LogicalPlan::Join(Join {
1659 left,
1660 right,
1661 on,
1662 filter,
1663 join_type,
1664 join_constraint,
1665 null_equality,
1666 null_aware,
1667 schema: _,
1669 }) => {
1670 let left: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1671 left.as_ref(),
1672 extension_codec,
1673 )?;
1674 let right: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1675 right.as_ref(),
1676 extension_codec,
1677 )?;
1678 let (left_join_key, right_join_key) = on
1679 .iter()
1680 .map(|(l, r)| {
1681 Ok((
1682 serialize_expr(l, extension_codec)?,
1683 serialize_expr(r, extension_codec)?,
1684 ))
1685 })
1686 .collect::<Result<Vec<_>, ToProtoError>>()?
1687 .into_iter()
1688 .unzip();
1689 let join_type = protobuf::JoinType::from(join_type.to_owned());
1690 let join_constraint =
1691 protobuf::JoinConstraint::from(join_constraint.to_owned());
1692 let null_equality =
1693 protobuf::NullEquality::from(null_equality.to_owned());
1694 let filter = filter
1695 .as_ref()
1696 .map(|e| serialize_expr(e, extension_codec).map(Box::new))
1697 .map_or(Ok(None), |v| v.map(Some))?;
1698 Ok(LogicalPlanNode {
1699 logical_plan_type: Some(LogicalPlanType::Join(Box::new(
1700 protobuf::JoinNode {
1701 left: Some(Box::new(left)),
1702 right: Some(Box::new(right)),
1703 join_type: join_type.into(),
1704 join_constraint: join_constraint.into(),
1705 left_join_key,
1706 right_join_key,
1707 null_equality: null_equality.into(),
1708 filter,
1709 null_aware: *null_aware,
1710 },
1711 ))),
1712 })
1713 }
1714 LogicalPlan::Subquery(subquery) => {
1715 LogicalPlanNode::try_from_logical_plan(
1719 &subquery.subquery,
1720 extension_codec,
1721 )
1722 }
1723 LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) => {
1724 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1725 input.as_ref(),
1726 extension_codec,
1727 )?;
1728 Ok(LogicalPlanNode {
1729 logical_plan_type: Some(LogicalPlanType::SubqueryAlias(Box::new(
1730 protobuf::SubqueryAliasNode {
1731 input: Some(Box::new(input)),
1732 alias: Some(protobuf::TableReference::from((*alias).clone())),
1733 },
1734 ))),
1735 })
1736 }
1737 LogicalPlan::Limit(limit) => {
1738 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1739 limit.input.as_ref(),
1740 extension_codec,
1741 )?;
1742 let SkipType::Literal(skip) = limit.get_skip_type()? else {
1743 return Err(proto_error(
1744 "LogicalPlan::Limit only supports literal skip values",
1745 ));
1746 };
1747 let FetchType::Literal(fetch) = limit.get_fetch_type()? else {
1748 return Err(proto_error(
1749 "LogicalPlan::Limit only supports literal fetch values",
1750 ));
1751 };
1752
1753 Ok(LogicalPlanNode {
1754 logical_plan_type: Some(LogicalPlanType::Limit(Box::new(
1755 protobuf::LimitNode {
1756 input: Some(Box::new(input)),
1757 skip: skip as i64,
1758 fetch: fetch.unwrap_or(i64::MAX as usize) as i64,
1759 },
1760 ))),
1761 })
1762 }
1763 LogicalPlan::Sort(Sort { input, expr, fetch }) => {
1764 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1765 input.as_ref(),
1766 extension_codec,
1767 )?;
1768 let sort_expr: Vec<protobuf::SortExprNode> =
1769 serialize_sorts(expr, extension_codec)?;
1770 Ok(LogicalPlanNode {
1771 logical_plan_type: Some(LogicalPlanType::Sort(Box::new(
1772 protobuf::SortNode {
1773 input: Some(Box::new(input)),
1774 expr: sort_expr,
1775 fetch: fetch.map(|f| f as i64).unwrap_or(-1i64),
1776 },
1777 ))),
1778 })
1779 }
1780 LogicalPlan::Repartition(Repartition {
1781 input,
1782 partitioning_scheme,
1783 }) => {
1784 use datafusion_expr::Partitioning;
1785 let input: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan(
1786 input.as_ref(),
1787 extension_codec,
1788 )?;
1789
1790 use protobuf::repartition_node::PartitionMethod;
1793
1794 let pb_partition_method = match partitioning_scheme {
1795 Partitioning::Hash(exprs, partition_count) => {
1796 PartitionMethod::Hash(protobuf::HashRepartition {
1797 hash_expr: serialize_exprs(exprs, extension_codec)?,
1798 partition_count: *partition_count as u64,
1799 })
1800 }
1801 Partitioning::RoundRobinBatch(partition_count) => {
1802 PartitionMethod::RoundRobin(*partition_count as u64)
1803 }
1804 Partitioning::Range(range_partitioning) => {
1805 let ordering = range_partitioning.ordering();
1806 let split_point = range_partitioning
1807 .split_points()
1808 .iter()
1809 .map(serialize_range_split_point)
1810 .collect::<Result<Vec<_>, _>>()?;
1811
1812 PartitionMethod::Range(protobuf::RangeRepartition {
1813 sort_expr: serialize_sorts(ordering, extension_codec)?,
1814 split_point,
1815 })
1816 }
1817 Partitioning::DistributeBy(_) => {
1818 return not_impl_err!("DistributeBy");
1819 }
1820 };
1821
1822 Ok(LogicalPlanNode {
1823 logical_plan_type: Some(LogicalPlanType::Repartition(Box::new(
1824 protobuf::RepartitionNode {
1825 input: Some(Box::new(input)),
1826 partition_method: Some(pb_partition_method),
1827 },
1828 ))),
1829 })
1830 }
1831 LogicalPlan::EmptyRelation(EmptyRelation {
1832 produce_one_row, ..
1833 }) => Ok(LogicalPlanNode {
1834 logical_plan_type: Some(LogicalPlanType::EmptyRelation(
1835 protobuf::EmptyRelationNode {
1836 produce_one_row: *produce_one_row,
1837 },
1838 )),
1839 }),
1840 LogicalPlan::Ddl(DdlStatement::CreateExternalTable(ce)) => {
1841 let CreateExternalTable {
1842 name,
1843 locations,
1844 file_type,
1845 schema: df_schema,
1846 table_partition_cols,
1847 if_not_exists,
1848 or_replace,
1849 definition,
1850 order_exprs,
1851 unbounded,
1852 options,
1853 constraints,
1854 column_defaults,
1855 temporary,
1856 } = ce.as_ref();
1857 let mut converted_order_exprs: Vec<SortExprNodeCollection> = vec![];
1858 for order in order_exprs {
1859 let temp = SortExprNodeCollection {
1860 sort_expr_nodes: serialize_sorts(order, extension_codec)?,
1861 };
1862 converted_order_exprs.push(temp);
1863 }
1864
1865 let mut converted_column_defaults =
1866 HashMap::with_capacity(column_defaults.len());
1867 for (col_name, expr) in column_defaults {
1868 converted_column_defaults
1869 .insert(col_name.clone(), serialize_expr(expr, extension_codec)?);
1870 }
1871 let (legacy_location, proto_locations) = match locations.as_slice() {
1872 [location] => (location.clone(), vec![]),
1873 _ => (String::new(), locations.clone()),
1874 };
1875
1876 Ok(LogicalPlanNode {
1877 logical_plan_type: Some(LogicalPlanType::CreateExternalTable(
1878 protobuf::CreateExternalTableNode {
1879 name: Some(protobuf::TableReference::from(name.clone())),
1880 location: legacy_location,
1881 locations: proto_locations,
1882 file_type: file_type.clone(),
1883 schema: Some(df_schema.try_into()?),
1884 table_partition_cols: table_partition_cols.clone(),
1885 if_not_exists: *if_not_exists,
1886 or_replace: *or_replace,
1887 temporary: *temporary,
1888 order_exprs: converted_order_exprs,
1889 definition: definition.clone().unwrap_or_default(),
1890 unbounded: *unbounded,
1891 options: options.clone(),
1892 constraints: Some(constraints.clone().into()),
1893 column_defaults: converted_column_defaults,
1894 },
1895 )),
1896 })
1897 }
1898 LogicalPlan::Ddl(DdlStatement::CreateView(CreateView {
1899 name,
1900 input,
1901 or_replace,
1902 definition,
1903 temporary,
1904 })) => Ok(LogicalPlanNode {
1905 logical_plan_type: Some(LogicalPlanType::CreateView(Box::new(
1906 protobuf::CreateViewNode {
1907 name: Some(protobuf::TableReference::from(name.clone())),
1908 input: Some(Box::new(LogicalPlanNode::try_from_logical_plan(
1909 input,
1910 extension_codec,
1911 )?)),
1912 or_replace: *or_replace,
1913 temporary: *temporary,
1914 definition: definition.clone().unwrap_or_default(),
1915 },
1916 ))),
1917 }),
1918 LogicalPlan::Ddl(DdlStatement::CreateCatalogSchema(
1919 CreateCatalogSchema {
1920 schema_name,
1921 if_not_exists,
1922 schema: df_schema,
1923 },
1924 )) => Ok(LogicalPlanNode {
1925 logical_plan_type: Some(LogicalPlanType::CreateCatalogSchema(
1926 protobuf::CreateCatalogSchemaNode {
1927 schema_name: schema_name.clone(),
1928 if_not_exists: *if_not_exists,
1929 schema: Some(df_schema.try_into()?),
1930 },
1931 )),
1932 }),
1933 LogicalPlan::Ddl(DdlStatement::CreateCatalog(CreateCatalog {
1934 catalog_name,
1935 if_not_exists,
1936 schema: df_schema,
1937 })) => Ok(LogicalPlanNode {
1938 logical_plan_type: Some(LogicalPlanType::CreateCatalog(
1939 protobuf::CreateCatalogNode {
1940 catalog_name: catalog_name.clone(),
1941 if_not_exists: *if_not_exists,
1942 schema: Some(df_schema.try_into()?),
1943 },
1944 )),
1945 }),
1946 LogicalPlan::Analyze(a) => {
1947 let input = LogicalPlanNode::try_from_logical_plan(
1948 a.input.as_ref(),
1949 extension_codec,
1950 )?;
1951 Ok(LogicalPlanNode {
1952 logical_plan_type: Some(LogicalPlanType::Analyze(Box::new(
1953 protobuf::AnalyzeNode {
1954 input: Some(Box::new(input)),
1955 verbose: a.verbose,
1956 analyze_level: a
1957 .analyze_level
1958 .map(|m| metric_type_to_proto(m) as i32),
1959 analyze_categories: a
1960 .analyze_categories
1961 .as_ref()
1962 .map(explain_analyze_categories_to_proto),
1963 format: match &a.format {
1964 ExplainFormat::Indent => protobuf::ExplainFormat::Indent,
1965 ExplainFormat::Tree => protobuf::ExplainFormat::Tree,
1966 ExplainFormat::PostgresJSON => {
1967 protobuf::ExplainFormat::Pgjson
1968 }
1969 ExplainFormat::Graphviz => {
1970 protobuf::ExplainFormat::Graphviz
1971 }
1972 } as i32,
1973 },
1974 ))),
1975 })
1976 }
1977 LogicalPlan::Explain(a) => {
1978 let input = LogicalPlanNode::try_from_logical_plan(
1979 a.plan.as_ref(),
1980 extension_codec,
1981 )?;
1982 Ok(LogicalPlanNode {
1983 logical_plan_type: Some(LogicalPlanType::Explain(Box::new(
1984 protobuf::ExplainNode {
1985 input: Some(Box::new(input)),
1986 verbose: a.verbose,
1987 format: match &a.explain_format {
1988 ExplainFormat::Indent => protobuf::ExplainFormat::Indent,
1989 ExplainFormat::Tree => protobuf::ExplainFormat::Tree,
1990 ExplainFormat::PostgresJSON => {
1991 protobuf::ExplainFormat::Pgjson
1992 }
1993 ExplainFormat::Graphviz => {
1994 protobuf::ExplainFormat::Graphviz
1995 }
1996 }
1997 .into(),
1998 show_statistics: a.show_statistics,
1999 },
2000 ))),
2001 })
2002 }
2003 LogicalPlan::Union(union) => {
2004 let inputs: Vec<LogicalPlanNode> = union
2005 .inputs
2006 .iter()
2007 .map(|i| LogicalPlanNode::try_from_logical_plan(i, extension_codec))
2008 .collect::<Result<_>>()?;
2009 Ok(LogicalPlanNode {
2010 logical_plan_type: Some(LogicalPlanType::Union(
2011 protobuf::UnionNode { inputs },
2012 )),
2013 })
2014 }
2015 LogicalPlan::Extension(extension) => {
2016 let mut buf: Vec<u8> = vec![];
2017 extension_codec.try_encode(extension, &mut buf)?;
2018
2019 let inputs: Vec<LogicalPlanNode> = extension
2020 .node
2021 .inputs()
2022 .iter()
2023 .map(|i| LogicalPlanNode::try_from_logical_plan(i, extension_codec))
2024 .collect::<Result<_>>()?;
2025
2026 Ok(LogicalPlanNode {
2027 logical_plan_type: Some(LogicalPlanType::Extension(
2028 LogicalExtensionNode { node: buf, inputs },
2029 )),
2030 })
2031 }
2032 LogicalPlan::Statement(Statement::Prepare(Prepare {
2033 name,
2034 fields,
2035 input,
2036 })) => {
2037 let input =
2038 LogicalPlanNode::try_from_logical_plan(input, extension_codec)?;
2039 Ok(LogicalPlanNode {
2040 logical_plan_type: Some(LogicalPlanType::Prepare(Box::new(
2041 protobuf::PrepareNode {
2042 name: name.clone(),
2043 input: Some(Box::new(input)),
2044 data_types: fields
2046 .iter()
2047 .map(|f| f.data_type().try_into())
2048 .collect::<Result<Vec<_>, _>>()?,
2049 fields: fields
2051 .iter()
2052 .map(|f| f.as_ref().try_into())
2053 .collect::<Result<Vec<_>, _>>()?,
2054 },
2055 ))),
2056 })
2057 }
2058 LogicalPlan::Unnest(Unnest {
2059 input,
2060 exec_columns,
2061 list_type_columns,
2062 struct_type_columns,
2063 dependency_indices,
2064 schema,
2065 options,
2066 }) => {
2067 let input =
2068 LogicalPlanNode::try_from_logical_plan(input, extension_codec)?;
2069 let proto_unnest_list_items = list_type_columns
2070 .iter()
2071 .map(|(index, ul)| ColumnUnnestListItem {
2072 input_index: *index as _,
2073 recursion: Some(ColumnUnnestListRecursion {
2074 output_column: Some(ul.output_column.to_owned().into()),
2075 depth: ul.depth as _,
2076 }),
2077 })
2078 .collect();
2079 Ok(LogicalPlanNode {
2080 logical_plan_type: Some(LogicalPlanType::Unnest(Box::new(
2081 protobuf::UnnestNode {
2082 input: Some(Box::new(input)),
2083 exec_columns: exec_columns
2084 .iter()
2085 .map(|col| col.into())
2086 .collect(),
2087 list_type_columns: proto_unnest_list_items,
2088 struct_type_columns: struct_type_columns
2089 .iter()
2090 .map(|c| *c as u64)
2091 .collect(),
2092 dependency_indices: dependency_indices
2093 .iter()
2094 .map(|c| *c as u64)
2095 .collect(),
2096 schema: Some(schema.try_into()?),
2097 options: Some(protobuf::UnnestOptions::from(options)),
2098 },
2099 ))),
2100 })
2101 }
2102 LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(_)) => Err(proto_error(
2103 "LogicalPlan serde is not yet implemented for CreateMemoryTable",
2104 )),
2105 LogicalPlan::Ddl(DdlStatement::CreateIndex(_)) => Err(proto_error(
2106 "LogicalPlan serde is not yet implemented for CreateIndex",
2107 )),
2108 LogicalPlan::Ddl(DdlStatement::DropTable(_)) => Err(proto_error(
2109 "LogicalPlan serde is not yet implemented for DropTable",
2110 )),
2111 LogicalPlan::Ddl(DdlStatement::DropView(DropView {
2112 name,
2113 if_exists,
2114 schema,
2115 })) => Ok(LogicalPlanNode {
2116 logical_plan_type: Some(LogicalPlanType::DropView(
2117 protobuf::DropViewNode {
2118 name: Some(protobuf::TableReference::from(name.clone())),
2119 if_exists: *if_exists,
2120 schema: Some(schema.try_into()?),
2121 },
2122 )),
2123 }),
2124 LogicalPlan::Ddl(DdlStatement::DropCatalogSchema(_)) => Err(proto_error(
2125 "LogicalPlan serde is not yet implemented for DropCatalogSchema",
2126 )),
2127 LogicalPlan::Ddl(DdlStatement::CreateFunction(_)) => Err(proto_error(
2128 "LogicalPlan serde is not yet implemented for CreateFunction",
2129 )),
2130 LogicalPlan::Ddl(DdlStatement::DropFunction(_)) => Err(proto_error(
2131 "LogicalPlan serde is not yet implemented for DropFunction",
2132 )),
2133 LogicalPlan::Statement(_) => Err(proto_error(
2134 "LogicalPlan serde is not yet implemented for Statement",
2135 )),
2136 LogicalPlan::Dml(DmlStatement {
2137 table_name,
2138 target,
2139 op,
2140 input,
2141 ..
2142 }) => {
2143 let input =
2144 LogicalPlanNode::try_from_logical_plan(input, extension_codec)?;
2145 let (dml_type, merge_into) = match op {
2146 WriteOp::Insert(InsertOp::Append) => {
2147 (dml_node::Type::InsertAppend, None)
2148 }
2149 WriteOp::Insert(InsertOp::Overwrite) => {
2150 (dml_node::Type::InsertOverwrite, None)
2151 }
2152 WriteOp::Insert(InsertOp::Replace) => {
2153 (dml_node::Type::InsertReplace, None)
2154 }
2155 WriteOp::Delete => (dml_node::Type::Delete, None),
2156 WriteOp::Update => (dml_node::Type::Update, None),
2157 WriteOp::Ctas => (dml_node::Type::Ctas, None),
2158 WriteOp::Truncate => (dml_node::Type::Truncate, None),
2159 WriteOp::MergeInto(merge_op) => (
2160 dml_node::Type::MergeInto,
2161 Some(Box::new(to_proto::serialize_merge_into_op(
2162 merge_op,
2163 extension_codec,
2164 )?)),
2165 ),
2166 other => {
2167 return Err(proto_error(format!(
2168 "WriteOp variant has no DmlNode encoding: {other}"
2169 )));
2170 }
2171 };
2172 Ok(LogicalPlanNode {
2173 logical_plan_type: Some(LogicalPlanType::Dml(Box::new(DmlNode {
2174 input: Some(Box::new(input)),
2175 target: Some(Box::new(from_table_source(
2176 table_name.clone(),
2177 Arc::clone(target),
2178 extension_codec,
2179 )?)),
2180 table_name: Some(protobuf::TableReference::from(
2181 table_name.clone(),
2182 )),
2183 dml_type: dml_type.into(),
2184 merge_into,
2185 }))),
2186 })
2187 }
2188 LogicalPlan::Copy(dml::CopyTo {
2189 input,
2190 output_url,
2191 file_type,
2192 partition_by,
2193 ..
2194 }) => {
2195 let input =
2196 LogicalPlanNode::try_from_logical_plan(input, extension_codec)?;
2197 let mut buf = Vec::new();
2198 extension_codec
2199 .try_encode_file_format(&mut buf, file_type_to_format(file_type)?)?;
2200
2201 Ok(LogicalPlanNode {
2202 logical_plan_type: Some(LogicalPlanType::CopyTo(Box::new(
2203 protobuf::CopyToNode {
2204 input: Some(Box::new(input)),
2205 output_url: output_url.to_string(),
2206 file_type: buf,
2207 partition_by: partition_by.clone(),
2208 },
2209 ))),
2210 })
2211 }
2212 LogicalPlan::DescribeTable(_) => Err(proto_error(
2213 "LogicalPlan serde is not yet implemented for DescribeTable",
2214 )),
2215 LogicalPlan::RecursiveQuery(recursive) => {
2216 let static_term = LogicalPlanNode::try_from_logical_plan(
2217 recursive.static_term.as_ref(),
2218 extension_codec,
2219 )?;
2220 let recursive_term = LogicalPlanNode::try_from_logical_plan(
2221 recursive.recursive_term.as_ref(),
2222 extension_codec,
2223 )?;
2224
2225 Ok(LogicalPlanNode {
2226 logical_plan_type: Some(LogicalPlanType::RecursiveQuery(Box::new(
2227 protobuf::RecursiveQueryNode {
2228 name: recursive.name.clone(),
2229 static_term: Some(Box::new(static_term)),
2230 recursive_term: Some(Box::new(recursive_term)),
2231 is_distinct: recursive.is_distinct,
2232 },
2233 ))),
2234 })
2235 }
2236 })
2237 }
2238}