diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 3f15b84784f16..8b8c5c463f79d 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -42,8 +42,9 @@ use crate::utils::{ expand_wildcard, find_valid_equijoin_key_pair, group_window_expr_by_sort_keys, }; use crate::{ - and, binary_expr, DmlStatement, Expr, ExprSchemable, Operator, RecursiveQuery, - TableProviderFilterPushDown, TableSource, WriteOp, + and, binary_expr, logical_plan::tree_node::unwrap_arc, DmlStatement, Expr, + ExprSchemable, Operator, RecursiveQuery, TableProviderFilterPushDown, TableSource, + WriteOp, }; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; @@ -1138,6 +1139,31 @@ impl LogicalPlanBuilder { )?)) } } + +/// Converts a `Arc` into `LogicalPlanBuilder` +/// fn employee_schema() -> Schema { +/// Schema::new(vec![ +/// Field::new("id", DataType::Int32, false), +/// Field::new("first_name", DataType::Utf8, false), +/// Field::new("last_name", DataType::Utf8, false), +/// Field::new("state", DataType::Utf8, false), +/// Field::new("salary", DataType::Int32, false), +/// ]) +/// } +/// let plan = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))? +/// .sort(vec![ +/// Expr::Sort(expr::Sort::new(Box::new(col("state")), true, true)), +/// Expr::Sort(expr::Sort::new(Box::new(col("salary")), false, false)), +/// ])? +/// .build()?; +/// let plan_builder: LogicalPlanBuilder = Arc::new(plan).into(); + +impl From> for LogicalPlanBuilder { + fn from(plan: Arc) -> Self { + LogicalPlanBuilder::from(unwrap_arc(plan)) + } +} + pub fn change_redundant_column(fields: &Fields) -> Vec { let mut name_map = HashMap::new(); fields @@ -2140,4 +2166,21 @@ mod tests { ); Ok(()) } + + #[test] + fn plan_builder_from_logical_plan() -> Result<()> { + let plan = + table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))? + .sort(vec![ + Expr::Sort(expr::Sort::new(Box::new(col("state")), true, true)), + Expr::Sort(expr::Sort::new(Box::new(col("salary")), false, false)), + ])? + .build()?; + + let plan_expected = format!("{plan:?}"); + let plan_builder: LogicalPlanBuilder = Arc::new(plan).into(); + assert_eq!(plan_expected, format!("{:?}", plan_builder.plan)); + + Ok(()) + } }