From 3ee5a71b2cef2f996fb912485fb10034237e0e0d Mon Sep 17 00:00:00 2001 From: AbrarNitk Date: Sun, 12 May 2024 15:47:28 +0530 Subject: [PATCH 1/3] implement From> for LogicalPlanBuilder --- datafusion/expr/src/logical_plan/builder.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 3f15b84784f16..40d2fafed436e 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -44,6 +44,7 @@ use crate::utils::{ use crate::{ and, binary_expr, DmlStatement, Expr, ExprSchemable, Operator, RecursiveQuery, TableProviderFilterPushDown, TableSource, WriteOp, + logical_plan::tree_node::unwrap_arc }; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; @@ -55,6 +56,7 @@ use datafusion_common::{ ToDFSchema, UnnestOptions, }; + /// Default table name for unnamed table pub const UNNAMED_TABLE: &str = "?table?"; @@ -1138,6 +1140,14 @@ impl LogicalPlanBuilder { )?)) } } + +/// Converts a `Arc` into `LogicalPlanBuilder` +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 From 9c9d1b61e890dedfe79a6a9d53cf9effea8184c0 Mon Sep 17 00:00:00 2001 From: AbrarNitk Date: Mon, 13 May 2024 17:45:29 +0530 Subject: [PATCH 2/3] make fmt happy --- datafusion/expr/src/logical_plan/builder.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 40d2fafed436e..ae5b8149fae08 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -42,9 +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, - logical_plan::tree_node::unwrap_arc + 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}; @@ -56,7 +56,6 @@ use datafusion_common::{ ToDFSchema, UnnestOptions, }; - /// Default table name for unnamed table pub const UNNAMED_TABLE: &str = "?table?"; From b6512725bc7e1c31b9f8e5e8579a534ad25915a6 Mon Sep 17 00:00:00 2001 From: AbrarNitk Date: Mon, 13 May 2024 18:08:04 +0530 Subject: [PATCH 3/3] added test case and doc comment --- datafusion/expr/src/logical_plan/builder.rs | 34 +++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index ae5b8149fae08..8b8c5c463f79d 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1141,6 +1141,23 @@ 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)) @@ -2149,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(()) + } }