rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-09 07:20:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(query): add runtime provider interface (#8386)

Browse Source 
* feat(query): add runtime provider interface

Signed-off-by: discord9 <discord9@163.com>

* docs(query): document runtime provider interface

Signed-off-by: discord9 <discord9@163.com>

* feat(query): pass runtime builder to provider

Signed-off-by: discord9 <discord9@163.com>

* fix(query): make runtime provider fallible

Signed-off-by: discord9 <discord9@163.com>

* fix(query): propagate runtime provider errors

Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-07-07 15:18:23 +08:00
committed by GitHub
parent 63cc232339
commit 9936a83384
6 changed files with 340 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/datanode/src/datanode.rs
View File

@@ -61,8 +61,8 @@ use tokio::sync::Notify;
use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig};
use crate::error::{
self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu,
GetMetadataSnafu, MissingCacheSnafu, MissingNodeIdSnafu, OpenLogStoreSnafu, Result,
ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu,
DataFusionSnafu, GetMetadataSnafu, MissingCacheSnafu, MissingNodeIdSnafu, OpenLogStoreSnafu,
Result, ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu,
};
use crate::event_listener::{
NoopRegionServerEventListener, RegionServerEventListenerRef, RegionServerEventReceiver,
@@ -432,7 +432,7 @@ impl DatanodeBuilder {
) -> Result<RegionServer> {
let opts: &DatanodeOptions = &self.opts;
let query_engine_factory = QueryEngineFactory::new_with_plugins(
let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
// query engine in datanode only executes plan with resolved table source.
DummyCatalogManager::arc(),
None,
@@ -443,7 +443,8 @@ impl DatanodeBuilder {
false,
self.plugins.clone(),
opts.query.clone(),
);
)
.context(DataFusionSnafu)?;
let query_engine = query_engine_factory.query_engine();
let table_provider_factory = self

11
src/flow/src/server.rs
View File

@@ -57,8 +57,8 @@ use crate::adapter::flownode_impl::{FlowDualEngine, FlowDualEngineRef};
use crate::adapter::{FlowStreamingEngineRef, create_worker};
use crate::batching_mode::engine::BatchingEngine;
use crate::error::{
CacheRequiredSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu, ShutdownServerSnafu,
StartServerSnafu, UnexpectedSnafu, to_status_with_last_err,
CacheRequiredSnafu, DatafusionSnafu, ExternalSnafu, ListFlowsSnafu, ParseAddrSnafu,
ShutdownServerSnafu, StartServerSnafu, UnexpectedSnafu, to_status_with_last_err,
};
use crate::heartbeat::HeartbeatTask;
use crate::metrics::{METRIC_FLOW_PROCESSING_TIME, METRIC_FLOW_ROWS};
@@ -384,7 +384,7 @@ impl FlownodeBuilder {
pub async fn build(mut self) -> Result<FlownodeInstance, Error> {
// TODO(discord9): does this query engine need those?
let query_engine_factory = QueryEngineFactory::new_with_plugins(
let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
// query engine in flownode is only used for translate plan with resolved table source.
self.catalog_manager.clone(),
None,
@@ -395,7 +395,10 @@ impl FlownodeBuilder {
false,
Default::default(),
self.opts.query.clone(),
);
)
.context(DatafusionSnafu {
context: "Failed to build query engine",
})?;
let manager = Arc::new(
self.build_manager(query_engine_factory.query_engine())
.await?,

5
src/frontend/src/instance/builder.rs
View File

@@ -45,7 +45,7 @@ use query::QueryEngineFactory;
use query::region_query::RegionQueryHandlerFactoryRef;
use snafu::{OptionExt, ResultExt};
use crate::error::{self, ExternalSnafu, Result};
use crate::error::{self, DataFusionSnafu, ExternalSnafu, Result};
use crate::events::EventHandlerImpl;
use crate::frontend::FrontendOptions;
use crate::heartbeat::frontend_peer_addr;
@@ -242,7 +242,7 @@ impl FrontendBuilder {
let mut query_options = self.options.query.clone();
query_options.enable_per_region_metrics = self.options.logging.enable_per_region_metrics;
let query_engine = QueryEngineFactory::new_with_plugins(
let query_engine = QueryEngineFactory::try_new_with_plugins(
self.catalog_manager.clone(),
Some(partition_manager.clone()),
Some(region_query_handler.clone()),
@@ -253,6 +253,7 @@ impl FrontendBuilder {
plugins.clone(),
query_options,
)
.context(DataFusionSnafu)?
.query_engine();
let frontend_peer_addr = frontend_peer_addr(&self.options);

59
src/query/src/query_engine.rs
View File

@@ -15,6 +15,7 @@
mod context;
mod default_serializer;
pub mod options;
pub mod runtime;
mod state;
use std::any::Any;
use std::sync::Arc;
@@ -41,6 +42,9 @@ use crate::error::Result;
use crate::options::QueryOptions;
use crate::planner::LogicalPlanner;
pub use crate::query_engine::context::QueryEngineContext;
pub use crate::query_engine::runtime::{
DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProvider, QueryRuntimeProviderRef,
};
pub use crate::query_engine::state::QueryEngineState;
use crate::region_query::RegionQueryHandlerRef;
@@ -113,7 +117,28 @@ impl QueryEngineFactory {
with_dist_planner: bool,
options: QueryOptions,
) -> Self {
Self::new_with_plugins(
Self::try_new(
catalog_manager,
region_query_handler,
table_mutation_handler,
procedure_service_handler,
flow_service_handler,
with_dist_planner,
options,
)
.expect("Failed to build query engine factory")
}
pub fn try_new(
catalog_manager: CatalogManagerRef,
region_query_handler: Option<RegionQueryHandlerRef>,
table_mutation_handler: Option<TableMutationHandlerRef>,
procedure_service_handler: Option<ProcedureServiceHandlerRef>,
flow_service_handler: Option<FlowServiceHandlerRef>,
with_dist_planner: bool,
options: QueryOptions,
) -> datafusion::error::Result<Self> {
Self::try_new_with_plugins(
catalog_manager,
None,
region_query_handler,
@@ -138,7 +163,33 @@ impl QueryEngineFactory {
plugins: Plugins,
options: QueryOptions,
) -> Self {
let state = Arc::new(QueryEngineState::new(
Self::try_new_with_plugins(
catalog_manager,
partition_rule_manager,
region_query_handler,
table_mutation_handler,
procedure_service_handler,
flow_service_handler,
with_dist_planner,
plugins,
options,
)
.expect("Failed to build query engine factory")
}
#[allow(clippy::too_many_arguments)]
pub fn try_new_with_plugins(
catalog_manager: CatalogManagerRef,
partition_rule_manager: Option<PartitionRuleManagerRef>,
region_query_handler: Option<RegionQueryHandlerRef>,
table_mutation_handler: Option<TableMutationHandlerRef>,
procedure_service_handler: Option<ProcedureServiceHandlerRef>,
flow_service_handler: Option<FlowServiceHandlerRef>,
with_dist_planner: bool,
plugins: Plugins,
options: QueryOptions,
) -> datafusion::error::Result<Self> {
let state = Arc::new(QueryEngineState::try_new(
catalog_manager,
partition_rule_manager,
region_query_handler,
@@ -148,10 +199,10 @@ impl QueryEngineFactory {
with_dist_planner,
plugins.clone(),
options,
));
)?);
let query_engine = Arc::new(DatafusionQueryEngine::new(state, plugins));
register_functions(&query_engine);
Self { query_engine }
Ok(Self { query_engine })
}
pub fn query_engine(&self) -> QueryEngineRef {

80
src/query/src/query_engine/runtime.rs Normal file
View File

@@ -0,0 +1,80 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::sync::Arc;
use datafusion::error::Result as DfResult;
use datafusion::execution::context::SessionConfig;
use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
use crate::options::QueryOptions;
use crate::query_engine::state::MetricsMemoryPool;
/// Reference-counted query runtime provider.
pub type QueryRuntimeProviderRef = Arc<dyn QueryRuntimeProvider>;
/// Context for building query runtime components.
#[derive(Clone, Copy)]
#[non_exhaustive]
pub struct QueryRuntimeContext<'a> {
/// Query options used by the query engine.
pub query_options: &'a QueryOptions,
/// Resolved memory pool size in bytes.
pub resolved_memory_pool_size: usize,
}
impl<'a> QueryRuntimeContext<'a> {
/// Creates a new query runtime context.
pub fn new(query_options: &'a QueryOptions, resolved_memory_pool_size: usize) -> Self {
Self {
query_options,
resolved_memory_pool_size,
}
}
}
/// Provides DataFusion session and runtime setup for the query engine.
pub trait QueryRuntimeProvider: Send + Sync + 'static {
/// Configures the DataFusion session config before building the session state.
fn configure_session_config(&self, _ctx: QueryRuntimeContext<'_>, _config: &mut SessionConfig) {
}
/// Builds the DataFusion runtime environment.
fn build_runtime_env(
&self,
_ctx: QueryRuntimeContext<'_>,
builder: RuntimeEnvBuilder,
) -> DfResult<Arc<RuntimeEnv>> {
builder.build().map(Arc::new)
}
}
/// Default query runtime provider.
#[derive(Debug, Default)]
pub struct DefaultQueryRuntimeProvider;
impl DefaultQueryRuntimeProvider {
/// Creates a default DataFusion runtime environment builder.
pub fn runtime_env_builder(ctx: QueryRuntimeContext<'_>) -> RuntimeEnvBuilder {
if ctx.resolved_memory_pool_size > 0 {
RuntimeEnvBuilder::new().with_memory_pool(Arc::new(MetricsMemoryPool::new(
ctx.resolved_memory_pool_size,
)))
} else {
RuntimeEnvBuilder::new()
}
}
}
impl QueryRuntimeProvider for DefaultQueryRuntimeProvider {}

205
src/query/src/query_engine/state.rs
View File

@@ -38,7 +38,6 @@ use datafusion::execution::memory_pool::{
GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
TrackConsumersPool,
};
use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_optimizer::optimizer::PhysicalOptimizer;
use datafusion::physical_optimizer::sanity_checker::SanityCheckPlan;
@@ -80,6 +79,9 @@ use crate::optimizer::windowed_sort::WindowedSortPhysicalRule;
use crate::options::QueryOptions as QueryOptionsNew;
use crate::query_engine::DefaultSerializer;
use crate::query_engine::options::QueryOptions;
use crate::query_engine::runtime::{
DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProviderRef,
};
use crate::range_select::planner::RangeSelectPlanner;
use crate::region_query::RegionQueryHandlerRef;
@@ -118,18 +120,37 @@ impl QueryEngineState {
plugins: Plugins,
options: QueryOptionsNew,
) -> Self {
Self::try_new(
catalog_list,
partition_rule_manager,
region_query_handler,
table_mutation_handler,
procedure_service_handler,
flow_service_handler,
with_dist_planner,
plugins,
options,
)
.expect("Failed to build query engine state")
}
#[allow(clippy::too_many_arguments)]
pub fn try_new(
catalog_list: CatalogManagerRef,
partition_rule_manager: Option<PartitionRuleManagerRef>,
region_query_handler: Option<RegionQueryHandlerRef>,
table_mutation_handler: Option<TableMutationHandlerRef>,
procedure_service_handler: Option<ProcedureServiceHandlerRef>,
flow_service_handler: Option<FlowServiceHandlerRef>,
with_dist_planner: bool,
plugins: Plugins,
options: QueryOptionsNew,
) -> DfResult<Self> {
let total_memory = get_total_memory_bytes().max(0) as u64;
let memory_pool_size = options.memory_pool_size.resolve(total_memory) as usize;
let runtime_env = if memory_pool_size > 0 {
Arc::new(
RuntimeEnvBuilder::new()
.with_memory_pool(Arc::new(MetricsMemoryPool::new(memory_pool_size)))
.build()
.expect("Failed to build RuntimeEnv"),
)
} else {
Arc::new(RuntimeEnv::default())
};
let runtime_provider = plugins
.get::<QueryRuntimeProviderRef>()
.unwrap_or_else(|| Arc::new(DefaultQueryRuntimeProvider));
let mut session_config = SessionConfig::new().with_create_default_catalog_and_schema(false);
if options.parallelism > 0 {
session_config = session_config.with_target_partitions(options.parallelism);
@@ -150,6 +171,11 @@ impl QueryEngineState {
.execution
.skip_physical_aggregate_schema_check = true;
let runtime_context = QueryRuntimeContext::new(&options, memory_pool_size);
runtime_provider.configure_session_config(runtime_context, &mut session_config);
let runtime_builder = DefaultQueryRuntimeProvider::runtime_env_builder(runtime_context);
let runtime_env = runtime_provider.build_runtime_env(runtime_context, runtime_builder)?;
// Apply extension rules
let mut extension_rules = Vec::new();
@@ -246,7 +272,7 @@ impl QueryEngineState {
let df_context = SessionContext::new_with_state(session_state);
register_function_aliases(&df_context);
Self {
Ok(Self {
df_context,
catalog_manager: catalog_list,
dyn_filter_registry_manager: Arc::new(DynFilterRegistryManager::default()),
@@ -260,7 +286,7 @@ impl QueryEngineState {
extension_rules,
plugins,
scalar_functions: Arc::new(RwLock::new(HashMap::new())),
}
})
}
fn remove_physical_optimizer_rule(
@@ -545,7 +571,7 @@ impl DfQueryPlanner {
/// This wrapper intercepts all memory pool operations and updates
/// Prometheus metrics for monitoring query memory usage and rejections.
#[derive(Debug)]
struct MetricsMemoryPool {
pub(super) struct MetricsMemoryPool {
inner: Arc<TrackConsumersPool<GreedyMemoryPool>>,
}
@@ -553,7 +579,7 @@ impl MetricsMemoryPool {
// Number of top memory consumers to report in OOM error messages
const TOP_CONSUMERS_TO_REPORT: usize = 5;
fn new(limit: usize) -> Self {
pub(super) fn new(limit: usize) -> Self {
Self {
inner: Arc::new(TrackConsumersPool::new(
GreedyMemoryPool::new(limit),
@@ -611,13 +637,25 @@ impl MemoryPool for MetricsMemoryPool {
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, Ordering};
use common_base::Plugins;
use common_base::memory_limit::MemoryLimit;
use common_base::readable_size::ReadableSize;
use datafusion::error::DataFusionError;
use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryLimit as DfMemoryLimit};
use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
use session::context::QueryContext;
use super::*;
use crate::options::QueryOptions;
use crate::query_engine::runtime::{QueryRuntimeProvider, QueryRuntimeProviderRef};
fn new_query_engine_state() -> QueryEngineState {
new_query_engine_state_with(Plugins::default(), QueryOptions::default())
}
fn new_query_engine_state_with(plugins: Plugins, options: QueryOptions) -> QueryEngineState {
QueryEngineState::new(
catalog::memory::new_memory_catalog_manager().unwrap(),
None,
@@ -626,9 +664,146 @@ mod tests {
None,
None,
false,
plugins,
options,
)
}
struct TestRuntimeProvider {
build_called: AtomicBool,
configure_called: AtomicBool,
}
impl TestRuntimeProvider {
fn new() -> Self {
Self {
build_called: AtomicBool::new(false),
configure_called: AtomicBool::new(false),
}
}
}
impl QueryRuntimeProvider for TestRuntimeProvider {
fn configure_session_config(
&self,
ctx: QueryRuntimeContext<'_>,
config: &mut SessionConfig,
) {
assert_eq!(ctx.resolved_memory_pool_size, 1024);
self.configure_called.store(true, Ordering::SeqCst);
*config = config.clone().with_target_partitions(7);
}
fn build_runtime_env(
&self,
ctx: QueryRuntimeContext<'_>,
builder: RuntimeEnvBuilder,
) -> DfResult<Arc<RuntimeEnv>> {
assert_eq!(ctx.resolved_memory_pool_size, 1024);
self.build_called.store(true, Ordering::SeqCst);
builder
.with_memory_pool(Arc::new(GreedyMemoryPool::new(2048)))
.build()
.map(Arc::new)
}
}
struct ErrorRuntimeProvider;
impl QueryRuntimeProvider for ErrorRuntimeProvider {
fn build_runtime_env(
&self,
_ctx: QueryRuntimeContext<'_>,
_builder: RuntimeEnvBuilder,
) -> DfResult<Arc<RuntimeEnv>> {
Err(DataFusionError::Execution("runtime provider error".into()))
}
}
#[test]
fn query_runtime_default_provider_keeps_bounded_memory_pool() {
let state = new_query_engine_state_with(
Plugins::default(),
QueryOptions {
memory_pool_size: MemoryLimit::Size(ReadableSize(1024)),
..Default::default()
},
);
assert!(matches!(
state
.session_state()
.runtime_env()
.memory_pool
.memory_limit(),
DfMemoryLimit::Finite(1024)
));
}
#[test]
fn query_runtime_provider_from_plugins_builds_runtime_env() {
let plugins = Plugins::default();
let provider = Arc::new(TestRuntimeProvider::new());
plugins.insert::<QueryRuntimeProviderRef>(provider.clone());
let state = new_query_engine_state_with(
plugins,
QueryOptions {
memory_pool_size: MemoryLimit::Size(ReadableSize(1024)),
..Default::default()
},
);
assert!(provider.build_called.load(Ordering::SeqCst));
assert!(matches!(
state
.session_state()
.runtime_env()
.memory_pool
.memory_limit(),
DfMemoryLimit::Finite(2048)
));
}
#[test]
fn query_runtime_provider_from_plugins_configures_session_config() {
let plugins = Plugins::default();
let provider = Arc::new(TestRuntimeProvider::new());
plugins.insert::<QueryRuntimeProviderRef>(provider.clone());
let state = new_query_engine_state_with(
plugins,
QueryOptions {
memory_pool_size: MemoryLimit::Size(ReadableSize(1024)),
..Default::default()
},
);
assert!(provider.configure_called.load(Ordering::SeqCst));
assert_eq!(7, state.session_state().config().target_partitions());
}
#[test]
fn query_runtime_provider_error_is_returned_by_try_new() {
let plugins = Plugins::default();
plugins.insert::<QueryRuntimeProviderRef>(Arc::new(ErrorRuntimeProvider));
let err = QueryEngineState::try_new(
catalog::memory::new_memory_catalog_manager().unwrap(),
None,
None,
None,
None,
None,
false,
plugins,
QueryOptions::default(),
)
.unwrap_err();
assert!(
matches!(err, DataFusionError::Execution(message) if message == "runtime provider error")
);
}
#[test]