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