feat(datanode): execute nested MergeScan plans on the datanode (#9206)

* feat(datanode): execute nested MergeScan plans on the datanode

Give a datanode the ability to decode and execute a plan whose MergeScan
carries a nested plan (e.g. an inner MergeScan over a build table), so a
distributed sub-plan pushed to a datanode can itself fan out to other
datanodes. This is the execution/codec capability only; it is not wired
to any user-facing switch and produces no plan by itself.

Codec (src/query/src/query_engine/default_serializer.rs):
- MergeScanAwareSerializer keeps the request's SessionState and the
  engine catalog manager, and decodes a MergeScan payload's tables by
  (catalog, schema, table) name through the engine catalog rather than
  the region-aware request catalog list (which binds every table to the
  request's region and breaks on mismatched columns).

Datanode (src/datanode):
- New DatanodeRegionQueryHandler (region_query.rs) serves region queries
  for the MergeScan nodes of a received plan, resolving region leaders
  via PartitionRuleManager and dispatching to peer datanodes.
- Wire KvBackendCatalogManager + PartitionRuleManager + NodeClients into
  the datanode's catalog/cache registry (with_dist_planner stays false:
  a datanode never rewrites its own plan).

Bug fixes uncovered while wiring this up:
- cache registry: register the table/table_info/table_name/table_route/
  partition_info caches the datanode's catalog manager and partition
  rule manager require (previously missing -> panic on lookup).
- default_serializer: decode_sub_plan bridge no longer drives a future
  on a current-thread runtime from within another executor (spawn a
  dedicated thread + runtime when not on a multi-thread runtime).

Tested by query_engine::default_serializer unit tests: nested MergeScan
round-trip and engine-catalog resolution with mismatched column names.
No user-facing behavior change.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(query): keep GreptimeDB function bindings after rebuilding decode state

SessionStateBuilder::new_from_existing collects the function registry
(HashMap) into a Vec and build() re-registers each entry; register_udf
registers aliases before the function's own name and the Vec order is
unstable. Rebuilding the session state to install MergeScanAwareSerializer
after the GreptimeDB functions were registered could therefore rebind a
name/alias (e.g. date_format, an alias of DataFusion's to_char) back to
the built-in function. This path runs for every decoded plan, not only
nested MergeScan ones.

Extract the GreptimeDB function registration into register_greptime_functions
and re-apply it after the final build of both the top-level decode state and
the MergeScan payload state. Add regression tests asserting date_format stays
bound to the GreptimeDB implementation on both paths.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(datanode): address review on nested MergeScan capability

Three follow-ups from review of the nested MergeScan execution capability:

1. Layer the datanode caches (P2b). table_cache derives from
   table_info_cache + table_name_cache, and partition_info_cache derives
   from table_route_cache, but they were all registered in a single flat
   CacheRegistry whose invalidate() runs concurrently with no ordering.
   A derived cache could be invalidated and refilled from a base cache
   that had not been invalidated yet, leaving stale metadata. Build the
   base caches and the derived caches as separate layers of a
   LayeredCacheRegistry (reusing the existing layering mechanism) so base
   caches finish invalidating before derived ones. Both the production
   path (cmd) and the test scaffolding now build the same layered
   registry; the flat build_datanode_cache_registry is removed.

2. Don't let a nested MergeScan's internal region query compete for the
   concurrency permit (P2a). An outer query holds its permit until its
   stream is fully consumed, and a nested MergeScan makes the datanode
   issue region queries to peer datanodes as execution stages of that
   same query. If such an inner stage also took a permit, the outer query
   waiting on it would deadlock against its own held permit (or time out
   when max_concurrent_queries > 0). Mark inner DN->DN region queries as
   internal stages: QueryRequest gains an  flag carried over
   the wire in the request header's query_context extensions (reserved
   key, not injectable via hints), and the region server skips acquiring
   a concurrency permit for internal-stage requests on both acquire
   sites. Outer queries still acquire normally.

3. Add an end-to-end capability test
   (tests-integration/.../nested_merge_scan_capability_test.rs): hand-build
   Join(<local probe region>, MergeScan(build regions on both datanodes))
   with mismatched column names and duplicate build keys, ship it to a
   datanode over a real TCP Flight endpoint, consume the stream fully and
   compare the row multiset against the frontend's reference join result,
   assert a real cross-datanode DoGet via an RPC counter, and assert an
   inner region failure fails the whole query. Also migrate the
   with_real_datanode_grpc_addr + DatanodeRpcStats test scaffolding.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(query): adapt date_format regression test to the updated DataFusion fork

ScalarUDFImpl's trait bounds changed so its impl can no longer be
downcast via as_any(); assert the function equals the GreptimeDB
implementation instead, which to_char can never satisfy.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(datanode): cover the internal stage under the concurrency limit

Add an end-to-end regression test for the nested MergeScan capability
under max_concurrent_queries = 1: both datanodes' only permit is held by
the outer queries' streams while their inner MergeScans fetch the build
regions from the same datanodes over real Flight connections. The result
consumption is bounded by a timeout, so a broken link anywhere in the
chain (marking the inner request internal, carrying the marker in the
request header, honoring it at the receiving acquire sites) fails the
test instead of deadlocking.

Also add a minimal with_datanode_options_override hook to the cluster
builder, and fix the test module doc: the outer request is handed to
RegionServer::handle_remote_read in process; only the inner region
queries go over the wire.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(datanode): cover cold route loads under unrelated invalidations

The table route cache uses InitStrategy::VersionChecked, whose version
counter belongs to the whole CacheContainer rather than to a single
table: any TableId invalidation bumps the shared version, so a cold
load of one table's route retries when an unrelated table is
invalidated concurrently. Document this as a known limitation on the
region query handler, and pin the behavior with an integration test:
run the nested plans while a background task invalidates an unrelated
table id on every datanode, assert the results still match the expected
multiset (retries only add work, they never return a stale route), and
report the observed load amplification via the cache miss counter
(2 loads without invalidations, 8 over three runs with 24 unrelated
invalidations in this run). Also expose the datanode cache registries
on the test cluster so a test can drive the local invalidation path
that the heartbeat handler drives after a metasrv broadcast.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* docs(query): trim nested plan decoder commentary

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor: remove redundant nested scan support plumbing

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(query): decode nested MergeScan payloads asynchronously

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): avoid unnecessary extension name allocation

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(query): retain extension decoding for manual session states

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor(datanode): keep nested scan query admission unchanged

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(datanode): honor query handler factories and trim nested scan fixtures

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(datanode): retain required cache error import

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor(query): reuse registered state for nested scan decoding

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): verify payload information schema table functions

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): retain nested scan runtime and cold-route guards

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor(query): address nested scan review feedback

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(session): preserve read preference across region query hops

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(session): verify strict context rejection at Flow ingress

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* test(query): remove redundant serializer error case

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

