From 9d35b8cad4df3f82324e57c2e0109d20edd01638 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Thu, 4 Dec 2025 17:38:41 +0800 Subject: [PATCH] refactor: remove datafusion data frame wrapper (#7347) Signed-off-by: Ruihang Xia --- src/datanode/src/tests.rs | 2 +- src/frontend/src/instance/promql.rs | 2 +- src/pipeline/src/manager/table.rs | 3 --- src/query/src/dataframe.rs | 31 ----------------------------- src/query/src/datafusion.rs | 16 +++++++-------- src/query/src/lib.rs | 1 - src/query/src/query_engine.rs | 2 +- src/query/src/sql.rs | 3 +-- src/servers/src/prom_store.rs | 8 +++----- src/servers/src/prometheus.rs | 3 +-- 10 files changed, 15 insertions(+), 56 deletions(-) delete mode 100644 src/query/src/dataframe.rs diff --git a/src/datanode/src/tests.rs b/src/datanode/src/tests.rs index 557a8c92ed1..8a011c1f233 100644 --- a/src/datanode/src/tests.rs +++ b/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}; diff --git a/src/frontend/src/instance/promql.rs b/src/frontend/src/instance/promql.rs index 4236474cf72..f527b18a71d 100644 --- a/src/frontend/src/instance/promql.rs +++ b/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)?; diff --git a/src/pipeline/src/manager/table.rs b/src/pipeline/src/manager/table.rs index d3e63e634ee..5369478e9a1 100644 --- a/src/pipeline/src/manager/table.rs +++ b/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 diff --git a/src/query/src/dataframe.rs b/src/query/src/dataframe.rs deleted file mode 100644 index ce630b99e79..00000000000 --- a/src/query/src/dataframe.rs +++ /dev/null @@ -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, - } - } -} diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index f22a209de77..9c81b8d524d 100644 --- a/src/query/src/datafusion.rs +++ b/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 { - 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() diff --git a/src/query/src/lib.rs b/src/query/src/lib.rs index 0c6ccaa41f7..2b159a91b1d 100644 --- a/src/query/src/lib.rs +++ b/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; diff --git a/src/query/src/query_engine.rs b/src/query/src/query_engine.rs index 34a4fee2099..53383cf905e 100644 --- a/src/query/src/query_engine.rs +++ b/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; diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index 39ef00e1fb0..56da1134cac 100644 --- a/src/query/src/sql.rs +++ b/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| { diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 24e8512ad0a..92c1815145c 100644 --- a/src/servers/src/prom_store.rs +++ b/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 /// Create a DataFrame from a remote Query #[tracing::instrument(skip_all)] pub fn query_to_plan(dataframe: DataFrame, q: &Query) -> Result { - 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(); diff --git a/src/servers/src/prometheus.rs b/src/servers/src/prometheus.rs index 0fbdb97c34c..64653570de2 100644 --- a/src/servers/src/prometheus.rs +++ b/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)?