Remove unused code and add test case for logical plan round trip with custom object store

Author: thinkharderdev

ballista/rust/core/src/serde/logical_plan/from_proto.rs

@@ -51,308 +51,6 @@ use std::{
     unimplemented,
 };
 
-impl TryInto<LogicalPlan> for &protobuf::LogicalPlanNode {
-    type Error = BallistaError;
-
-    fn try_into(self) -> Result<LogicalPlan, Self::Error> {
-        let plan = self.logical_plan_type.as_ref().ok_or_else(|| {
-            proto_error(format!(
-                "logical_plan::from_proto() Unsupported logical plan '{:?}'",
-                self
-            ))
-        })?;
-        match plan {
-            LogicalPlanType::Values(values) => {
-                let n_cols = values.n_cols as usize;
-                let values: Vec<Vec<Expr>> = if values.values_list.is_empty() {
-                    Ok(Vec::new())
-                } else if values.values_list.len() % n_cols != 0 {
-                    Err(BallistaError::General(format!(
-                        "Invalid values list length, expect {} to be divisible by {}",
-                        values.values_list.len(),
-                        n_cols
-                    )))
-                } else {
-                    values
-                        .values_list
-                        .chunks_exact(n_cols)
-                        .map(|r| {
-                            r.iter()
-                                .map(|v| v.try_into())
-                                .collect::<Result<Vec<_>, _>>()
-                        })
-                        .collect::<Result<Vec<_>, _>>()
-                }?;
-                LogicalPlanBuilder::values(values)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::Projection(projection) => {
-                let input: LogicalPlan = convert_box_required!(projection.input)?;
-                let x: Vec<Expr> = projection
-                    .expr
-                    .iter()
-                    .map(|expr| expr.try_into())
-                    .collect::<Result<Vec<_>, _>>()?;
-                LogicalPlanBuilder::from(input)
-                    .project_with_alias(
-                        x,
-                        projection.optional_alias.as_ref().map(|a| match a {
-                            protobuf::projection_node::OptionalAlias::Alias(alias) => {
-                                alias.clone()
-                            }
-                        }),
-                    )?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::Selection(selection) => {
-                let input: LogicalPlan = convert_box_required!(selection.input)?;
-                let expr: Expr = selection
-                    .expr
-                    .as_ref()
-                    .ok_or_else(|| {
-                        BallistaError::General("expression required".to_string())
-                    })?
-                    .try_into()?;
-                LogicalPlanBuilder::from(input)
-                    .filter(expr)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::Window(window) => {
-                let input: LogicalPlan = convert_box_required!(window.input)?;
-                let window_expr = window
-                    .window_expr
-                    .iter()
-                    .map(|expr| expr.try_into())
-                    .collect::<Result<Vec<Expr>, _>>()?;
-                LogicalPlanBuilder::from(input)
-                    .window(window_expr)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::Aggregate(aggregate) => {
-                let input: LogicalPlan = convert_box_required!(aggregate.input)?;
-                let group_expr = aggregate
-                    .group_expr
-                    .iter()
-                    .map(|expr| expr.try_into())
-                    .collect::<Result<Vec<Expr>, _>>()?;
-                let aggr_expr = aggregate
-                    .aggr_expr
-                    .iter()
-                    .map(|expr| expr.try_into())
-                    .collect::<Result<Vec<Expr>, _>>()?;
-                LogicalPlanBuilder::from(input)
-                    .aggregate(group_expr, aggr_expr)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::ListingScan(scan) => {
-                let schema: Schema = convert_required!(scan.schema)?;
-
-                let mut projection = None;
-                if let Some(columns) = &scan.projection {
-                    let column_indices = columns
-                        .columns
-                        .iter()
-                        .map(|name| schema.index_of(name))
-                        .collect::<Result<Vec<usize>, _>>()?;
-                    projection = Some(column_indices);
-                }
-
-                let filters = scan
-                    .filters
-                    .iter()
-                    .map(|e| e.try_into())
-                    .collect::<Result<Vec<_>, _>>()?;
-
-                let file_format: Arc<dyn FileFormat> =
-                    match scan.file_format_type.as_ref().ok_or_else(|| {
-                        proto_error(format!(
-                            "logical_plan::from_proto() Unsupported file format '{:?}'",
-                            self
-                        ))
-                    })? {
-                        &FileFormatType::Parquet(protobuf::ParquetFormat {
-                            enable_pruning,
-                        }) => Arc::new(
-                            ParquetFormat::default().with_enable_pruning(enable_pruning),
-                        ),
-                        FileFormatType::Csv(protobuf::CsvFormat {
-                            has_header,
-                            delimiter,
-                        }) => Arc::new(
-                            CsvFormat::default()
-                                .with_has_header(*has_header)
-                                .with_delimiter(str_to_byte(delimiter)?),
-                        ),
-                        FileFormatType::Avro(..) => Arc::new(AvroFormat::default()),
-                    };
-
-                let options = ListingOptions {
-                    file_extension: scan.file_extension.clone(),
-                    format: file_format,
-                    table_partition_cols: scan.table_partition_cols.clone(),
-                    collect_stat: scan.collect_stat,
-                    target_partitions: scan.target_partitions as usize,
-                };
-
-                let provider = ListingTable::new(
-                    Arc::new(LocalFileSystem {}),
-                    scan.path.clone(),
-                    Arc::new(schema),
-                    options,
-                );
-
-                LogicalPlanBuilder::scan_with_filters(
-                    &scan.table_name,
-                    Arc::new(provider),
-                    projection,
-                    filters,
-                )?
-                .build()
-                .map_err(|e| e.into())
-            }
-            LogicalPlanType::Sort(sort) => {
-                let input: LogicalPlan = convert_box_required!(sort.input)?;
-                let sort_expr: Vec<Expr> = sort
-                    .expr
-                    .iter()
-                    .map(|expr| expr.try_into())
-                    .collect::<Result<Vec<Expr>, _>>()?;
-                LogicalPlanBuilder::from(input)
-                    .sort(sort_expr)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::Repartition(repartition) => {
-                use datafusion::logical_plan::Partitioning;
-                let input: LogicalPlan = convert_box_required!(repartition.input)?;
-                use protobuf::repartition_node::PartitionMethod;
-                let pb_partition_method = repartition.partition_method.clone().ok_or_else(|| {
-                    BallistaError::General(String::from(
-                        "Protobuf deserialization error, RepartitionNode was missing required field 'partition_method'",
-                    ))
-                })?;
-
-                let partitioning_scheme = match pb_partition_method {
-                    PartitionMethod::Hash(protobuf::HashRepartition {
-                        hash_expr: pb_hash_expr,
-                        partition_count,
-                    }) => Partitioning::Hash(
-                        pb_hash_expr
-                            .iter()
-                            .map(|pb_expr| pb_expr.try_into())
-                            .collect::<Result<Vec<_>, _>>()?,
-                        partition_count as usize,
-                    ),
-                    PartitionMethod::RoundRobin(partition_count) => {
-                        Partitioning::RoundRobinBatch(partition_count as usize)
-                    }
-                };
-
-                LogicalPlanBuilder::from(input)
-                    .repartition(partitioning_scheme)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::EmptyRelation(empty_relation) => {
-                LogicalPlanBuilder::empty(empty_relation.produce_one_row)
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::CreateExternalTable(create_extern_table) => {
-                let pb_schema = (create_extern_table.schema.clone()).ok_or_else(|| {
-                    BallistaError::General(String::from(
-                        "Protobuf deserialization error, CreateExternalTableNode was missing required field schema.",
-                    ))
-                })?;
-
-                let pb_file_type: protobuf::FileType =
-                    create_extern_table.file_type.try_into()?;
-
-                Ok(LogicalPlan::CreateExternalTable(CreateExternalTable {
-                    schema: pb_schema.try_into()?,
-                    name: create_extern_table.name.clone(),
-                    location: create_extern_table.location.clone(),
-                    file_type: pb_file_type.into(),
-                    has_header: create_extern_table.has_header,
-                }))
-            }
-            LogicalPlanType::Analyze(analyze) => {
-                let input: LogicalPlan = convert_box_required!(analyze.input)?;
-                LogicalPlanBuilder::from(input)
-                    .explain(analyze.verbose, true)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::Explain(explain) => {
-                let input: LogicalPlan = convert_box_required!(explain.input)?;
-                LogicalPlanBuilder::from(input)
-                    .explain(explain.verbose, false)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::Limit(limit) => {
-                let input: LogicalPlan = convert_box_required!(limit.input)?;
-                LogicalPlanBuilder::from(input)
-                    .limit(limit.limit as usize)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-            LogicalPlanType::Join(join) => {
-                let left_keys: Vec<Column> =
-                    join.left_join_column.iter().map(|i| i.into()).collect();
-                let right_keys: Vec<Column> =
-                    join.right_join_column.iter().map(|i| i.into()).collect();
-                let join_type =
-                    protobuf::JoinType::from_i32(join.join_type).ok_or_else(|| {
-                        proto_error(format!(
-                            "Received a JoinNode message with unknown JoinType {}",
-                            join.join_type
-                        ))
-                    })?;
-                let join_constraint = protobuf::JoinConstraint::from_i32(
-                    join.join_constraint,
-                )
-                .ok_or_else(|| {
-                    proto_error(format!(
-                        "Received a JoinNode message with unknown JoinConstraint {}",
-                        join.join_constraint
-                    ))
-                })?;
-
-                let builder = LogicalPlanBuilder::from(convert_box_required!(join.left)?);
-                let builder = match join_constraint.into() {
-                    JoinConstraint::On => builder.join(
-                        &convert_box_required!(join.right)?,
-                        join_type.into(),
-                        (left_keys, right_keys),
-                    )?,
-                    JoinConstraint::Using => builder.join_using(
-                        &convert_box_required!(join.right)?,
-                        join_type.into(),
-                        left_keys,
-                    )?,
-                };
-
-                builder.build().map_err(|e| e.into())
-            }
-            LogicalPlanType::CrossJoin(crossjoin) => {
-                let left = convert_box_required!(crossjoin.left)?;
-                let right = convert_box_required!(crossjoin.right)?;
-
-                LogicalPlanBuilder::from(left)
-                    .cross_join(&right)?
-                    .build()
-                    .map_err(|e| e.into())
-            }
-        }
-    }
-}
-
 impl From<&protobuf::Column> for Column {
     fn from(c: &protobuf::Column) -> Column {
         let c = c.clone();