---------

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/cache/src/lib.rs

@@ -20,8 +20,8 @@ use std::time::Duration;
 
 use catalog::kvbackend::new_table_cache;
 use common_meta::cache::{
-    CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistryBuilder, new_schema_cache,
-    new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache,
+    CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder,
+    new_schema_cache, new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache,
     new_table_route_cache, new_table_schema_cache, new_view_info_cache,
 };
 use common_meta::kv_backend::KvBackendRef;
@@ -53,11 +53,11 @@ pub const TABLE_FLOWNODE_SET_CACHE_NAME: &str = "table_flownode_set_cache";
 pub const TABLE_ROUTE_CACHE_NAME: &str = "table_route_cache";
 pub const PARTITION_INFO_CACHE_NAME: &str = "partition_info_cache";
 
-/// Builds cache registry for datanode, including:
-/// - Schema cache.
-/// - Table id to schema name cache.
-pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry {
-    // Builds table id schema name cache that never expires.
+/// Builds the datanode cache registry with base metadata caches before derived caches.
+///
+/// Invalidating the base layer first prevents a concurrent query from refilling a derived cache
+/// from stale metadata after that derived cache has been invalidated.
+pub fn build_datanode_layered_cache_registry(kv_backend: KvBackendRef) -> LayeredCacheRegistry {
     let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY).build();
     let table_id_schema_cache = Arc::new(new_table_schema_cache(
         TABLE_SCHEMA_NAME_CACHE_NAME.to_string(),
@@ -65,7 +65,6 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry
         kv_backend.clone(),
     ));
 
-    // Builds schema cache
     let cache = default_cache();
     let schema_cache = Arc::new(new_schema_cache(
         SCHEMA_CACHE_NAME.to_string(),
@@ -73,9 +72,57 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry
         kv_backend.clone(),
     ));
 
-    CacheRegistryBuilder::default()
+    let cache = default_cache();
+    let table_info_cache = Arc::new(new_table_info_cache(
+        TABLE_INFO_CACHE_NAME.to_string(),
+        cache,
+        kv_backend.clone(),
+    ));
+
+    let cache = default_cache();
+    let table_name_cache = Arc::new(new_table_name_cache(
+        TABLE_NAME_CACHE_NAME.to_string(),
+        cache,
+        kv_backend.clone(),
+    ));
+
+    let cache = default_cache();
+    let table_route_cache = Arc::new(new_table_route_cache(
+        TABLE_ROUTE_CACHE_NAME.to_string(),
+        cache,
+        kv_backend.clone(),
+    ));
+
+    let cache = default_cache();
+    let table_cache = Arc::new(new_table_cache(
+        TABLE_CACHE_NAME.to_string(),
+        cache,
+        table_info_cache.clone(),
+        table_name_cache.clone(),
+    ));
+
+    let cache = default_cache();
+    let partition_info_cache = Arc::new(new_partition_info_cache(
+        PARTITION_INFO_CACHE_NAME.to_string(),
+        cache,
+        table_route_cache.clone(),
+    ));
+
+    let base_registry = CacheRegistryBuilder::default()
         .add_cache(table_id_schema_cache)
         .add_cache(schema_cache)
+        .add_cache(table_info_cache)
+        .add_cache(table_name_cache)
+        .add_cache(table_route_cache)
+        .build();
+    let derived_registry = CacheRegistryBuilder::default()
+        .add_cache(table_cache)
+        .add_cache(partition_info_cache)
+        .build();
+
+    LayeredCacheRegistryBuilder::default()
+        .add_cache_registry(base_registry)
+        .add_cache_registry(derived_registry)
         .build()
 }
 

src/cmd/src/datanode/builder.rs

@@ -14,10 +14,9 @@
 
 use std::sync::Arc;
 
-use cache::build_datanode_cache_registry;
+use cache::build_datanode_layered_cache_registry;
 use catalog::kvbackend::new_read_only_meta_kv_backend;
 use common_base::Plugins;
