feat: parallelism hint in grpc

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

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::<u64>() {
+                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)
     }
 

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",
 ];