refactor: Restructure project layout and clean up dependencies (#30)

This commit introduces a significant reorganization of the project structure to improve maintainability and clarity.

Key changes include:
- Adjusted the directory layout for a more logical module organization.
- Removed unused crate dependencies, reducing the overall project size and potentially speeding up build times.
- Updated import paths and configuration files to reflect the structural changes.
This commit is contained in:
houseme
2025-07-02 19:33:12 +08:00
committed by GitHub
parent 0be4264eb1
commit 5826396cd0
322 changed files with 977 additions and 1542 deletions
@@ -0,0 +1,15 @@
// Copyright 2024 RustFS 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.
pub mod table_source;
@@ -0,0 +1,152 @@
// Copyright 2024 RustFS 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 std::any::Any;
use std::borrow::Cow;
use std::fmt::Display;
use std::sync::Arc;
use std::write;
use async_trait::async_trait;
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::common::Result as DFResult;
use datafusion::datasource::listing::ListingTable;
use datafusion::datasource::{TableProvider, provider_as_source};
use datafusion::error::DataFusionError;
use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, TableProviderFilterPushDown, TableSource};
use datafusion::prelude::Expr;
use datafusion::sql::TableReference;
use tracing::debug;
pub const TEMP_LOCATION_TABLE_NAME: &str = "external_location_table";
pub struct TableSourceAdapter {
database_name: String,
table_name: String,
table_handle: TableHandle,
plan: LogicalPlan,
}
impl TableSourceAdapter {
pub fn try_new(
table_ref: impl Into<TableReference>,
table_name: impl Into<String>,
table_handle: impl Into<TableHandle>,
) -> Result<Self, DataFusionError> {
let table_name: String = table_name.into();
let table_handle = table_handle.into();
let plan = match &table_handle {
// TableScan
TableHandle::External(t) => {
let table_source = provider_as_source(t.clone());
LogicalPlanBuilder::scan(table_ref, table_source, None)?.build()?
}
// TableScan
TableHandle::TableProvider(t) => {
let table_source = provider_as_source(t.clone());
if let Some(plan) = table_source.get_logical_plan() {
LogicalPlanBuilder::from(plan.into_owned()).build()?
} else {
LogicalPlanBuilder::scan(table_ref, table_source, None)?.build()?
}
}
};
debug!("Table source logical plan node of {}:\n{}", table_name, plan.display_indent_schema());
Ok(Self {
database_name: "default_db".to_string(),
table_name,
table_handle,
plan,
})
}
pub fn database_name(&self) -> &str {
&self.database_name
}
pub fn table_name(&self) -> &str {
&self.table_name
}
pub fn table_handle(&self) -> &TableHandle {
&self.table_handle
}
}
#[async_trait]
impl TableSource for TableSourceAdapter {
fn as_any(&self) -> &dyn Any {
self
}
fn schema(&self) -> SchemaRef {
self.table_handle.schema()
}
fn supports_filters_pushdown(&self, filter: &[&Expr]) -> DFResult<Vec<TableProviderFilterPushDown>> {
self.table_handle.supports_filters_pushdown(filter)
}
/// Called by [`InlineTableScan`]
fn get_logical_plan(&self) -> Option<Cow<LogicalPlan>> {
Some(Cow::Owned(self.plan.clone()))
}
}
#[derive(Clone)]
pub enum TableHandle {
TableProvider(Arc<dyn TableProvider>),
External(Arc<ListingTable>),
}
impl TableHandle {
pub fn schema(&self) -> SchemaRef {
match self {
Self::External(t) => t.schema(),
Self::TableProvider(t) => t.schema(),
}
}
pub fn supports_filters_pushdown(&self, filter: &[&Expr]) -> DFResult<Vec<TableProviderFilterPushDown>> {
match self {
Self::External(t) => t.supports_filters_pushdown(filter),
Self::TableProvider(t) => t.supports_filters_pushdown(filter),
}
}
}
impl From<Arc<dyn TableProvider>> for TableHandle {
fn from(value: Arc<dyn TableProvider>) -> Self {
TableHandle::TableProvider(value)
}
}
impl From<Arc<ListingTable>> for TableHandle {
fn from(value: Arc<ListingTable>) -> Self {
TableHandle::External(value)
}
}
impl Display for TableHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::External(e) => write!(f, "External({:?})", e.table_paths()),
Self::TableProvider(_) => write!(f, "TableProvider"),
}
}
}
@@ -0,0 +1,367 @@
// Copyright 2024 RustFS 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 std::{
ops::Deref,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
use async_trait::async_trait;
use datafusion::{
arrow::{
datatypes::{Schema, SchemaRef},
record_batch::RecordBatch,
},
datasource::{
file_format::{csv::CsvFormat, json::JsonFormat, parquet::ParquetFormat},
listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl},
},
error::Result as DFResult,
execution::{RecordBatchStream, SendableRecordBatchStream},
};
use futures::{Stream, StreamExt};
use lazy_static::lazy_static;
use rustfs_s3select_api::{
QueryError, QueryResult,
query::{
Query,
ast::ExtStatement,
dispatcher::QueryDispatcher,
execution::{Output, QueryStateMachine},
function::FuncMetaManagerRef,
logical_planner::{LogicalPlanner, Plan},
parser::Parser,
session::{SessionCtx, SessionCtxFactory},
},
};
use s3s::dto::{FileHeaderInfo, SelectObjectContentInput};
use crate::{
execution::factory::QueryExecutionFactoryRef,
metadata::{ContextProviderExtension, MetadataProvider, TableHandleProviderRef, base_table::BaseTableProvider},
sql::logical::planner::DefaultLogicalPlanner,
};
lazy_static! {
static ref IGNORE: FileHeaderInfo = FileHeaderInfo::from_static(FileHeaderInfo::IGNORE);
static ref NONE: FileHeaderInfo = FileHeaderInfo::from_static(FileHeaderInfo::NONE);
static ref USE: FileHeaderInfo = FileHeaderInfo::from_static(FileHeaderInfo::USE);
}
#[derive(Clone)]
pub struct SimpleQueryDispatcher {
input: Arc<SelectObjectContentInput>,
// client for default tenant
_default_table_provider: TableHandleProviderRef,
session_factory: Arc<SessionCtxFactory>,
// parser
parser: Arc<dyn Parser + Send + Sync>,
// get query execution factory
query_execution_factory: QueryExecutionFactoryRef,
func_manager: FuncMetaManagerRef,
}
#[async_trait]
impl QueryDispatcher for SimpleQueryDispatcher {
async fn execute_query(&self, query: &Query) -> QueryResult<Output> {
let query_state_machine = { self.build_query_state_machine(query.clone()).await? };
let logical_plan = self.build_logical_plan(query_state_machine.clone()).await?;
let logical_plan = match logical_plan {
Some(plan) => plan,
None => return Ok(Output::Nil(())),
};
let result = self.execute_logical_plan(logical_plan, query_state_machine).await?;
Ok(result)
}
async fn build_logical_plan(&self, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Option<Plan>> {
let session = &query_state_machine.session;
let query = &query_state_machine.query;
let scheme_provider = self.build_scheme_provider(session).await?;
let logical_planner = DefaultLogicalPlanner::new(&scheme_provider);
let statements = self.parser.parse(query.content())?;
// not allow multi statement
if statements.len() > 1 {
return Err(QueryError::MultiStatement {
num: statements.len(),
sql: query_state_machine.query.content().to_string(),
});
}
let stmt = match statements.front() {
Some(stmt) => stmt.clone(),
None => return Ok(None),
};
let logical_plan = self
.statement_to_logical_plan(stmt, &logical_planner, query_state_machine)
.await?;
Ok(Some(logical_plan))
}
async fn execute_logical_plan(&self, logical_plan: Plan, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Output> {
self.execute_logical_plan(logical_plan, query_state_machine).await
}
async fn build_query_state_machine(&self, query: Query) -> QueryResult<Arc<QueryStateMachine>> {
let session = self.session_factory.create_session_ctx(query.context()).await?;
let query_state_machine = Arc::new(QueryStateMachine::begin(query, session));
Ok(query_state_machine)
}
}
impl SimpleQueryDispatcher {
async fn statement_to_logical_plan<S: ContextProviderExtension + Send + Sync>(
&self,
stmt: ExtStatement,
logical_planner: &DefaultLogicalPlanner<'_, S>,
query_state_machine: Arc<QueryStateMachine>,
) -> QueryResult<Plan> {
// begin analyze
query_state_machine.begin_analyze();
let logical_plan = logical_planner
.create_logical_plan(stmt, &query_state_machine.session)
.await?;
query_state_machine.end_analyze();
Ok(logical_plan)
}
async fn execute_logical_plan(&self, logical_plan: Plan, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Output> {
let execution = self
.query_execution_factory
.create_query_execution(logical_plan, query_state_machine.clone())
.await?;
match execution.start().await {
Ok(Output::StreamData(stream)) => Ok(Output::StreamData(Box::pin(TrackedRecordBatchStream { inner: stream }))),
Ok(nil @ Output::Nil(_)) => Ok(nil),
Err(err) => Err(err),
}
}
async fn build_scheme_provider(&self, session: &SessionCtx) -> QueryResult<MetadataProvider> {
let path = format!("s3://{}/{}", self.input.bucket, self.input.key);
let table_path = ListingTableUrl::parse(path)?;
let (listing_options, need_rename_volume_name, need_ignore_volume_name) =
if let Some(csv) = self.input.request.input_serialization.csv.as_ref() {
let mut need_rename_volume_name = false;
let mut need_ignore_volume_name = false;
let mut file_format = CsvFormat::default()
.with_comment(
csv.comments
.clone()
.map(|c| c.as_bytes().first().copied().unwrap_or_default()),
)
.with_escape(
csv.quote_escape_character
.clone()
.map(|e| e.as_bytes().first().copied().unwrap_or_default()),
);
if let Some(delimiter) = csv.field_delimiter.as_ref() {
if delimiter.len() == 1 {
file_format = file_format.with_delimiter(delimiter.as_bytes()[0]);
}
}
// TODO waiting for processing @junxiang Mu
// if csv.file_header_info.is_some() {}
match csv.file_header_info.as_ref() {
Some(info) => {
if *info == *NONE {
file_format = file_format.with_has_header(false);
need_rename_volume_name = true;
} else if *info == *IGNORE {
file_format = file_format.with_has_header(true);
need_rename_volume_name = true;
need_ignore_volume_name = true;
} else if *info == *USE {
file_format = file_format.with_has_header(true);
} else {
return Err(QueryError::NotImplemented {
err: "unsupported FileHeaderInfo".to_string(),
});
}
}
_ => {
return Err(QueryError::NotImplemented {
err: "unsupported FileHeaderInfo".to_string(),
});
}
}
if let Some(quote) = csv.quote_character.as_ref() {
file_format = file_format.with_quote(quote.as_bytes().first().copied().unwrap_or_default());
}
(
ListingOptions::new(Arc::new(file_format)).with_file_extension(".csv"),
need_rename_volume_name,
need_ignore_volume_name,
)
} else if self.input.request.input_serialization.parquet.is_some() {
let file_format = ParquetFormat::new();
(ListingOptions::new(Arc::new(file_format)).with_file_extension(".parquet"), false, false)
} else if self.input.request.input_serialization.json.is_some() {
let file_format = JsonFormat::default();
(ListingOptions::new(Arc::new(file_format)).with_file_extension(".json"), false, false)
} else {
return Err(QueryError::NotImplemented {
err: "not support this file type".to_string(),
});
};
let resolve_schema = listing_options.infer_schema(session.inner(), &table_path).await?;
let config = if need_rename_volume_name {
let mut new_fields = Vec::new();
for (i, field) in resolve_schema.fields().iter().enumerate() {
let f_name = field.name();
let mut_field = field.deref().clone();
if f_name.starts_with("column_") {
let re_name = f_name.replace("column_", "_");
new_fields.push(mut_field.with_name(re_name));
} else if need_ignore_volume_name {
let re_name = format!("_{}", i + 1);
new_fields.push(mut_field.with_name(re_name));
} else {
new_fields.push(mut_field);
}
}
let new_schema = Arc::new(Schema::new(new_fields).with_metadata(resolve_schema.metadata().clone()));
ListingTableConfig::new(table_path)
.with_listing_options(listing_options)
.with_schema(new_schema)
} else {
ListingTableConfig::new(table_path)
.with_listing_options(listing_options)
.with_schema(resolve_schema)
};
// rename default
let provider = Arc::new(ListingTable::try_new(config)?);
let current_session_table_provider = self.build_table_handle_provider()?;
let metadata_provider =
MetadataProvider::new(provider, current_session_table_provider, self.func_manager.clone(), session.clone());
Ok(metadata_provider)
}
fn build_table_handle_provider(&self) -> QueryResult<TableHandleProviderRef> {
let current_session_table_provider: Arc<BaseTableProvider> = Arc::new(BaseTableProvider::default());
Ok(current_session_table_provider)
}
}
pub struct TrackedRecordBatchStream {
inner: SendableRecordBatchStream,
}
impl RecordBatchStream for TrackedRecordBatchStream {
fn schema(&self) -> SchemaRef {
self.inner.schema()
}
}
impl Stream for TrackedRecordBatchStream {
type Item = DFResult<RecordBatch>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.inner.poll_next_unpin(cx)
}
}
#[derive(Default, Clone)]
pub struct SimpleQueryDispatcherBuilder {
input: Option<Arc<SelectObjectContentInput>>,
default_table_provider: Option<TableHandleProviderRef>,
session_factory: Option<Arc<SessionCtxFactory>>,
parser: Option<Arc<dyn Parser + Send + Sync>>,
query_execution_factory: Option<QueryExecutionFactoryRef>,
func_manager: Option<FuncMetaManagerRef>,
}
impl SimpleQueryDispatcherBuilder {
pub fn with_input(mut self, input: Arc<SelectObjectContentInput>) -> Self {
self.input = Some(input);
self
}
pub fn with_default_table_provider(mut self, default_table_provider: TableHandleProviderRef) -> Self {
self.default_table_provider = Some(default_table_provider);
self
}
pub fn with_session_factory(mut self, session_factory: Arc<SessionCtxFactory>) -> Self {
self.session_factory = Some(session_factory);
self
}
pub fn with_parser(mut self, parser: Arc<dyn Parser + Send + Sync>) -> Self {
self.parser = Some(parser);
self
}
pub fn with_query_execution_factory(mut self, query_execution_factory: QueryExecutionFactoryRef) -> Self {
self.query_execution_factory = Some(query_execution_factory);
self
}
pub fn with_func_manager(mut self, func_manager: FuncMetaManagerRef) -> Self {
self.func_manager = Some(func_manager);
self
}
pub fn build(self) -> QueryResult<Arc<SimpleQueryDispatcher>> {
let input = self.input.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of input".to_string(),
})?;
let session_factory = self.session_factory.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of session_factory".to_string(),
})?;
let parser = self.parser.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of parser".to_string(),
})?;
let query_execution_factory = self.query_execution_factory.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of query_execution_factory".to_string(),
})?;
let func_manager = self.func_manager.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of func_manager".to_string(),
})?;
let default_table_provider = self.default_table_provider.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of default_table_provider".to_string(),
})?;
let dispatcher = Arc::new(SimpleQueryDispatcher {
input,
_default_table_provider: default_table_provider,
session_factory,
parser,
query_execution_factory,
func_manager,
});
Ok(dispatcher)
}
}
@@ -0,0 +1,15 @@
// Copyright 2024 RustFS 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.
pub mod manager;
@@ -0,0 +1,60 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use async_trait::async_trait;
use rustfs_s3select_api::{
QueryError,
query::{
execution::{QueryExecutionFactory, QueryExecutionRef, QueryStateMachineRef},
logical_planner::Plan,
optimizer::Optimizer,
scheduler::SchedulerRef,
},
};
use super::query::SqlQueryExecution;
pub type QueryExecutionFactoryRef = Arc<dyn QueryExecutionFactory + Send + Sync>;
pub struct SqlQueryExecutionFactory {
optimizer: Arc<dyn Optimizer + Send + Sync>,
scheduler: SchedulerRef,
}
impl SqlQueryExecutionFactory {
#[inline(always)]
pub fn new(optimizer: Arc<dyn Optimizer + Send + Sync>, scheduler: SchedulerRef) -> Self {
Self { optimizer, scheduler }
}
}
#[async_trait]
impl QueryExecutionFactory for SqlQueryExecutionFactory {
async fn create_query_execution(
&self,
plan: Plan,
state_machine: QueryStateMachineRef,
) -> Result<QueryExecutionRef, QueryError> {
match plan {
Plan::Query(query_plan) => Ok(Arc::new(SqlQueryExecution::new(
state_machine,
query_plan,
self.optimizer.clone(),
self.scheduler.clone(),
))),
}
}
}
@@ -0,0 +1,17 @@
// Copyright 2024 RustFS 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.
pub mod factory;
pub mod query;
pub mod scheduler;
@@ -0,0 +1,106 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use async_trait::async_trait;
use futures::stream::AbortHandle;
use parking_lot::Mutex;
use rustfs_s3select_api::query::execution::{Output, QueryExecution, QueryStateMachineRef};
use rustfs_s3select_api::query::logical_planner::QueryPlan;
use rustfs_s3select_api::query::optimizer::Optimizer;
use rustfs_s3select_api::query::scheduler::SchedulerRef;
use rustfs_s3select_api::{QueryError, QueryResult};
use tracing::debug;
pub struct SqlQueryExecution {
query_state_machine: QueryStateMachineRef,
plan: QueryPlan,
optimizer: Arc<dyn Optimizer + Send + Sync>,
scheduler: SchedulerRef,
abort_handle: Mutex<Option<AbortHandle>>,
}
impl SqlQueryExecution {
pub fn new(
query_state_machine: QueryStateMachineRef,
plan: QueryPlan,
optimizer: Arc<dyn Optimizer + Send + Sync>,
scheduler: SchedulerRef,
) -> Self {
Self {
query_state_machine,
plan,
optimizer,
scheduler,
abort_handle: Mutex::new(None),
}
}
async fn start(&self) -> QueryResult<Output> {
// begin optimize
self.query_state_machine.begin_optimize();
let physical_plan = self.optimizer.optimize(&self.plan, &self.query_state_machine.session).await?;
self.query_state_machine.end_optimize();
// begin schedule
self.query_state_machine.begin_schedule();
let stream = self
.scheduler
.schedule(physical_plan.clone(), self.query_state_machine.session.inner().task_ctx())
.await?
.stream();
debug!("Success build result stream.");
self.query_state_machine.end_schedule();
Ok(Output::StreamData(stream))
}
}
#[async_trait]
impl QueryExecution for SqlQueryExecution {
async fn start(&self) -> QueryResult<Output> {
let (task, abort_handle) = futures::future::abortable(self.start());
{
*self.abort_handle.lock() = Some(abort_handle);
}
task.await.map_err(|_| QueryError::Cancel)?
}
fn cancel(&self) -> QueryResult<()> {
debug!(
"cancel sql query execution: sql: {}, state: {:?}",
self.query_state_machine.query.content(),
self.query_state_machine.state()
);
// change state
self.query_state_machine.cancel();
// stop future task
if let Some(e) = self.abort_handle.lock().as_ref() {
e.abort()
};
debug!(
"canceled sql query execution: sql: {}, state: {:?}",
self.query_state_machine.query.content(),
self.query_state_machine.state()
);
Ok(())
}
}
@@ -0,0 +1,36 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use async_trait::async_trait;
use datafusion::error::DataFusionError;
use datafusion::execution::context::TaskContext;
use datafusion::physical_plan::{ExecutionPlan, execute_stream};
use rustfs_s3select_api::query::scheduler::{ExecutionResults, Scheduler};
pub struct LocalScheduler {}
#[async_trait]
impl Scheduler for LocalScheduler {
async fn schedule(
&self,
plan: Arc<dyn ExecutionPlan>,
context: Arc<TaskContext>,
) -> Result<ExecutionResults, DataFusionError> {
let stream = execute_stream(plan, context)?;
Ok(ExecutionResults::new(stream))
}
}
@@ -0,0 +1,15 @@
// Copyright 2024 RustFS 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.
pub mod local;
+15
View File
@@ -0,0 +1,15 @@
// Copyright 2024 RustFS 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.
pub mod simple_func_manager;
@@ -0,0 +1,119 @@
// Copyright 2024 RustFS 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 std::collections::HashMap;
use std::sync::Arc;
use datafusion::execution::SessionStateDefaults;
use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
use rustfs_s3select_api::query::function::FunctionMetadataManager;
use rustfs_s3select_api::{QueryError, QueryResult};
use tracing::debug;
pub type SimpleFunctionMetadataManagerRef = Arc<SimpleFunctionMetadataManager>;
#[derive(Debug)]
pub struct SimpleFunctionMetadataManager {
/// Scalar functions that are registered with the context
pub scalar_functions: HashMap<String, Arc<ScalarUDF>>,
/// Aggregate functions registered in the context
pub aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
/// Window functions registered in the context
pub window_functions: HashMap<String, Arc<WindowUDF>>,
}
impl Default for SimpleFunctionMetadataManager {
fn default() -> Self {
let mut func_meta_manager = Self {
scalar_functions: Default::default(),
aggregate_functions: Default::default(),
window_functions: Default::default(),
};
SessionStateDefaults::default_scalar_functions().into_iter().for_each(|udf| {
let existing_udf = func_meta_manager.register_udf(udf.clone());
if let Ok(()) = existing_udf {
debug!("Overwrote an existing UDF: {}", udf.name());
}
});
SessionStateDefaults::default_aggregate_functions()
.into_iter()
.for_each(|udaf| {
let existing_udaf = func_meta_manager.register_udaf(udaf.clone());
if let Ok(()) = existing_udaf {
debug!("Overwrote an existing UDAF: {}", udaf.name());
}
});
SessionStateDefaults::default_window_functions().into_iter().for_each(|udwf| {
let existing_udwf = func_meta_manager.register_udwf(udwf.clone());
if let Ok(()) = existing_udwf {
debug!("Overwrote an existing UDWF: {}", udwf.name());
}
});
func_meta_manager
}
}
impl FunctionMetadataManager for SimpleFunctionMetadataManager {
fn register_udf(&mut self, f: Arc<ScalarUDF>) -> QueryResult<()> {
self.scalar_functions.insert(f.inner().name().to_uppercase(), f);
Ok(())
}
fn register_udaf(&mut self, f: Arc<AggregateUDF>) -> QueryResult<()> {
self.aggregate_functions.insert(f.inner().name().to_uppercase(), f);
Ok(())
}
fn register_udwf(&mut self, f: Arc<WindowUDF>) -> QueryResult<()> {
self.window_functions.insert(f.inner().name().to_uppercase(), f);
Ok(())
}
fn udf(&self, name: &str) -> QueryResult<Arc<ScalarUDF>> {
let result = self.scalar_functions.get(&name.to_uppercase());
result
.cloned()
.ok_or_else(|| QueryError::FunctionExists { name: name.to_string() })
}
fn udaf(&self, name: &str) -> QueryResult<Arc<AggregateUDF>> {
let result = self.aggregate_functions.get(&name.to_uppercase());
result
.cloned()
.ok_or_else(|| QueryError::FunctionNotExists { name: name.to_string() })
}
fn udwf(&self, name: &str) -> QueryResult<Arc<WindowUDF>> {
let result = self.window_functions.get(&name.to_uppercase());
result
.cloned()
.ok_or_else(|| QueryError::FunctionNotExists { name: name.to_string() })
}
fn udfs(&self) -> Vec<String> {
self.scalar_functions.keys().cloned().collect()
}
fn udafs(&self) -> Vec<String> {
self.aggregate_functions.keys().cloned().collect()
}
fn udwfs(&self) -> Vec<String> {
self.window_functions.keys().cloned().collect()
}
}
+222
View File
@@ -0,0 +1,222 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use async_trait::async_trait;
use derive_builder::Builder;
use rustfs_s3select_api::{
QueryResult,
query::{
Query, dispatcher::QueryDispatcher, execution::QueryStateMachineRef, logical_planner::Plan, session::SessionCtxFactory,
},
server::dbms::{DatabaseManagerSystem, QueryHandle},
};
use s3s::dto::SelectObjectContentInput;
use crate::{
dispatcher::manager::SimpleQueryDispatcherBuilder,
execution::{factory::SqlQueryExecutionFactory, scheduler::local::LocalScheduler},
function::simple_func_manager::SimpleFunctionMetadataManager,
metadata::base_table::BaseTableProvider,
sql::{optimizer::CascadeOptimizerBuilder, parser::DefaultParser},
};
#[derive(Builder)]
pub struct RustFSms<D: QueryDispatcher> {
// query dispatcher & query execution
query_dispatcher: Arc<D>,
}
#[async_trait]
impl<D> DatabaseManagerSystem for RustFSms<D>
where
D: QueryDispatcher,
{
async fn execute(&self, query: &Query) -> QueryResult<QueryHandle> {
let result = self.query_dispatcher.execute_query(query).await?;
Ok(QueryHandle::new(query.clone(), result))
}
async fn build_query_state_machine(&self, query: Query) -> QueryResult<QueryStateMachineRef> {
let query_state_machine = self.query_dispatcher.build_query_state_machine(query).await?;
Ok(query_state_machine)
}
async fn build_logical_plan(&self, query_state_machine: QueryStateMachineRef) -> QueryResult<Option<Plan>> {
let logical_plan = self.query_dispatcher.build_logical_plan(query_state_machine).await?;
Ok(logical_plan)
}
async fn execute_logical_plan(
&self,
logical_plan: Plan,
query_state_machine: QueryStateMachineRef,
) -> QueryResult<QueryHandle> {
let query = query_state_machine.query.clone();
let result = self
.query_dispatcher
.execute_logical_plan(logical_plan, query_state_machine)
.await?;
Ok(QueryHandle::new(query.clone(), result))
}
}
pub async fn make_rustfsms(input: Arc<SelectObjectContentInput>, is_test: bool) -> QueryResult<impl DatabaseManagerSystem> {
// init Function Manager, we can define some UDF if need
let func_manager = SimpleFunctionMetadataManager::default();
// TODO session config need load global system config
let session_factory = Arc::new(SessionCtxFactory { is_test });
let parser = Arc::new(DefaultParser::default());
let optimizer = Arc::new(CascadeOptimizerBuilder::default().build());
// TODO wrap, and num_threads configurable
let scheduler = Arc::new(LocalScheduler {});
let query_execution_factory = Arc::new(SqlQueryExecutionFactory::new(optimizer, scheduler));
let default_table_provider = Arc::new(BaseTableProvider::default());
let query_dispatcher = SimpleQueryDispatcherBuilder::default()
.with_input(input)
.with_func_manager(Arc::new(func_manager))
.with_default_table_provider(default_table_provider)
.with_session_factory(session_factory)
.with_parser(parser)
.with_query_execution_factory(query_execution_factory)
.build()?;
let mut builder = RustFSmsBuilder::default();
let db_server = builder.query_dispatcher(query_dispatcher).build().expect("build db server");
Ok(db_server)
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use datafusion::{arrow::util::pretty, assert_batches_eq};
use rustfs_s3select_api::{
query::{Context, Query},
server::dbms::DatabaseManagerSystem,
};
use s3s::dto::{
CSVInput, CSVOutput, ExpressionType, FieldDelimiter, FileHeaderInfo, InputSerialization, OutputSerialization,
RecordDelimiter, SelectObjectContentInput, SelectObjectContentRequest,
};
use crate::instance::make_rustfsms;
#[tokio::test]
#[ignore]
async fn test_simple_sql() {
let sql = "select * from S3Object";
let input = Arc::new(SelectObjectContentInput {
bucket: "dandan".to_string(),
expected_bucket_owner: None,
key: "test.csv".to_string(),
sse_customer_algorithm: None,
sse_customer_key: None,
sse_customer_key_md5: None,
request: SelectObjectContentRequest {
expression: sql.to_string(),
expression_type: ExpressionType::from_static("SQL"),
input_serialization: InputSerialization {
csv: Some(CSVInput {
file_header_info: Some(FileHeaderInfo::from_static(FileHeaderInfo::USE)),
..Default::default()
}),
..Default::default()
},
output_serialization: OutputSerialization {
csv: Some(CSVOutput::default()),
..Default::default()
},
request_progress: None,
scan_range: None,
},
});
let db = make_rustfsms(input.clone(), true).await.unwrap();
let query = Query::new(Context { input }, sql.to_string());
let result = db.execute(&query).await.unwrap();
let results = result.result().chunk_result().await.unwrap().to_vec();
let expected = [
"+----------------+---------+-----+------------+--------+",
"| id | name | age | department | salary |",
"+----------------+---------+-----+------------+--------+",
"| 1 | Alice | 25 | HR | 5000 |",
"| 2 | Bob | 30 | IT | 6000 |",
"| 3 | Charlie | 35 | Finance | 7000 |",
"| 4 | Diana | 22 | Marketing | 4500 |",
"| 5 | Eve | 28 | IT | 5500 |",
"| 6 | Frank | 40 | Finance | 8000 |",
"| 7 | Grace | 26 | HR | 5200 |",
"| 8 | Henry | 32 | IT | 6200 |",
"| 9 | Ivy | 24 | Marketing | 4800 |",
"| 10 | Jack | 38 | Finance | 7500 |",
"+----------------+---------+-----+------------+--------+",
];
assert_batches_eq!(expected, &results);
pretty::print_batches(&results).unwrap();
}
#[tokio::test]
#[ignore]
async fn test_func_sql() {
let sql = "SELECT * FROM S3Object s";
let input = Arc::new(SelectObjectContentInput {
bucket: "dandan".to_string(),
expected_bucket_owner: None,
key: "test.csv".to_string(),
sse_customer_algorithm: None,
sse_customer_key: None,
sse_customer_key_md5: None,
request: SelectObjectContentRequest {
expression: sql.to_string(),
expression_type: ExpressionType::from_static("SQL"),
input_serialization: InputSerialization {
csv: Some(CSVInput {
file_header_info: Some(FileHeaderInfo::from_static(FileHeaderInfo::IGNORE)),
field_delimiter: Some(FieldDelimiter::from("")),
record_delimiter: Some(RecordDelimiter::from("\n")),
..Default::default()
}),
..Default::default()
},
output_serialization: OutputSerialization {
csv: Some(CSVOutput::default()),
..Default::default()
},
request_progress: None,
scan_range: None,
},
});
let db = make_rustfsms(input.clone(), true).await.unwrap();
let query = Query::new(Context { input }, sql.to_string());
let result = db.execute(&query).await.unwrap();
let results = result.result().chunk_result().await.unwrap().to_vec();
pretty::print_batches(&results).unwrap();
}
}
+21
View File
@@ -0,0 +1,21 @@
// Copyright 2024 RustFS 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.
pub mod data_source;
pub mod dispatcher;
pub mod execution;
pub mod function;
pub mod instance;
pub mod metadata;
pub mod sql;
@@ -0,0 +1,31 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use datafusion::common::Result as DFResult;
use datafusion::datasource::listing::ListingTable;
use crate::data_source::table_source::TableHandle;
use super::TableHandleProvider;
#[derive(Default)]
pub struct BaseTableProvider {}
impl TableHandleProvider for BaseTableProvider {
fn build_table_handle(&self, provider: Arc<ListingTable>) -> DFResult<TableHandle> {
Ok(TableHandle::External(provider))
}
}
+140
View File
@@ -0,0 +1,140 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use async_trait::async_trait;
use datafusion::arrow::datatypes::DataType;
use datafusion::common::Result as DFResult;
use datafusion::datasource::listing::ListingTable;
use datafusion::logical_expr::var_provider::is_system_variables;
use datafusion::logical_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF};
use datafusion::variable::VarType;
use datafusion::{
config::ConfigOptions,
sql::{TableReference, planner::ContextProvider},
};
use rustfs_s3select_api::query::{function::FuncMetaManagerRef, session::SessionCtx};
use crate::data_source::table_source::{TableHandle, TableSourceAdapter};
pub mod base_table;
#[async_trait]
pub trait ContextProviderExtension: ContextProvider {
fn get_table_source_(&self, name: TableReference) -> datafusion::common::Result<Arc<TableSourceAdapter>>;
}
pub type TableHandleProviderRef = Arc<dyn TableHandleProvider + Send + Sync>;
pub trait TableHandleProvider {
fn build_table_handle(&self, provider: Arc<ListingTable>) -> DFResult<TableHandle>;
}
pub struct MetadataProvider {
provider: Arc<ListingTable>,
session: SessionCtx,
config_options: ConfigOptions,
func_manager: FuncMetaManagerRef,
current_session_table_provider: TableHandleProviderRef,
}
impl MetadataProvider {
#[allow(clippy::too_many_arguments)]
pub fn new(
provider: Arc<ListingTable>,
current_session_table_provider: TableHandleProviderRef,
func_manager: FuncMetaManagerRef,
session: SessionCtx,
) -> Self {
Self {
provider,
current_session_table_provider,
config_options: session.inner().config_options().clone(),
session,
func_manager,
}
}
fn build_table_handle(&self) -> datafusion::common::Result<TableHandle> {
self.current_session_table_provider.build_table_handle(self.provider.clone())
}
}
impl ContextProviderExtension for MetadataProvider {
fn get_table_source_(&self, table_ref: TableReference) -> datafusion::common::Result<Arc<TableSourceAdapter>> {
let name = table_ref.clone().resolve("", "");
let table_name = &*name.table;
let table_handle = self.build_table_handle()?;
Ok(Arc::new(TableSourceAdapter::try_new(table_ref.clone(), table_name, table_handle)?))
}
}
impl ContextProvider for MetadataProvider {
fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
self.func_manager
.udf(name)
.ok()
.or(self.session.inner().scalar_functions().get(name).cloned())
}
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
self.func_manager.udaf(name).ok()
}
fn get_variable_type(&self, variable_names: &[String]) -> Option<DataType> {
if variable_names.is_empty() {
return None;
}
let var_type = if is_system_variables(variable_names) {
VarType::System
} else {
VarType::UserDefined
};
self.session
.inner()
.execution_props()
.get_var_provider(var_type)
.and_then(|p| p.get_type(variable_names))
}
fn options(&self) -> &ConfigOptions {
// TODO refactor
&self.config_options
}
fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
self.func_manager.udwf(name).ok()
}
fn get_table_source(&self, name: TableReference) -> DFResult<Arc<dyn TableSource>> {
Ok(self.get_table_source_(name)?)
}
fn udf_names(&self) -> Vec<String> {
self.func_manager.udfs()
}
fn udaf_names(&self) -> Vec<String> {
self.func_manager.udafs()
}
fn udwf_names(&self) -> Vec<String> {
self.func_manager.udwfs()
}
}
+47
View File
@@ -0,0 +1,47 @@
// Copyright 2024 RustFS 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::logical_expr::LogicalPlan;
use datafusion::optimizer::analyzer::Analyzer as DFAnalyzer;
use rustfs_s3select_api::QueryResult;
use rustfs_s3select_api::query::analyzer::Analyzer;
use rustfs_s3select_api::query::session::SessionCtx;
pub struct DefaultAnalyzer {
inner: DFAnalyzer,
}
impl DefaultAnalyzer {
pub fn new() -> Self {
let analyzer = DFAnalyzer::default();
// we can add analyzer rule at here
Self { inner: analyzer }
}
}
impl Default for DefaultAnalyzer {
fn default() -> Self {
Self::new()
}
}
impl Analyzer for DefaultAnalyzer {
fn analyze(&self, plan: &LogicalPlan, session: &SessionCtx) -> QueryResult<LogicalPlan> {
let plan = self
.inner
.execute_and_check(plan.to_owned(), session.inner().config_options(), |_, _| {})?;
Ok(plan)
}
}
+345
View File
@@ -0,0 +1,345 @@
// Copyright 2024 RustFS 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::sql::sqlparser::dialect::Dialect;
#[derive(Debug, Default)]
pub struct RustFsDialect;
impl Dialect for RustFsDialect {
fn is_identifier_start(&self, ch: char) -> bool {
ch.is_alphabetic() || ch == '_' || ch == '#' || ch == '@'
}
fn is_identifier_part(&self, ch: char) -> bool {
ch.is_alphabetic() || ch.is_ascii_digit() || ch == '@' || ch == '$' || ch == '#' || ch == '_'
}
fn supports_group_by_expr(&self) -> bool {
true
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_rustfs_dialect_creation() {
let _dialect = RustFsDialect;
// Test that dialect can be created successfully
assert!(std::mem::size_of::<RustFsDialect>() == 0, "Dialect should be zero-sized");
}
#[test]
fn test_rustfs_dialect_debug() {
let dialect = RustFsDialect;
let debug_str = format!("{dialect:?}");
assert!(!debug_str.is_empty(), "Debug output should not be empty");
assert!(debug_str.contains("RustFsDialect"), "Debug output should contain dialect name");
}
#[test]
fn test_is_identifier_start_alphabetic() {
let dialect = RustFsDialect;
// Test alphabetic characters
assert!(dialect.is_identifier_start('a'), "Lowercase letter should be valid identifier start");
assert!(dialect.is_identifier_start('A'), "Uppercase letter should be valid identifier start");
assert!(dialect.is_identifier_start('z'), "Last lowercase letter should be valid identifier start");
assert!(dialect.is_identifier_start('Z'), "Last uppercase letter should be valid identifier start");
// Test Unicode alphabetic characters
assert!(dialect.is_identifier_start('α'), "Greek letter should be valid identifier start");
assert!(dialect.is_identifier_start('中'), "Chinese character should be valid identifier start");
assert!(dialect.is_identifier_start('ñ'), "Accented letter should be valid identifier start");
}
#[test]
fn test_is_identifier_start_special_chars() {
let dialect = RustFsDialect;
// Test special characters that are allowed
assert!(dialect.is_identifier_start('_'), "Underscore should be valid identifier start");
assert!(dialect.is_identifier_start('#'), "Hash should be valid identifier start");
assert!(dialect.is_identifier_start('@'), "At symbol should be valid identifier start");
}
#[test]
fn test_is_identifier_start_invalid_chars() {
let dialect = RustFsDialect;
// Test characters that should not be valid identifier starts
assert!(!dialect.is_identifier_start('0'), "Digit should not be valid identifier start");
assert!(!dialect.is_identifier_start('9'), "Digit should not be valid identifier start");
assert!(!dialect.is_identifier_start('$'), "Dollar sign should not be valid identifier start");
assert!(!dialect.is_identifier_start(' '), "Space should not be valid identifier start");
assert!(!dialect.is_identifier_start('\t'), "Tab should not be valid identifier start");
assert!(!dialect.is_identifier_start('\n'), "Newline should not be valid identifier start");
assert!(!dialect.is_identifier_start('.'), "Dot should not be valid identifier start");
assert!(!dialect.is_identifier_start(','), "Comma should not be valid identifier start");
assert!(!dialect.is_identifier_start(';'), "Semicolon should not be valid identifier start");
assert!(!dialect.is_identifier_start('('), "Left paren should not be valid identifier start");
assert!(!dialect.is_identifier_start(')'), "Right paren should not be valid identifier start");
assert!(!dialect.is_identifier_start('['), "Left bracket should not be valid identifier start");
assert!(!dialect.is_identifier_start(']'), "Right bracket should not be valid identifier start");
assert!(!dialect.is_identifier_start('{'), "Left brace should not be valid identifier start");
assert!(!dialect.is_identifier_start('}'), "Right brace should not be valid identifier start");
assert!(!dialect.is_identifier_start('='), "Equals should not be valid identifier start");
assert!(!dialect.is_identifier_start('+'), "Plus should not be valid identifier start");
assert!(!dialect.is_identifier_start('-'), "Minus should not be valid identifier start");
assert!(!dialect.is_identifier_start('*'), "Asterisk should not be valid identifier start");
assert!(!dialect.is_identifier_start('/'), "Slash should not be valid identifier start");
assert!(!dialect.is_identifier_start('%'), "Percent should not be valid identifier start");
assert!(!dialect.is_identifier_start('<'), "Less than should not be valid identifier start");
assert!(!dialect.is_identifier_start('>'), "Greater than should not be valid identifier start");
assert!(!dialect.is_identifier_start('!'), "Exclamation should not be valid identifier start");
assert!(!dialect.is_identifier_start('?'), "Question mark should not be valid identifier start");
assert!(!dialect.is_identifier_start('&'), "Ampersand should not be valid identifier start");
assert!(!dialect.is_identifier_start('|'), "Pipe should not be valid identifier start");
assert!(!dialect.is_identifier_start('^'), "Caret should not be valid identifier start");
assert!(!dialect.is_identifier_start('~'), "Tilde should not be valid identifier start");
assert!(!dialect.is_identifier_start('`'), "Backtick should not be valid identifier start");
assert!(!dialect.is_identifier_start('"'), "Double quote should not be valid identifier start");
assert!(!dialect.is_identifier_start('\''), "Single quote should not be valid identifier start");
}
#[test]
fn test_is_identifier_part_alphabetic() {
let dialect = RustFsDialect;
// Test alphabetic characters
assert!(dialect.is_identifier_part('a'), "Lowercase letter should be valid identifier part");
assert!(dialect.is_identifier_part('A'), "Uppercase letter should be valid identifier part");
assert!(dialect.is_identifier_part('z'), "Last lowercase letter should be valid identifier part");
assert!(dialect.is_identifier_part('Z'), "Last uppercase letter should be valid identifier part");
// Test Unicode alphabetic characters
assert!(dialect.is_identifier_part('α'), "Greek letter should be valid identifier part");
assert!(dialect.is_identifier_part('中'), "Chinese character should be valid identifier part");
assert!(dialect.is_identifier_part('ñ'), "Accented letter should be valid identifier part");
}
#[test]
fn test_is_identifier_part_digits() {
let dialect = RustFsDialect;
// Test ASCII digits
assert!(dialect.is_identifier_part('0'), "Digit 0 should be valid identifier part");
assert!(dialect.is_identifier_part('1'), "Digit 1 should be valid identifier part");
assert!(dialect.is_identifier_part('5'), "Digit 5 should be valid identifier part");
assert!(dialect.is_identifier_part('9'), "Digit 9 should be valid identifier part");
}
#[test]
fn test_is_identifier_part_special_chars() {
let dialect = RustFsDialect;
// Test special characters that are allowed
assert!(dialect.is_identifier_part('_'), "Underscore should be valid identifier part");
assert!(dialect.is_identifier_part('#'), "Hash should be valid identifier part");
assert!(dialect.is_identifier_part('@'), "At symbol should be valid identifier part");
assert!(dialect.is_identifier_part('$'), "Dollar sign should be valid identifier part");
}
#[test]
fn test_is_identifier_part_invalid_chars() {
let dialect = RustFsDialect;
// Test characters that should not be valid identifier parts
assert!(!dialect.is_identifier_part(' '), "Space should not be valid identifier part");
assert!(!dialect.is_identifier_part('\t'), "Tab should not be valid identifier part");
assert!(!dialect.is_identifier_part('\n'), "Newline should not be valid identifier part");
assert!(!dialect.is_identifier_part('.'), "Dot should not be valid identifier part");
assert!(!dialect.is_identifier_part(','), "Comma should not be valid identifier part");
assert!(!dialect.is_identifier_part(';'), "Semicolon should not be valid identifier part");
assert!(!dialect.is_identifier_part('('), "Left paren should not be valid identifier part");
assert!(!dialect.is_identifier_part(')'), "Right paren should not be valid identifier part");
assert!(!dialect.is_identifier_part('['), "Left bracket should not be valid identifier part");
assert!(!dialect.is_identifier_part(']'), "Right bracket should not be valid identifier part");
assert!(!dialect.is_identifier_part('{'), "Left brace should not be valid identifier part");
assert!(!dialect.is_identifier_part('}'), "Right brace should not be valid identifier part");
assert!(!dialect.is_identifier_part('='), "Equals should not be valid identifier part");
assert!(!dialect.is_identifier_part('+'), "Plus should not be valid identifier part");
assert!(!dialect.is_identifier_part('-'), "Minus should not be valid identifier part");
assert!(!dialect.is_identifier_part('*'), "Asterisk should not be valid identifier part");
assert!(!dialect.is_identifier_part('/'), "Slash should not be valid identifier part");
assert!(!dialect.is_identifier_part('%'), "Percent should not be valid identifier part");
assert!(!dialect.is_identifier_part('<'), "Less than should not be valid identifier part");
assert!(!dialect.is_identifier_part('>'), "Greater than should not be valid identifier part");
assert!(!dialect.is_identifier_part('!'), "Exclamation should not be valid identifier part");
assert!(!dialect.is_identifier_part('?'), "Question mark should not be valid identifier part");
assert!(!dialect.is_identifier_part('&'), "Ampersand should not be valid identifier part");
assert!(!dialect.is_identifier_part('|'), "Pipe should not be valid identifier part");
assert!(!dialect.is_identifier_part('^'), "Caret should not be valid identifier part");
assert!(!dialect.is_identifier_part('~'), "Tilde should not be valid identifier part");
assert!(!dialect.is_identifier_part('`'), "Backtick should not be valid identifier part");
assert!(!dialect.is_identifier_part('"'), "Double quote should not be valid identifier part");
assert!(!dialect.is_identifier_part('\''), "Single quote should not be valid identifier part");
}
#[test]
fn test_supports_group_by_expr() {
let dialect = RustFsDialect;
assert!(dialect.supports_group_by_expr(), "RustFsDialect should support GROUP BY expressions");
}
#[test]
fn test_identifier_validation_comprehensive() {
let dialect = RustFsDialect;
// Test valid identifier patterns
let valid_starts = ['a', 'A', 'z', 'Z', '_', '#', '@', 'α', '中'];
let valid_parts = ['a', 'A', '0', '9', '_', '#', '@', '$', 'α', '中'];
for start_char in valid_starts {
assert!(
dialect.is_identifier_start(start_char),
"Character '{start_char}' should be valid identifier start"
);
for part_char in valid_parts {
assert!(
dialect.is_identifier_part(part_char),
"Character '{part_char}' should be valid identifier part"
);
}
}
}
#[test]
fn test_identifier_edge_cases() {
let dialect = RustFsDialect;
// Test edge cases with control characters
assert!(!dialect.is_identifier_start('\0'), "Null character should not be valid identifier start");
assert!(!dialect.is_identifier_part('\0'), "Null character should not be valid identifier part");
assert!(
!dialect.is_identifier_start('\x01'),
"Control character should not be valid identifier start"
);
assert!(
!dialect.is_identifier_part('\x01'),
"Control character should not be valid identifier part"
);
assert!(!dialect.is_identifier_start('\x7F'), "DEL character should not be valid identifier start");
assert!(!dialect.is_identifier_part('\x7F'), "DEL character should not be valid identifier part");
}
#[test]
fn test_identifier_unicode_support() {
let dialect = RustFsDialect;
// Test various Unicode categories
let unicode_letters = ['α', 'β', 'γ', 'Α', 'Β', 'Γ', '中', '文', '日', '本', 'ñ', 'ü', 'ç'];
for ch in unicode_letters {
assert!(dialect.is_identifier_start(ch), "Unicode letter '{ch}' should be valid identifier start");
assert!(dialect.is_identifier_part(ch), "Unicode letter '{ch}' should be valid identifier part");
}
}
#[test]
fn test_identifier_ascii_digits() {
let dialect = RustFsDialect;
// Test all ASCII digits
for digit in '0'..='9' {
assert!(
!dialect.is_identifier_start(digit),
"ASCII digit '{digit}' should not be valid identifier start"
);
assert!(dialect.is_identifier_part(digit), "ASCII digit '{digit}' should be valid identifier part");
}
}
#[test]
fn test_dialect_consistency() {
let dialect = RustFsDialect;
// Test that all valid identifier starts are also valid identifier parts
let test_chars = [
'a', 'A', 'z', 'Z', '_', '#', '@', 'α', '中', 'ñ', '0', '9', '$', ' ', '.', ',', ';', '(', ')', '=', '+', '-',
];
for ch in test_chars {
if dialect.is_identifier_start(ch) {
assert!(
dialect.is_identifier_part(ch),
"Character '{ch}' that is valid identifier start should also be valid identifier part"
);
}
}
}
#[test]
fn test_dialect_memory_efficiency() {
let dialect = RustFsDialect;
// Test that dialect doesn't use excessive memory
let dialect_size = std::mem::size_of_val(&dialect);
assert!(dialect_size < 100, "Dialect should not use excessive memory");
}
#[test]
fn test_dialect_trait_implementation() {
let dialect = RustFsDialect;
// Test that dialect properly implements the Dialect trait
let dialect_ref: &dyn Dialect = &dialect;
// Test basic functionality through trait
assert!(dialect_ref.is_identifier_start('a'), "Trait method should work for valid start");
assert!(!dialect_ref.is_identifier_start('0'), "Trait method should work for invalid start");
assert!(dialect_ref.is_identifier_part('a'), "Trait method should work for valid part");
assert!(dialect_ref.is_identifier_part('0'), "Trait method should work for digit part");
assert!(
dialect_ref.supports_group_by_expr(),
"Trait method should return true for GROUP BY support"
);
}
#[test]
fn test_dialect_clone_and_default() {
let dialect1 = RustFsDialect;
let dialect2 = RustFsDialect;
// Test that multiple instances behave the same
let test_chars = ['a', 'A', '0', '_', '#', '@', '$', ' ', '.'];
for ch in test_chars {
assert_eq!(
dialect1.is_identifier_start(ch),
dialect2.is_identifier_start(ch),
"Different instances should behave the same for is_identifier_start"
);
assert_eq!(
dialect1.is_identifier_part(ch),
dialect2.is_identifier_part(ch),
"Different instances should behave the same for is_identifier_part"
);
}
assert_eq!(
dialect1.supports_group_by_expr(),
dialect2.supports_group_by_expr(),
"Different instances should behave the same for supports_group_by_expr"
);
}
}
@@ -0,0 +1,16 @@
// Copyright 2024 RustFS 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.
pub mod optimizer;
pub mod planner;
@@ -0,0 +1,125 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use datafusion::{
execution::SessionStateBuilder,
logical_expr::LogicalPlan,
optimizer::{
OptimizerRule, common_subexpr_eliminate::CommonSubexprEliminate,
decorrelate_predicate_subquery::DecorrelatePredicateSubquery, eliminate_cross_join::EliminateCrossJoin,
eliminate_duplicated_expr::EliminateDuplicatedExpr, eliminate_filter::EliminateFilter, eliminate_join::EliminateJoin,
eliminate_limit::EliminateLimit, eliminate_outer_join::EliminateOuterJoin,
extract_equijoin_predicate::ExtractEquijoinPredicate, filter_null_join_keys::FilterNullJoinKeys,
propagate_empty_relation::PropagateEmptyRelation, push_down_filter::PushDownFilter, push_down_limit::PushDownLimit,
replace_distinct_aggregate::ReplaceDistinctWithAggregate, scalar_subquery_to_join::ScalarSubqueryToJoin,
simplify_expressions::SimplifyExpressions, single_distinct_to_groupby::SingleDistinctToGroupBy,
unwrap_cast_in_comparison::UnwrapCastInComparison,
},
};
use rustfs_s3select_api::{
QueryResult,
query::{analyzer::AnalyzerRef, logical_planner::QueryPlan, session::SessionCtx},
};
use tracing::debug;
use crate::sql::analyzer::DefaultAnalyzer;
pub trait LogicalOptimizer: Send + Sync {
fn optimize(&self, plan: &QueryPlan, session: &SessionCtx) -> QueryResult<LogicalPlan>;
fn inject_optimizer_rule(&mut self, optimizer_rule: Arc<dyn OptimizerRule + Send + Sync>);
}
pub struct DefaultLogicalOptimizer {
// fit datafusion
// TODO refactor
analyzer: AnalyzerRef,
rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>,
}
impl DefaultLogicalOptimizer {
#[allow(dead_code)]
fn with_optimizer_rules(mut self, rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>) -> Self {
self.rules = rules;
self
}
}
impl Default for DefaultLogicalOptimizer {
fn default() -> Self {
let analyzer = Arc::new(DefaultAnalyzer::default());
// additional optimizer rule
let rules: Vec<Arc<dyn OptimizerRule + Send + Sync>> = vec![
// df default rules start
Arc::new(SimplifyExpressions::new()),
Arc::new(UnwrapCastInComparison::new()),
Arc::new(ReplaceDistinctWithAggregate::new()),
Arc::new(EliminateJoin::new()),
Arc::new(DecorrelatePredicateSubquery::new()),
Arc::new(ScalarSubqueryToJoin::new()),
Arc::new(ExtractEquijoinPredicate::new()),
// simplify expressions does not simplify expressions in subqueries, so we
// run it again after running the optimizations that potentially converted
// subqueries to joins
Arc::new(SimplifyExpressions::new()),
Arc::new(EliminateDuplicatedExpr::new()),
Arc::new(EliminateFilter::new()),
Arc::new(EliminateCrossJoin::new()),
Arc::new(CommonSubexprEliminate::new()),
Arc::new(EliminateLimit::new()),
Arc::new(PropagateEmptyRelation::new()),
Arc::new(FilterNullJoinKeys::default()),
Arc::new(EliminateOuterJoin::new()),
// Filters can't be pushed down past Limits, we should do PushDownFilter after PushDownLimit
Arc::new(PushDownLimit::new()),
Arc::new(PushDownFilter::new()),
Arc::new(SingleDistinctToGroupBy::new()),
// The previous optimizations added expressions and projections,
// that might benefit from the following rules
Arc::new(SimplifyExpressions::new()),
Arc::new(UnwrapCastInComparison::new()),
Arc::new(CommonSubexprEliminate::new()),
// PushDownProjection can pushdown Projections through Limits, do PushDownLimit again.
Arc::new(PushDownLimit::new()),
// df default rules end
// custom rules can add at here
];
Self { analyzer, rules }
}
}
impl LogicalOptimizer for DefaultLogicalOptimizer {
fn optimize(&self, plan: &QueryPlan, session: &SessionCtx) -> QueryResult<LogicalPlan> {
let analyzed_plan = { self.analyzer.analyze(&plan.df_plan, session)? };
debug!("Analyzed logical plan:\n{}\n", plan.df_plan.display_indent_schema(),);
let optimizeed_plan = {
SessionStateBuilder::new_from_existing(session.inner().clone())
.with_optimizer_rules(self.rules.clone())
.build()
.optimize(&analyzed_plan)?
};
Ok(optimizeed_plan)
}
fn inject_optimizer_rule(&mut self, optimizer_rule: Arc<dyn OptimizerRule + Send + Sync>) {
self.rules.push(optimizer_rule);
}
}
@@ -0,0 +1,17 @@
// Copyright 2024 RustFS 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 crate::sql::planner::SqlPlanner;
pub type DefaultLogicalPlanner<'a, S> = SqlPlanner<'a, S>;
+21
View File
@@ -0,0 +1,21 @@
// Copyright 2024 RustFS 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.
pub mod analyzer;
pub mod dialect;
pub mod logical;
pub mod optimizer;
pub mod parser;
pub mod physical;
pub mod planner;
+199
View File
@@ -0,0 +1,199 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use async_trait::async_trait;
use datafusion::physical_plan::{ExecutionPlan, displayable};
use rustfs_s3select_api::{
QueryResult,
query::{logical_planner::QueryPlan, optimizer::Optimizer, physical_planner::PhysicalPlanner, session::SessionCtx},
};
use tracing::debug;
use super::{
logical::optimizer::{DefaultLogicalOptimizer, LogicalOptimizer},
physical::{optimizer::PhysicalOptimizer, planner::DefaultPhysicalPlanner},
};
pub struct CascadeOptimizer {
logical_optimizer: Arc<dyn LogicalOptimizer + Send + Sync>,
physical_planner: Arc<dyn PhysicalPlanner + Send + Sync>,
physical_optimizer: Arc<dyn PhysicalOptimizer + Send + Sync>,
}
#[async_trait]
impl Optimizer for CascadeOptimizer {
async fn optimize(&self, plan: &QueryPlan, session: &SessionCtx) -> QueryResult<Arc<dyn ExecutionPlan>> {
debug!("Original logical plan:\n{}\n", plan.df_plan.display_indent_schema(),);
let optimized_logical_plan = self.logical_optimizer.optimize(plan, session)?;
debug!("Final logical plan:\n{}\n", optimized_logical_plan.display_indent_schema(),);
let physical_plan = {
self.physical_planner
.create_physical_plan(&optimized_logical_plan, session)
.await?
};
debug!("Original physical plan:\n{}\n", displayable(physical_plan.as_ref()).indent(false));
let optimized_physical_plan = { self.physical_optimizer.optimize(physical_plan, session)? };
Ok(optimized_physical_plan)
}
}
#[derive(Default)]
pub struct CascadeOptimizerBuilder {
logical_optimizer: Option<Arc<dyn LogicalOptimizer + Send + Sync>>,
physical_planner: Option<Arc<dyn PhysicalPlanner + Send + Sync>>,
physical_optimizer: Option<Arc<dyn PhysicalOptimizer + Send + Sync>>,
}
impl CascadeOptimizerBuilder {
pub fn with_logical_optimizer(mut self, logical_optimizer: Arc<dyn LogicalOptimizer + Send + Sync>) -> Self {
self.logical_optimizer = Some(logical_optimizer);
self
}
pub fn with_physical_planner(mut self, physical_planner: Arc<dyn PhysicalPlanner + Send + Sync>) -> Self {
self.physical_planner = Some(physical_planner);
self
}
pub fn with_physical_optimizer(mut self, physical_optimizer: Arc<dyn PhysicalOptimizer + Send + Sync>) -> Self {
self.physical_optimizer = Some(physical_optimizer);
self
}
pub fn build(self) -> CascadeOptimizer {
let default_logical_optimizer = Arc::new(DefaultLogicalOptimizer::default());
let default_physical_planner = Arc::new(DefaultPhysicalPlanner::default());
let logical_optimizer = self.logical_optimizer.unwrap_or(default_logical_optimizer);
let physical_planner = self.physical_planner.unwrap_or_else(|| default_physical_planner.clone());
let physical_optimizer = self.physical_optimizer.unwrap_or(default_physical_planner);
CascadeOptimizer {
logical_optimizer,
physical_planner,
physical_optimizer,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cascade_optimizer_builder_default() {
let _builder = CascadeOptimizerBuilder::default();
// Test that builder can be created successfully
assert!(
std::mem::size_of::<CascadeOptimizerBuilder>() > 0,
"Builder should be created successfully"
);
}
#[test]
fn test_cascade_optimizer_builder_build_with_defaults() {
let _builder = CascadeOptimizerBuilder::default();
let optimizer = _builder.build();
// Test that optimizer can be built with default components
assert!(std::mem::size_of_val(&optimizer) > 0, "Optimizer should be built successfully");
}
#[test]
fn test_cascade_optimizer_builder_basic_functionality() {
// Test that builder methods can be called and return self
let _builder = CascadeOptimizerBuilder::default();
// Test that we can call builder methods (even if we don't have mock implementations)
// This tests the builder pattern itself
assert!(
std::mem::size_of::<CascadeOptimizerBuilder>() > 0,
"Builder should be created successfully"
);
}
#[test]
fn test_cascade_optimizer_builder_memory_efficiency() {
let _builder = CascadeOptimizerBuilder::default();
// Test that builder doesn't use excessive memory
let builder_size = std::mem::size_of_val(&_builder);
assert!(builder_size < 1000, "Builder should not use excessive memory");
let optimizer = _builder.build();
let optimizer_size = std::mem::size_of_val(&optimizer);
assert!(optimizer_size < 1000, "Optimizer should not use excessive memory");
}
#[test]
fn test_cascade_optimizer_builder_multiple_builds() {
let _builder = CascadeOptimizerBuilder::default();
// Test that we can build multiple optimizers from the same configuration
let optimizer1 = _builder.build();
assert!(std::mem::size_of_val(&optimizer1) > 0, "First optimizer should be built successfully");
// Note: builder is consumed by build(), so we can't build again from the same instance
// This is the expected behavior
}
#[test]
fn test_cascade_optimizer_builder_default_fallbacks() {
let _builder = CascadeOptimizerBuilder::default();
let optimizer = _builder.build();
// Test that default components are used when none are specified
// We can't directly access the internal components, but we can verify the optimizer was built
assert!(std::mem::size_of_val(&optimizer) > 0, "Optimizer should use default components");
}
#[test]
fn test_cascade_optimizer_component_types() {
let optimizer = CascadeOptimizerBuilder::default().build();
// Test that optimizer contains the expected component types
// We can't directly access the components, but we can verify the optimizer structure
assert!(std::mem::size_of_val(&optimizer) > 0, "Optimizer should contain components");
// The optimizer should have three Arc fields for the components
// This is a basic structural test
}
#[test]
fn test_cascade_optimizer_builder_consistency() {
// Test that multiple builders with the same configuration produce equivalent optimizers
let optimizer1 = CascadeOptimizerBuilder::default().build();
let optimizer2 = CascadeOptimizerBuilder::default().build();
// Both optimizers should be built successfully
assert!(std::mem::size_of_val(&optimizer1) > 0, "First optimizer should be built");
assert!(std::mem::size_of_val(&optimizer2) > 0, "Second optimizer should be built");
// They should have the same memory footprint (same structure)
assert_eq!(
std::mem::size_of_val(&optimizer1),
std::mem::size_of_val(&optimizer2),
"Optimizers with same configuration should have same size"
);
}
}
+449
View File
@@ -0,0 +1,449 @@
// Copyright 2024 RustFS 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 std::{collections::VecDeque, fmt::Display};
use datafusion::sql::sqlparser::{
dialect::Dialect,
parser::{Parser, ParserError},
tokenizer::{Token, Tokenizer},
};
use rustfs_s3select_api::{
ParserSnafu,
query::{ast::ExtStatement, parser::Parser as RustFsParser},
};
use snafu::ResultExt;
use super::dialect::RustFsDialect;
pub type Result<T, E = ParserError> = std::result::Result<T, E>;
// Use `Parser::expected` instead, if possible
macro_rules! parser_err {
($MSG:expr) => {
Err(ParserError::ParserError($MSG.to_string()))
};
}
#[derive(Default)]
pub struct DefaultParser {}
impl RustFsParser for DefaultParser {
fn parse(&self, sql: &str) -> rustfs_s3select_api::QueryResult<VecDeque<ExtStatement>> {
ExtParser::parse_sql(sql).context(ParserSnafu)
}
}
/// SQL Parser
pub struct ExtParser<'a> {
parser: Parser<'a>,
}
impl<'a> ExtParser<'a> {
/// Parse the specified tokens with dialect
fn new_with_dialect(sql: &str, dialect: &'a dyn Dialect) -> Result<Self> {
let mut tokenizer = Tokenizer::new(dialect, sql);
let tokens = tokenizer.tokenize()?;
Ok(ExtParser {
parser: Parser::new(dialect).with_tokens(tokens),
})
}
/// Parse a SQL statement and produce a set of statements
pub fn parse_sql(sql: &str) -> Result<VecDeque<ExtStatement>> {
let dialect = &RustFsDialect {};
ExtParser::parse_sql_with_dialect(sql, dialect)
}
/// Parse a SQL statement and produce a set of statements
pub fn parse_sql_with_dialect(sql: &str, dialect: &dyn Dialect) -> Result<VecDeque<ExtStatement>> {
let mut parser = ExtParser::new_with_dialect(sql, dialect)?;
let mut stmts = VecDeque::new();
let mut expecting_statement_delimiter = false;
loop {
// ignore empty statements (between successive statement delimiters)
while parser.parser.consume_token(&Token::SemiColon) {
expecting_statement_delimiter = false;
}
if parser.parser.peek_token() == Token::EOF {
break;
}
if expecting_statement_delimiter {
return parser.expected("end of statement", parser.parser.peek_token());
}
let statement = parser.parse_statement()?;
stmts.push_back(statement);
expecting_statement_delimiter = true;
}
// debug!("Parser sql: {}, stmts: {:#?}", sql, stmts);
Ok(stmts)
}
/// Parse a new expression
fn parse_statement(&mut self) -> Result<ExtStatement> {
Ok(ExtStatement::SqlStatement(Box::new(self.parser.parse_statement()?)))
}
// Report unexpected token
fn expected<T>(&self, expected: &str, found: impl Display) -> Result<T> {
parser_err!(format!("Expected {}, found: {}", expected, found))
}
}
#[cfg(test)]
mod tests {
use super::*;
use rustfs_s3select_api::query::ast::ExtStatement;
#[test]
fn test_default_parser_creation() {
let _parser = DefaultParser::default();
// Test that parser can be created successfully
assert!(std::mem::size_of::<DefaultParser>() == 0, "Parser should be zero-sized");
}
#[test]
fn test_default_parser_simple_select() {
let parser = DefaultParser::default();
let sql = "SELECT * FROM S3Object";
let result = parser.parse(sql);
assert!(result.is_ok(), "Simple SELECT should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
// Just verify we get a SQL statement without diving into AST details
match &statements[0] {
ExtStatement::SqlStatement(_) => {
// Successfully parsed as SQL statement
}
}
}
#[test]
fn test_default_parser_select_with_columns() {
let parser = DefaultParser::default();
let sql = "SELECT id, name, age FROM S3Object";
let result = parser.parse(sql);
assert!(result.is_ok(), "SELECT with columns should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
match &statements[0] {
ExtStatement::SqlStatement(_) => {
// Successfully parsed as SQL statement
}
}
}
#[test]
fn test_default_parser_select_with_where() {
let parser = DefaultParser::default();
let sql = "SELECT * FROM S3Object WHERE age > 25";
let result = parser.parse(sql);
assert!(result.is_ok(), "SELECT with WHERE should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
match &statements[0] {
ExtStatement::SqlStatement(_) => {
// Successfully parsed as SQL statement
}
}
}
#[test]
fn test_default_parser_multiple_statements() {
let parser = DefaultParser::default();
let sql = "SELECT * FROM S3Object; SELECT id FROM S3Object;";
let result = parser.parse(sql);
assert!(result.is_ok(), "Multiple statements should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 2, "Should have exactly two statements");
}
#[test]
fn test_default_parser_empty_statements() {
let parser = DefaultParser::default();
let sql = ";;; SELECT * FROM S3Object; ;;;";
let result = parser.parse(sql);
assert!(result.is_ok(), "Empty statements should be ignored");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one non-empty statement");
}
#[test]
fn test_default_parser_invalid_sql() {
let parser = DefaultParser::default();
let sql = "INVALID SQL SYNTAX";
let result = parser.parse(sql);
assert!(result.is_err(), "Invalid SQL should return error");
}
#[test]
fn test_default_parser_empty_sql() {
let parser = DefaultParser::default();
let sql = "";
let result = parser.parse(sql);
assert!(result.is_ok(), "Empty SQL should parse successfully");
let statements = result.unwrap();
assert!(statements.is_empty(), "Should have no statements");
}
#[test]
fn test_default_parser_whitespace_only() {
let parser = DefaultParser::default();
let sql = " \n\t ";
let result = parser.parse(sql);
assert!(result.is_ok(), "Whitespace-only SQL should parse successfully");
let statements = result.unwrap();
assert!(statements.is_empty(), "Should have no statements");
}
#[test]
fn test_ext_parser_parse_sql() {
let sql = "SELECT * FROM S3Object";
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "ExtParser::parse_sql should work");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_ext_parser_parse_sql_with_dialect() {
let sql = "SELECT * FROM S3Object";
let dialect = &RustFsDialect;
let result = ExtParser::parse_sql_with_dialect(sql, dialect);
assert!(result.is_ok(), "ExtParser::parse_sql_with_dialect should work");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_ext_parser_new_with_dialect() {
let sql = "SELECT * FROM S3Object";
let dialect = &RustFsDialect;
let result = ExtParser::new_with_dialect(sql, dialect);
assert!(result.is_ok(), "ExtParser::new_with_dialect should work");
}
#[test]
fn test_ext_parser_complex_query() {
let sql = "SELECT id, name, age FROM S3Object WHERE age > 25 AND department = 'IT' ORDER BY age DESC LIMIT 10";
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "Complex query should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
match &statements[0] {
ExtStatement::SqlStatement(_) => {
// Successfully parsed as SQL statement
}
}
}
#[test]
fn test_ext_parser_aggregate_functions() {
let sql = "SELECT COUNT(*), AVG(age), MAX(salary) FROM S3Object GROUP BY department";
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "Aggregate functions should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
match &statements[0] {
ExtStatement::SqlStatement(_) => {
// Successfully parsed as SQL statement
}
}
}
#[test]
fn test_ext_parser_join_query() {
let sql = "SELECT s1.id, s2.name FROM S3Object s1 JOIN S3Object s2 ON s1.id = s2.id";
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "JOIN query should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_ext_parser_subquery() {
let sql = "SELECT * FROM S3Object WHERE id IN (SELECT id FROM S3Object WHERE age > 30)";
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "Subquery should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_ext_parser_case_insensitive() {
let sql = "select * from s3object where age > 25";
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "Case insensitive SQL should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_ext_parser_quoted_identifiers() {
let sql = r#"SELECT "id", "name" FROM "S3Object" WHERE "age" > 25"#;
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "Quoted identifiers should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_ext_parser_string_literals() {
let sql = "SELECT * FROM S3Object WHERE name = 'John Doe' AND department = 'IT'";
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "String literals should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_ext_parser_numeric_literals() {
let sql = "SELECT * FROM S3Object WHERE age = 25 AND salary = 50000.50";
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "Numeric literals should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_ext_parser_error_handling() {
let invalid_sqls = vec![
"SELECT FROM", // Missing column list
"SELECT * FROM", // Missing table name
"SELECT * FROM S3Object WHERE", // Incomplete WHERE clause
"SELECT * FROM S3Object GROUP", // Incomplete GROUP BY
"SELECT * FROM S3Object ORDER", // Incomplete ORDER BY
];
for sql in invalid_sqls {
let result = ExtParser::parse_sql(sql);
assert!(result.is_err(), "Invalid SQL '{sql}' should return error");
}
}
#[test]
fn test_ext_parser_memory_efficiency() {
let sql = "SELECT * FROM S3Object";
// Test that parser doesn't use excessive memory
let result = ExtParser::parse_sql(sql);
assert!(result.is_ok(), "Parser should work efficiently");
let statements = result.unwrap();
let memory_size = std::mem::size_of_val(&statements);
assert!(memory_size < 10000, "Parsed statements should not use excessive memory");
}
#[test]
fn test_ext_parser_large_query() {
// Test with a reasonably large query
let mut sql = String::from("SELECT ");
for i in 0..100 {
if i > 0 {
sql.push_str(", ");
}
sql.push_str(&format!("col{i}"));
}
sql.push_str(" FROM S3Object WHERE ");
for i in 0..50 {
if i > 0 {
sql.push_str(" AND ");
}
sql.push_str(&format!("col{i} > {i}"));
}
let result = ExtParser::parse_sql(&sql);
assert!(result.is_ok(), "Large query should parse successfully");
let statements = result.unwrap();
assert_eq!(statements.len(), 1, "Should have exactly one statement");
}
#[test]
fn test_parser_err_macro() {
let error: Result<()> = parser_err!("Test error message");
assert!(error.is_err(), "parser_err! macro should create error");
match error {
Err(ParserError::ParserError(msg)) => {
assert_eq!(msg, "Test error message", "Error message should match");
}
_ => panic!("Expected ParserError::ParserError"),
}
}
#[test]
fn test_ext_parser_expected_method() {
let sql = "SELECT * FROM S3Object";
let dialect = &RustFsDialect;
let parser = ExtParser::new_with_dialect(sql, dialect).unwrap();
let result: Result<()> = parser.expected("test token", "found token");
assert!(result.is_err(), "expected method should return error");
match result {
Err(ParserError::ParserError(msg)) => {
assert!(msg.contains("Expected test token"), "Error should contain expected message");
assert!(msg.contains("found: found token"), "Error should contain found message");
}
_ => panic!("Expected ParserError::ParserError"),
}
}
}
@@ -0,0 +1,16 @@
// Copyright 2024 RustFS 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.
pub mod optimizer;
pub mod planner;
@@ -0,0 +1,26 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;
use rustfs_s3select_api::QueryResult;
use rustfs_s3select_api::query::session::SessionCtx;
pub trait PhysicalOptimizer {
fn optimize(&self, plan: Arc<dyn ExecutionPlan>, session: &SessionCtx) -> QueryResult<Arc<dyn ExecutionPlan>>;
fn inject_optimizer_rule(&mut self, optimizer_rule: Arc<dyn PhysicalOptimizerRule + Send + Sync>);
}
@@ -0,0 +1,118 @@
// Copyright 2024 RustFS 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 std::sync::Arc;
use async_trait::async_trait;
use datafusion::execution::SessionStateBuilder;
use datafusion::logical_expr::LogicalPlan;
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_optimizer::aggregate_statistics::AggregateStatistics;
use datafusion::physical_optimizer::coalesce_batches::CoalesceBatches;
use datafusion::physical_optimizer::join_selection::JoinSelection;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_planner::{
DefaultPhysicalPlanner as DFDefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner as DFPhysicalPlanner,
};
use rustfs_s3select_api::QueryResult;
use rustfs_s3select_api::query::physical_planner::PhysicalPlanner;
use rustfs_s3select_api::query::session::SessionCtx;
use super::optimizer::PhysicalOptimizer;
pub struct DefaultPhysicalPlanner {
ext_physical_transform_rules: Vec<Arc<dyn ExtensionPlanner + Send + Sync>>,
/// Responsible for optimizing a physical execution plan
ext_physical_optimizer_rules: Vec<Arc<dyn PhysicalOptimizerRule + Send + Sync>>,
}
impl DefaultPhysicalPlanner {
#[allow(dead_code)]
fn with_physical_transform_rules(mut self, rules: Vec<Arc<dyn ExtensionPlanner + Send + Sync>>) -> Self {
self.ext_physical_transform_rules = rules;
self
}
}
impl DefaultPhysicalPlanner {
#[allow(dead_code)]
fn with_optimizer_rules(mut self, rules: Vec<Arc<dyn PhysicalOptimizerRule + Send + Sync>>) -> Self {
self.ext_physical_optimizer_rules = rules;
self
}
}
impl Default for DefaultPhysicalPlanner {
fn default() -> Self {
let ext_physical_transform_rules: Vec<Arc<dyn ExtensionPlanner + Send + Sync>> = vec![
// can add rules at here
];
// We need to take care of the rule ordering. They may influence each other.
let ext_physical_optimizer_rules: Vec<Arc<dyn PhysicalOptimizerRule + Sync + Send>> = vec![
Arc::new(AggregateStatistics::new()),
// Statistics-based join selection will change the Auto mode to a real join implementation,
// like collect left, or hash join, or future sort merge join, which will influence the
// EnforceDistribution and EnforceSorting rules as they decide whether to add additional
// repartitioning and local sorting steps to meet distribution and ordering requirements.
// Therefore, it should run before EnforceDistribution and EnforceSorting.
Arc::new(JoinSelection::new()),
// The CoalesceBatches rule will not influence the distribution and ordering of the
// whole plan tree. Therefore, to avoid influencing other rules, it should run last.
Arc::new(CoalesceBatches::new()),
];
Self {
ext_physical_transform_rules,
ext_physical_optimizer_rules,
}
}
}
#[async_trait]
impl PhysicalPlanner for DefaultPhysicalPlanner {
async fn create_physical_plan(
&self,
logical_plan: &LogicalPlan,
session: &SessionCtx,
) -> QueryResult<Arc<dyn ExecutionPlan>> {
// 将扩展的物理计划优化规则注入 df 的 session state
let new_state = SessionStateBuilder::new_from_existing(session.inner().clone())
.with_physical_optimizer_rules(self.ext_physical_optimizer_rules.clone())
.build();
// 通过扩展的物理计划转换规则构造 df 的 Physical Planner
let planner = DFDefaultPhysicalPlanner::with_extension_planners(self.ext_physical_transform_rules.clone());
// 执行 df 的物理计划规划及优化
planner
.create_physical_plan(logical_plan, &new_state)
.await
.map_err(|e| e.into())
}
fn inject_physical_transform_rule(&mut self, rule: Arc<dyn ExtensionPlanner + Send + Sync>) {
self.ext_physical_transform_rules.push(rule)
}
}
impl PhysicalOptimizer for DefaultPhysicalPlanner {
fn optimize(&self, plan: Arc<dyn ExecutionPlan>, _session: &SessionCtx) -> QueryResult<Arc<dyn ExecutionPlan>> {
Ok(plan)
}
fn inject_optimizer_rule(&mut self, optimizer_rule: Arc<dyn PhysicalOptimizerRule + Send + Sync>) {
self.ext_physical_optimizer_rules.push(optimizer_rule);
}
}
+74
View File
@@ -0,0 +1,74 @@
// Copyright 2024 RustFS 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 async_recursion::async_recursion;
use async_trait::async_trait;
use datafusion::sql::{planner::SqlToRel, sqlparser::ast::Statement};
use rustfs_s3select_api::{
QueryError, QueryResult,
query::{
ast::ExtStatement,
logical_planner::{LogicalPlanner, Plan, QueryPlan},
session::SessionCtx,
},
};
use crate::metadata::ContextProviderExtension;
pub struct SqlPlanner<'a, S: ContextProviderExtension> {
_schema_provider: &'a S,
df_planner: SqlToRel<'a, S>,
}
#[async_trait]
impl<S: ContextProviderExtension + Send + Sync> LogicalPlanner for SqlPlanner<'_, S> {
async fn create_logical_plan(&self, statement: ExtStatement, session: &SessionCtx) -> QueryResult<Plan> {
let plan = { self.statement_to_plan(statement, session).await? };
Ok(plan)
}
}
impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> {
/// Create a new query planner
pub fn new(schema_provider: &'a S) -> Self {
SqlPlanner {
_schema_provider: schema_provider,
df_planner: SqlToRel::new(schema_provider),
}
}
/// Generate a logical plan from an Extent SQL statement
#[async_recursion]
pub(crate) async fn statement_to_plan(&self, statement: ExtStatement, session: &SessionCtx) -> QueryResult<Plan> {
match statement {
ExtStatement::SqlStatement(stmt) => self.df_sql_to_plan(*stmt, session).await,
}
}
async fn df_sql_to_plan(&self, stmt: Statement, _session: &SessionCtx) -> QueryResult<Plan> {
match stmt {
Statement::Query(_) => {
let df_plan = self.df_planner.sql_statement_to_plan(stmt)?;
let plan = Plan::Query(QueryPlan {
df_plan,
is_tag_scan: false,
});
Ok(plan)
}
_ => Err(QueryError::NotImplemented { err: stmt.to_string() }),
}
}
}