-use common_meta::cache::LayeredCacheRegistryBuilder;
 use common_telemetry::info;
 use common_version::{short_version, verbose_version};
 use datanode::datanode::DatanodeBuilder;
@@ -111,11 +110,9 @@ impl InstanceBuilder {
         let backend = new_read_only_meta_kv_backend(client.clone());
         let mut builder = DatanodeBuilder::new(dn_opts.clone(), plugins.clone(), backend.clone());
 
-        let registry = Arc::new(
-            LayeredCacheRegistryBuilder::default()
-                .add_cache_registry(build_datanode_cache_registry(backend))
-                .build(),
-        );
+        // The derived caches (table/partition caches) are put into the layer after the caches
+        // they are derived from, so they are invalidated only after them.
+        let registry = Arc::new(build_datanode_layered_cache_registry(backend));
         builder
             .with_cache_registry(registry)
             .with_meta_client(client.clone());

src/common/substrait/src/df_substrait.rs

@@ -25,7 +25,9 @@ use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::optimizer::AnalyzerRule;
 use datafusion::prelude::SessionConfig;
 use datafusion_expr::LogicalPlan;
-use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
+use datafusion_substrait::logical_plan::consumer::{
+    SubstraitConsumer, from_substrait_plan, from_substrait_plan_with_consumer,
+};
 use datafusion_substrait::logical_plan::producer::to_substrait_plan;
 use datafusion_substrait::substrait::proto::Plan;
 use prost::Message;
@@ -74,6 +76,21 @@ impl SubstraitPlan for DFLogicalSubstraitConvertor {
 }
 
 impl DFLogicalSubstraitConvertor {
+    /// Decode with a custom consumer, retaining Greptime's state-aggregate ordering correction.
+    pub async fn decode_with_consumer(
+        &self,
+        plan: &Plan,
+        state: &SessionState,
+        consumer: &impl SubstraitConsumer,
+    ) -> Result<LogicalPlan, Error> {
+        let df_plan = from_substrait_plan_with_consumer(consumer, plan)
+            .await
+            .context(DecodeDfPlanSnafu)?;
+        FixStateUdafOrderingAnalyzer {}
+            .analyze(df_plan, state.config_options())
+            .context(DecodeDfPlanSnafu)
+    }
+
     pub fn to_sub_plan(
         &self,
         plan: &LogicalPlan,

src/common/substrait/src/lib.rs

@@ -20,6 +20,7 @@ use async_trait::async_trait;
 use bytes::{Buf, Bytes};
 use datafusion::execution::context::SessionState;
 pub use datafusion::execution::registry::SerializerRegistry;
+pub use datafusion_substrait::extensions::Extensions;
 /// Re-export the Substrait module of datafusion,
 pub use datafusion_substrait::substrait as substrait_proto_df;
 pub use datafusion_substrait::{logical_plan as df_logical_plan, variation_const};

src/datanode/Cargo.toml

@@ -16,6 +16,8 @@ api.workspace = true
 arrow-flight.workspace = true
 async-trait.workspace = true
 bytes.workspace = true
+catalog.workspace = true
+client.workspace = true
 common-base.workspace = true
 common-catalog.workspace = true
 common-config.workspace = true
@@ -23,6 +25,7 @@ common-datasource.workspace = true
 common-error.workspace = true
 common-function.workspace = true
 common-greptimedb-telemetry.workspace = true
+common-grpc.workspace = true
 common-macro.workspace = true
 common-meta.workspace = true
 common-options.workspace = true
@@ -54,6 +57,7 @@ metric-engine.workspace = true
 mito2.workspace = true
 num_cpus.workspace = true
 object-store.workspace = true
+partition.workspace = true
 prometheus.workspace = true
 prost.workspace = true
 query.workspace = true
@@ -78,5 +82,4 @@ common-query.workspace = true
 common-test-util.workspace = true
 datafusion-common.workspace = true
 mito2 = { workspace = true, features = ["test"] }
-partition.workspace = true
 session.workspace = true

src/datanode/src/datanode.rs

@@ -19,17 +19,22 @@ use std::path::Path;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
+use catalog::information_schema::NoopInformationExtension;
+use catalog::kvbackend::KvBackendCatalogManagerBuilder;
+use client::client_manager::NodeClients;
 use common_base::Plugins;
 use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE};
 use common_datasource::object_store::LocalFileAccess;
 use common_error::ext::BoxedError;
 use common_greptimedb_telemetry::GreptimeDBTelemetryTask;
+use common_grpc::channel_manager::ChannelConfig;
 use common_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef};
 use common_meta::cache_invalidator::CacheInvalidatorRef;
 use common_meta::datanode::TopicStatsReporter;
 use common_meta::key::runtime_switch::RuntimeSwitchManager;
 use common_meta::key::{SchemaMetadataManager, SchemaMetadataManagerRef};
 use common_meta::kv_backend::KvBackendRef;
+use common_meta::node_manager::NodeManagerRef;
 pub use common_procedure::options::ProcedureConfig;
 use common_query::prelude::set_default_prefix;
 use common_stat::ResourceStatImpl;
@@ -54,7 +59,8 @@ use object_store::ObjectStore;
 use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef};
 use object_store::util::normalize_dir;
 use query::QueryEngineFactory;
-use query::dummy_catalog::{DummyCatalogManager, TableProviderFactoryRef};
+use query::dummy_catalog::TableProviderFactoryRef;
+use query::region_query::RegionQueryHandlerFactoryRef;
 use servers::server::ServerHandlers;
 use snafu::{OptionExt, ResultExt, ensure};
 use store_api::logstore::LogStore;
