refactor: remove DfPlan wrapper (#4733)

* refactor: remove DfPlan wrapper Signed-off-by: Ruihang Xia <[email protected]> * clean up Signed-off-by: Ruihang Xia <[email protected]> * remove unused errors Signed-off-by: Ruihang Xia <[email protected]> * fix test assertion Signed-off-by: Ruihang Xia <[email protected]> --------- Signed-off-by: Ruihang Xia <[email protected]>
GreptimeTeam · Sep 19, 2024 · f5cf25b · f5cf25b
1 parent 1acda74
commit f5cf25b
Show file tree

Hide file tree

Showing 38 changed files with 194 additions and 293 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/cmd/src/cli/repl.rs b/src/cmd/src/cli/repl.rs
@@ -35,7 +35,6 @@ use either::Either;
 use meta_client::client::MetaClientBuilder;
 use query::datafusion::DatafusionQueryEngine;
 use query::parser::QueryLanguageParser;
-use query::plan::LogicalPlan;
 use query::query_engine::{DefaultSerializer, QueryEngineState};
 use query::QueryEngine;
 use rustyline::error::ReadlineError;
@@ -179,7 +178,7 @@ impl Repl {
                 .await
                 .context(PlanStatementSnafu)?;
 
-            let LogicalPlan::DfPlan(plan) = query_engine
+            let plan = query_engine
                 .optimize(&query_engine.engine_context(query_ctx), &plan)
                 .context(PlanStatementSnafu)?;
 

diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs
@@ -843,7 +843,7 @@ impl RegionServerInner {
 
         let result = self
             .query_engine
-            .execute(request.plan.into(), query_ctx)
+            .execute(request.plan, query_ctx)
             .await
             .context(ExecuteLogicalPlanSnafu)?;
 

diff --git a/src/datanode/src/tests.rs b/src/datanode/src/tests.rs
@@ -24,8 +24,8 @@ use common_function::scalars::aggregate::AggregateFunctionMetaRef;
 use common_query::prelude::ScalarUdf;
 use common_query::Output;
 use common_runtime::Runtime;
+use datafusion_expr::LogicalPlan;
 use query::dataframe::DataFrame;
-use query::plan::LogicalPlan;
 use query::planner::LogicalPlanner;
 use query::query_engine::{DescribeResult, QueryEngineState};
 use query::{QueryEngine, QueryEngineContext};

diff --git a/src/flow/src/df_optimizer.rs b/src/flow/src/df_optimizer.rs
@@ -41,7 +41,6 @@ use datafusion_expr::{
     BinaryExpr, Expr, Operator, Projection, ScalarUDFImpl, Signature, TypeSignature, Volatility,
 };
 use query::parser::QueryLanguageParser;
-use query::plan::LogicalPlan;
 use query::query_engine::DefaultSerializer;
 use query::QueryEngine;
 use snafu::ResultExt;
@@ -111,7 +110,6 @@ pub async fn sql_to_flow_plan(
         .await
         .map_err(BoxedError::new)
         .context(ExternalSnafu)?;
-    let LogicalPlan::DfPlan(plan) = plan;
 
     let opted_plan = apply_df_optimizer(plan).await?;
 

diff --git a/src/flow/src/transform.rs b/src/flow/src/transform.rs
@@ -163,7 +163,6 @@ mod test {
     use itertools::Itertools;
     use prost::Message;
     use query::parser::QueryLanguageParser;
-    use query::plan::LogicalPlan;
     use query::query_engine::DefaultSerializer;
     use query::QueryEngine;
     use session::context::QueryContext;
@@ -274,7 +273,6 @@ mod test {
             .plan(stmt, QueryContext::arc())
             .await
             .unwrap();
-        let LogicalPlan::DfPlan(plan) = plan;
         let plan = apply_df_optimizer(plan).await.unwrap();
 
         // encode then decode so to rely on the impl of conversion from logical plan to substrait plan
@@ -297,7 +295,6 @@ mod test {
             .plan(stmt, QueryContext::arc())
             .await
             .unwrap();
-        let LogicalPlan::DfPlan(plan) = plan;
         let plan = apply_df_optimizer(plan).await;
 
         assert!(plan.is_err());

diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml
@@ -37,6 +37,7 @@ common-runtime.workspace = true
 common-telemetry.workspace = true
 common-time.workspace = true
 common-version.workspace = true
+datafusion-expr.workspace = true
 datanode.workspace = true
 humantime-serde.workspace = true
 lazy_static.workspace = true

diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs
@@ -40,6 +40,7 @@ use common_procedure::options::ProcedureConfig;
 use common_procedure::ProcedureManagerRef;
 use common_query::Output;
 use common_telemetry::{debug, error, tracing};
+use datafusion_expr::LogicalPlan;
 use log_store::raft_engine::RaftEngineBackend;
 use operator::delete::DeleterRef;
 use operator::insert::InserterRef;
@@ -48,7 +49,6 @@ use pipeline::pipeline_operator::PipelineOperator;
 use prometheus::HistogramTimer;
 use query::metrics::OnDone;
 use query::parser::{PromQuery, QueryLanguageParser, QueryStatement};
-use query::plan::LogicalPlan;
 use query::query_engine::options::{validate_catalog_and_schema, QueryOptions};
 use query::query_engine::DescribeResult;
 use query::QueryEngineRef;

diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs
@@ -133,13 +133,6 @@ pub enum Error {
         source: query::error::Error,
     },
 
-    #[snafu(display("Failed to get schema from logical plan"))]
-    GetSchema {
-        #[snafu(implicit)]
-        location: Location,
-        source: query::error::Error,
-    },
-
     #[snafu(display("Column datatype error"))]
     ColumnDataType {
         #[snafu(implicit)]
@@ -184,6 +177,13 @@ pub enum Error {
         source: datatypes::error::Error,
     },
 
+    #[snafu(display("Failed to convert datafusion schema"))]
+    ConvertSchema {
+        source: datatypes::error::Error,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Failed to convert expr to struct"))]
     InvalidExpr {
         #[snafu(implicit)]
@@ -795,6 +795,7 @@ impl ErrorExt for Error {
             | Error::PrepareFileTable { .. }
             | Error::InferFileTableSchema { .. }
             | Error::SchemaIncompatible { .. }
+            | Error::ConvertSchema { .. }
             | Error::UnsupportedRegionRequest { .. }
             | Error::InvalidTableName { .. }
             | Error::InvalidViewName { .. }
@@ -872,7 +873,6 @@ impl ErrorExt for Error {
             | Error::FindNewColumnsOnInsertion { source, .. } => source.status_code(),
 
             Error::ExecuteStatement { source, .. }
-            | Error::GetSchema { source, .. }
             | Error::ExtractTableNames { source, .. }
             | Error::PlanStatement { source, .. }
             | Error::ParseQuery { source, .. }

diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs
@@ -42,9 +42,9 @@ use common_query::Output;
 use common_telemetry::tracing;
 use common_time::range::TimestampRange;
 use common_time::Timestamp;
+use datafusion_expr::LogicalPlan;
 use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
 use query::parser::QueryStatement;
-use query::plan::LogicalPlan;
 use query::QueryEngineRef;
 use session::context::QueryContextRef;
 use session::table_name::table_idents_to_full_name;

diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs
@@ -29,7 +29,6 @@ use datafusion::datasource::DefaultTableSource;
 use datafusion_common::TableReference as DfTableReference;
 use datafusion_expr::LogicalPlanBuilder;
 use object_store::ObjectStore;
-use query::plan::LogicalPlan;
 use session::context::QueryContextRef;
 use snafu::{OptionExt, ResultExt};
 use table::requests::CopyTableRequest;
@@ -133,7 +132,7 @@ impl StatementExecutor {
 
         let output = self
             .query_engine
-            .execute(LogicalPlan::DfPlan(plan), query_ctx)
+            .execute(plan, query_ctx)
             .await
             .context(ExecLogicalPlanSnafu)?;
         let stream = match output.data {

diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs
@@ -37,7 +37,7 @@ use common_query::Output;
 use common_telemetry::{debug, info, tracing};
 use common_time::Timezone;
 use datatypes::prelude::ConcreteDataType;
-use datatypes::schema::RawSchema;
+use datatypes::schema::{RawSchema, Schema};
 use datatypes::value::Value;
 use lazy_static::lazy_static;
 use partition::expr::{Operand, PartitionExpr, RestrictedOp};
@@ -69,11 +69,11 @@ use table::TableRef;
 use super::StatementExecutor;
 use crate::error::{
     self, AlterExprToRequestSnafu, CatalogSnafu, ColumnDataTypeSnafu, ColumnNotFoundSnafu,
-    CreateLogicalTablesSnafu, CreateTableInfoSnafu, DeserializePartitionSnafu, EmptyDdlExprSnafu,
-    ExtractTableNamesSnafu, FlowNotFoundSnafu, InvalidPartitionRuleSnafu, InvalidPartitionSnafu,
-    InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu, ParseSqlValueSnafu, Result,
-    SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu, SubstraitCodecSnafu,
-    TableAlreadyExistsSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu,
+    ConvertSchemaSnafu, CreateLogicalTablesSnafu, CreateTableInfoSnafu, DeserializePartitionSnafu,
+    EmptyDdlExprSnafu, ExtractTableNamesSnafu, FlowNotFoundSnafu, InvalidPartitionRuleSnafu,
+    InvalidPartitionSnafu, InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu,
+    ParseSqlValueSnafu, Result, SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu,
+    SubstraitCodecSnafu, TableAlreadyExistsSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu,
     UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu,
 };
 use crate::expr_factory;
@@ -406,9 +406,12 @@ impl StatementExecutor {
 
         // Save the columns in plan, it may changed when the schemas of tables in plan
         // are altered.
-        let plan_columns: Vec<_> = logical_plan
+        let schema: Schema = logical_plan
             .schema()
-            .context(error::GetSchemaSnafu)?
+            .clone()
+            .try_into()
+            .context(ConvertSchemaSnafu)?;
+        let plan_columns: Vec<_> = schema
             .column_schemas()
             .iter()
             .map(|c| c.name.clone())
@@ -434,9 +437,8 @@ impl StatementExecutor {
 
         // Extract the table names from the original plan
         // and rewrite them as fully qualified names.
-        let (table_names, plan) =
-            extract_and_rewrite_full_table_names(logical_plan.unwrap_df_plan(), ctx.clone())
-                .context(ExtractTableNamesSnafu)?;
+        let (table_names, plan) = extract_and_rewrite_full_table_names(logical_plan, ctx.clone())
+            .context(ExtractTableNamesSnafu)?;
 
         let table_names = table_names.into_iter().map(|t| t.into()).collect();
 

diff --git a/src/operator/src/statement/tql.rs b/src/operator/src/statement/tql.rs
@@ -16,11 +16,11 @@ use std::collections::HashMap;
 
 use common_query::Output;
 use common_telemetry::tracing;
+use datafusion_expr::LogicalPlan;
 use query::parser::{
     PromQuery, QueryLanguageParser, ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME,
     DEFAULT_LOOKBACK_STRING, EXPLAIN_NODE_NAME, EXPLAIN_VERBOSE_NODE_NAME,
 };
-use query::plan::LogicalPlan;
 use session::context::QueryContextRef;
 use snafu::ResultExt;
 use sql::statements::tql::Tql;

diff --git a/src/pipeline/src/manager/table.rs b/src/pipeline/src/manager/table.rs
@@ -26,15 +26,14 @@ use common_telemetry::{debug, info};
 use common_time::timestamp::{TimeUnit, Timestamp};
 use datafusion::logical_expr::col;
 use datafusion_common::{TableReference, ToDFSchema};
-use datafusion_expr::{DmlStatement, LogicalPlan as DfLogicalPlan};
+use datafusion_expr::{DmlStatement, LogicalPlan};
 use datatypes::prelude::ScalarVector;
 use datatypes::timestamp::TimestampNanosecond;
 use datatypes::vectors::{StringVector, TimestampNanosecondVector, Vector};
 use moka::sync::Cache;
 use operator::insert::InserterRef;
 use operator::statement::StatementExecutorRef;
 use query::dataframe::DataFrame;
-use query::plan::LogicalPlan;
 use query::QueryEngineRef;
 use session::context::{QueryContextBuilder, QueryContextRef};
 use snafu::{ensure, OptionExt, ResultExt};
@@ -373,7 +372,7 @@ impl PipelineTable {
             Arc::new(dataframe.into_parts().1),
         );
 
-        let plan = LogicalPlan::DfPlan(DfLogicalPlan::Dml(stmt));
+        let plan = LogicalPlan::Dml(stmt);
 
         // 4. execute dml stmt
         let output = self
@@ -427,7 +426,7 @@ impl PipelineTable {
             .limit(0, Some(1))
             .context(BuildDfLogicalPlanSnafu)?;
 
-        let plan = LogicalPlan::DfPlan(dataframe.into_parts().1);
+        let plan = dataframe.into_parts().1;
 
         let table_info = self.table.table_info();