diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 685f27c355..c74ed4b68e 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -478,6 +478,17 @@ impl QueryEngine for DatafusionQueryEngine { fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext { let mut state = self.state.session_state(); state.config_mut().set_extension(query_ctx.clone()); + if let Some(parallelism) = query_ctx.extension("query_parallelism") { + if let Ok(n) = parallelism.parse::() { + let new_cfg = state.config().clone().with_target_partitions(n as usize); + *state.config_mut() = new_cfg; + } else { + common_telemetry::warn!( + "Failed to parse query_parallelism: {}, using default value", + parallelism + ); + } + } QueryEngineContext::new(state, query_ctx) } diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index 4c0f359b84..327061e4e4 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -18,7 +18,7 @@ pub const HINTS_KEY_PREFIX: &str = "x-greptime-hint-"; pub const READ_PREFERENCE_HINT: &str = "read_preference"; -pub const HINT_KEYS: [&str; 7] = [ +pub const HINT_KEYS: [&str; 8] = [ "x-greptime-hint-auto_create_table", "x-greptime-hint-ttl", "x-greptime-hint-append_mode", @@ -26,4 +26,6 @@ pub const HINT_KEYS: [&str; 7] = [ "x-greptime-hint-physical_table", "x-greptime-hint-skip_wal", "x-greptime-hint-read_preference", + // same as `query.parallelism` in query options, but only for this query + "x-greptime-hint-query_parallelism", ];