@@ -80,6 +86,7 @@ use crate::event_listener::{
 use crate::greptimedb_telemetry::get_greptimedb_telemetry_task;
 use crate::heartbeat::HeartbeatTask;
 use crate::partition_expr_fetcher::MetaPartitionExprFetcher;
+use crate::region_query::DatanodeRegionQueryHandler;
 use crate::region_server::{DummyTableProviderFactory, RegionServer};
 use crate::store::{self, new_object_store_without_cache};
 use crate::utils::{RegionOpenRequests, build_region_open_requests};
@@ -357,6 +364,7 @@ impl DatanodeBuilder {
                 schema_metadata_manager,
                 region_event_listener,
                 file_ref_manager,
+                cache_registry.clone(),
             )
             .await?;
 
@@ -497,14 +505,41 @@ impl DatanodeBuilder {
         schema_metadata_manager: SchemaMetadataManagerRef,
         event_listener: RegionServerEventListenerRef,
         file_ref_manager: FileReferenceManagerRef,
+        cache_registry: Arc<LayeredCacheRegistry>,
     ) -> Result<RegionServer> {
         let opts: &DatanodeOptions = &self.opts;
 
+        // Keep `with_dist_planner` false: datanodes execute received MergeScan plans.
+        let catalog_manager = KvBackendCatalogManagerBuilder::new(
+            Arc::new(NoopInformationExtension),
+            self.kv_backend.clone(),
+            cache_registry.clone(),
+        )
+        .build();
+        let partition_manager = catalog_manager.partition_manager();
+
+        // MergeScan region requests stream until completion, so they have no client timeout.
+        let mut channel_config = ChannelConfig {
+            timeout: None,
+            ..Default::default()
+        };
+        if opts.grpc.flight_compression.transport_compression() {
+            channel_config.accept_compression = true;
+            channel_config.send_compression = true;
+        }
+        let node_manager: NodeManagerRef = Arc::new(NodeClients::new(channel_config));
+        let region_query_handler =
+            if let Some(factory) = self.plugins.get::<RegionQueryHandlerFactoryRef>() {
+                factory.build(partition_manager.clone(), node_manager)
+            } else {
+                DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager)
+            };
+
         let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
             // query engine in datanode only executes plan with resolved table source.
-            DummyCatalogManager::arc(),
-            None,
-            None,
+            catalog_manager,
+            Some(partition_manager),
+            Some(region_query_handler),
             None,
             None,
             None,
@@ -1060,17 +1095,17 @@ mod tests {
     use std::sync::atomic::{AtomicUsize, Ordering};
     use std::time::Duration;
 
-    use cache::build_datanode_cache_registry;
+    use cache::build_datanode_layered_cache_registry;
     use common_base::Plugins;
     use common_base::readable_size::ReadableSize;
     use common_config::Configurable;
     use common_error::ext::ErrorExt;
     use common_error::status_code::StatusCode;
-    use common_meta::cache::LayeredCacheRegistryBuilder;
     use common_meta::key::RegionRoleSet;
     use common_meta::key::datanode_table::DatanodeTableManager;
     use common_meta::kv_backend::KvBackendRef;
     use common_meta::kv_backend::memory::MemoryKvBackend;
+    use common_meta::node_manager::NodeManagerRef;
     use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir};
     use common_wal::config::DatanodeWalConfig;
     use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION};
@@ -1080,6 +1115,8 @@ mod tests {
     use mito2::engine::MITO_ENGINE_NAME;
     use object_store::ObjectStore;
     use object_store::manager::ObjectStoreManager;
+    use partition::manager::PartitionRuleManagerRef;
+    use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef};
     use store_api::logstore::LogStore;
     use store_api::region_request::RegionRequest;
     use store_api::storage::RegionId;
@@ -1090,6 +1127,7 @@ mod tests {
         wal_object_store,
     };
     use crate::error::{self, Error};
