refactor: remove datafusion data frame wrapper (#7347)

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

src/datanode/src/tests.rs

@@ -24,8 +24,8 @@ use common_query::Output;
 use common_runtime::Runtime;
 use common_runtime::runtime::{BuilderBuild, RuntimeTrait};
 use datafusion::catalog::TableFunction;
+use datafusion::dataframe::DataFrame;
 use datafusion_expr::{AggregateUDF, LogicalPlan};
-use query::dataframe::DataFrame;
 use query::planner::LogicalPlanner;
 use query::query_engine::{DescribeResult, QueryEngineState};
 use query::{QueryEngine, QueryEngineContext};

src/frontend/src/instance/promql.rs

@@ -136,7 +136,7 @@ impl Instance {
                 table_name: format_full_table_name(ctx.current_catalog(), &table_schema, &metric),
             })?;
 
-        let scan_plan = dataframe.into_logical_plan();
+        let scan_plan = dataframe.into_unoptimized_plan();
         let filter_conditions =
             PromPlanner::matchers_to_expr(Matchers::new(matchers), scan_plan.schema())
                 .context(PrometheusLabelValuesQueryPlanSnafu)?;

src/pipeline/src/manager/table.rs

@@ -34,7 +34,6 @@ use itertools::Itertools;
 use operator::insert::InserterRef;
 use operator::statement::StatementExecutorRef;
 use query::QueryEngineRef;
-use query::dataframe::DataFrame;
 use session::context::{QueryContextBuilder, QueryContextRef};
 use snafu::{OptionExt, ResultExt, ensure};
 use table::TableRef;
