gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-08-11 15:46:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

tmp2

Browse Source 
Signed-off-by: junxiang Mu <1948535941@qq.com>
This commit is contained in:
junxiang Mu
2025-03-13 09:43:53 +00:00
parent 63ef986bac
commit 0b270bf0cc
41 changed files with 4402 additions and 1256 deletions
Download Patch File Download Diff File Expand all files Collapse all files
s3select/query/src/dispatcher/manager.rs
+85 -140
View File
@@ -1,51 +1,62 @@
use std::{
collections::HashMap,
pin::Pin,
sync::{Arc, Mutex},
task::{Context, Poll},
};
use api::{
query::{
ast::ExtStatement,
dispatcher::QueryDispatcher,
execution::{Output, QueryStateMachine},
function::FuncMetaManagerRef,
logical_planner::Plan,
logical_planner::{LogicalPlanner, Plan},
parser::Parser,
session::{SessionCtx, SessionCtxFactory},
Query,
},
QueryError, QueryResult,
};
use async_trait::async_trait;
use datafusion::{
arrow::{datatypes::SchemaRef, record_batch::RecordBatch},
config::CsvOptions,
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 s3s::dto::SelectObjectContentInput;
use tokio::task::JoinHandle;
use crate::metadata::TableHandleProviderRef;
use crate::{
execution::factory::QueryExecutionFactoryRef,
metadata::{
base_table::BaseTableProvider,
ContextProviderExtension, MetadataProvider, TableHandleProviderRef,
},
sql::logical::planner::DefaultLogicalPlanner,
};
#[derive(Clone)]
pub struct SimpleQueryDispatcher {
input: SelectObjectContentInput,
// client for default tenant
default_table_provider: TableHandleProviderRef,
// memory pool
// memory_pool: MemoryPoolRef,
// query tracker
session_factory: Arc<SessionCtxFactory>,
// parser
parser: Arc<dyn Parser + Send + Sync>,
// get query execution factory
query_execution_factory: QueryExecutionFactoryRef,
func_manager: FuncMetaManagerRef,
async_task_joinhandle: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
failed_task_joinhandle: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
}
#[async_trait]
impl QueryDispatcher for SimpleQueryDispatcher {
async fn start(&self) -> QueryResult<()> {
self.execute_persister_query(self.coord.node_id()).await
}
fn stop(&self) {
// TODO
}
async fn execute_query(&self, query: &Query) -> QueryResult<Output> {
let query_state_machine = { self.build_query_state_machine(query.clone()).await? };
@@ -66,7 +77,6 @@ impl QueryDispatcher for SimpleQueryDispatcher {
let logical_planner = DefaultLogicalPlanner::new(&scheme_provider);
let span_recorder = session.get_child_span("parse sql");
let statements = self.parser.parse(query.content())?;
// not allow multi statement
@@ -82,8 +92,6 @@ impl QueryDispatcher for SimpleQueryDispatcher {
None => return Ok(None),
};
drop(span_recorder);
let logical_plan = self
.statement_to_logical_plan(stmt, &logical_planner, query_state_machine)
.await?;
@@ -95,9 +103,7 @@ impl QueryDispatcher for SimpleQueryDispatcher {
}
async fn build_query_state_machine(&self, query: Query) -> QueryResult<Arc<QueryStateMachine>> {
let session = self
.session_factory
.create_session_ctx(query.context(), self.memory_pool.clone(), self.coord.clone())?;
let session = self.session_factory.create_session_ctx(query.context())?;
let query_state_machine = Arc::new(QueryStateMachine::begin(query, session));
Ok(query_state_machine)
@@ -114,7 +120,7 @@ impl SimpleQueryDispatcher {
// begin analyze
query_state_machine.begin_analyze();
let logical_plan = logical_planner
.create_logical_plan(stmt, &query_state_machine.session, self.coord.get_config().query.auth_enabled)
.create_logical_plan(stmt, &query_state_machine.session)
.await?;
query_state_machine.end_analyze();
@@ -127,79 +133,85 @@ impl SimpleQueryDispatcher {
.create_query_execution(logical_plan, query_state_machine.clone())
.await?;
// TrackedQuery.drop() is called implicitly when the value goes out of scope,
self.query_tracker
.try_track_query(query_state_machine.query_id, execution)
.await?
.start()
.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 meta_client = self.build_current_session_meta_client(session).await?;
let current_session_table_provider = self.build_table_handle_provider(meta_client.clone())?;
let metadata_provider = MetadataProvider::new(
self.coord.clone(),
meta_client,
current_session_table_provider,
self.default_table_provider.clone(),
self.func_manager.clone(),
self.query_tracker.clone(),
session.clone(),
);
let path = format!("s3://{}/{}", self.input.bucket, self.input.key);
let table_path = ListingTableUrl::parse(path)?;
let listing_options = if self.input.request.input_serialization.csv.is_some() {
let file_format = CsvFormat::default().with_options(CsvOptions::default().with_has_header(false));
ListingOptions::new(Arc::new(file_format)).with_file_extension(".csv")
} 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")
} 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")
} 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 = ListingTableConfig::new(table_path)
.with_listing_options(listing_options)
.with_schema(resolve_schema);
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)
}
async fn build_current_session_meta_client(&self, session: &SessionCtx) -> QueryResult<MetaClientRef> {
let meta_client = self
.coord
.tenant_meta(session.tenant())
.await
.ok_or_else(|| MetaError::TenantNotFound {
tenant: session.tenant().to_string(),
})
.context(MetaSnafu)?;
Ok(meta_client)
}
fn build_table_handle_provider(&self, meta_client: MetaClientRef) -> QueryResult<TableHandleProviderRef> {
let current_session_table_provider: Arc<BaseTableProvider> = Arc::new(BaseTableProvider::new(
self.coord.clone(),
self.split_manager.clone(),
meta_client,
self.stream_provider_manager.clone(),
));
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 {
coord: Option<CoordinatorRef>,
input: Option<SelectObjectContentInput>,
default_table_provider: Option<TableHandleProviderRef>,
split_manager: Option<SplitManagerRef>,
session_factory: Option<Arc<SessionCtxFactory>>,
parser: Option<Arc<dyn Parser + Send + Sync>>,
query_execution_factory: Option<QueryExecutionFactoryRef>,
query_tracker: Option<Arc<QueryTracker>>,
memory_pool: Option<MemoryPoolRef>, // memory
func_manager: Option<FuncMetaManagerRef>,
stream_provider_manager: Option<StreamProviderManagerRef>,
span_ctx: Option<SpanContext>,
auth_cache: Option<Arc<AuthCache<AuthCacheKey, User>>>,
}
impl SimpleQueryDispatcherBuilder {
pub fn with_coord(mut self, coord: CoordinatorRef) -> Self {
self.coord = Some(coord);
pub fn with_input(mut self, input: 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
@@ -210,11 +222,6 @@ impl SimpleQueryDispatcherBuilder {
self
}
pub fn with_split_manager(mut self, split_manager: SplitManagerRef) -> Self {
self.split_manager = Some(split_manager);
self
}
pub fn with_parser(mut self, parser: Arc<dyn Parser + Send + Sync>) -> Self {
self.parser = Some(parser);
self
@@ -225,43 +232,14 @@ impl SimpleQueryDispatcherBuilder {
self
}
pub fn with_query_tracker(mut self, query_tracker: Arc<QueryTracker>) -> Self {
self.query_tracker = Some(query_tracker);
self
}
pub fn with_memory_pool(mut self, memory_pool: MemoryPoolRef) -> Self {
self.memory_pool = Some(memory_pool);
self
}
pub fn with_func_manager(mut self, func_manager: FuncMetaManagerRef) -> Self {
self.func_manager = Some(func_manager);
self
}
pub fn with_stream_provider_manager(mut self, stream_provider_manager: StreamProviderManagerRef) -> Self {
self.stream_provider_manager = Some(stream_provider_manager);
self
}
pub fn with_span_ctx(mut self, span_ctx: Option<SpanContext>) -> Self {
self.span_ctx = span_ctx;
self
}
pub fn with_auth_cache(mut self, auth_cache: Arc<AuthCache<AuthCacheKey, User>>) -> Self {
self.auth_cache = Some(auth_cache);
self
}
pub fn build(self) -> QueryResult<Arc<SimpleQueryDispatcher>> {
let coord = self.coord.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of coord".to_string(),
})?;
let split_manager = self.split_manager.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of split manager".to_string(),
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 {
@@ -276,56 +254,23 @@ impl SimpleQueryDispatcherBuilder {
err: "lost of query_execution_factory".to_string(),
})?;
let query_tracker = self.query_tracker.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of query_tracker".to_string(),
})?;
let func_manager = self.func_manager.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of func_manager".to_string(),
})?;
let stream_provider_manager = self.stream_provider_manager.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of stream_provider_manager".to_string(),
})?;
let memory_pool = self.memory_pool.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of memory pool".to_string(),
})?;
let default_table_provider = self.default_table_provider.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of default_table_provider".to_string(),
})?;
let span_ctx = self.span_ctx;
let auth_cache = self.auth_cache.ok_or_else(|| QueryError::BuildQueryDispatcher {
err: "lost of auth_cache".to_string(),
})?;
let dispatcher = Arc::new(SimpleQueryDispatcher {
coord,
input,
default_table_provider,
split_manager,
session_factory,
memory_pool,
parser,
query_execution_factory,
query_tracker,
func_manager,
stream_provider_manager,
span_ctx,
async_task_joinhandle: Arc::new(Mutex::new(HashMap::new())),
failed_task_joinhandle: Arc::new(Mutex::new(HashMap::new())),
auth_cache,
});
let meta_task_receiver = dispatcher
.coord
.meta_manager()
.take_resourceinfo_rx()
.expect("meta resource channel only has one consumer");
tokio::spawn(SimpleQueryDispatcher::recv_meta_modify(dispatcher.clone(), meta_task_receiver));
Ok(dispatcher)
}
}