+    use crate::region_query::DatanodeRegionQueryHandler;
     use crate::tests::{MockRegionEngine, mock_region_server};
 
     async fn setup_table_datanode(kv: &KvBackendRef) {
@@ -1118,11 +1156,8 @@ mod tests {
         mock_region_server.register_engine(mock_region.clone());
 
         let kv_backend = Arc::new(MemoryKvBackend::new());
-        let layered_cache_registry = Arc::new(
-            LayeredCacheRegistryBuilder::default()
-                .add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
-                .build(),
-        );
+        let layered_cache_registry =
+            Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
 
         let mut builder = DatanodeBuilder::new(
             DatanodeOptions {
@@ -1164,11 +1199,8 @@ mod tests {
         wal: DatanodeWalConfig,
         kv_backend: KvBackendRef,
     ) -> DatanodeBuilder {
-        let layered_cache_registry = Arc::new(
-            LayeredCacheRegistryBuilder::default()
-                .add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
-                .build(),
-        );
+        let layered_cache_registry =
+            Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
         let mut opts = DatanodeOptions {
             node_id: Some(0),
             wal,
@@ -1446,6 +1478,38 @@ mod tests {
         );
     }
 
+    struct CountingRegionQueryHandlerFactory(Arc<AtomicUsize>);
+
+    impl RegionQueryHandlerFactory for CountingRegionQueryHandlerFactory {
+        fn build(
+            &self,
+            partition_manager: PartitionRuleManagerRef,
+            node_manager: NodeManagerRef,
+        ) -> query::region_query::RegionQueryHandlerRef {
+            self.0.fetch_add(1, Ordering::SeqCst);
+            DatanodeRegionQueryHandler::arc(partition_manager, node_manager)
+        }
+    }
+
+    #[tokio::test]
+    async fn test_build_uses_installed_region_query_handler_factory() {
+        let data_home = create_temp_dir("region-query-handler-factory");
+        let builds = Arc::new(AtomicUsize::new(0));
+        let plugins = Plugins::new();
+        plugins.insert(Arc::new(CountingRegionQueryHandlerFactory(builds.clone()))
+            as RegionQueryHandlerFactoryRef);
+        let mut builder = datanode_builder(
+            data_home.path().to_str().unwrap(),
+            DatanodeWalConfig::default(),
+            Arc::new(MemoryKvBackend::new()),
+        );
+        builder.set_plugins(plugins);
+
+        let mut datanode = builder.build().await.unwrap();
+        assert_eq!(builds.load(Ordering::SeqCst), 1);
+        datanode.shutdown().await.unwrap();
+    }
+
     #[tokio::test]
     async fn test_build_and_shutdown_with_raft_engine_wal() {
         let data_home = create_temp_dir("raft-engine-wal");

src/datanode/src/error.rs

@@ -145,6 +145,13 @@ pub enum Error {
         source: Box<log_store::error::Error>,
     },
 
+    #[snafu(display("Invalid query context: {}", reason))]
+    InvalidQueryContext {
+        reason: String,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Invalid SQL, error: {}", msg))]
     InvalidSql { msg: String },
 
@@ -485,7 +492,8 @@ impl ErrorExt for Error {
 
             Delete { source, .. } => source.status_code(),
 
-            InvalidSql { .. }
+            InvalidQueryContext { .. }
+            | InvalidSql { .. }
             | IllegalPrimaryKeysDef { .. }
             | MissingTimestampColumn { .. }
             | SchemaNotFound { .. }

src/datanode/src/lib.rs

@@ -22,6 +22,7 @@ pub mod heartbeat;
 pub mod metrics;
 mod partition_expr_fetcher;
 pub mod query_stream;
+pub mod region_query;
 pub mod region_server;
 pub mod service;
 pub mod store;

src/datanode/src/region_query.rs

@@ -0,0 +1,109 @@
+// 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.
+
+//! The datanode's [`RegionQueryHandler`] implementation for querying remote regions in a
+//! `MergeScan` plan.
+
+use std::sync::Arc;
+
+use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
+use async_trait::async_trait;
+use common_error::ext::BoxedError;
+use common_meta::node_manager::NodeManagerRef;
+use common_query::request::QueryRequest;
+use common_recordbatch::SendableRecordBatchStream;
+use partition::manager::PartitionRuleManagerRef;
+use query::error::{RegionQuerySnafu, Result as QueryResult, UnimplementedSnafu};
+use query::region_query::{RegionQueryHandler, RegionQueryTarget};
+use session::ReadPreference;
+use snafu::ResultExt;
+use store_api::storage::RegionId;
+
+/// Serves region queries by resolving their leader and forwarding requests to its datanode.
+pub(crate) struct DatanodeRegionQueryHandler {
+    partition_manager: PartitionRuleManagerRef,
+    node_manager: NodeManagerRef,
+}
+
+impl DatanodeRegionQueryHandler {
+    pub fn arc(
+        partition_manager: PartitionRuleManagerRef,
+        node_manager: NodeManagerRef,
+    ) -> Arc<Self> {
+        Arc::new(Self {
+            partition_manager,
+            node_manager,
+        })
+    }
+}
+
+#[async_trait]
+impl RegionQueryHandler for DatanodeRegionQueryHandler {
+    async fn select_target(
+        &self,
+        _read_preference: ReadPreference,
+        region_id: RegionId,
+    ) -> QueryResult<RegionQueryTarget> {
+        // The leader of a region is the only peer that can serve a region query.
+        //
+        let peer = self
+            .partition_manager
+            .find_region_leader(region_id)
+            .await
+            .map_err(BoxedError::new)
+            .context(RegionQuerySnafu)?;
+
+        Ok(RegionQueryTarget::new(peer))
+    }
+
+    async fn do_get(
+        &self,
+        target: &RegionQueryTarget,
+        request: QueryRequest,
+    ) -> QueryResult<SendableRecordBatchStream> {
+        self.node_manager
+            .datanode(target.peer())
+            .await
+            .handle_query(request)
+            .await
+            .map_err(BoxedError::new)
+            .context(RegionQuerySnafu)
+    }
+
+    /// Datanodes forward received region queries and do not produce remote dynamic filters.
+    async fn handle_remote_dyn_filter_update(
+        &self,
+        _target: &RegionQueryTarget,
+        _query_id: String,
+        _update: RemoteDynFilterUpdate,
+    ) -> QueryResult<()> {
+        UnimplementedSnafu {
+            operation: "remote dynamic filter update from a datanode",
+        }
+        .fail()
+    }
+
+    /// See [`RegionQueryHandler::handle_remote_dyn_filter_update`].
+    async fn handle_remote_dyn_filter_unregister(
+        &self,
+        _target: &RegionQueryTarget,
+        _query_id: String,
+        _unregister: RemoteDynFilterUnregister,
+    ) -> QueryResult<()> {
+        UnimplementedSnafu {
+            operation: "remote dynamic filter unregister from a datanode",
+        }
+        .fail()
+    }
+}

src/datanode/src/region_server.rs

@@ -354,7 +354,17 @@ impl RegionServer {
             None
         };
 
-        let ctx = request.header.as_ref().map(|h| h.into());
+        let ctx = request
+            .header
+            .as_ref()
+            .map(QueryContext::try_from)
+            .transpose()
+            .map_err(|err| {
+                crate::error::InvalidQueryContextSnafu {
+                    reason: err.to_string(),
+                }
+                .build()
+            })?;
         let query_ctx = Arc::new(ctx.unwrap_or_else(|| QueryContextBuilder::default().build()));
 
         let region_id = request.region_id;
@@ -977,7 +987,15 @@ impl FlightCraft for RegionServer {
         let query_ctx = request
             .header
             .as_ref()
-            .map(|h| Arc::new(QueryContext::from(h)))
+            .map(QueryContext::try_from)
+            .transpose()
+            .map_err(|err| {
+                servers_error::InvalidParameterSnafu {
+                    reason: err.to_string(),
+                }
+                .build()
+            })?
+            .map(Arc::new)
             .unwrap_or(QueryContext::arc());
 
         let region_server = self.clone();

src/flow/src/adapter/flownode_impl.rs

@@ -826,7 +826,16 @@ impl common_meta::node_manager::Flownode for FlowDualEngine {
         let query_ctx = request
             .header
             .and_then(|h| h.query_context)
-            .map(|ctx| ctx.into());
+            .map(session::context::QueryContext::try_from)
+            .transpose()
+            .map_err(|err| {
+                to_meta_err(snafu::location!())(
+                    crate::error::InvalidQuerySnafu {
+                        reason: err.to_string(),
+                    }
+                    .build(),
+                )
+            })?;
         match request.body {
             Some(flow_request::Body::Create(CreateRequest {
                 flow_id: Some(task_id),

src/servers/src/grpc/greptime_handler.rs

@@ -322,7 +322,8 @@ mod tests {
     use common_time::Timezone;
     use query::options::FLOW_SCHEDULED_TIME_MILLIS;
     use session::hints::{
-        INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
+        INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
+        REMOTE_QUERY_ID_EXTENSION_KEY,
     };
     use snafu::ResultExt;
     use tonic::Code;
@@ -418,6 +419,10 @@ mod tests {
             vec![
                 ("auto_create_table".to_string(), "true".to_string()),
                 ("read_preference".to_string(), "leader".to_string()),
+                (
+                    READ_PREFERENCE_EXTENSION_KEY.to_string(),
+                    "follower".to_string(),
+                ),
                 (
                     REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
                     "spoofed".to_string(),
@@ -446,6 +451,8 @@ mod tests {
             ReadPreference::Leader
         ));
         assert_eq!(query_context.extension("auto_create_table"), Some("true"));
+        assert_eq!(query_context.extension(READ_PREFERENCE_HINT), None);
+        assert_eq!(query_context.extension(READ_PREFERENCE_EXTENSION_KEY), None);
         assert_ne!(query_context.remote_query_id(), Some("spoofed"));
         assert!(
             query_context

src/servers/src/http/hints.rs

@@ -49,7 +49,8 @@ mod tests {
     use query::options::FLOW_SCHEDULED_TIME_MILLIS;
     use session::context::{QueryContextBuilder, generate_remote_query_id};
     use session::hints::{
-        INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
+        INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
+        REMOTE_QUERY_ID_EXTENSION_KEY,
     };
 
     use super::apply_hints;
@@ -75,6 +76,10 @@ mod tests {
                     INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(),
                     "spoofed-regs".to_string(),
                 ),
+                (
+                    READ_PREFERENCE_EXTENSION_KEY.to_string(),
+                    "follower".to_string(),
+                ),
                 (
                     FLOW_SCHEDULED_TIME_MILLIS.to_string(),
                     "1700000000000".to_string(),
@@ -92,6 +97,8 @@ mod tests {
                 .extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
                 .is_none()
         );
+        assert_eq!(query_ctx.read_preference(), session::ReadPreference::Leader);
+        assert!(query_ctx.extension(READ_PREFERENCE_EXTENSION_KEY).is_none());
         assert_eq!(
             query_ctx.extension(FLOW_SCHEDULED_TIME_MILLIS),
             Some("1700000000000")

src/session/Cargo.toml

@@ -27,6 +27,9 @@ snafu.workspace = true
 sql.workspace = true
 uuid.workspace = true
 
+[dev-dependencies]
+prost.workspace = true
+
 [package.metadata.cargo-shear]
 # `common-error` is required by `common-macro`'s `stack_trace_debug`: the
 # generated code references `::common_error` paths, so it must stay a direct

src/session/src/context.rs

@@ -15,6 +15,7 @@
 use std::collections::HashMap;
 use std::fmt::{Display, Formatter};
 use std::net::SocketAddr;
+use std::str::FromStr;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 
@@ -34,12 +35,14 @@ use derive_builder::Builder;
 use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect};
 
 pub use crate::hints::{
-    LIVE_ANALYZE_METRICS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
-    SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
+    LIVE_ANALYZE_METRICS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
+    REMOTE_QUERY_ID_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
 };
 use crate::protocol_ctx::ProtocolCtx;
 use crate::query_id::QueryId;
-use crate::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle};
+use crate::session_config::{
+    InvalidConfigValueSnafu, PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle,
+};
 use crate::{MutableInner, ReadPreference};
 
 pub type QueryContextRef = Arc<QueryContext>;
@@ -180,25 +183,41 @@ impl QueryContextBuilder {
     }
 }
 
-impl From<&RegionRequestHeader> for QueryContext {
-    fn from(value: &RegionRequestHeader) -> Self {
+impl TryFrom<&RegionRequestHeader> for QueryContext {
+    type Error = crate::session_config::Error;
+
+    fn try_from(value: &RegionRequestHeader) -> Result<Self, Self::Error> {
         if let Some(ctx) = &value.query_context {
-            ctx.clone().into()
+            ctx.clone().try_into()
         } else {
-            QueryContextBuilder::default()
+            Ok(QueryContextBuilder::default()
                 .set_extension(
                     REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
                     generate_remote_query_id(),
                 )
-                .build()
+                .build())
         }
     }
 }
 
-impl From<api::v1::QueryContext> for QueryContext {
-    fn from(ctx: api::v1::QueryContext) -> Self {
+impl TryFrom<api::v1::QueryContext> for QueryContext {
+    type Error = crate::session_config::Error;
+
+    fn try_from(mut ctx: api::v1::QueryContext) -> Result<Self, Self::Error> {
+        let read_preference = ctx.extensions.remove(READ_PREFERENCE_EXTENSION_KEY);
+        let read_preference = match read_preference {
+            Some(value) => ReadPreference::from_str(&value).map_err(|_| {
+                InvalidConfigValueSnafu {
+                    name: "read_preference",
+                    value,
+                    hint: "Expected a supported read preference",
+                }
+                .build()
+            })?,
+            None => ReadPreference::default(),
+        };
         let sequences = ctx.snapshot_seqs.as_ref();
-        QueryContextBuilder::default()
+        Ok(QueryContextBuilder::default()
             .current_catalog(ctx.current_catalog)
             .current_schema(ctx.current_schema)
             .timezone(parse_timezone(Some(&ctx.timezone)))
@@ -215,7 +234,8 @@ impl From<api::v1::QueryContext> for QueryContext {
                     .unwrap_or_default(),
             )))
             .explain_options(ctx.explain)
-            .build()
+            .read_preference(read_preference)
+            .build())
     }
 }
 
@@ -234,6 +254,17 @@ impl From<QueryContext> for api::v1::QueryContext {
     ) -> Self {
         let explain = mutable_query_context_data.read().unwrap().explain_options;
         let mutable_inner = mutable_inner.read().unwrap();
+        let mut extensions = extensions;
+        extensions.remove(READ_PREFERENCE_EXTENSION_KEY);
+        if mutable_inner.read_preference != ReadPreference::Leader {
+            extensions.insert(
+                READ_PREFERENCE_EXTENSION_KEY.to_string(),
+                mutable_inner
+                    .read_preference
+                    .to_string()
+                    .to_ascii_lowercase(),
+            );
+        }
         api::v1::QueryContext {
             current_catalog,
             current_schema: mutable_inner.schema.clone(),
@@ -843,7 +874,7 @@ mod test {
             0
         );
         let wire: api::v1::QueryContext = admitted.into();
-        let restored = QueryContext::from(wire);
+        let restored = QueryContext::try_from(wire).unwrap();
         assert_eq!(restored.write_rows_to_admit("greptime", "public", 100), 100);
     }
 
@@ -856,7 +887,7 @@ mod test {
                 .extensions
                 .contains_key(crate::hints::INSERT_SKIP_WAL_HINT)
         );
-        let restored: QueryContext = api_context.into();
+        let restored: QueryContext = api_context.try_into().unwrap();
         assert!(!restored.skip_wal());
     }
 
@@ -875,7 +906,7 @@ mod test {
             explain: None,
         };
 
-        let session_ctx: QueryContext = api_ctx.clone().into();
+        let session_ctx: QueryContext = api_ctx.clone().try_into().unwrap();
         let roundtrip_api: api::v1::QueryContext = session_ctx.into();
 
         assert_eq!(roundtrip_api.current_catalog, api_ctx.current_catalog);
@@ -910,7 +941,7 @@ mod test {
         assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id);
 
         let proto: api::v1::QueryContext = (&ctx).into();
-        let restored = QueryContext::from(proto);
+        let restored = QueryContext::try_from(proto).unwrap();
         assert_eq!(restored.remote_query_id(), Some(query_id));
         assert_eq!(
             restored.remote_query_id_value().unwrap().to_string(),
@@ -918,6 +949,117 @@ mod test {
         );
     }
 
+    #[test]
+    fn test_query_context_read_preference_wire_values() {
+        let wire = |value: Option<&str>| api::v1::QueryContext {
+            current_catalog: DEFAULT_CATALOG_NAME.to_string(),
+            current_schema: DEFAULT_SCHEMA_NAME.to_string(),
+            timezone: "UTC".to_string(),
+            extensions: value
+                .map(|value| {
+                    HashMap::from([(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string())])
+                })
+                .unwrap_or_default(),
+            channel: Channel::Internal as u32,
+            snapshot_seqs: None,
+            explain: None,
+        };
+
+        let header = RegionRequestHeader::default();
+        let header_context = QueryContext::try_from(&header).unwrap();
+        assert_eq!(header_context.read_preference(), ReadPreference::Leader);
+
+        for value in [None, Some("leader"), Some("LEADER")] {
+            let ctx = QueryContext::try_from(wire(value)).unwrap();
+            assert_eq!(ctx.read_preference(), ReadPreference::Leader);
+            assert!(!ctx.extensions().contains_key(READ_PREFERENCE_EXTENSION_KEY));
+        }
+        for value in ["", "unknown", "follower", "follower_preferred", "random"] {
+            let error = QueryContext::try_from(wire(Some(value))).unwrap_err();
+            assert!(matches!(
+                error,
+                crate::session_config::Error::InvalidConfigValue {
+                    name,
+                    value: actual,
+                    ..
+                } if name == "read_preference" && actual == value
+            ));
+        }
+    }
+
+    #[test]
+    fn test_legacy_query_context_wire_defaults_to_leader() {
+        use prost::Message;
+
+        let legacy = api::v1::QueryContext::decode(
+            include_bytes!(
+                "../../../tests/compatibility/wire/read_preference/legacy_query_context.pb"
+            )
+            .as_slice(),
+        )
+        .unwrap();
+        assert_eq!(legacy.current_catalog, "c1");
+        assert_eq!(legacy.current_schema, "s1");
+        assert_eq!(legacy.timezone, "UTC");
+        assert_eq!(legacy.channel, Channel::Internal as u32);
+        assert_eq!(
+            legacy
+                .extensions
+                .get("flow.return_region_seq")
+                .map(String::as_str),
+            Some("true")
+        );
+        assert!(
+            !legacy
+                .extensions
+                .contains_key(READ_PREFERENCE_EXTENSION_KEY)
+        );
+
+        let restored = QueryContext::try_from(legacy).unwrap();
+        assert_eq!(restored.read_preference(), ReadPreference::Leader);
+        assert_eq!(restored.current_catalog(), "c1");
+        assert_eq!(restored.current_schema(), "s1");
+        assert_eq!(restored.timezone().to_string(), "UTC");
+        assert_eq!(restored.channel(), Channel::Internal);
+        assert_eq!(restored.extension("flow.return_region_seq"), Some("true"));
+
+        let outbound: api::v1::QueryContext = restored.into();
+        let roundtrip = api::v1::QueryContext::decode(outbound.encode_to_vec().as_slice()).unwrap();
+        assert_eq!(roundtrip.current_catalog, "c1");
+        assert_eq!(roundtrip.current_schema, "s1");
+        assert_eq!(roundtrip.timezone, "UTC");
+        assert_eq!(roundtrip.channel, Channel::Internal as u32);
+        assert_eq!(
+            roundtrip
+                .extensions
+                .get("flow.return_region_seq")
+                .map(String::as_str),
+            Some("true")
+        );
+        assert!(
+            !roundtrip
+                .extensions
+                .contains_key(READ_PREFERENCE_EXTENSION_KEY)
+        );
+    }
+
+    #[test]
+    fn test_outbound_read_preference_clears_stale_extension() {
+        let context = QueryContextBuilder::default()
+            .set_extension(
+                READ_PREFERENCE_EXTENSION_KEY.to_string(),
+                "stale".to_string(),
+            )
+            .set_extension("unrelated".to_string(), "preserved".to_string())
+            .build();
+        let wire: api::v1::QueryContext = context.into();
+        assert!(!wire.extensions.contains_key(READ_PREFERENCE_EXTENSION_KEY));
+        assert_eq!(
+            wire.extensions.get("unrelated").map(String::as_str),
+            Some("preserved")
+        );
+    }
+
     #[test]
     fn test_live_analyze_metrics_requires_matching_remote_query_id() {
         let mut ctx = QueryContext::arc().as_ref().clone();
@@ -944,7 +1086,7 @@ mod test {
         assert!(ctx.clone().batching_enabled());
         assert!(ctx.fork().batching_enabled());
         let wire: api::v1::QueryContext = ctx.into();
-        let restored = QueryContext::from(wire);
+        let restored = QueryContext::try_from(wire).unwrap();
         assert!(!restored.batching_enabled());
         assert!(!restored.logical_batching_enabled());
         let ctx = QueryContextBuilder::default()

src/session/src/hints.rs

@@ -27,11 +27,13 @@ pub const LIVE_ANALYZE_METRICS_EXTENSION_KEY: &str = "query.live_analyze_metrics
 pub const INSERT_SKIP_WAL_HINT: &str = "insert_skip_wal";
 
 pub const READ_PREFERENCE_HINT: &str = "read_preference";
-pub const RESERVED_EXTENSION_KEYS: [&str; 4] = [
+pub const READ_PREFERENCE_EXTENSION_KEY: &str = "query.read_preference";
+pub const RESERVED_EXTENSION_KEYS: [&str; 5] = [
     REMOTE_QUERY_ID_EXTENSION_KEY,
     INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
     SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
     LIVE_ANALYZE_METRICS_EXTENSION_KEY,
+    READ_PREFERENCE_EXTENSION_KEY,
 ];
 
 /// Deprecated, use `HINTS_KEY` instead.
@@ -65,6 +67,7 @@ mod tests {
         assert!(is_reserved_extension_key(
             LIVE_ANALYZE_METRICS_EXTENSION_KEY
         ));
+        assert!(is_reserved_extension_key(READ_PREFERENCE_EXTENSION_KEY));
         assert!(!is_reserved_extension_key(READ_PREFERENCE_HINT));
     }
 }