@@ -413,7 +412,6 @@ impl PipelineTable {
             .query_engine
             .read_table(self.table.clone())
             .context(DataFrameSnafu)?;
-        let DataFrame::DataFusion(dataframe) = dataframe;
 
         let dataframe = dataframe
             .filter(prepare_dataframe_conditions(name, version))
@@ -474,7 +472,6 @@ impl PipelineTable {
             .query_engine
             .read_table(self.table.clone())
             .context(DataFrameSnafu)?;
-        let DataFrame::DataFusion(dataframe) = dataframe;
 
         // select all pipelines with name and version
         let dataframe = dataframe

src/query/src/dataframe.rs

@@ -1,31 +0,0 @@
-// Copyright 2023 Greptime Team
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use datafusion::dataframe::DataFrame as DfDataFrame;
-use datafusion_expr::LogicalPlan;
-
-/// DataFrame represents a logical set of rows with the same named columns.
-/// Similar to a Pandas DataFrame or Spark DataFrame
-#[derive(Clone)]
-pub enum DataFrame {
-    DataFusion(DfDataFrame),
-}
-
-impl DataFrame {
-    pub fn into_logical_plan(self) -> LogicalPlan {
-        match self {
-            Self::DataFusion(dataframe) => dataframe.into_parts().1,
-        }
-    }
-}

src/query/src/datafusion.rs

@@ -32,6 +32,7 @@ use common_recordbatch::adapter::RecordBatchStreamAdapter;
 use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream};
 use common_telemetry::tracing;
 use datafusion::catalog::TableFunction;
+use datafusion::dataframe::DataFrame;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_plan::analyze::AnalyzeExec;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
@@ -49,7 +50,6 @@ use table::TableRef;
 use table::requests::{DeleteRequest, InsertRequest};
 
 use crate::analyze::DistAnalyzeExec;
-use crate::dataframe::DataFrame;
 pub use crate::datafusion::planner::DfContextProviderAdapter;
 use crate::dist_plan::{DistPlannerOptions, MergeScanLogicalPlan};
 use crate::error::{
@@ -515,13 +515,11 @@ impl QueryEngine for DatafusionQueryEngine {
     }
 
     fn read_table(&self, table: TableRef) -> Result<DataFrame> {
-        Ok(DataFrame::DataFusion(
-            self.state
-                .read_table(table)
-                .context(error::DatafusionSnafu)
-                .map_err(BoxedError::new)
-                .context(QueryExecutionSnafu)?,
-        ))
+        self.state
+            .read_table(table)
+            .context(error::DatafusionSnafu)
+            .map_err(BoxedError::new)
+            .context(QueryExecutionSnafu)
     }
 
     fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext {
@@ -799,7 +797,7 @@ mod tests {
             .await
             .unwrap();
 
-        let DataFrame::DataFusion(df) = engine.read_table(table).unwrap();
+        let df = engine.read_table(table).unwrap();
         let df = df
             .select_columns(&["number"])
             .unwrap()

src/query/src/lib.rs

@@ -19,7 +19,6 @@
 #![feature(box_patterns)]
 
 mod analyze;
-pub mod dataframe;
 pub mod datafusion;
 pub mod dist_plan;
 pub mod dummy_catalog;

src/query/src/query_engine.rs

@@ -29,6 +29,7 @@ use common_function::handlers::{
 };
 use common_query::Output;
 use datafusion::catalog::TableFunction;
+use datafusion::dataframe::DataFrame;
 use datafusion_expr::{AggregateUDF, LogicalPlan};
 use datatypes::schema::Schema;
 pub use default_serializer::{DefaultPlanDecoder, DefaultSerializer};
@@ -36,7 +37,6 @@ use partition::manager::PartitionRuleManagerRef;
 use session::context::QueryContextRef;
 use table::TableRef;
 
-use crate::dataframe::DataFrame;
 use crate::datafusion::DatafusionQueryEngine;
 use crate::error::Result;
 use crate::options::QueryOptions;

src/query/src/sql.rs

@@ -68,7 +68,6 @@ use table::TableRef;
 use table::requests::{FILE_TABLE_LOCATION_KEY, FILE_TABLE_PATTERN_KEY};
 
 use crate::QueryEngineRef;
-use crate::dataframe::DataFrame;
 use crate::error::{self, Result, UnsupportedVariableSnafu};
 use crate::planner::DfLogicalPlanner;
 
@@ -270,7 +269,7 @@ async fn query_from_information_schema_table(
             ),
         })?;
 
-    let DataFrame::DataFusion(dataframe) = query_engine.read_table(table)?;
+    let dataframe = query_engine.read_table(table)?;
 
     // Apply filters
     let dataframe = filters.into_iter().try_fold(dataframe, |df, expr| {

src/servers/src/prom_store.rs

@@ -27,11 +27,11 @@ use common_grpc::precision::Precision;
 use common_query::prelude::{greptime_timestamp, greptime_value};
 use common_recordbatch::{RecordBatch, RecordBatches};
 use common_telemetry::tracing;
+use datafusion::dataframe::DataFrame;
 use datafusion::prelude::{Expr, col, lit, regexp_match};
 use datafusion_common::ScalarValue;
 use datafusion_expr::LogicalPlan;
 use openmetrics_parser::{MetricsExposition, PrometheusType, PrometheusValue};
-use query::dataframe::DataFrame;
 use snafu::{OptionExt, ResultExt};
 use snap::raw::{Decoder, Encoder};
 
@@ -102,8 +102,6 @@ pub fn extract_schema_from_read_request(request: &ReadRequest) -> Option<String>
 /// Create a DataFrame from a remote Query
 #[tracing::instrument(skip_all)]
 pub fn query_to_plan(dataframe: DataFrame, q: &Query) -> Result<LogicalPlan> {
-    let DataFrame::DataFusion(dataframe) = dataframe;
-
     let start_timestamp_ms = q.start_timestamp_ms;
     let end_timestamp_ms = q.end_timestamp_ms;
 
@@ -654,7 +652,7 @@ mod tests {
         let table_provider = Arc::new(DfTableProviderAdapter::new(table));
 
         let dataframe = ctx.read_table(table_provider.clone()).unwrap();
-        let plan = query_to_plan(DataFrame::DataFusion(dataframe), &q).unwrap();
+        let plan = query_to_plan(dataframe, &q).unwrap();
         let display_string = format!("{}", plan.display_indent());
 
         let ts_col = greptime_timestamp();
@@ -688,7 +686,7 @@ mod tests {
         };
 
         let dataframe = ctx.read_table(table_provider).unwrap();
-        let plan = query_to_plan(DataFrame::DataFusion(dataframe), &q).unwrap();
+        let plan = query_to_plan(dataframe, &q).unwrap();
         let display_string = format!("{}", plan.display_indent());
 
         let ts_col = greptime_timestamp();

src/servers/src/prometheus.rs

@@ -16,10 +16,10 @@ use catalog::system_schema::information_schema::tables::{
     ENGINE as TABLE_ENGINE, TABLE_CATALOG, TABLE_NAME, TABLE_SCHEMA,
 };
 use common_telemetry::tracing;
+use datafusion::dataframe::DataFrame;
 use datafusion::prelude::{Expr, col, lit, regexp_match};
 use datafusion_expr::LogicalPlan;
 use promql_parser::label::{MatchOp, Matcher};
-use query::dataframe::DataFrame;
 use session::context::QueryContextRef;
 use snafu::ResultExt;
 
@@ -71,7 +71,6 @@ pub fn metric_name_matchers_to_plan(
     // Safety: conditions MUST not be empty, reduce always return Some(expr).
     let conditions = conditions.into_iter().reduce(Expr::and).unwrap();
 
-    let DataFrame::DataFusion(dataframe) = dataframe;
     let dataframe = dataframe
         .filter(conditions)
         .context(error::DataFrameSnafu)?