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();