From fc07f6d723bdbffbd075224e90c910c4cfe3cdd8 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:15:00 +0800 Subject: [PATCH 01/37] 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> --- Cargo.lock | 3 + src/cache/src/lib.rs | 55 +++ src/datanode/Cargo.toml | 5 +- src/datanode/src/datanode.rs | 69 ++- src/datanode/src/lib.rs | 1 + src/datanode/src/region_query.rs | 115 +++++ .../src/query_engine/default_serializer.rs | 467 +++++++++++++++++- 7 files changed, 707 insertions(+), 8 deletions(-) create mode 100644 src/datanode/src/region_query.rs diff --git a/Cargo.lock b/Cargo.lock index e985f483efbf..c16572c2fcc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4480,6 +4480,8 @@ dependencies = [ "async-trait", "bytes", "cache", + "catalog", + "client", "common-base", "common-catalog", "common-config", @@ -4487,6 +4489,7 @@ dependencies = [ "common-error", "common-function", "common-greptimedb-telemetry", + "common-grpc", "common-macro", "common-meta", "common-options", diff --git a/src/cache/src/lib.rs b/src/cache/src/lib.rs index 1a234aa802e3..bbcf68b30ff6 100644 --- a/src/cache/src/lib.rs +++ b/src/cache/src/lib.rs @@ -56,6 +56,15 @@ pub const PARTITION_INFO_CACHE_NAME: &str = "partition_info_cache"; /// Builds cache registry for datanode, including: /// - Schema cache. /// - Table id to schema name cache. +/// - Table info cache. +/// - Table name cache. +/// - Table cache. +/// - Table route cache. +/// - Partition info cache. +/// +/// The table/route caches are needed by the catalog manager and by the partition rule manager of +/// the datanode, which are used to plan the `MergeScan` nodes of the plans received from the +/// frontend (e.g. a nested merge scan that queries another datanode). pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry { // Builds table id schema name cache that never expires. let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY).build(); @@ -73,9 +82,55 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry kv_backend.clone(), )); + // Builds table info cache + let cache = default_cache(); + let table_info_cache = Arc::new(new_table_info_cache( + TABLE_INFO_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + // Builds table name cache + let cache = default_cache(); + let table_name_cache = Arc::new(new_table_name_cache( + TABLE_NAME_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + // Builds table cache + 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(), + )); + + // Builds table route cache + let cache = default_cache(); + let table_route_cache = Arc::new(new_table_route_cache( + TABLE_ROUTE_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + // Builds partition info cache + 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(), + )); + 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_cache) + .add_cache(table_route_cache) + .add_cache(partition_info_cache) .build() } diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 7c1267647234..4d4e470e6653 100644 --- a/src/datanode/Cargo.toml +++ b/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 diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 8a80083aef3e..2c3dc93c8f21 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -19,17 +19,25 @@ use std::path::Path; use std::sync::Arc; use std::time::{Duration, Instant}; +use catalog::CatalogManagerRef; +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_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef}; +use common_grpc::channel_manager::ChannelConfig; +use common_meta::cache::{ + LayeredCacheRegistry, SchemaCacheRef, TableRouteCacheRef, 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; @@ -53,8 +61,10 @@ use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef}; use object_store::ObjectStore; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::util::normalize_dir; +use partition::cache::PartitionInfoCacheRef; +use partition::manager::PartitionRuleManager; use query::QueryEngineFactory; -use query::dummy_catalog::{DummyCatalogManager, TableProviderFactoryRef}; +use query::dummy_catalog::TableProviderFactoryRef; use servers::server::ServerHandlers; use snafu::{OptionExt, ResultExt, ensure}; use store_api::logstore::LogStore; @@ -80,6 +90,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}; @@ -350,6 +361,7 @@ impl DatanodeBuilder { schema_metadata_manager, region_event_listener, file_ref_manager, + cache_registry.clone(), ) .await?; @@ -490,14 +502,61 @@ impl DatanodeBuilder { schema_metadata_manager: SchemaMetadataManagerRef, event_listener: RegionServerEventListenerRef, file_ref_manager: FileReferenceManagerRef, + cache_registry: Arc, ) -> Result { let opts: &DatanodeOptions = &self.opts; + // The datanode executes the plans of the region queries it receives. Such a plan may + // contain a `MergeScan` node: the dist planner of the datanode plans it into a + // `MergeScanExec`, which resolves the metadata of the target table from the catalog and the + // routes of its regions from the partition rule manager, then queries their leaders through + // the datanode client. + // + // Notes: + // - The `DistExtensionPlanner` is registered as long as both the partition rule manager and + // the region query handler are given (see `DfQueryPlanner::new`), it doesn't depend on + // `with_dist_planner`. + // - `with_dist_planner` stays `false` on purpose: the datanode only executes the plans it + // receives, it must not rewrite its own plans into distributed ones. + let catalog_manager: CatalogManagerRef = KvBackendCatalogManagerBuilder::new( + Arc::new(NoopInformationExtension), + self.kv_backend.clone(), + cache_registry.clone(), + ) + .build(); + + let table_route_cache: TableRouteCacheRef = + cache_registry.get().context(MissingCacheSnafu)?; + let partition_info_cache: PartitionInfoCacheRef = + cache_registry.get().context(MissingCacheSnafu)?; + let partition_manager = Arc::new(PartitionRuleManager::new( + self.kv_backend.clone(), + table_route_cache, + partition_info_cache, + )); + + // Clients to the other datanodes. They are only used when a plan contains a `MergeScan` + // node, i.e. they stay idle for the plans that the datanode executes locally. + // + // A region query is a streaming request, so it must not be cut off by the request timeout: + // same as the frontend to datanode client. + 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 = + 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, diff --git a/src/datanode/src/lib.rs b/src/datanode/src/lib.rs index c14ec886facc..5a8948e2c846 100644 --- a/src/datanode/src/lib.rs +++ b/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; diff --git a/src/datanode/src/region_query.rs b/src/datanode/src/region_query.rs new file mode 100644 index 000000000000..84cd18f8fc61 --- /dev/null +++ b/src/datanode/src/region_query.rs @@ -0,0 +1,115 @@ +// 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. + +//! A [`RegionQueryHandler`] implementation for the datanode. +//! +//! A datanode executes the plans it receives through region queries. Such a plan may contain a +//! `MergeScan` node (see [`query::dist_plan::MergeScanLogicalPlan`]) that is planned into a +//! `MergeScanExec` by the dist planner of the datanode, which then has to query the regions of +//! that `MergeScan` from the datanodes that host them. This handler provides the datanode with the +//! client side of those region queries. + +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 the region queries issued by the datanode itself. +pub struct DatanodeRegionQueryHandler { + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, +} + +impl DatanodeRegionQueryHandler { + pub fn arc( + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + ) -> Arc { + 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 { + // 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 { + self.node_manager + .datanode(target.peer()) + .await + .handle_query(request) + .await + .map_err(BoxedError::new) + .context(RegionQuerySnafu) + } + + /// The remote dynamic filter of a nested merge scan is not supported yet (PoC): a datanode only + /// dispatches region queries for the `MergeScan` nodes of the plans it receives, it doesn't + /// produce remote dynamic filters of its own. + 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() + } +} diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index f4d6b29fa2cd..cd341b2a21a7 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -12,8 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::fmt; use std::sync::Arc; +use bytes::Bytes; +use catalog::CatalogManagerRef; use common_error::ext::BoxedError; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; @@ -44,12 +47,13 @@ use promql::functions::{ Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use prost::Message; -use session::context::QueryContextRef; +use session::context::{QueryContext, QueryContextRef}; use snafu::ResultExt; use substrait::extension_serializer::ExtensionSerializer; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use crate::dist_plan::MergeScanLogicalPlan; +use crate::query_engine::QueryEngineState; /// Extended [`substrait::extension_serializer::ExtensionSerializer`] but supports [`MergeScanLogicalPlan`] serialization. #[derive(Debug)] @@ -88,6 +92,7 @@ impl SerializerRegistry for DefaultSerializer { if name == MergeScanLogicalPlan::name() { // TODO(dennis): missing `session_state` to decode the logical plan in `MergeScanLogicalPlan`, // so we only save the unoptimized logical plan for view currently. + // Plan decoding paths that need to decode a `MergeScan` use [`MergeScanAwareSerializer`]. Err(DataFusionError::Substrait(format!( "Unsupported plan node: {name}" ))) @@ -97,6 +102,214 @@ impl SerializerRegistry for DefaultSerializer { } } +/// Extended [`substrait::extension_serializer::ExtensionSerializer`] that supports both +/// serialization and deserialization of [`MergeScanLogicalPlan`]. +/// +/// [`DefaultSerializer`] is a unit struct that is shared by every encoding path, hence it cannot +/// decode a `MergeScan` node: the plan embedded in the `MergeScan` payload is decoded with a +/// [`SessionState`] which is only known at the decoding site. This registry keeps such a session +/// state and is installed (instead of `DefaultSerializer`) into the session state built by +/// [`DefaultPlanDecoder::decode`], so that plans containing a `MergeScan` can be decoded. +struct MergeScanAwareSerializer { + /// The session state used to decode the plans embedded in `MergeScan` payloads: the state of + /// the request being decoded, with the functions and the catalog list given to + /// [`DefaultPlanDecoder::decode`]. + session_state: SessionState, + /// The catalog of the query engine, which knows every table of the cluster. It resolves the + /// tables of a `MergeScan` payload, see [`Self::decode_payload`]. + /// + /// It is `None` when the session state doesn't carry the engine state (a session state built + /// by hand, e.g. in a unit test): the payload is then decoded with the catalog list of the + /// request, which is the behavior of every payload before this field existed. + catalog_manager: Option, +} + +impl fmt::Debug for MergeScanAwareSerializer { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // `dyn CatalogManager` is not `Debug`. + f.debug_struct("MergeScanAwareSerializer") + .field("session_state", &self.session_state) + .field("has_catalog_manager", &self.catalog_manager.is_some()) + .finish() + } +} + +impl MergeScanAwareSerializer { + /// Returns the query context of the session state, or an empty one if it doesn't carry any. + fn query_ctx(&self) -> QueryContextRef { + self.session_state + .config() + .get_extension() + .unwrap_or_else(QueryContext::arc) + } + + /// Returns the session state used to decode a `MergeScan` payload whose tables are resolved + /// with `catalog_list`. + /// + /// The state is rebuilt with this registry, so that a `MergeScan` nested in the payload (e.g. + /// the plan sent to a datanode that has to query another datanode) is decoded recursively. + /// + /// The default catalog and schema of the state are the ones of the request, because the table + /// references of a payload may be bare (a table name as it was written in the query): they are + /// resolved against that pair, which is how they refer to a `(catalog, schema, table)` of the + /// catalog. Decoding the payload with the catalog list of a region aware request used to make + /// that pair irrelevant (that list resolves every catalog and schema to the region); a catalog + /// that resolves the tables of the cluster needs it. + fn payload_state(&self, catalog_list: Arc) -> SessionState { + let query_ctx = self.query_ctx(); + let mut config = self.session_state.config().clone(); + { + let catalog_options = &mut config.options_mut().catalog; + catalog_options.default_catalog = query_ctx.current_catalog().to_string(); + catalog_options.default_schema = query_ctx.current_schema(); + } + + SessionStateBuilder::new_from_existing(self.session_state.clone()) + .with_config(config) + .with_serializer_registry(Arc::new(Self { + session_state: self.session_state.clone(), + catalog_manager: self.catalog_manager.clone(), + })) + .with_catalog_list(catalog_list) + .build() + } + + /// Decodes `payload`, the plan embedded in a `MergeScan`. + /// + /// The tables of the payload are the tables that the merge scan reads, i.e. the tables of that + /// merge scan and not the region that the enclosing plan is decoded for: they are resolved by + /// `(catalog, schema, table)` name through the catalog of the query engine, which knows every + /// table of the cluster. + /// + /// The catalog list of the request is deliberately not used for the payload: on a datanode that + /// list is the region aware one (`NameAwareCatalogList`), which resolves *every* table name to + /// the region of the request. A nested `MergeScan` reads another table than the enclosing plan + /// (the build side of a nested join, e.g.), and binding it to the region of the request decodes + /// it with the schema of the table of that region (the columns don't even match when the two + /// tables don't share their column names). The tables of a payload are read by the datanodes + /// that own their regions: `MergeScanExec` sends the plan of the payload to them as is, and + /// they decode it against their own regions. + /// + /// When the session state carries a catalog manager (the query engine does), the payload is + /// decoded with it and a failure is returned as is: falling back to the catalog list of the + /// request would decode the payload with a *different* table binding semantics (on a datanode + /// it binds every table to the region of the request), masking a real catalog, metadata or + /// decoding error behind a wrong schema. The catalog list of the request is only used when no + /// catalog manager is available (e.g. a hand-built session state of a unit test). + fn decode_payload(&self, payload: Vec) -> Result { + if let Some(catalog_manager) = &self.catalog_manager { + let engine_catalog = Arc::new( + catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( + catalog_manager.clone(), + self.query_ctx(), + ), + ); + return decode_sub_plan(payload, self.payload_state(engine_catalog)); + } + + decode_sub_plan( + payload, + self.payload_state(self.session_state.catalog_list().clone()), + ) + } +} + +impl SerializerRegistry for MergeScanAwareSerializer { + fn serialize_logical_plan(&self, node: &dyn UserDefinedLogicalNode) -> Result> { + // Encoding doesn't need any session state, delegate to the default serializer. + DefaultSerializer.serialize_logical_plan(node) + } + + fn deserialize_logical_plan( + &self, + name: &str, + bytes: &[u8], + ) -> Result> { + if name != MergeScanLogicalPlan::name() { + return DefaultSerializer.deserialize_logical_plan(name, bytes); + } + + let merge_scan = PbMergeScan::decode(bytes).map_err(|e| { + DataFusionError::Substrait(format!("Failed to decode the MergeScan plan node: {e}")) + })?; + + let input = self.decode_payload(merge_scan.input)?; + + // `PbMergeScan` doesn't carry `partition_cols`, so the decoded node maps no partition + // column to its aliases (an empty `AliasMapping`). `MergeScanExec` uses it to declare the + // hash partitioning of its output and to keep a hash partitioned input of an upstream + // operator (see `MergeScanExec::try_with_new_distribution`), so an empty mapping only + // loses that optimization (an extra repartition may be inserted by the planner), it + // doesn't change the results. Carrying `partition_cols` in the payload requires a proto + // change and is left to a follow-up. + Ok(Arc::new(MergeScanLogicalPlan::new( + input, + merge_scan.is_placeholder, + Default::default(), + ))) + } +} + +/// Decodes `sub_plan`, the plan embedded in a `MergeScan` payload. +/// +/// This sync/async bridge exists because the two sides of it have incompatible shapes (a PoC +/// simplification): +/// +/// * `SerializerRegistry::deserialize_logical_plan` is synchronous — it is called from the +/// synchronous `consume_extension_leaf` of the datafusion fork — while decoding a substrait plan +/// is asynchronous. +/// * `consume_extension_leaf` does not treat the input of a `MergeScan` as a subtree it recurses +/// into, so a nested `MergeScan` is only decoded by recursing from here, in our own registry. +/// Each recursion level therefore re-enters this synchronous bridge *from inside the future it +/// is currently driving*, which a plain `block_on` cannot do: the nested executor aborts with +/// `cannot execute LocalPool executor from within another executor: EnterError`. +/// +/// The future must thus be driven on a runtime other than the one the caller is already running +/// in. Driving it at all needs a runtime — not a bare polling loop — because resolving the tables +/// of the payload may read the metadata through the catalog (a remote read on a datanode), which +/// needs a reactor. +/// +/// * On a multi-threaded runtime (the datanode production path) `block_in_place` hands the other +/// tasks of the current worker thread over to another worker, then drives the future on the +/// current runtime; the remaining workers keep serving the I/O the future awaits. +/// * Otherwise (a current-thread runtime, e.g. the default of `#[tokio::test]`, or no runtime at +/// all) there is no other worker to hand the tasks over to, so driving the future here would +/// both starve the runtime and panic as soon as this function is re-entered recursively. The +/// future is moved to a fresh thread that builds its own current-thread runtime instead. That +/// thread is unrelated to whatever the caller runs in, so a recursive call simply spawns another +/// fresh thread; the caller thread only joins and never re-enters its own executor. +/// +/// Scope of that current-thread branch (best effort, not a general guarantee): the production +/// datanode runs on a multi-threaded runtime and uses the `block_in_place` branch above. The +/// current-thread branch exists for environments such as `#[tokio::test]`. It assumes the +/// catalog resolution the future awaits does *not* depend on tasks running on the caller's +/// runtime (the PoC uses in-memory / KV catalogs, for which that holds). `Send`-ness of the +/// future alone does not establish that independence, so this branch makes no safety claim for +/// arbitrary runtime setups; it also offers no cancellation of the synchronous `join`. +fn decode_sub_plan(sub_plan: Vec, session_state: SessionState) -> Result { + let decode = async move { + DFLogicalSubstraitConvertor + .decode(Bytes::from(sub_plan), session_state) + .await + .map_err(|e| DataFusionError::External(Box::new(e))) + }; + + match tokio::runtime::Handle::try_current() { + Ok(handle) if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread => { + tokio::task::block_in_place(|| handle.block_on(decode)) + } + _ => std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|e| DataFusionError::External(Box::new(e)))?; + runtime.block_on(decode) + }) + .join() + .unwrap_or_else(|panic| std::panic::resume_unwind(panic)), + } +} + /// The datafusion `[LogicalPlan]` decoder. pub struct DefaultPlanDecoder { session_state: SessionState, @@ -123,7 +336,8 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { catalog_list: Arc, optimize: bool, ) -> common_query::error::Result { - // The session_state already has the `DefaultSerialzier` as `SerializerRegistry`. + // The session_state already has the `DefaultSerialzier` as `SerializerRegistry`. It is + // replaced by `MergeScanAwareSerializer` below, once this state is fully built. let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) .with_catalog_list(catalog_list) .build(); @@ -223,6 +437,26 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { let _ = session_state.register_udaf(Arc::new(udaf)); } + // Install a registry that is also able to decode `MergeScan` nodes (the plans sent to a + // datanode may contain one) after all the functions above are registered, so that the + // sub-plans of a `MergeScan` see the same functions. + // + // The registry resolves the tables of a `MergeScan` payload with the catalog of the query + // engine, which it retrieves from the engine state the query engine leaves in the session + // config (`QueryEngine::engine_context`). A session state that carries no engine state has + // no catalog manager either, in which case the payloads are decoded with the catalog list + // of the request, see `MergeScanAwareSerializer::decode_payload`. + let catalog_manager = session_state + .config() + .get_extension::() + .map(|engine_state| engine_state.catalog_manager().clone()); + let session_state = SessionStateBuilder::new_from_existing(session_state.clone()) + .with_serializer_registry(Arc::new(MergeScanAwareSerializer { + session_state, + catalog_manager, + })) + .build(); + let logical_plan = DFLogicalSubstraitConvertor .decode(message, session_state) .await @@ -241,6 +475,8 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { #[cfg(test)] mod tests { + use catalog::RegisterTableRequest; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; use common_query::native_histogram::native_histogram_value_type; use datafusion::catalog::TableProvider; use datafusion::datasource::MemTable; @@ -253,6 +489,7 @@ mod tests { use datatypes::data_type::DataType; use promql::extension_plan::RangeManipulate; use session::context::QueryContext; + use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable}; use super::*; use crate::QueryEngineFactory; @@ -488,4 +725,230 @@ mod tests { assert!(decoded.contains("prom_mixed_range_float")); assert!(decoded.contains("prom_mixed_range_histogram")); } + + #[tokio::test] + async fn test_serializer_decode_merge_scan() { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + // The payload scans `numbers`: it has to be resolvable through the catalog of the query + // engine, which is what a datanode uses to decode the payload of a `MergeScan` (the + // catalog list of the request is not a fallback for the engine catalog). + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + let factory = QueryEngineFactory::new( + catalog_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + + // A `MergeScan` embeds its input plan into its payload, so decoding a `MergeScan` has to + // decode that plan recursively. + let input = LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .build() + .unwrap(); + let plan = + MergeScanLogicalPlan::new(input.clone(), true, Default::default()).into_logical_plan(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(QueryContext::arc()) + .new_plan_decoder() + .unwrap(); + // The catalog list of the request is deliberately unrelated to the payload: the payload + // is resolved through the catalog of the query engine, not through this list. + let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))); + + let decoded = plan_decoder + .decode(bytes, catalog_list, false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + assert!(merge_scan.is_placeholder()); + assert_eq!(merge_scan.input().to_string(), input.to_string()); + // `PbMergeScan` doesn't carry `partition_cols`, so decoded nodes have an empty mapping. + assert!(merge_scan.partition_cols().is_empty()); + } + + /// The plan embedded in a `MergeScan` is decoded with the catalog of the query engine, which + /// resolves its tables by `(catalog, schema, table)` name, and not with the catalog list of the + /// request. + /// + /// On a datanode the catalog list of the request is the region aware one + /// (`NameAwareCatalogList`), which resolves *every* table name to the region of the request: a + /// payload reading another table is then decoded with the schema of the table of that region, + /// which breaks the plan as soon as the two tables don't share their columns. This test covers + /// that binding with a `numbers` table whose column names differ from the request region's. + #[tokio::test] + async fn test_serializer_decode_merge_scan_with_engine_catalog() { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + // The table of the payload: `numbers`, with a single `number` column. + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + let factory = QueryEngineFactory::new( + catalog_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + + // The payload scans `numbers` by name, so the catalog that decodes it has to resolve a + // bare table name too. + let input = LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .build() + .unwrap(); + let plan = MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + + // The catalog list of the request resolves every table name to a provider whose schema is + // not the one of `numbers` (the region aware list of a datanode does the same). + let table_provider = Arc::new(mock_table_provider(1.into())); + let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider)); + + let plan_decoder = engine + .engine_context(QueryContext::arc()) + .new_plan_decoder() + .unwrap(); + let decoded = plan_decoder + .decode(bytes, catalog_list, false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + let LogicalPlan::TableScan(scan) = merge_scan.input() else { + panic!("Expect a table scan, got: {}", merge_scan.input()); + }; + // The scan is the `numbers` table of the catalog of the query engine, not the table of the + // request catalog list. + assert_eq!(scan.table_name.table(), NUMBERS_TABLE_NAME); + assert_eq!(scan.source.schema().fields().len(), 1); + assert_eq!(scan.source.schema().field(0).name(), "number"); + } + + /// A `MergeScan` nested in the payload of another `MergeScan` is decoded recursively: this is + /// the shape of a datanode querying another datanode. + #[tokio::test] + async fn test_serializer_decode_nested_merge_scan() { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + // The payload of both `MergeScan` levels scans `numbers`, which has to be resolvable + // through the catalog of the query engine (the request catalog list is not a fallback). + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + let factory = QueryEngineFactory::new( + catalog_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + + let numbers_scan = || { + LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .build() + .unwrap() + }; + let inner = MergeScanLogicalPlan::new(numbers_scan(), false, Default::default()) + .into_logical_plan(); + let input = LogicalPlanBuilder::from(inner) + .filter(col("number").lt(lit(10u32))) + .unwrap() + .build() + .unwrap(); + let plan = + MergeScanLogicalPlan::new(input.clone(), false, Default::default()).into_logical_plan(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(QueryContext::arc()) + .new_plan_decoder() + .unwrap(); + // The catalog list of the request is deliberately unrelated to the payload. + let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))); + + let decoded = plan_decoder + .decode(bytes, catalog_list, false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + // Both levels are `MergeScan` nodes. + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + } } From 60d5713cd87c1cd1501ea698c75b1898a418eac5 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:34:53 +0800 Subject: [PATCH 02/37] 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> --- .../src/query_engine/default_serializer.rs | 477 ++++++++++++++---- 1 file changed, 372 insertions(+), 105 deletions(-) diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index cd341b2a21a7..e0d41674b24e 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -155,7 +155,12 @@ impl MergeScanAwareSerializer { /// catalog. Decoding the payload with the catalog list of a region aware request used to make /// that pair irrelevant (that list resolves every catalog and schema to the region); a catalog /// that resolves the tables of the cluster needs it. - fn payload_state(&self, catalog_list: Arc) -> SessionState { + /// + /// The functions of the state are registered once more after it is built, see + /// [`register_greptime_functions`]: building a state may bind a function back to a built-in + /// DataFusion function of the same name or alias, so the payload has to be decoded with a state + /// that holds the GreptimeDB functions, like the state that decodes the enclosing plan. + fn payload_state(&self, catalog_list: Arc) -> Result { let query_ctx = self.query_ctx(); let mut config = self.session_state.config().clone(); { @@ -164,14 +169,17 @@ impl MergeScanAwareSerializer { catalog_options.default_schema = query_ctx.current_schema(); } - SessionStateBuilder::new_from_existing(self.session_state.clone()) + let mut state = SessionStateBuilder::new_from_existing(self.session_state.clone()) .with_config(config) .with_serializer_registry(Arc::new(Self { session_state: self.session_state.clone(), catalog_manager: self.catalog_manager.clone(), })) .with_catalog_list(catalog_list) - .build() + .build(); + register_greptime_functions(&mut state, &query_ctx).map_err(DataFusionError::from)?; + + Ok(state) } /// Decodes `payload`, the plan embedded in a `MergeScan`. @@ -204,12 +212,12 @@ impl MergeScanAwareSerializer { self.query_ctx(), ), ); - return decode_sub_plan(payload, self.payload_state(engine_catalog)); + return decode_sub_plan(payload, self.payload_state(engine_catalog)?); } decode_sub_plan( payload, - self.payload_state(self.session_state.catalog_list().clone()), + self.payload_state(self.session_state.catalog_list().clone())?, ) } } @@ -328,6 +336,123 @@ impl DefaultPlanDecoder { } } +/// Registers every function of GreptimeDB into `session_state`, overwriting any built-in +/// function of the same name or alias. +/// +/// The sub strait decoder looks the functions of a plan up in the session state, and some of the +/// GreptimeDB functions collide with a built-in DataFusion function or one of its aliases: the +/// built-in `to_char` declares `date_format` as an alias, while GreptimeDB has a `date_format` +/// function of its own. +/// +/// That is why the functions have to be registered *after* the session state is built. Building a +/// state re-registers the functions of the state it was given — `SessionStateBuilder::build` +/// takes them from `new_from_existing`, which turns the function maps into vectors — in the +/// unspecified order in which a hash map yields them, and registering a function inserts its +/// aliases *before* its own name. A state that holds both the built-in `to_char` and the GreptimeDB +/// `date_format` therefore binds `date_format` back to the built-in `to_char` whenever it +/// re-registers `to_char` last. +/// +/// The query context must be passed to the functions to set the timezone. This function is +/// idempotent, and it must be applied to the state that actually decodes a plan, i.e. after *every* +/// `SessionStateBuilder::build`. +fn register_greptime_functions( + session_state: &mut SessionState, + query_ctx: &QueryContextRef, +) -> common_query::error::Result<()> { + for func in FUNCTION_REGISTRY.scalar_functions() { + let udf = func.provide(FunctionContext { + query_ctx: query_ctx.clone(), + state: Default::default(), + }); + session_state + .register_udf(Arc::new(udf)) + .context(RegisterUdfSnafu { name: func.name() })?; + } + + for func in FUNCTION_REGISTRY.aggregate_functions() { + let name = func.name().to_string(); + session_state + .register_udaf(Arc::new(func)) + .context(RegisterUdfSnafu { name })?; + } + + let _ = session_state.register_udaf(quantile_udaf()); + + let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); + let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Rate::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Increase::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Delta::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Resets::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Changes::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf())); + let _ = session_state.register_udf(Arc::new(Round::scalar_udf())); + let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf())); + let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf())); + let double_exponential_smoothing_udf = + DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]); + let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf)); + + for udf in [ + NativeHistogramAbsentOverTime::scalar_udf(), + NativeHistogramAdd::scalar_udf(), + NativeHistogramAvg::scalar_udf(), + NativeHistogramAvgOverTime::scalar_udf(), + NativeHistogramChanges::scalar_udf(), + NativeHistogramCount::scalar_udf(), + NativeHistogramCountOverTime::scalar_udf(), + NativeHistogramDelta::scalar_udf(), + NativeHistogramDivScalar::scalar_udf(), + NativeHistogramDrop::bool_false_udf(String::new(), None), + NativeHistogramDrop::bool_true_udf(String::new(), None), + NativeHistogramDrop::float_null_udf(String::new(), None), + NativeHistogramEq::scalar_udf(), + NativeHistogramFraction::scalar_udf(), + NativeHistogramIDelta::scalar_udf(), + NativeHistogramIRate::scalar_udf(), + NativeHistogramIncrease::scalar_udf(), + NativeHistogramLastOverTime::scalar_udf(), + MixedRange::float_udf(None), + MixedRange::histogram_udf(None), + NativeHistogramMulScalar::scalar_udf(), + NativeHistogramNeg::scalar_udf(), + NativeHistogramNotEq::scalar_udf(), + NativeHistogramPresentOverTime::scalar_udf(), + NativeHistogramQuantile::scalar_udf(), + NativeHistogramRate::scalar_udf(), + NativeHistogramResets::scalar_udf(), + NativeHistogramScalarMul::scalar_udf(), + NativeHistogramStddev::scalar_udf(), + NativeHistogramStdvar::scalar_udf(), + NativeHistogramSub::scalar_udf(), + NativeHistogramSum::scalar_udf(), + NativeHistogramSumOverTime::scalar_udf(), + NativeHistogramToString::scalar_udf(), + PromqlFloatToString::scalar_udf(), + ] { + let _ = session_state.register_udf(Arc::new(udf)); + } + + for udaf in [ + NativeHistogramAggAvg::aggregate_udf(), + NativeHistogramAggSum::aggregate_udf(), + ] { + let _ = session_state.register_udaf(Arc::new(udaf)); + } + + Ok(()) +} + #[async_trait::async_trait] impl SubstraitPlanDecoder for DefaultPlanDecoder { async fn decode( @@ -341,105 +466,10 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) .with_catalog_list(catalog_list) .build(); - // Substrait decoder will look up the UDFs in SessionState, so we need to register them - // Note: the query context must be passed to set the timezone - // We MUST register the UDFs after we build the session state, otherwise the UDFs will be lost - // if they have the same name as the default UDFs or their alias. - // e.g. The default UDF `to_char()` has an alias `date_format()`, if we register a UDF with the name `date_format()` - // before we build the session state, the UDF will be lost. - for func in FUNCTION_REGISTRY.scalar_functions() { - let udf = func.provide(FunctionContext { - query_ctx: self.query_ctx.clone(), - state: Default::default(), - }); - session_state - .register_udf(Arc::new(udf)) - .context(RegisterUdfSnafu { name: func.name() })?; - } - - for func in FUNCTION_REGISTRY.aggregate_functions() { - let name = func.name().to_string(); - session_state - .register_udaf(Arc::new(func)) - .context(RegisterUdfSnafu { name })?; - } - - let _ = session_state.register_udaf(quantile_udaf()); - - let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); - let _ = session_state.register_udf(Arc::new(IDelta::::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Rate::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Increase::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Delta::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Resets::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Changes::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf())); - let _ = session_state.register_udf(Arc::new(Round::scalar_udf())); - let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf())); - let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf())); - let double_exponential_smoothing_udf = - DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]); - let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf)); - - for udf in [ - NativeHistogramAbsentOverTime::scalar_udf(), - NativeHistogramAdd::scalar_udf(), - NativeHistogramAvg::scalar_udf(), - NativeHistogramAvgOverTime::scalar_udf(), - NativeHistogramChanges::scalar_udf(), - NativeHistogramCount::scalar_udf(), - NativeHistogramCountOverTime::scalar_udf(), - NativeHistogramDelta::scalar_udf(), - NativeHistogramDivScalar::scalar_udf(), - NativeHistogramDrop::bool_false_udf(String::new(), None), - NativeHistogramDrop::bool_true_udf(String::new(), None), - NativeHistogramDrop::float_null_udf(String::new(), None), - NativeHistogramEq::scalar_udf(), - NativeHistogramFraction::scalar_udf(), - NativeHistogramIDelta::scalar_udf(), - NativeHistogramIRate::scalar_udf(), - NativeHistogramIncrease::scalar_udf(), - NativeHistogramLastOverTime::scalar_udf(), - MixedRange::float_udf(None), - MixedRange::histogram_udf(None), - NativeHistogramMulScalar::scalar_udf(), - NativeHistogramNeg::scalar_udf(), - NativeHistogramNotEq::scalar_udf(), - NativeHistogramPresentOverTime::scalar_udf(), - NativeHistogramQuantile::scalar_udf(), - NativeHistogramRate::scalar_udf(), - NativeHistogramResets::scalar_udf(), - NativeHistogramScalarMul::scalar_udf(), - NativeHistogramStddev::scalar_udf(), - NativeHistogramStdvar::scalar_udf(), - NativeHistogramSub::scalar_udf(), - NativeHistogramSum::scalar_udf(), - NativeHistogramSumOverTime::scalar_udf(), - NativeHistogramToString::scalar_udf(), - PromqlFloatToString::scalar_udf(), - ] { - let _ = session_state.register_udf(Arc::new(udf)); - } - for udaf in [ - NativeHistogramAggAvg::aggregate_udf(), - NativeHistogramAggSum::aggregate_udf(), - ] { - let _ = session_state.register_udaf(Arc::new(udaf)); - } + register_greptime_functions(&mut session_state, &self.query_ctx)?; // Install a registry that is also able to decode `MergeScan` nodes (the plans sent to a - // datanode may contain one) after all the functions above are registered, so that the - // sub-plans of a `MergeScan` see the same functions. + // datanode may contain one), so that the sub-plans of a `MergeScan` see the same functions. // // The registry resolves the tables of a `MergeScan` payload with the catalog of the query // engine, which it retrieves from the engine state the query engine leaves in the session @@ -450,12 +480,17 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { .config() .get_extension::() .map(|engine_state| engine_state.catalog_manager().clone()); - let session_state = SessionStateBuilder::new_from_existing(session_state.clone()) + let mut session_state = SessionStateBuilder::new_from_existing(session_state.clone()) .with_serializer_registry(Arc::new(MergeScanAwareSerializer { session_state, catalog_manager, })) .build(); + // Building the state above re-registered the functions of the previous state in hash order, + // which can bind a function back to a built-in DataFusion function of the same name or + // alias: register the GreptimeDB functions once more, so that the state that decodes the + // plan below holds the GreptimeDB bindings. + register_greptime_functions(&mut session_state, &self.query_ctx)?; let logical_plan = DFLogicalSubstraitConvertor .decode(message, session_state) @@ -478,11 +513,13 @@ mod tests { use catalog::RegisterTableRequest; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; use common_query::native_histogram::native_histogram_value_type; + use common_time::Timezone; use datafusion::catalog::TableProvider; use datafusion::datasource::MemTable; + use datafusion::functions::datetime::to_char::ToCharFunc; use datafusion::logical_expr::Extension; - use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, col, lit}; + use datafusion_expr::expr::{Cast, ScalarFunction}; + use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, ScalarUDF, col, lit}; use datatypes::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, }; @@ -951,4 +988,234 @@ mod tests { // Both levels are `MergeScan` nodes. assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); } + + /// Returns the `date_format` implementation of GreptimeDB: the implementation a decoded plan + /// has to use. + fn greptime_date_format_udf(query_ctx: &QueryContextRef) -> Arc { + FUNCTION_REGISTRY + .get_function("date_format") + .expect("`date_format` is a function of GreptimeDB") + .provide(FunctionContext { + query_ctx: query_ctx.clone(), + state: Default::default(), + }) + .into() + } + + /// The query context of the tests below: the timezone is not the default one, so that the + /// implementation bound to `date_format` is observable (only the GreptimeDB implementation + /// formats in the timezone of the query context). + fn date_format_query_ctx() -> QueryContextRef { + let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()); + Arc::new(query_ctx) + } + + /// A `date_format(timestamp, format)` expression that uses the GreptimeDB implementation. + fn greptime_date_format_expr(query_ctx: &QueryContextRef, timestamp: Expr) -> Expr { + Expr::ScalarFunction(ScalarFunction { + func: greptime_date_format_udf(query_ctx), + args: vec![timestamp, lit("%Y-%m-%d %H:%i:%S")], + }) + } + + /// Returns the scalar functions of every expression of `plan`, including the expressions of its + /// inputs and the ones of the plan embedded in a `MergeScan` (the payload of a `MergeScan` is + /// not an input of the node, see `MergeScanLogicalPlan::inputs`). + fn scalar_functions_of(plan: &LogicalPlan) -> Vec> { + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + + let mut udfs = Vec::new(); + for expr in plan.expressions() { + expr.apply(|expr| { + if let Expr::ScalarFunction(func) = expr { + udfs.push(Arc::clone(&func.func)); + } + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + } + + if let LogicalPlan::Extension(extension) = plan + && let Some(merge_scan) = extension + .node + .as_any() + .downcast_ref::() + { + udfs.extend(scalar_functions_of(merge_scan.input())); + } + + for input in plan.inputs() { + udfs.extend(scalar_functions_of(input)); + } + + udfs + } + + /// Asserts that the only function of `plan` is the `date_format` of GreptimeDB. + /// + /// The built-in `to_char` of DataFusion declares `date_format` as an alias, so a plan where + /// `date_format` was bound to it holds a `to_char` function, which the substrait encoder + /// renders as `to_char(...) AS date_format(...)`: the binding itself is asserted here, and not + /// the way the expression is printed. + fn assert_only_function_is_greptime_date_format( + plan: &LogicalPlan, + query_ctx: &QueryContextRef, + ) { + let udfs = scalar_functions_of(plan); + assert_eq!(udfs.len(), 1, "expected one function in: {plan}"); + + let udf = &udfs[0]; + assert_eq!(udf.name(), "date_format", "got: {plan}"); + assert!( + udf.inner().as_any().downcast_ref::().is_none(), + "`date_format` is bound to the built-in `to_char` of DataFusion: {plan}" + ); + assert_eq!( + udf.as_ref(), + greptime_date_format_udf(query_ctx).as_ref(), + "`date_format` is not bound to the implementation of GreptimeDB: {plan}" + ); + } + + /// The built-in `to_char` of DataFusion declares `date_format` as an alias, and + /// [`SessionStateBuilder`] re-registers the functions of a state in the (unspecified) order in + /// which a hash map yields them, so building the session state can bind `date_format` back to + /// the built-in `to_char` instead of the GreptimeDB implementation. Decoding a plan has to keep + /// the GreptimeDB binding, on the plain path (no `MergeScan` to decode) as well. + #[tokio::test] + async fn test_serializer_decode_keeps_greptime_date_format() { + // The functions of a session state are collected from a hash map, and the map order differs + // between the states (and between the decodings, which rebuild the state): decode the same + // plan several times, with a fresh engine each time, so that an order dependent binding is + // caught whatever the order of a single state is. + for _ in 0..8 { + let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); + let factory = QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + + let table_provider = Arc::new(mock_table_provider(1.into())); + let query_ctx = date_format_query_ctx(); + let plan = LogicalPlanBuilder::scan( + "devices", + Arc::new(LogicalTableSource::new(table_provider.schema().clone())), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))]) + .unwrap() + .build() + .unwrap(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .new_plan_decoder() + .unwrap(); + let decoded = plan_decoder + .decode( + bytes, + Arc::new(DummyCatalogList::with_table_provider(table_provider)), + false, + ) + .await + .unwrap(); + + assert_only_function_is_greptime_date_format(&decoded, &query_ctx); + } + } + + /// Same as [`test_serializer_decode_keeps_greptime_date_format`], for a `date_format` inside the + /// payload of a `MergeScan`: the payload is decoded with the session state built by + /// [`MergeScanAwareSerializer::payload_state`], which is rebuilt as well and thus has to keep + /// the GreptimeDB functions too. + #[tokio::test] + async fn test_serializer_decode_merge_scan_payload_keeps_greptime_date_format() { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + // The payload scans `numbers`, which is resolved through the catalog of the query engine. + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + + for _ in 0..8 { + let factory = QueryEngineFactory::new( + catalog_manager.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + let query_ctx = date_format_query_ctx(); + + // `numbers` has a single `number` column of type `u32`: cast it to a timestamp, which + // is what `date_format` formats. + let input = LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr( + &query_ctx, + Expr::Cast(Cast::new( + Box::new(col("number")), + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + )), + )]) + .unwrap() + .build() + .unwrap(); + let plan = + MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .new_plan_decoder() + .unwrap(); + // The catalog list of the request is deliberately unrelated to the payload. + let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))); + + let decoded = plan_decoder + .decode(bytes, catalog_list, false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); + + assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx); + } + } } From ae4edafeca0645c648598d47bd32bd418aefd288 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:56:58 +0800 Subject: [PATCH 03/37] 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(, 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> --- src/cache/src/lib.rs | 247 ++++-- src/catalog/src/information_extension.rs | 1 + src/client/src/region.rs | 163 +++- src/cmd/src/datanode/builder.rs | 11 +- src/common/query/src/request.rs | 58 ++ src/datanode/src/datanode.rs | 19 +- src/datanode/src/region_query.rs | 11 +- src/datanode/src/region_server.rs | 161 +++- src/query/src/dist_plan/merge_scan.rs | 1 + src/session/src/hints.rs | 30 +- src/standalone/src/information_extension.rs | 1 + tests-integration/src/cluster.rs | 216 ++++- tests-integration/src/standalone.rs | 11 +- tests-integration/src/tests.rs | 1 + .../nested_merge_scan_capability_test.rs | 818 ++++++++++++++++++ 15 files changed, 1620 insertions(+), 129 deletions(-) create mode 100644 tests-integration/src/tests/nested_merge_scan_capability_test.rs diff --git a/src/cache/src/lib.rs b/src/cache/src/lib.rs index bbcf68b30ff6..4fd9ae5fcc13 100644 --- a/src/cache/src/lib.rs +++ b/src/cache/src/lib.rs @@ -18,15 +18,16 @@ use std::hash::Hash; use std::sync::Arc; use std::time::Duration; -use catalog::kvbackend::new_table_cache; +use catalog::kvbackend::{TableCacheRef, 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, + SchemaCacheRef, TableInfoCacheRef, TableNameCacheRef, TableRouteCacheRef, TableSchemaCacheRef, + 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; use moka::future::{Cache, CacheBuilder}; -use partition::cache::new_partition_info_cache; +use partition::cache::{PartitionInfoCacheRef, new_partition_info_cache}; use snafu::OptionExt; use crate::error::Result; @@ -53,7 +54,117 @@ 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: +/// The caches of a datanode. +/// +/// The `table_cache` is derived from the `table_info_cache` and the `table_name_cache`, and the +/// `partition_info_cache` is derived from the `table_route_cache`. +struct DatanodeCaches { + table_id_schema_cache: TableSchemaCacheRef, + schema_cache: SchemaCacheRef, + table_info_cache: TableInfoCacheRef, + table_name_cache: TableNameCacheRef, + table_route_cache: TableRouteCacheRef, + table_cache: TableCacheRef, + partition_info_cache: PartitionInfoCacheRef, +} + +impl DatanodeCaches { + /// Builds the caches of a datanode. + /// + /// The table/route caches are needed by the catalog manager and by the partition rule manager + /// of the datanode, which are used to plan the `MergeScan` nodes of the plans received from the + /// frontend (e.g. a nested merge scan that queries another datanode). + fn new(kv_backend: KvBackendRef) -> Self { + // Builds table id schema name cache that never expires. + 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(), + cache, + kv_backend.clone(), + )); + + // Builds schema cache + let cache = default_cache(); + let schema_cache = Arc::new(new_schema_cache( + SCHEMA_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + // Builds table info cache + let cache = default_cache(); + let table_info_cache = Arc::new(new_table_info_cache( + TABLE_INFO_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + // Builds table name cache + let cache = default_cache(); + let table_name_cache = Arc::new(new_table_name_cache( + TABLE_NAME_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + // Builds table cache + 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(), + )); + + // Builds table route cache + let cache = default_cache(); + let table_route_cache = Arc::new(new_table_route_cache( + TABLE_ROUTE_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + // Builds partition info cache + 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(), + )); + + Self { + table_id_schema_cache, + schema_cache, + table_info_cache, + table_name_cache, + table_route_cache, + table_cache, + partition_info_cache, + } + } + + /// Splits the caches into the layer of the caches that are read from the kv backend, and the + /// layer of the caches derived from them. + fn into_layers(self) -> (CacheRegistry, CacheRegistry) { + let base_registry = CacheRegistryBuilder::default() + .add_cache(self.table_id_schema_cache) + .add_cache(self.schema_cache) + .add_cache(self.table_info_cache) + .add_cache(self.table_name_cache) + .add_cache(self.table_route_cache) + .build(); + let derived_registry = CacheRegistryBuilder::default() + .add_cache(self.table_cache) + .add_cache(self.partition_info_cache) + .build(); + + (base_registry, derived_registry) + } +} + +/// Builds the layered cache registry for datanode. +/// +/// The registry holds every cache a datanode needs: /// - Schema cache. /// - Table id to schema name cache. /// - Table info cache. @@ -62,75 +173,24 @@ pub const PARTITION_INFO_CACHE_NAME: &str = "partition_info_cache"; /// - Table route cache. /// - Partition info cache. /// -/// The table/route caches are needed by the catalog manager and by the partition rule manager of -/// the datanode, which are used to plan the `MergeScan` nodes of the plans received from the -/// frontend (e.g. a nested merge scan that queries another datanode). -pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry { - // Builds table id schema name cache that never expires. - 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(), - cache, - kv_backend.clone(), - )); - - // Builds schema cache - let cache = default_cache(); - let schema_cache = Arc::new(new_schema_cache( - SCHEMA_CACHE_NAME.to_string(), - cache, - kv_backend.clone(), - )); - - // Builds table info cache - let cache = default_cache(); - let table_info_cache = Arc::new(new_table_info_cache( - TABLE_INFO_CACHE_NAME.to_string(), - cache, - kv_backend.clone(), - )); - - // Builds table name cache - let cache = default_cache(); - let table_name_cache = Arc::new(new_table_name_cache( - TABLE_NAME_CACHE_NAME.to_string(), - cache, - kv_backend.clone(), - )); - - // Builds table cache - 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(), - )); - - // Builds table route cache - let cache = default_cache(); - let table_route_cache = Arc::new(new_table_route_cache( - TABLE_ROUTE_CACHE_NAME.to_string(), - cache, - kv_backend.clone(), - )); - - // Builds partition info cache - 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(), - )); +/// The first layer holds the caches the datanode reads from the kv backend: the schema cache, the +/// table id to schema name cache, the table info cache, the table name cache and the table route +/// cache. The second layer holds the caches derived from them: the table cache (from the table info +/// and the table name caches) and the partition info cache (from the table route cache). +/// +/// [LayeredCacheRegistry] invalidates a layer only after the previous layer finished invalidating +/// (see [LayeredCacheRegistryBuilder::add_cache_registry]), so the derived caches are invalidated +/// after the caches they are derived from. A flat [CacheRegistry] invalidates its caches +/// concurrently: a derived cache could be invalidated first and then refilled by a concurrent query +/// from a base cache that still holds the old value, and the stale metadata would stay in the +/// derived cache until the entry expires. This is the only datanode cache registry: production code +/// and test infrastructure build the same one. +pub fn build_datanode_layered_cache_registry(kv_backend: KvBackendRef) -> LayeredCacheRegistry { + let (base_registry, derived_registry) = DatanodeCaches::new(kv_backend).into_layers(); - 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_cache) - .add_cache(table_route_cache) - .add_cache(partition_info_cache) + LayeredCacheRegistryBuilder::default() + .add_cache_registry(base_registry) + .add_cache_registry(derived_registry) .build() } @@ -242,3 +302,44 @@ pub fn with_default_composite_cache_registry( Ok(builder.add_cache_registry(registry)) } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use catalog::kvbackend::TableCacheRef; + use common_meta::cache::{ + SchemaCacheRef, TableInfoCacheRef, TableNameCacheRef, TableRouteCacheRef, + TableSchemaCacheRef, + }; + use common_meta::kv_backend::KvBackendRef; + use common_meta::kv_backend::memory::MemoryKvBackend; + use partition::cache::PartitionInfoCacheRef; + + use super::*; + + fn memory_backend() -> KvBackendRef { + Arc::new(MemoryKvBackend::::new()) + } + + /// Asserts that the registry holds every cache the catalog manager and the partition rule + /// manager of a datanode look up. + macro_rules! assert_datanode_caches { + ($registry:expr) => { + assert!($registry.get::().is_some()); + assert!($registry.get::().is_some()); + assert!($registry.get::().is_some()); + assert!($registry.get::().is_some()); + assert!($registry.get::().is_some()); + assert!($registry.get::().is_some()); + assert!($registry.get::().is_some()); + }; + } + + #[test] + fn test_datanode_layered_cache_registry_holds_all_caches() { + let registry = build_datanode_layered_cache_registry(memory_backend()); + + assert_datanode_caches!(registry); + } +} diff --git a/src/catalog/src/information_extension.rs b/src/catalog/src/information_extension.rs index 9d55f9f7a1ee..b9b0baa0a585 100644 --- a/src/catalog/src/information_extension.rs +++ b/src/catalog/src/information_extension.rs @@ -82,6 +82,7 @@ impl DistributedInformationExtension { plan, region_id: RegionId::default(), header: None, + internal: false, }) .await .context(crate::error::HandleQuerySnafu) diff --git a/src/client/src/region.rs b/src/client/src/region.rs index fb2e4d10ca16..e357e18a4209 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -30,7 +30,9 @@ use common_error::status_code::StatusCode; use common_grpc::flight::{FlightDecoder, FlightMessage}; use common_meta::error::{self as meta_error, Result as MetaResult}; use common_meta::node_manager::Datanode; -use common_query::request::QueryRequest; +use common_query::request::{ + QUERY_INTERNAL_STAGE_EXTENSION_KEY, QueryRequest, mark_query_internal_stage, +}; use common_recordbatch::error::ExternalSnafu; use common_recordbatch::{RecordBatch, RecordBatchStreamWrapper, SendableRecordBatchStream}; use common_telemetry::error; @@ -85,7 +87,7 @@ impl Datanode for RegionRequester { .context(meta_error::ExternalSnafu)? .to_vec(); let request = api::v1::region::QueryRequest { - header: request.header, + header: encode_region_request_header(request.header, request.internal), region_id: request.region_id.as_u64(), plan, }; @@ -100,6 +102,50 @@ impl Datanode for RegionRequester { } } +/// Encodes the header of a region query, carrying [`QueryRequest::internal`] to the peer datanode. +/// +/// The wire format of a region query (`api::v1::region::QueryRequest`) has no field for the marker, +/// so an internal request (an execution stage of the plan the sender received) is marked with the +/// [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] extension of the query context of the header, which the +/// peer reads back in `RegionServer::handle_remote_read_inner`. Without it the peer can't tell an +/// inner stage of another query from an independent query and takes a concurrency permit for it, +/// which makes the outer query wait for the inner stage while the inner stage waits for the permit +/// held by the outer query. +/// +/// A header-less request keeps no header unless it is internal. An internal request that has no +/// header is encoded with a header that carries the marker only, so the peer builds the query +/// context of the request from that header, i.e. without a remote query id (the peer generates a +/// fresh query context, with an id, only for a request that has no header at all). That only +/// disables the remote dynamic filter cleanup of the request. +fn encode_region_request_header( + header: Option, + internal: bool, +) -> Option { + let mut header = match header { + Some(header) => header, + // A request that has no header is an independent query: it has no marker to carry. + None if !internal => return None, + None => RegionRequestHeader::default(), + }; + if internal { + // Creates the query context if the caller didn't set one: the peer honors the marker only + // in the query context of the header. + mark_query_internal_stage( + &mut header + .query_context + .get_or_insert_with(api::v1::QueryContext::default) + .extensions, + ); + } else if let Some(query_context) = header.query_context.as_mut() { + // [`QueryRequest::internal`] owns the marker: a header that carries it while the request + // claims to be an independent query must not make the peer skip its concurrency limiter. + query_context + .extensions + .remove(QUERY_INTERNAL_STAGE_EXTENSION_KEY); + } + Some(header) +} + impl RegionRequester { pub fn new(client: Client, send_compression: bool, accept_compression: bool) -> Self { Self { @@ -838,3 +884,116 @@ mod test { assert!(recordbatches.next().await.is_none()); } } + +#[cfg(test)] +mod internal_stage_test { + use std::collections::HashMap; + + use common_query::request::{QUERY_INTERNAL_STAGE_EXTENSION_VALUE, is_internal_stage}; + + use super::*; + + /// Reads the marker of a header the way the receiver does: `RegionServer::do_get` builds the + /// query context of the request from the header, `session::context::QueryContext` copies the + /// extensions of that context verbatim. + fn extensions_of(header: &RegionRequestHeader) -> HashMap { + header + .query_context + .as_ref() + .map(|ctx| ctx.extensions.clone()) + .unwrap_or_default() + } + + #[test] + fn test_encode_region_request_header_marks_internal_request() { + // An internal request carries the marker even when its caller set no header at all. + let header = encode_region_request_header(None, true).unwrap(); + + assert!(is_internal_stage(&extensions_of(&header))); + } + + #[test] + fn test_encode_region_request_header_keeps_context_of_internal_request() { + let header = encode_region_request_header( + Some(RegionRequestHeader { + dbname: "greptime".to_string(), + query_context: Some(api::v1::QueryContext { + current_schema: "public".to_string(), + ..Default::default() + }), + ..Default::default() + }), + true, + ) + .unwrap(); + + assert_eq!(header.dbname, "greptime"); + assert_eq!( + header + .query_context + .as_ref() + .map(|ctx| ctx.current_schema.as_str()), + Some("public") + ); + assert!(is_internal_stage(&extensions_of(&header))); + } + + #[test] + fn test_encode_region_request_header_does_not_mark_independent_request() { + // An independent request with no header keeps no header at all. + assert!(encode_region_request_header(None, false).is_none()); + + // A marker that a caller left in the header of an independent request is dropped: the + // marker is owned by `QueryRequest::internal`. + let header = encode_region_request_header( + Some(RegionRequestHeader { + query_context: Some(api::v1::QueryContext { + extensions: HashMap::from([( + QUERY_INTERNAL_STAGE_EXTENSION_KEY.to_string(), + QUERY_INTERNAL_STAGE_EXTENSION_VALUE.to_string(), + )]), + ..Default::default() + }), + ..Default::default() + }), + false, + ) + .unwrap(); + + assert!(!is_internal_stage(&extensions_of(&header))); + } + + #[test] + fn test_internal_marker_survives_the_wire_format() { + let request = api::v1::region::QueryRequest { + header: encode_region_request_header( + Some(RegionRequestHeader { + query_context: Some(api::v1::QueryContext { + current_schema: "public".to_string(), + extensions: HashMap::from([( + "remote_query_id".to_string(), + "1".to_string(), + )]), + ..Default::default() + }), + ..Default::default() + }), + true, + ), + region_id: 1, + plan: Vec::new(), + }; + let bytes = request.encode_to_vec(); + + // The receiver decodes the ticket and builds the query context of the request from the + // header of the decoded request. + let decoded = api::v1::region::QueryRequest::decode(bytes.as_slice()).unwrap(); + let extensions = extensions_of(decoded.header.as_ref().unwrap()); + + assert!(is_internal_stage(&extensions)); + assert_eq!( + extensions.get("remote_query_id").map(String::as_str), + Some("1") + ); + } +} diff --git a/src/cmd/src/datanode/builder.rs b/src/cmd/src/datanode/builder.rs index 810dfeedc7be..6df7ac728e8a 100644 --- a/src/cmd/src/datanode/builder.rs +++ b/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()); diff --git a/src/common/query/src/request.rs b/src/common/query/src/request.rs index 86c6758aa255..223500b21d8f 100644 --- a/src/common/query/src/request.rs +++ b/src/common/query/src/request.rs @@ -16,6 +16,7 @@ mod base64_serde; mod initial_remote_dyn_filter_reg; use std::cell::Cell; +use std::collections::HashMap; use std::sync::Arc; use api::v1::region::RegionRequestHeader; @@ -339,6 +340,24 @@ impl DynFilterUpdate { } } +/// The extension key of the marker that a region query is an execution stage of another query, +/// instead of an independent query. +/// +/// The wire format of a region query (`api::v1::region::QueryRequest`) has no field for +/// [`QueryRequest::internal`], so the marker travels to the peer datanode as an extension of the +/// query context of the [`RegionRequestHeader`] of the request (the only map-shaped part of the +/// header that survives the round trip). +/// +/// The key is reserved (see `session::hints::RESERVED_EXTENSION_KEYS`), so the hint path +/// (`x-greptime-hints`) can't set it. The marker is trusted on the region query port of a datanode, +/// like the other extensions of the header and the plan itself. +pub const QUERY_INTERNAL_STAGE_EXTENSION_KEY: &str = "query.internal_stage"; + +/// The value of [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] that marks a region query as an execution +/// stage of another query. Any other value, including an absent one, leaves the query subject to +/// the concurrency limiter. +pub const QUERY_INTERNAL_STAGE_EXTENSION_VALUE: &str = "true"; + /// The query request to be handled by the RegionServer (Datanode). #[derive(Clone, Debug)] pub struct QueryRequest { @@ -350,6 +369,45 @@ pub struct QueryRequest { /// The form of the query: a logical plan. pub plan: LogicalPlan, + + /// Whether this request is an execution stage of another query, instead of an independent + /// query. + /// + /// A datanode executes the `MergeScan` nodes of the plan it received by querying the regions + /// of those nodes from the datanodes hosting them. Such an inner region query is a stage of + /// the outer query: the outer query waits for it while holding its own concurrency permit, so + /// the inner query must not acquire a permit of its own, otherwise the outer query waits for + /// the inner stage while the inner stage waits for the permit held by the outer query. + /// + /// `false` means an ordinary, independent query, which is subject to the concurrency limiter. + /// + /// A request that crosses the wire carries the marker as the + /// [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] extension of the query context of its header (the + /// sender sets it, the receiver reads it back into this field), so the marker also survives a + /// dispatch to a peer datanode. + pub internal: bool, +} + +/// Marks the query context of a region query as an execution stage of another query: sets the +/// [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] extension to [`QUERY_INTERNAL_STAGE_EXTENSION_VALUE`]. +/// +/// The sender side of a region query uses it to carry [`QueryRequest::internal`] to a peer +/// datanode; [`is_internal_stage`] reads the marker back on the receiver side. +pub fn mark_query_internal_stage(extensions: &mut HashMap) { + extensions.insert( + QUERY_INTERNAL_STAGE_EXTENSION_KEY.to_string(), + QUERY_INTERNAL_STAGE_EXTENSION_VALUE.to_string(), + ); +} + +/// Returns whether `extensions` marks a region query as an execution stage of another query, i.e. +/// whether the sender of the query set [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] to +/// [`QUERY_INTERNAL_STAGE_EXTENSION_VALUE`]. +pub fn is_internal_stage(extensions: &HashMap) -> bool { + extensions + .get(QUERY_INTERNAL_STAGE_EXTENSION_KEY) + .map(String::as_str) + == Some(QUERY_INTERNAL_STAGE_EXTENSION_VALUE) } #[cfg(test)] diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 2c3dc93c8f21..4fef618b90dd 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -1112,13 +1112,12 @@ 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; @@ -1170,11 +1169,10 @@ 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(), - ); + // Use the same layered registry as the datanode of a running instance, so the derived + // caches of the test are invalidated after the caches they are derived from. + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut builder = DatanodeBuilder::new( DatanodeOptions { @@ -1216,11 +1214,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, diff --git a/src/datanode/src/region_query.rs b/src/datanode/src/region_query.rs index 84cd18f8fc61..6b0fae6678d2 100644 --- a/src/datanode/src/region_query.rs +++ b/src/datanode/src/region_query.rs @@ -74,8 +74,17 @@ impl RegionQueryHandler for DatanodeRegionQueryHandler { async fn do_get( &self, target: &RegionQueryTarget, - request: QueryRequest, + mut request: QueryRequest, ) -> QueryResult { + // A datanode serves the region queries of the `MergeScan` nodes of the plan it received. + // Such a region query is an execution stage of that plan, not an independent query, so it + // must leave the concurrency permits to the query that dispatched the plan. The marker is + // carried to the peer with the request: the client encodes it into the query context of the + // header (`common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_KEY`), where the peer reads + // it back and admits the stage without a permit of its own (see + // `RegionServer::handle_remote_read_inner`). + request.internal = true; + self.node_manager .datanode(target.peer()) .await diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 7d5218f30b75..d8ef5e123189 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -74,6 +74,7 @@ use servers::grpc::region_server::RegionServerHandler; use session::context::{ FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef, }; +use session::hints::{QUERY_INTERNAL_STAGE_EXTENSION_KEY, QUERY_INTERNAL_STAGE_EXTENSION_VALUE}; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::{ FILE_ENGINE_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, @@ -290,11 +291,13 @@ impl RegionServer { request: api::v1::region::QueryRequest, query_ctx: QueryContextRef, ) -> Result { - let permit = if let Some(p) = &self.inner.parallelism { - Some(p.acquire().await?) - } else { - None - }; + // A peer datanode that executes the `MergeScan` nodes of the plan it received queries the + // regions of those nodes from this datanode and marks the region query as an execution + // stage of its query (see [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`]). Such a stage must not + // take a permit of its own: the outer query holds one until its stream is fully consumed + // and waits for the stage to return. + let internal = is_internal_stage_query(&query_ctx); + let permit = acquire_query_permit(&self.inner.parallelism, internal).await?; let region_id = RegionId::from_u64(request.region_id); let catalog_list = Arc::new(NameAwareCatalogList::new( @@ -328,6 +331,10 @@ impl RegionServer { header: request.header, region_id, plan, + // The stage marker is passed on: the plan of this read is executed as part of + // the stage this remote read was admitted for, so nothing on that path may + // take a permit of its own either. + internal, }, query_ctx.clone(), ) @@ -348,11 +355,10 @@ impl RegionServer { } async fn handle_read_inner(&self, request: QueryRequest) -> Result { - let permit = if let Some(p) = &self.inner.parallelism { - Some(p.acquire().await?) - } else { - None - }; + // The datanode executes the `MergeScan` nodes of the plan it received by querying regions + // from the datanodes hosting them. Those inner region queries are marked as internal + // stages of the query that dispatched the plan and must not take a permit of their own. + let permit = acquire_query_permit(&self.inner.parallelism, request.internal).await?; let ctx = request.header.as_ref().map(|h| h.into()); let query_ctx = Arc::new(ctx.unwrap_or_else(|| QueryContextBuilder::default().build())); @@ -1086,6 +1092,39 @@ impl RegionServerParallelism { } } +/// Returns whether the query context of a region query received over the wire marks the query as +/// an execution stage of another query, instead of an independent query. +/// +/// The sender of the query sets the [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] extension of the query +/// context to [`QUERY_INTERNAL_STAGE_EXTENSION_VALUE`] when the query is an execution stage of the +/// plan it received (see [`QueryRequest::internal`]), i.e. when the query is dispatched by a +/// datanode for the `MergeScan` nodes of that plan. +fn is_internal_stage_query(query_ctx: &QueryContextRef) -> bool { + query_ctx.extension(QUERY_INTERNAL_STAGE_EXTENSION_KEY) + == Some(QUERY_INTERNAL_STAGE_EXTENSION_VALUE) +} + +/// Acquires a concurrency permit for a request that is an independent query, if the limiter is +/// enabled. +/// +/// A request that is an execution stage of another query (see [`QueryRequest::internal`]) doesn't +/// acquire a permit: the query that dispatched it already holds one, holds it until its stream is +/// fully consumed and waits for the stage to return. Letting the stage acquire a permit of its own +/// makes the outer query wait for the inner stage while the inner stage waits for the permit held +/// by the outer query. +async fn acquire_query_permit( + parallelism: &Option, + internal: bool, +) -> Result> { + if internal { + return Ok(None); + } + match parallelism { + Some(p) => Ok(Some(p.acquire().await?)), + None => Ok(None), + } +} + /// Wraps a record batch stream and holds a concurrency permit until the stream is /// fully consumed (dropped), so `max_concurrent_queries` bounds the number of /// in-flight read streams, not just query planning. @@ -2786,6 +2825,108 @@ mod tests { assert!(forth_query.is_ok()); } + #[tokio::test] + async fn test_internal_query_skips_concurrency_permit() { + let parallelism = + Some(RegionServerParallelism::from_opts(1, Duration::from_millis(1)).unwrap()); + + // The query that dispatched the nested plan holds the only permit of the server for as + // long as its stream is alive. + let outer_query = acquire_query_permit(¶llelism, false).await.unwrap(); + assert!(outer_query.is_some()); + + // An inner stage of that query must neither wait for nor consume a permit: it runs while + // the outer query waits for it to return. + let inner_stage = acquire_query_permit(¶llelism, true).await.unwrap(); + assert!(inner_stage.is_none()); + + // An independent query is still subject to the limiter. + assert!(acquire_query_permit(¶llelism, false).await.is_err()); + + // The inner stage reserved nothing: the permit goes back to the limiter with the outer + // query's stream. + drop(outer_query); + assert!( + acquire_query_permit(¶llelism, false) + .await + .unwrap() + .is_some() + ); + + // The limiter stays disabled when `max_concurrent_queries` is 0. + assert!(acquire_query_permit(&None, false).await.unwrap().is_none()); + } + + #[test] + fn test_internal_stage_extension_keys_match_between_crates() { + // `session` can't depend on `common-query`: the sender writes the marker with the constant + // of `common-query`, the receiver reads it with the constant of `session`. + assert_eq!( + QUERY_INTERNAL_STAGE_EXTENSION_KEY, + common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_KEY + ); + assert_eq!( + QUERY_INTERNAL_STAGE_EXTENSION_VALUE, + common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_VALUE + ); + } + + /// Asserts that a region query received over the wire is admitted without a permit only when the + /// query context of its header marks it as an execution stage of another query. + #[tokio::test] + async fn test_internal_stage_query_skips_concurrency_permit() { + let parallelism = + Some(RegionServerParallelism::from_opts(1, Duration::from_millis(1)).unwrap()); + + // A datanode that fans a nested `MergeScan` out to its peer marks the region query it sends + // as an execution stage of its own query. + let mut extensions = std::collections::HashMap::new(); + common_query::request::mark_query_internal_stage(&mut extensions); + let internal_ctx: QueryContextRef = Arc::new( + QueryContextBuilder::default() + .extensions(extensions) + .build(), + ); + assert!(is_internal_stage_query(&internal_ctx)); + + // The query of the peer itself holds the only permit of the server. + let outer_query = acquire_query_permit(¶llelism, false).await.unwrap(); + assert!(outer_query.is_some()); + + // The stage is admitted while the outer query waits for it to return. + let stage_permit = + acquire_query_permit(¶llelism, is_internal_stage_query(&internal_ctx)) + .await + .unwrap(); + assert!(stage_permit.is_none()); + + // A query with no marker, a marker with another value or an empty query context is an + // independent query and is subject to the limiter. + for query_ctx in [ + QueryContext::arc(), + Arc::new( + QueryContextBuilder::default() + .set_extension( + QUERY_INTERNAL_STAGE_EXTENSION_KEY.to_string(), + "false".to_string(), + ) + .build(), + ), + Arc::new( + QueryContextBuilder::default() + .set_extension("query.internal_stage.other".to_string(), "true".to_string()) + .build(), + ), + ] { + assert!(!is_internal_stage_query(&query_ctx)); + assert!( + acquire_query_permit(¶llelism, is_internal_stage_query(&query_ctx)) + .await + .is_err() + ); + } + } + fn mock_region_metadata(region_id: RegionId) -> RegionMetadata { let mut metadata_builder = RegionMetadataBuilder::new(region_id); metadata_builder.push_column_metadata(ColumnMetadata { diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index ee4e9574280c..7993e2b9e1b5 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -728,6 +728,7 @@ impl MergeScanExec { }), region_id, plan: plan.clone(), + internal: false, }; if explain_verbose { common_telemetry::info!( diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index 971426f4cc03..b97d1cf94cda 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -27,11 +27,36 @@ 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] = [ +/// The extension key of the marker that a query is an execution stage of another query, instead of +/// an independent query. +/// +/// A datanode executes the `MergeScan` nodes of the plan it received by querying the regions of +/// those nodes from the datanodes hosting them. Such an inner region query is a stage of the outer +/// query: the outer query waits for it while holding its own concurrency permit, so the inner query +/// must not acquire a permit of its own, otherwise the outer query waits for the inner stage while +/// the inner stage waits for the permit held by the outer query. The marker travels to the peer +/// datanode as an extension of the query context of the `RegionRequestHeader` of the request. +/// +/// The key is reserved (see [`RESERVED_EXTENSION_KEYS`]), so the hint path (`x-greptime-hints`) +/// can't set it: a SQL/PromQL client of the frontend can't make its own query bypass the +/// concurrency limiter of a datanode. The marker is trusted on the region query port of a datanode, +/// like the other extensions of the header and the plan itself. +/// +/// Keep in sync with `common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_KEY` +/// (`session` can't depend on `common-query`). +pub const QUERY_INTERNAL_STAGE_EXTENSION_KEY: &str = "query.internal_stage"; +/// The value of [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] that marks an execution stage of another +/// query. Any other value, including an absent one, leaves the query subject to the concurrency +/// limiter. +/// +/// Keep in sync with `common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_VALUE`. +pub const QUERY_INTERNAL_STAGE_EXTENSION_VALUE: &str = "true"; +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, + QUERY_INTERNAL_STAGE_EXTENSION_KEY, ]; /// Deprecated, use `HINTS_KEY` instead. @@ -65,6 +90,9 @@ mod tests { assert!(is_reserved_extension_key( LIVE_ANALYZE_METRICS_EXTENSION_KEY )); + assert!(is_reserved_extension_key( + QUERY_INTERNAL_STAGE_EXTENSION_KEY + )); assert!(!is_reserved_extension_key(READ_PREFERENCE_HINT)); } } diff --git a/src/standalone/src/information_extension.rs b/src/standalone/src/information_extension.rs index 8d52eb12ccc1..b465d1b1270d 100644 --- a/src/standalone/src/information_extension.rs +++ b/src/standalone/src/information_extension.rs @@ -169,6 +169,7 @@ impl InformationExtension for StandaloneInformationExtension { .context(catalog::error::DatafusionSnafu)?, region_id: RegionId::default(), header: None, + internal: false, }; self.region_server diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 4f762ba709b8..681a32a77590 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -16,14 +16,16 @@ use std::collections::{BTreeSet, HashMap}; use std::env; use std::net::TcpListener; use std::ops::RangeInclusive; -use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; use std::time::Duration; use api::v1::health_check_server::HealthCheckServer; use api::v1::region::region_server::RegionServer; use arrow_flight::flight_service_server::FlightServiceServer; use cache::{ - build_datanode_cache_registry, build_fundamental_cache_registry, + build_datanode_layered_cache_registry, build_fundamental_cache_registry, with_default_composite_cache_registry, }; use catalog::information_extension::DistributedInformationExtension; @@ -74,9 +76,11 @@ use servers::grpc::{GrpcOptions, HealthCheckHandler}; use servers::server::ServerHandlers; use store_api::storage::RegionId; use tempfile::TempDir; +use tokio::io::{AsyncRead, AsyncWrite}; use tonic::codec::CompressionEncoding; use tonic::transport::Server; -use tower::service_fn; +use tonic::transport::server::Connected; +use tower::{Layer, Service, service_fn}; use uuid::Uuid; use crate::test_util::{ @@ -84,11 +88,39 @@ use crate::test_util::{ create_tmp_dir_and_datanode_opts, }; +/// The gRPC requests another node sent to a datanode over its network address (see +/// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]). +/// +/// The frontend and the datanodes inside the test process talk to a datanode through an +/// in-process client, so a request counted here comes from a remote node, e.g. the region query +/// that a datanode issues for the regions of another datanode. +#[derive(Debug, Default)] +pub struct DatanodeRpcStats { + requests: AtomicUsize, + paths: Mutex>, +} + +impl DatanodeRpcStats { + /// The number of gRPC requests another node sent to the datanode. + pub fn requests(&self) -> usize { + self.requests.load(Ordering::Relaxed) + } + + /// The gRPC paths of those requests, e.g. + /// `/arrow.flight.protocol.FlightService/DoGet` for a region query. + pub fn paths(&self) -> BTreeSet { + self.paths.lock().unwrap().clone() + } +} + pub struct GreptimeDbCluster { pub guards: Vec, pub datanode_options: Vec, pub datanode_instances: HashMap, + /// The stats of the network address of every datanode, empty unless the cluster serves the + /// datanodes at their own addresses. + pub datanode_rpc_stats: HashMap>, pub kv_backend: KvBackendRef, pub metasrv: Arc, pub frontend: Arc, @@ -99,6 +131,13 @@ impl GreptimeDbCluster { &self.frontend.instance } + /// The stats of the network address of the datanode `datanode_id`, if the cluster serves the + /// datanodes at their own addresses (see + /// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]). + pub fn datanode_rpc_stats(&self, datanode_id: DatanodeId) -> Option<&Arc> { + self.datanode_rpc_stats.get(&datanode_id) + } + /// List all SST files from all datanodes. pub async fn list_sst_files_from_all_datanodes(&self) -> BTreeSet { let mut sst_files = BTreeSet::new(); @@ -176,6 +215,7 @@ pub struct GreptimeDbClusterBuilder { meta_selector: Option, local_file_access: LocalFileAccess, event_recorder_options: EventRecorderOptions, + real_datanode_grpc_addr: bool, } impl GreptimeDbClusterBuilder { @@ -212,6 +252,7 @@ impl GreptimeDbClusterBuilder { meta_selector: None, local_file_access: LocalFileAccess::default(), event_recorder_options: EventRecorderOptions::default(), + real_datanode_grpc_addr: false, } } @@ -290,6 +331,21 @@ impl GreptimeDbClusterBuilder { self } + /// Serves every datanode on its own real localhost TCP port and registers that + /// address in metasrv, instead of the default in-process duplex placeholder address + /// that all datanodes share. + /// + /// Datanodes serve the region queries of the plans they execute by querying the + /// leader of every region of the plan (see `DatanodeRegionQueryHandler`), so a test + /// that exercises a datanode-to-datanode region query needs each datanode to be + /// reachable at the address recorded in the table route, i.e. the address the + /// datanode registers with its heartbeat. + #[must_use] + pub fn with_real_datanode_grpc_addr(mut self, real_datanode_grpc_addr: bool) -> Self { + self.real_datanode_grpc_addr = real_datanode_grpc_addr; + self + } + pub async fn build_with( &self, datanode_options: Vec, @@ -335,7 +391,22 @@ impl GreptimeDbClusterBuilder { .build_datanodes_with_options(&metasrv, &datanode_options) .await; - build_datanode_clients(datanode_clients.clone(), &datanode_instances).await; + let datanode_rpc_stats = if self.real_datanode_grpc_addr { + datanode_options + .iter() + .map(|opts| (opts.node_id.unwrap(), Arc::new(DatanodeRpcStats::default()))) + .collect::>() + } else { + HashMap::new() + }; + + build_datanode_clients( + datanode_clients.clone(), + &datanode_instances, + &datanode_options, + &datanode_rpc_stats, + ) + .await; self.wait_datanodes_alive(metasrv.metasrv.meta_peer_client(), datanodes) .await; @@ -354,6 +425,7 @@ impl GreptimeDbClusterBuilder { datanode_options, guards, datanode_instances, + datanode_rpc_stats, kv_backend: self.kv_backend.clone(), metasrv: metasrv.metasrv, frontend: Arc::new(frontend), @@ -410,12 +482,29 @@ impl GreptimeDbClusterBuilder { opts }; opts.node_id = Some(datanode_id); + if self.real_datanode_grpc_addr { + // Give every datanode its own address, so that the other datanodes can + // reach it (see `with_real_datanode_grpc_addr`). + let addr = self.choose_datanode_grpc_addr(); + opts.grpc.bind_addr = addr.clone(); + opts.grpc.server_addr = addr; + } options.push(opts); } (options, guards) } + /// Chooses a free localhost port for a datanode. + /// + /// Uses a port range disjoint from the one of the frontend servers + /// (see [`Self::build_frontend_options`]). + fn choose_datanode_grpc_addr(&self) -> String { + let localhost = "127.0.0.1"; + let port = self.choose_random_unused_port(24000..=34000, 20, localhost); + format!("{}:{}", localhost, port) + } + async fn build_datanodes_with_options( &self, metasrv: &MockInfo, @@ -463,11 +552,10 @@ impl GreptimeDbClusterBuilder { let meta_backend = new_read_only_meta_kv_backend(meta_client.clone()); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(meta_backend.clone())) - .build(), - ); + // Use the same layered registry as the datanode of a running instance, so the derived + // caches of the test are invalidated after the caches they are derived from. + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(meta_backend.clone())); let mut builder = DatanodeBuilder::new(opts.clone(), Plugins::default(), meta_backend); builder @@ -619,18 +707,34 @@ impl GreptimeDbClusterBuilder { async fn build_datanode_clients( clients: Arc, instances: &HashMap, + options: &[DatanodeOptions], + rpc_stats: &HashMap>, ) { - for (&datanode_id, instance) in instances { - let (addr, client) = create_datanode_client(instance).await; + for opts in options { + let datanode_id = opts.node_id.unwrap(); + let instance = &instances[&datanode_id]; + let (addr, client) = create_datanode_client( + instance, + &opts.grpc.server_addr, + rpc_stats.get(&datanode_id).cloned(), + ) + .await; clients .insert_client(Peer::new(datanode_id, addr), client) .await; } } -async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { - let (client, server) = tokio::io::duplex(1024); - +/// Runs a gRPC server that serves the region server of `datanode` over `incoming` until the +/// cluster is dropped. +fn serve_datanode_region_server( + datanode: &Datanode, + incoming: I, + rpc_stats: Option>, +) where + I: futures::Stream> + Send + 'static, + IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static, +{ let runtime = RuntimeBuilder::default() .worker_threads(2) .thread_name("grpc-handlers") @@ -641,8 +745,9 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { let region_server_handler = RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime); - let _handle = tokio::spawn(async move { + tokio::spawn(async move { Server::builder() + .layer(CountRequestsLayer { rpc_stats }) .add_service( FlightServiceServer::new(flight_handler) .accept_compressed(CompressionEncoding::Gzip) @@ -658,12 +763,40 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { .send_compressed(CompressionEncoding::Zstd), ) .add_service(HealthCheckServer::new(HealthCheckHandler)) - .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)])) + .serve_with_incoming(incoming) .await }); +} + +/// Serves the region server of `datanode` and returns the address and the client of the address. +/// +/// The returned client is always an in-process client: the frontend and the datanodes of the test +/// process reach the region server through a duplex stream, no matter which address they use. This +/// keeps the in-process traffic out of `rpc_stats`. +/// +/// [`test_util::PEER_PLACEHOLDER_ADDR`] is not a real address, so it cannot be served over the +/// network at all. Any other address is a real one: the region server is additionally served on +/// its local TCP port (`rpc_stats` counts the requests of the remote nodes that reach it, e.g. the +/// region query that a datanode sends for the regions of another datanode). +async fn create_datanode_client( + datanode: &Datanode, + addr: &str, + rpc_stats: Option>, +) -> (String, Client) { + let (client, server) = tokio::io::duplex(1024); + serve_datanode_region_server( + datanode, + futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]), + None, + ); + + if addr != test_util::PEER_PLACEHOLDER_ADDR { + let listener = tokio::net::TcpListener::bind(addr).await.unwrap(); + let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener); + serve_datanode_region_server(datanode, incoming, rpc_stats); + } let mut client = Some(client); - let addr = test_util::PEER_PLACEHOLDER_ADDR; let channel_manager = ChannelManager::new(); channel_manager .reset_with_connector( @@ -688,6 +821,55 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { ) } +/// Counts the gRPC requests that the wrapped server receives. +#[derive(Clone)] +struct CountRequests { + inner: S, + rpc_stats: Option>, +} + +impl Service> for CountRequests +where + S: Service>, +{ + type Response = S::Response; + type Error = S::Error; + type Future = S::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, request: http::Request) -> Self::Future { + if let Some(rpc_stats) = &self.rpc_stats { + rpc_stats.requests.fetch_add(1, Ordering::Relaxed); + rpc_stats + .paths + .lock() + .unwrap() + .insert(request.uri().path().to_string()); + } + + self.inner.call(request) + } +} + +#[derive(Clone)] +struct CountRequestsLayer { + rpc_stats: Option>, +} + +impl Layer for CountRequestsLayer { + type Service = CountRequests; + + fn layer(&self, inner: S) -> Self::Service { + CountRequests { + inner, + rpc_stats: self.rpc_stats.clone(), + } + } +} + // Mock connectors are registered on a shared channel manager. #[cfg(test)] #[allow(deprecated)] diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 3fa1d3944893..3feaca375f9c 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use std::time::Duration; use cache::{ - build_datanode_cache_registry, build_fundamental_cache_registry, + build_datanode_layered_cache_registry, build_fundamental_cache_registry, with_default_composite_cache_registry, }; use catalog::information_schema::NoopInformationExtension; @@ -216,11 +216,10 @@ impl GreptimeDbStandaloneBuilder { ) -> GreptimeDbStandalone { let plugins = self.plugin.clone().unwrap_or_default(); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + // Use the same layered registry as the datanode of a running instance, so the derived + // caches of the test are invalidated after the caches they are derived from. + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut builder = DatanodeBuilder::new(opts.datanode_options(), plugins.clone(), kv_backend.clone()); diff --git a/tests-integration/src/tests.rs b/tests-integration/src/tests.rs index d70af021651a..4a5b2f0c6c5f 100644 --- a/tests-integration/src/tests.rs +++ b/tests-integration/src/tests.rs @@ -16,6 +16,7 @@ mod gc; mod instance_kafka_wal_test; mod instance_noop_wal_test; mod instance_test; +mod nested_merge_scan_capability_test; mod promql_test; mod reconcile_table; mod remote_dyn_filter_test; diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs new file mode 100644 index 000000000000..2f4b02c86e46 --- /dev/null +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -0,0 +1,818 @@ +// 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. + +//! Capability test of the nested `MergeScan` execution of a datanode. +//! +//! A datanode executes the plans it receives. Such a plan may carry a nested plan in a `MergeScan` +//! node, i.e. a `MergeScan` that is not a placeholder, and the datanode then plans that node into +//! a `MergeScanExec` which queries the regions of the *inner* table from the datanodes that own +//! them (see `DatanodeRegionQueryHandler`). That is what lets a distributed sub-plan of a join run +//! on the datanode that owns a region of the probe table while the build side of the join is read +//! from the datanodes that own the regions of the build table. +//! +//! The plans below are built **by hand** and sent to a datanode directly, so these tests exercise +//! the execution capability of the datanode only: no rewriter, hint or switch of this repository +//! takes part in producing the plan. +//! +//! The plan of one probe region is `Join(, MergeScan())`: +//! +//! * the probe side is decoded with the region aware catalog list of the request +//! (`NameAwareCatalogList`), so it reads the region of the probe table that the request names, +//! * the `MergeScan` payload is decoded with the catalog of the query engine (see +//! `query::query_engine::default_serializer`), so its table scan refers to the *build table* of +//! the cluster and not to the region of the request. The two tables of the test share no column +//! name besides their join key on purpose: a payload that was bound to the region of the request +//! would decode with a schema that does not even match the payload. +//! +//! The checks are: +//! +//! * every region of the probe table is asked to execute the nested plan and the union of the +//! results is the same multiset of rows as the result of the equivalent join SQL on the frontend +//! (the build table has duplicate join keys, so a probe row joins several build rows: a lost row +//! and a row that was counted twice both fail the comparison), +//! * the datanode really queried the *other* datanode for the regions of the build table that it +//! owns (see [`assert_cross_datanode_region_query`]), +//! * a nested plan whose inner relation does not exist, and a nested plan with an inner region that +//! is not served any more, both make the whole query fail: the datanode never returns the rows of +//! the inner regions that happen to succeed. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use api::v1::region::{QueryRequest, RegionRequestHeader}; +use common_query::Output; +use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; +use common_telemetry::info; +use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; +use datanode::region_server::RegionServer; +use frontend::instance::Instance; +use query::datafusion::QUERY_PARALLELISM_HINT; +use query::dist_plan::MergeScanLogicalPlan; +use query::parser::QueryLanguageParser; +use query::query_engine::DefaultSerializer; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::{QueryContext, QueryContextRef}; +use store_api::region_request::{RegionCloseRequest, RegionRequest}; +use store_api::storage::RegionId; +use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; + +use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; +use crate::test_util::execute_sql; + +/// The probe table of the nested join: partitioned, so that its regions are spread over both +/// datanodes. +const PROBE_TABLE: &str = "nested_cap_probe"; +/// The build table of the nested join: partitioned as well, and its join keys are duplicated. +const BUILD_TABLE: &str = "nested_cap_build"; +/// The build table of the error path: no such table exists, so the inner relation of the nested +/// plan cannot be resolved. +const MISSING_INNER_TABLE: &str = "nested_cap_missing_build"; + +/// `(a_id, probe_key, probe_v)` of the probe table: `a_id` is the primary key and the partition +/// column, `probe_key` is the join key and `probe_v` is the value of the row. +/// +/// The probe table and the build table share no column name besides their join key on purpose, see +/// the module documentation. +const PROBE_ROWS: &[(i32, i32, i32)] = &[ + (1, 10, 100), + (2, 20, 200), + // The second partition of the probe table. + (120, 30, 300), + (121, 10, 400), +]; + +/// `(b_id, build_key, build_v)` of the build table: the join key `10` appears twice and `40` +/// matches no row of the probe table. +const BUILD_ROWS: &[(i32, i32, i32)] = &[ + (1, 10, 1), + (2, 10, 2), + (3, 20, 3), + // The second partition of the build table. + (120, 30, 4), + (121, 40, 5), +]; + +/// The datanodes that have to serve a region query of the nested plan of a probe region. +#[derive(Debug, Default)] +struct CrossDatanodeExpectation { + /// Datanode id -> the number of (probe region, build region) pairs that the nested plans have + /// to read from that datanode over the network. + remote_datanodes: BTreeMap, +} + +/// The capability of a datanode to execute a hand-built nested plan: every region of the probe +/// table is executed with +/// `Join(, MergeScan())`, and +/// the datanodes fetch the regions of the build table from the datanode that owns them. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_join_across_datanodes() { + common_telemetry::init_default_ut_logging(); + + let cluster = build_cluster("test_nested_merge_scan_capability_join_across_datanodes").await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + info!("region leaders of {PROBE_TABLE}: {probe_leaders:?}"); + info!("region leaders of {BUILD_TABLE}: {build_leaders:?}"); + // Every table is partitioned into two regions. Both tables span both datanodes, so a nested + // plan of a probe region always has to read a region of the build table from the other + // datanode over the network. + assert_eq!( + 2, + datanodes(&probe_leaders).len(), + "expected the regions of {PROBE_TABLE} to be spread over both datanodes, actual region \ + leaders: {probe_leaders:?}" + ); + assert_eq!( + 2, + datanodes(&build_leaders).len(), + "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual region \ + leaders: {build_leaders:?}" + ); + + // The reference: the frontend executes the same join with its default distributed plan, i.e. + // both sides of the join are read with a `MergeScan` and the join runs on the frontend. + let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; + info!("join result of the frontend:\n{reference}"); + assert_eq!( + expected_join_rows(), + multiset(table_cells(&reference)), + "unexpected result of the join on the frontend:\n{reference}" + ); + + // The nested plan is the same for every probe region: the region of the request decides which + // rows of the probe table the plan reads, the inner `MergeScan` reads all regions of the build + // table. + let query_ctx = query_ctx(); + let nested_plan = nested_plan(&frontend, &query_ctx).await; + let plan = encode(&nested_plan); + + // The requests of the remote nodes are counted per datanode: a nested plan of a probe region + // that a datanode `d` owns has to query every datanode that owns a region of the build table + // *other* than `d`. + let expectation = cross_datanode_expectation(&probe_leaders, &build_leaders); + let baseline_requests = expectation + .remote_datanodes + .keys() + .map(|datanode| (*datanode, rpc_requests(&cluster, *datanode))) + .collect::>(); + + let mut rows = Vec::new(); + for (probe_region, (probe_datanode, _)) in &probe_leaders { + let stream = region_server(&cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(&query_ctx)), + region_id: *probe_region, + plan: plan.clone(), + }, + query_ctx.clone(), + ) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region}: {e}" + ) + }); + let batches = RecordBatches::try_collect(stream) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region} completely: {e}" + ) + }); + let actual = batches + .pretty_print() + .expect("the result of the datanode must be printable"); + info!("join result of the nested plan of the region {probe_region}:\n{actual}"); + rows.extend(table_cells(&actual)); + } + + // The nested plan of a probe region reads that region of the probe table, so the union of the + // results of all probe regions is the result of the whole join. The join key `10` appears twice + // in the build table, so the probe rows of that key appear twice as well: comparing the rows as + // a multiset catches both a lost row and a row that was counted twice. + assert_eq!( + expected_join_rows(), + multiset(rows.clone()), + "the nested plans of the datanodes returned an unexpected multiset of rows: {rows:?}" + ); + assert_eq!( + multiset(table_cells(&reference)), + multiset(rows), + "the nested plans of the datanodes returned a different multiset of rows than the \ + frontend:\nfrontend:\n{reference}" + ); + + assert_cross_datanode_region_query(&cluster, &expectation, &baseline_requests); +} + +/// A failing inner region makes the whole query fail: the datanode must not return the rows of the +/// inner regions that succeed. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_inner_failure_fails_query() { + common_telemetry::init_default_ut_logging(); + + let cluster = + build_cluster("test_nested_merge_scan_capability_inner_failure_fails_query").await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + info!("region leaders of {PROBE_TABLE}: {probe_leaders:?}"); + info!("region leaders of {BUILD_TABLE}: {build_leaders:?}"); + + // The request goes to a region of the probe table that the datanode `1` owns, and the nested + // plan reads the regions of the build table, so one of them comes from the other datanode. + let (probe_region, (probe_datanode, _)) = probe_leaders + .iter() + .find(|(_, (datanode, _))| *datanode == 1) + .unwrap_or_else(|| { + panic!( + "expected a region of {PROBE_TABLE} on the datanode 1, actual region leaders: \ + {probe_leaders:?}" + ) + }); + let (build_region, (build_datanode, build_addr)) = build_leaders + .iter() + .find(|(_, (datanode, _))| *datanode != *probe_datanode) + .unwrap_or_else(|| { + panic!( + "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual \ + region leaders: {build_leaders:?}" + ) + }); + + // 1. The inner relation of the nested plan does not exist. The payload of the `MergeScan` is + // resolved with the catalog of the datanode, so the datanode fails to decode the plan instead + // of binding the inner table scan to the region of the request. + let query_ctx = query_ctx(); + let missing_inner_plan = + nested_plan_with_inner_table(&frontend, &query_ctx, MISSING_INNER_TABLE).await; + let plan = encode(&missing_inner_plan); + let error = region_server(&cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(&query_ctx)), + region_id: *probe_region, + plan, + }, + query_ctx.clone(), + ) + .await + .err() + .expect("a nested plan whose inner relation does not exist must not be executed"); + // The top level error of the datanode only names the failing step (`Failed to decode logical + // plan`): the reason is the source of it, which names the table that could not be resolved. + let chain = error_chain(&error); + info!("the datanode rejected the nested plan with the missing inner relation: {chain}"); + assert!( + chain.contains(MISSING_INNER_TABLE), + "expected the datanode to fail while resolving the inner relation {MISSING_INNER_TABLE}, \ + actual error: {chain}" + ); + + // 2. An inner region of the nested plan fails: close the region of the build table that the + // other datanode owns. The nested plan of the probe region reads every region of the build + // table, so it has to query that datanode, which then cannot serve the region any more. The + // query of the probe region must fail as a whole: a partial result would hide the failure of + // the inner region, because the region of the build table that the datanode owns keeps serving + // its rows. + let baseline = rpc_requests(&cluster, *build_datanode); + let closed = region_server(&cluster, *build_datanode) + .handle_request( + RegionId::from_u64(*build_region), + RegionRequest::Close(RegionCloseRequest { + flush_on_close: false, + }), + ) + .await + .expect("the datanode must close the region of the build table"); + info!( + "closed the region {build_region} of {BUILD_TABLE} on the datanode {build_datanode} \ + ({closed:?}), which is reachable at {build_addr}" + ); + + let nested_plan = nested_plan(&frontend, &query_ctx).await; + let plan = encode(&nested_plan); + let stream = region_server(&cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(&query_ctx)), + region_id: *probe_region, + plan, + }, + query_ctx.clone(), + ) + .await + .expect("the request itself must be accepted by the datanode"); + let error = collect_error(stream) + .await + .expect("the query must fail when one of its inner regions fails"); + info!("the nested plan failed as a whole: {error}"); + + // The failure comes from the region query of the *other* datanode: it reached that datanode + // over the network, i.e. the query failed while reading a region of the inner table, not while + // planning the local datanode. + let stats = cluster + .datanode_rpc_stats(*build_datanode) + .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {build_datanode}")); + info!( + "the datanode {build_datanode} served {} remote gRPC requests, paths: {:?}", + stats.requests(), + stats.paths() + ); + assert!( + stats.requests() > baseline, + "expected the failed nested plan to query the datanode {build_datanode} for the region \ + {build_region} it owns, but the requests of the datanode did not grow" + ); + assert!( + stats.paths().iter().any(|path| path.contains("DoGet")), + "expected the datanode {build_datanode} to serve the region query of the nested plan, \ + actual remote gRPC paths: {:?}", + stats.paths() + ); +} + +async fn build_cluster(test_name: &str) -> GreptimeDbCluster { + GreptimeDbClusterBuilder::new(test_name) + .await + // A datanode serves the inner `MergeScan` of a nested plan by querying the datanodes that + // own the regions of the inner table, so every datanode must be reachable at the address it + // registers in metasrv, i.e. the address of the table routes. The cluster then also counts + // the requests of the remote nodes, see [`assert_cross_datanode_region_query`]. + .with_real_datanode_grpc_addr(true) + .with_datanodes(2) + .build(false) + .await +} + +/// The `DoGet` requests that the nested plans of the probe regions have to send to the datanodes +/// that own the regions of the build table. +fn cross_datanode_expectation( + probe_leaders: &BTreeMap, + build_leaders: &BTreeMap, +) -> CrossDatanodeExpectation { + let mut expectation = CrossDatanodeExpectation::default(); + for (probe_region, (probe_datanode, probe_addr)) in probe_leaders { + for (build_region, (build_datanode, build_addr)) in build_leaders { + if probe_datanode == build_datanode { + continue; + } + *expectation + .remote_datanodes + .entry(*build_datanode) + .or_default() += 1; + info!( + "the datanode {probe_datanode} ({probe_addr}) executing the region {probe_region} \ + of {PROBE_TABLE} queries the region {build_region} of {BUILD_TABLE} from the \ + datanode {build_datanode} ({build_addr})" + ); + } + } + expectation +} + +/// Asserts that the datanodes issued a real region query to every datanode that owns a region of +/// the build table. +/// +/// The cluster serves every datanode at its own address, while the frontend and the local +/// bookkeeping of the test process reach a datanode through an in-process client, so a request +/// recorded at the address of a datanode comes from a remote node: a request of the `DoGet` (a +/// region query) or of the region service proves the datanode-to-datanode region query. +fn assert_cross_datanode_region_query( + cluster: &GreptimeDbCluster, + expectation: &CrossDatanodeExpectation, + baseline_requests: &BTreeMap, +) { + assert!( + !expectation.remote_datanodes.is_empty(), + "expected the nested plans of the probe regions to read a region of the build table from \ + another datanode, but the regions of {BUILD_TABLE} are not spread over both datanodes" + ); + + for (datanode, probe_regions) in &expectation.remote_datanodes { + let stats = cluster + .datanode_rpc_stats(*datanode) + .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {datanode}")); + let baseline = baseline_requests.get(datanode).copied().unwrap_or_default(); + info!( + "the datanode {datanode} served {} remote gRPC requests (baseline {baseline}), \ + paths: {:?}", + stats.requests(), + stats.paths() + ); + assert!( + stats.requests() > baseline, + "expected the nested plans of {PROBE_TABLE} to query the datanode {datanode} \ + {probe_regions} times for the regions of {BUILD_TABLE} it owns, but the datanode \ + served no additional remote request ({} requests, baseline {baseline})", + stats.requests() + ); + assert!( + stats + .paths() + .iter() + .any(|path| path.contains("DoGet") || path.contains("Region")), + "expected the datanode {datanode} to serve the region query of the nested plan, actual \ + remote gRPC paths: {:?}", + stats.paths() + ); + } +} + +/// The number of gRPC requests that the remote nodes sent to the datanode so far. +fn rpc_requests(cluster: &GreptimeDbCluster, datanode_id: u64) -> usize { + cluster + .datanode_rpc_stats(datanode_id) + .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {datanode_id}")) + .requests() +} + +/// Returns the region server of the datanode `datanode_id`. +fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer { + cluster + .datanode_instances + .get(&datanode_id) + .unwrap_or_else(|| panic!("expected a datanode {datanode_id}")) + .region_server() +} + +/// Creates the probe and build tables of the tests and inserts the rows. +async fn prepare_tables(frontend: &Arc) { + execute_sql( + frontend, + &format!( + "CREATE TABLE {PROBE_TABLE} ( + a_id INT, + probe_key INT, + probe_v INT, + a_ts TIMESTAMP, + TIME INDEX (a_ts), + PRIMARY KEY(a_id) + ) + PARTITION ON COLUMNS (a_id) ( + a_id < 100, + a_id >= 100 + ) + engine=mito" + ), + ) + .await; + execute_sql( + frontend, + &format!( + "CREATE TABLE {BUILD_TABLE} ( + b_id INT, + build_key INT, + build_v INT, + b_ts TIMESTAMP, + TIME INDEX (b_ts), + PRIMARY KEY(b_id) + ) + PARTITION ON COLUMNS (b_id) ( + b_id < 100, + b_id >= 100 + ) + engine=mito" + ), + ) + .await; + + execute_sql( + frontend, + &format!( + "INSERT INTO {PROBE_TABLE}(a_id, probe_key, probe_v, a_ts) VALUES {}", + PROBE_ROWS + .iter() + .enumerate() + .map(|(i, (a_id, probe_key, probe_v))| { + format!("({a_id}, {probe_key}, {probe_v}, {}000)", i + 1) + }) + .collect::>() + .join(",") + ), + ) + .await; + execute_sql( + frontend, + &format!( + "INSERT INTO {BUILD_TABLE}(b_id, build_key, build_v, b_ts) VALUES {}", + BUILD_ROWS + .iter() + .enumerate() + .map(|(i, (b_id, build_key, build_v))| { + format!("({b_id}, {build_key}, {build_v}, {}000)", i + 1) + }) + .collect::>() + .join(",") + ), + ) + .await; +} + +/// Encodes a plan the way a frontend sends it to a datanode. +fn encode(plan: &LogicalPlan) -> Vec { + DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .expect("the plan must be encodable") + .to_vec() +} + +/// Builds the nested plan of the tests by hand: +/// `Join(, MergeScan())`. +/// +/// The plan is the remote input of a `MergeScan` that a frontend sends to the datanode that owns a +/// region of the probe table: the probe side is decoded with the region aware catalog list of the +/// request, so it reads the region of the request, and the `MergeScan` is not a placeholder, so the +/// datanode plans it into a `MergeScanExec` over the regions of the build table. +/// +/// The table scans are planned by the frontend (a `SELECT` of the columns of the join), which gives +/// them the schema of the tables of the cluster. The join and the nested `MergeScan` around the +/// build side are assembled here: no rewrite of the frontend is involved. +async fn nested_plan(frontend: &Arc, query_ctx: &QueryContextRef) -> LogicalPlan { + let probe = plan_sql( + frontend, + &format!("SELECT a_id, probe_key, probe_v FROM {PROBE_TABLE}"), + query_ctx, + ) + .await; + let build = plan_sql( + frontend, + &format!("SELECT build_key, build_v FROM {BUILD_TABLE}"), + query_ctx, + ) + .await; + + nested_plan_of(&probe, build) +} + +/// Same as [`nested_plan`], but the inner `MergeScan` scans `inner_table` instead of +/// [`BUILD_TABLE`]. +/// +/// The build side of the nested plan is the build table with the name of its scan replaced: the +/// scan keeps the schema and the source of the build table, so the payload of the `MergeScan` looks +/// exactly like the payload of [`nested_plan`] to the encoder, while the table it names does not +/// exist. +async fn nested_plan_with_inner_table( + frontend: &Arc, + query_ctx: &QueryContextRef, + inner_table: &str, +) -> LogicalPlan { + let probe = plan_sql( + frontend, + &format!("SELECT a_id, probe_key, probe_v FROM {PROBE_TABLE}"), + query_ctx, + ) + .await; + let build = plan_sql( + frontend, + &format!("SELECT build_key, build_v FROM {BUILD_TABLE}"), + query_ctx, + ) + .await; + let build = rename_scans(build, inner_table); + + nested_plan_of(&probe, build) +} + +/// Assembles `Join(probe, MergeScan(build))` and logs both sides. +fn nested_plan_of(probe: &LogicalPlan, build: LogicalPlan) -> LogicalPlan { + info!("probe side of the nested plan:\n{probe}"); + info!("build side of the nested plan:\n{build}"); + + let inner = MergeScanLogicalPlan::new(build, false, Default::default()).into_logical_plan(); + let plan = LogicalPlanBuilder::from(probe.clone()) + .join( + inner, + JoinType::Inner, + (vec!["probe_key"], vec!["build_key"]), + None, + ) + .expect("the join keys must be resolvable in the sides of the join") + .build() + .expect("the nested plan must be a valid logical plan"); + info!("nested plan sent to the datanode:\n{plan}"); + + plan +} + +/// Returns `plan` with every table scan renamed to `table_name`, keeping the schema and the source +/// of the scan. +fn rename_scans(plan: LogicalPlan, table_name: &str) -> LogicalPlan { + use datafusion::common::TableReference; + use datafusion::common::tree_node::{Transformed, TreeNode}; + + plan.transform_up(|plan| match plan { + LogicalPlan::TableScan(mut scan) => { + info!( + "renaming the scan of the nested plan from {} to {table_name}", + scan.table_name + ); + scan.table_name = TableReference::bare(table_name); + Ok(Transformed::yes(LogicalPlan::TableScan(scan))) + } + other => Ok(Transformed::no(other)), + }) + .expect("the plan must be rewritable") + .data +} + +/// Plans `sql` on the frontend and returns the logical plan. +async fn plan_sql(frontend: &Arc, sql: &str, query_ctx: &QueryContextRef) -> LogicalPlan { + let stmt = QueryLanguageParser::parse_sql(sql, query_ctx) + .unwrap_or_else(|e| panic!("failed to parse `{sql}`: {e}")); + frontend + .statement_executor() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap_or_else(|e| panic!("failed to plan `{sql}`: {e}")) +} + +/// The join of the tests, over the columns of the nested plan of [`nested_plan`] in order: +/// `(a_id, probe_key, probe_v, build_key, build_v)`. +fn join_sql() -> String { + format!( + "SELECT p.a_id, p.probe_key, p.probe_v, b.build_key, b.build_v + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b ON p.probe_key = b.build_key" + ) +} + +/// The rows of [`join_sql`] as a multiset: every probe row joins every build row of its join key. +fn expected_join_rows() -> Vec> { + let mut rows = Vec::new(); + for (a_id, probe_key, probe_v) in PROBE_ROWS { + for (_, build_key, build_v) in BUILD_ROWS { + if probe_key == build_key { + rows.push(vec![ + a_id.to_string(), + probe_key.to_string(), + probe_v.to_string(), + build_key.to_string(), + build_v.to_string(), + ]); + } + } + } + multiset(rows) +} + +/// Returns the region leaders of `table`: region id -> (datanode id, datanode address). +async fn region_leaders(frontend: &Arc, table: &str) -> BTreeMap { + use datatypes::arrow::array::AsArray; + use datatypes::arrow::datatypes::UInt64Type; + + let sql = format!( + "SELECT region_id, peer_id, peer_addr FROM information_schema.region_peers + WHERE table_schema = 'public' AND table_name = '{table}' AND is_leader = 'Yes' + ORDER BY region_id" + ); + let batches = output_batches(run_sql(frontend, &sql, query_ctx()).await).await; + + let mut leaders = BTreeMap::new(); + for batch in batches.iter() { + let batch = batch.df_record_batch(); + let region_ids = batch + .column_by_name("region_id") + .expect("region_id column") + .as_primitive::(); + let peer_ids = batch + .column_by_name("peer_id") + .expect("peer_id column") + .as_primitive::(); + let peer_addrs = batch + .column_by_name("peer_addr") + .expect("peer_addr column") + .as_string::(); + for row in 0..batch.num_rows() { + leaders.insert( + region_ids.value(row), + (peer_ids.value(row), peer_addrs.value(row).to_string()), + ); + } + } + + leaders +} + +/// The datanode ids of the region leaders. +fn datanodes(leaders: &BTreeMap) -> BTreeSet { + leaders.values().map(|(id, _)| *id).collect() +} + +/// Collects the record batches of a query output. +async fn output_batches(output: Output) -> RecordBatches { + match output.data { + common_query::OutputData::Stream(stream) => RecordBatches::try_collect(stream) + .await + .expect("the output stream must be collectable"), + common_query::OutputData::RecordBatches(batches) => batches, + common_query::OutputData::AffectedRows(rows) => { + panic!("expected a query output, got {rows} affected rows") + } + } +} + +/// Consumes `stream` and returns the error it failed with, or `None` if it returned completely. +async fn collect_error(stream: SendableRecordBatchStream) -> Option { + use futures::TryStreamExt; + + let result = stream.try_collect::>().await; + match result { + Ok(rows) => { + info!( + "the query returned {} batches instead of failing", + rows.len() + ); + None + } + Err(error) => Some(error_chain(&error)), + } +} + +/// Concatenates the message of every error of the source chain of `error`, so that the failure of +/// an outer step (e.g. decoding a plan) does not hide its reason (e.g. an unknown table). +fn error_chain(error: &(dyn std::error::Error + 'static)) -> String { + let mut messages = vec![error.to_string()]; + let mut source = error.source(); + while let Some(inner) = source { + messages.push(inner.to_string()); + source = inner.source(); + } + messages.join(": ") +} + +/// Runs `sql` on the frontend and returns the pretty printed result. +async fn query_pretty(frontend: &Arc, sql: &str, query_ctx: QueryContextRef) -> String { + run_sql(frontend, sql, query_ctx) + .await + .data + .pretty_print() + .await +} + +async fn run_sql(frontend: &Arc, sql: &str, query_ctx: QueryContextRef) -> Output { + SqlQueryHandler::do_query(frontend.as_ref(), sql, query_ctx) + .await + .remove(0) + .unwrap() +} + +/// The query context of a new session, with a single target partition so that the plans of the +/// tests are stable. The datanode builds the query context of its own execution from the header of +/// the request, so the request of the tests carries the same query context. +fn query_ctx() -> QueryContextRef { + let mut query_ctx = QueryContext::with_db_name(None); + query_ctx.set_extension(QUERY_PARALLELISM_HINT, "1"); + Arc::new(query_ctx) +} + +/// The header of the region query of the tests, carrying the query context of the session of the +/// plan. +fn region_request_header(query_ctx: &QueryContextRef) -> RegionRequestHeader { + RegionRequestHeader { + query_context: Some(query_ctx.as_ref().into()), + ..Default::default() + } +} + +/// Parses the cells of the data rows of a pretty printed table. +fn table_cells(pretty: &str) -> Vec> { + let mut lines = pretty.lines().filter(|line| line.starts_with('|')); + // The first `|` line is the header of the table. + lines.next(); + + lines + .map(|line| { + line.trim_matches('|') + .split('|') + .map(|cell| cell.trim().to_string()) + .collect() + }) + .collect() +} + +/// Sorts the rows of a result, so that two results can be compared as multisets. +fn multiset(mut rows: Vec>) -> Vec> { + rows.sort(); + rows +} From a68d3798d6a368828b500b2984614097e9b6cf53 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:24:38 +0800 Subject: [PATCH 04/37] 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> --- src/query/src/query_engine/default_serializer.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index e0d41674b24e..271961d6ebff 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -516,7 +516,6 @@ mod tests { use common_time::Timezone; use datafusion::catalog::TableProvider; use datafusion::datasource::MemTable; - use datafusion::functions::datetime::to_char::ToCharFunc; use datafusion::logical_expr::Extension; use datafusion_expr::expr::{Cast, ScalarFunction}; use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, ScalarUDF, col, lit}; @@ -1067,10 +1066,10 @@ mod tests { let udf = &udfs[0]; assert_eq!(udf.name(), "date_format", "got: {plan}"); - assert!( - udf.inner().as_any().downcast_ref::().is_none(), - "`date_format` is bound to the built-in `to_char` of DataFusion: {plan}" - ); + // The built-in `to_char` of DataFusion declares `date_format` as an alias, so a plan + // where `date_format` was bound to it holds a `to_char` function (rendered + // `to_char(...) AS date_format(...)` by the substrait encoder). Its instance can never + // equal the implementation of GreptimeDB, which the assertion below checks. assert_eq!( udf.as_ref(), greptime_date_format_udf(query_ctx).as_ref(), From 238bd0075aec7c83df314a57a1636e1fe8eb7d20 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:24:43 +0800 Subject: [PATCH 05/37] 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> --- tests-integration/src/cluster.rs | 24 ++ .../nested_merge_scan_capability_test.rs | 216 +++++++++++++++++- 2 files changed, 236 insertions(+), 4 deletions(-) diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 681a32a77590..22731d2d2e29 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -200,6 +200,11 @@ impl GreptimeDbCluster { } } +/// Adjusts the options of a datanode of the cluster before the datanode is created, see +/// [`GreptimeDbClusterBuilder::with_datanode_options_override`]. The first argument is the id of +/// the datanode, so a test can vary the options per datanode. +type DatanodeOptionsOverride = Box; + pub struct GreptimeDbClusterBuilder { cluster_name: String, kv_backend: KvBackendRef, @@ -216,6 +221,7 @@ pub struct GreptimeDbClusterBuilder { local_file_access: LocalFileAccess, event_recorder_options: EventRecorderOptions, real_datanode_grpc_addr: bool, + datanode_options_override: Option, } impl GreptimeDbClusterBuilder { @@ -253,6 +259,7 @@ impl GreptimeDbClusterBuilder { local_file_access: LocalFileAccess::default(), event_recorder_options: EventRecorderOptions::default(), real_datanode_grpc_addr: false, + datanode_options_override: None, } } @@ -346,6 +353,20 @@ impl GreptimeDbClusterBuilder { self } + /// Adjusts the options of every datanode before the datanode is created, e.g. to set the + /// concurrency limit of the region server of the datanode (`max_concurrent_queries`). + /// + /// The hook runs after the id and the gRPC address of the datanode have been set, and it + /// receives the id of the datanode, so a test can vary the options per datanode. + #[must_use] + pub fn with_datanode_options_override( + mut self, + override_options: impl Fn(u64, &mut DatanodeOptions) + Send + Sync + 'static, + ) -> Self { + self.datanode_options_override = Some(Box::new(override_options)); + self + } + pub async fn build_with( &self, datanode_options: Vec, @@ -489,6 +510,9 @@ impl GreptimeDbClusterBuilder { opts.grpc.bind_addr = addr.clone(); opts.grpc.server_addr = addr; } + if let Some(override_options) = &self.datanode_options_override { + override_options(datanode_id, &mut opts); + } options.push(opts); } diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 2f4b02c86e46..bb58e34142ec 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -21,9 +21,14 @@ //! on the datanode that owns a region of the probe table while the build side of the join is read //! from the datanodes that own the regions of the build table. //! -//! The plans below are built **by hand** and sent to a datanode directly, so these tests exercise -//! the execution capability of the datanode only: no rewriter, hint or switch of this repository -//! takes part in producing the plan. +//! The plans below are built **by hand** and handed to a datanode by calling +//! `RegionServer::handle_remote_read` in process, so these tests exercise the execution capability +//! of the datanode only: no rewriter, hint or switch of this repository takes part in producing the +//! plan, and the outer request itself never goes over the wire. The region queries that the +//! datanode then issues for the regions of the inner table do: the tests build the cluster with +//! `GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`, so every datanode is reachable at the +//! address it registered and the inner `MergeScan` reaches the datanodes that own the regions of +//! the inner table over a real Flight connection. //! //! The plan of one probe region is `Join(, MergeScan())`: //! @@ -45,10 +50,15 @@ //! owns (see [`assert_cross_datanode_region_query`]), //! * a nested plan whose inner relation does not exist, and a nested plan with an inner region that //! is not served any more, both make the whole query fail: the datanode never returns the rows of -//! the inner regions that happen to succeed. +//! the inner regions that happen to succeed, +//! * with the concurrency limit of every datanode set to a single query, every probe region still +//! executes: the inner region queries of the nested plans are admitted as execution stages of the +//! outer queries, so they don't wait for the permit that their own query holds (see +//! [`test_nested_merge_scan_capability_internal_stage_under_concurrency_limit`]). use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; +use std::time::Duration; use api::v1::region::{QueryRequest, RegionRequestHeader}; use common_query::Output; @@ -353,6 +363,180 @@ async fn test_nested_merge_scan_capability_inner_failure_fails_query() { ); } +/// The nested plans execute under the concurrency limit of their datanodes: a datanode admits a +/// single query at a time, the outer query of a probe region holds the only permit of its datanode +/// until its stream is fully consumed, and the nested plan then queries the regions of the build +/// table from the datanode itself and from the other datanode. +/// +/// Every inner region query has to be admitted as an execution stage of the outer query, without a +/// permit of its own. The marker of the stage travels from the region query handler of the datanode +/// (`request.internal = true`), through the client that encodes it into the header of the request, +/// to the region server of the peer that reads it back and skips its limiter. If any one of them is +/// missing, the inner region query of the *same* datanode waits for the permit held by the outer +/// query and fails after `concurrent_query_limiter_timeout`, and the inner region query of the +/// *other* datanode waits for the outer query of that datanode, which waits for its own nested +/// plan: the consumption of the outer queries below is bounded by a timeout, so both cases fail the +/// test instead of hanging it. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_internal_stage_under_concurrency_limit() { + common_telemetry::init_default_ut_logging(); + + let cluster = build_cluster_with_query_limit( + "test_nested_merge_scan_capability_internal_stage_under_concurrency_limit", + ) + .await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + info!("region leaders of {PROBE_TABLE}: {probe_leaders:?}"); + info!("region leaders of {BUILD_TABLE}: {build_leaders:?}"); + // Every table is partitioned into two regions. Both tables span both datanodes, so a nested + // plan of a probe region always has to read a region of the build table from the other + // datanode, and the two probe regions sit on different datanodes: asking both datanodes for + // their probe region exhausts the single permit of every datanode of the cluster. + assert_eq!( + 2, + datanodes(&probe_leaders).len(), + "expected the regions of {PROBE_TABLE} to be spread over both datanodes, actual region \ + leaders: {probe_leaders:?}" + ); + assert_eq!( + 2, + datanodes(&build_leaders).len(), + "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual region \ + leaders: {build_leaders:?}" + ); + + // The reference: the frontend executes the same join with its default distributed plan. + let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; + info!("join result of the frontend:\n{reference}"); + assert_eq!( + expected_join_rows(), + multiset(table_cells(&reference)), + "unexpected result of the join on the frontend:\n{reference}" + ); + + let query_ctx = query_ctx(); + let nested_plan = nested_plan(&frontend, &query_ctx).await; + let plan = encode(&nested_plan); + + let expectation = cross_datanode_expectation(&probe_leaders, &build_leaders); + let baseline_requests = expectation + .remote_datanodes + .keys() + .map(|datanode| (*datanode, rpc_requests(&cluster, *datanode))) + .collect::>(); + + // Ask every datanode for its probe region before consuming anything: the region server takes + // the permit of the datanode when it admits the request and holds it until the stream of the + // request is consumed, so no datanode has a free permit left while the nested plans run. + let mut outer_queries = Vec::with_capacity(probe_leaders.len()); + for (probe_region, (probe_datanode, _)) in &probe_leaders { + let stream = region_server(&cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(&query_ctx)), + region_id: *probe_region, + plan: plan.clone(), + }, + query_ctx.clone(), + ) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must admit the nested plan of the region \ + {probe_region} while it serves a single query at a time: {e}" + ) + }); + outer_queries.push((*probe_region, *probe_datanode, stream)); + } + + // Consume the outer queries of both datanodes at the same time, so that every nested plan + // queries the regions of the build table while the only permit of every datanode is held by the + // outer query of that datanode. An inner region query that misses the stage marker waits for + // that permit up to `concurrent_query_limiter_timeout` and fails, or the two outer queries wait + // for each other: the timeout turns the second case into a failure instead of a hang. + let collected = tokio::time::timeout( + Duration::from_secs(30), + futures::future::join_all(outer_queries.into_iter().map( + |(probe_region, probe_datanode, stream)| async move { + ( + probe_region, + probe_datanode, + RecordBatches::try_collect(stream).await, + ) + }, + )), + ) + .await + .unwrap_or_else(|_| { + panic!( + "the nested plans of the probe regions did not return within the timeout while every \ + datanode served a single query at a time: an inner region query is waiting for the \ + concurrency permit held by the outer query that dispatched it, i.e. the stage marker \ + of the inner region query was lost" + ) + }); + + let mut rows = Vec::new(); + for (probe_region, probe_datanode, result) in collected { + let batches = result.unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region} completely while it serves a single query at a time: {e}" + ) + }); + let actual = batches + .pretty_print() + .expect("the result of the datanode must be printable"); + info!("join result of the nested plan of the region {probe_region}:\n{actual}"); + rows.extend(table_cells(&actual)); + } + + // The nested plan of a probe region reads that region of the probe table, so the union of the + // results of all probe regions is the result of the whole join. Comparing the rows as a + // multiset catches both a lost row and a row that was counted twice. + assert_eq!( + expected_join_rows(), + multiset(rows.clone()), + "the nested plans of the datanodes returned an unexpected multiset of rows under the \ + concurrency limit: {rows:?}" + ); + assert_eq!( + multiset(table_cells(&reference)), + multiset(rows), + "the nested plans of the datanodes returned a different multiset of rows than the \ + frontend:\nfrontend:\n{reference}" + ); + + assert_cross_datanode_region_query(&cluster, &expectation, &baseline_requests); + + // The nested plan really reached the other datanode over the network: the region query of a + // `MergeScan` is a Flight `DoGet` (see `RegionRequester::handle_query`), and the remote + // requests of a datanode are counted at the address it registered, so a `DoGet` recorded there + // was sent by another datanode of the cluster. + for datanode in expectation.remote_datanodes.keys() { + let stats = cluster + .datanode_rpc_stats(*datanode) + .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {datanode}")); + info!( + "the datanode {datanode} served {} remote gRPC requests, paths: {:?}", + stats.requests(), + stats.paths() + ); + assert!( + stats.paths().iter().any(|path| path.contains("DoGet")), + "expected the nested plans of {PROBE_TABLE} to query the datanode {datanode} for the \ + regions of {BUILD_TABLE} it owns with a Flight `DoGet`, actual remote gRPC paths: \ + {:?}", + stats.paths() + ); + } +} + async fn build_cluster(test_name: &str) -> GreptimeDbCluster { GreptimeDbClusterBuilder::new(test_name) .await @@ -366,6 +550,30 @@ async fn build_cluster(test_name: &str) -> GreptimeDbCluster { .await } +/// Same as [`build_cluster`], but every datanode admits a single query at a time. +/// +/// The outer region query of a probe region holds the only permit of its datanode until its stream +/// is fully consumed, so a region query that the nested plan of that query issues to the same +/// datanode can only succeed if it is admitted as an execution stage of the query instead of +/// taking a permit of its own (see +/// [`test_nested_merge_scan_capability_internal_stage_under_concurrency_limit`]). +async fn build_cluster_with_query_limit(test_name: &str) -> GreptimeDbCluster { + GreptimeDbClusterBuilder::new(test_name) + .await + .with_real_datanode_grpc_addr(true) + .with_datanodes(2) + .with_datanode_options_override(|_, opts| { + opts.max_concurrent_queries = 1; + // Wider than the default: the test must fail because the stage marker of an inner + // region query is lost, not because a slow machine made an admitted query wait for a + // permit for too long. The timeout around the consumption of the outer queries is much + // larger than this. + opts.concurrent_query_limiter_timeout = Duration::from_secs(1); + }) + .build(false) + .await +} + /// The `DoGet` requests that the nested plans of the probe regions have to send to the datanodes /// that own the regions of the build table. fn cross_datanode_expectation( From 652be4519312c62d55b5f74385c1ba954183dfba Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:23:59 +0800 Subject: [PATCH 06/37] 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> --- src/datanode/src/region_query.rs | 12 + tests-integration/src/cluster.rs | 38 +- .../nested_merge_scan_capability_test.rs | 337 +++++++++++++++++- 3 files changed, 376 insertions(+), 11 deletions(-) diff --git a/src/datanode/src/region_query.rs b/src/datanode/src/region_query.rs index 6b0fae6678d2..d22a93457b9e 100644 --- a/src/datanode/src/region_query.rs +++ b/src/datanode/src/region_query.rs @@ -36,6 +36,15 @@ use snafu::ResultExt; use store_api::storage::RegionId; /// Serves the region queries issued by the datanode itself. +/// +/// Known limitation: [`RegionQueryHandler::select_target`] resolves the leader of a region through +/// the table route cache, which uses `InitStrategy::VersionChecked`. The version counter of that +/// strategy belongs to the `CacheContainer` and not to a single table: any `TableId` invalidation +/// bumps the shared version, so a cold load of one table's route retries (and loads again) when an +/// *unrelated* table is invalidated concurrently. Clusters with frequent DDL or region movement can +/// therefore see amplified cold-load retries on the datanode; the tail latency of cold loads in +/// such clusters is worth verifying. The retries only add work, they keep the loaded route correct +/// (a stale route is never returned). pub struct DatanodeRegionQueryHandler { partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, @@ -61,6 +70,9 @@ impl RegionQueryHandler for DatanodeRegionQueryHandler { region_id: RegionId, ) -> QueryResult { // The leader of a region is the only peer that can serve a region query. + // + // This lookup goes through the version checked table route cache: see the known limitation + // of `DatanodeRegionQueryHandler` about retries of unrelated invalidations. let peer = self .partition_manager .find_region_leader(region_id) diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 22731d2d2e29..965ee867e1e3 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -41,7 +41,7 @@ use common_datasource::object_store::LocalFileAccess; use common_event_recorder::EventRecorderOptions; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::DatanodeId; -use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; +use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder}; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::chroot::ChrootKvBackend; use common_meta::kv_backend::etcd::EtcdStore; @@ -121,6 +121,11 @@ pub struct GreptimeDbCluster { /// The stats of the network address of every datanode, empty unless the cluster serves the /// datanodes at their own addresses. pub datanode_rpc_stats: HashMap>, + /// The layered cache registry of every datanode, i.e. the registry that the datanode + /// invalidates locally after it received an invalidation instruction from metasrv. A test can + /// invalidate a cache ident (e.g. an unrelated `TableId`) exactly the way the datanode does + /// when it handles the broadcast of metasrv. + pub datanode_cache_registries: HashMap>, pub kv_backend: KvBackendRef, pub metasrv: Arc, pub frontend: Arc, @@ -408,7 +413,7 @@ impl GreptimeDbClusterBuilder { ) .await; - let datanode_instances = self + let (datanode_instances, datanode_cache_registries) = self .build_datanodes_with_options(&metasrv, &datanode_options) .await; @@ -447,6 +452,7 @@ impl GreptimeDbClusterBuilder { guards, datanode_instances, datanode_rpc_stats, + datanode_cache_registries, kv_backend: self.kv_backend.clone(), metasrv: metasrv.metasrv, frontend: Arc::new(frontend), @@ -533,15 +539,22 @@ impl GreptimeDbClusterBuilder { &self, metasrv: &MockInfo, options: &[DatanodeOptions], - ) -> HashMap { + ) -> ( + HashMap, + HashMap>, + ) { let mut instances = HashMap::with_capacity(options.len()); + let mut cache_registries = HashMap::with_capacity(options.len()); for opts in options { - let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await; - instances.insert(opts.node_id.unwrap(), datanode); + let (datanode, cache_registry) = + self.create_datanode(opts.clone(), metasrv.clone()).await; + let datanode_id = opts.node_id.unwrap(); + instances.insert(datanode_id, datanode); + cache_registries.insert(datanode_id, cache_registry); } - instances + (instances, cache_registries) } async fn wait_datanodes_alive( @@ -567,7 +580,11 @@ impl GreptimeDbClusterBuilder { panic!("Some Datanodes are not alive in 10 seconds!") } - async fn create_datanode(&self, opts: DatanodeOptions, metasrv: MockInfo) -> Datanode { + async fn create_datanode( + &self, + opts: DatanodeOptions, + metasrv: MockInfo, + ) -> (Datanode, Arc) { let mut meta_client = MetaClientBuilder::datanode_default_options(opts.node_id.unwrap()) .channel_manager(metasrv.channel_manager) .build(); @@ -583,13 +600,16 @@ impl GreptimeDbClusterBuilder { let mut builder = DatanodeBuilder::new(opts.clone(), Plugins::default(), meta_backend); builder - .with_cache_registry(layered_cache_registry) + .with_cache_registry(layered_cache_registry.clone()) .with_meta_client(meta_client); let mut datanode = builder.build().await.unwrap(); datanode.start_heartbeat().await.unwrap(); - datanode + // The datanode invalidates this registry locally when it receives an invalidation + // instruction, so a test can drive the same path: see + // [`GreptimeDbCluster::datanode_cache_registries`]. + (datanode, layered_cache_registry) } async fn build_frontend( diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index bb58e34142ec..f2e92d6428cd 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -54,13 +54,22 @@ //! * with the concurrency limit of every datanode set to a single query, every probe region still //! executes: the inner region queries of the nested plans are admitted as execution stages of the //! outer queries, so they don't wait for the permit that their own query holds (see -//! [`test_nested_merge_scan_capability_internal_stage_under_concurrency_limit`]). +//! [`test_nested_merge_scan_capability_internal_stage_under_concurrency_limit`]), +//! * the inner `MergeScan` resolves the leader of the regions of the build table through the table +//! route cache of the datanode, whose version counter is shared by every table of the container: +//! a cold load of the route of the build table still returns the rows of the join while an +//! unrelated table id is invalidated in a loop (see +//! [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]). use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; use api::v1::region::{QueryRequest, RegionRequestHeader}; +use common_meta::cache::{LayeredCacheRegistry, TableRouteCacheRef}; +use common_meta::cache_invalidator::{CacheInvalidator, Context}; +use common_meta::instruction::CacheIdent; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::info; @@ -74,7 +83,7 @@ use query::query_engine::DefaultSerializer; use servers::query_handler::sql::SqlQueryHandler; use session::context::{QueryContext, QueryContextRef}; use store_api::region_request::{RegionCloseRequest, RegionRequest}; -use store_api::storage::RegionId; +use store_api::storage::{RegionId, TableId}; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; @@ -89,6 +98,21 @@ const BUILD_TABLE: &str = "nested_cap_build"; /// plan cannot be resolved. const MISSING_INNER_TABLE: &str = "nested_cap_missing_build"; +/// The table id that no table of the cluster uses. The test of the cold route load invalidates this +/// table id while it cold-loads the route of [`BUILD_TABLE`]: the version counter of the version +/// checked table route cache belongs to the cache container and not to a single table, so this +/// unrelated invalidation bumps the version that the cold load compares. +const UNRELATED_TABLE_ID: TableId = 999_999; + +/// The interval of the unrelated invalidation loop of +/// [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]. +/// +/// Short enough to hit an in-flight load of the table route cache, long enough to leave the retry +/// loop of the version checked load room to finish (a retry sleeps for at least `10ms`): the test +/// must fail because a route was stale or belonged to another table, not because the retry loop of +/// the cache gave up. +const UNRELATED_INVALIDATION_INTERVAL: Duration = Duration::from_millis(20); + /// `(a_id, probe_key, probe_v)` of the probe table: `a_id` is the primary key and the partition /// column, `probe_key` is the join key and `probe_v` is the value of the row. /// @@ -233,6 +257,171 @@ async fn test_nested_merge_scan_capability_join_across_datanodes() { assert_cross_datanode_region_query(&cluster, &expectation, &baseline_requests); } +/// The route of the build table is loaded cold while an *unrelated* table id is invalidated: the +/// version counter of the version checked table route cache belongs to the cache container and not +/// to a single table, so an invalidation of any `TableId` bumps the version that a cold load +/// compares, and a load that runs while an unrelated table is invalidated retries (retrying is what +/// keeps the loaded route correct; see the known limitation of `DatanodeRegionQueryHandler`). +/// +/// The invalidation of the unrelated table id is injected into the layered cache registry of every +/// datanode, which is exactly the local handling path of the invalidation instruction that a +/// datanode receives from metasrv; the broadcast to the datanode is the only part not exercised +/// here. +/// +/// The test pins the correctness of the retry path. The route of the build table is dropped on +/// every datanode (so the nested plans really load it from metasrv, the assertion below checks the +/// cache is cold), then every probe region executes the nested plan while a background task +/// invalidates the unrelated table id in a loop. The rows must be the rows of the join in every +/// round: a retried load must neither return the route of another table nor a stale route. +/// +/// The test also reports how often the table route caches loaded a route with and without the +/// unrelated invalidation. A version checked load counts every attempt, so the amplification of the +/// unrelated invalidation is visible in the counter; the test logs the numbers instead of asserting +/// them, because the unrelated invalidation only hits an in-flight load with some probability. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation() { + common_telemetry::init_default_ut_logging(); + + let cluster = + build_cluster("test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation") + .await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + info!("region leaders of {PROBE_TABLE}: {probe_leaders:?}"); + info!("region leaders of {BUILD_TABLE}: {build_leaders:?}"); + // Both datanodes own a probe region, so both of them resolve the leader of the regions of the + // build table when they execute the nested plan of their probe region. + assert_eq!( + 2, + datanodes(&probe_leaders).len(), + "expected the regions of {PROBE_TABLE} to be spread over both datanodes, actual region \ + leaders: {probe_leaders:?}" + ); + assert_eq!( + 2, + datanodes(&build_leaders).len(), + "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual region \ + leaders: {build_leaders:?}" + ); + + let build_table_id = table_id(&build_leaders); + info!( + "the regions of {BUILD_TABLE} belong to the table id {build_table_id}; the test \ + invalidates the unrelated table id {UNRELATED_TABLE_ID} while the datanodes load the \ + route of {BUILD_TABLE}" + ); + + // The reference join of the frontend and the nested plan, as in the other capability tests. + let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; + info!("join result of the frontend:\n{reference}"); + assert_eq!( + expected_join_rows(), + multiset(table_cells(&reference)), + "unexpected result of the join on the frontend:\n{reference}" + ); + + let query_ctx = query_ctx(); + let plan = encode(&nested_plan(&frontend, &query_ctx).await); + + // 1. The baseline: the route of the build table is dropped on every datanode and nothing + // invalidates a table while the nested plans run, so both datanodes load the route of the build + // table from metasrv exactly once. + invalidate_table_ids(&cluster, &[build_table_id]).await; + assert_route_cold(&cluster, build_table_id).await; + + let loads_before = table_route_cache_loads(); + let rows_cold = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; + let cold_loads = table_route_cache_loads() - loads_before; + assert_eq!( + expected_join_rows(), + multiset(rows_cold.clone()), + "the nested plans returned an unexpected multiset of rows after the route of {BUILD_TABLE} \ + was dropped: {rows_cold:?}" + ); + // Every datanode owns a probe region, so every datanode resolved the leader of the regions of + // the build table with a cold route cache: the baseline really started without the route. + assert!( + cold_loads >= 2, + "expected the datanodes to load the dropped route of {BUILD_TABLE} from metasrv, but the \ + table route cache only loaded a route {cold_loads} time(s)" + ); + + // 2. The unrelated invalidation: a background task invalidates the table id of a table that the + // cluster does not have, the way the datanodes handle the invalidation instruction of a DDL + // that touched some other table. Every round drops the route of the build table again, so the + // nested plans load a cold route while the unrelated invalidation runs. + let stop = Arc::new(AtomicBool::new(false)); + let unrelated_invalidations = Arc::new(AtomicUsize::new(0)); + let invalidation_task = tokio::spawn(invalidate_unrelated_table_id( + cluster + .datanode_cache_registries + .values() + .cloned() + .collect::>>(), + stop.clone(), + unrelated_invalidations.clone(), + )); + + let mut loads_with_invalidation = 0; + for round in 0..3 { + invalidate_table_ids(&cluster, &[build_table_id]).await; + assert_route_cold(&cluster, build_table_id).await; + + let loads_before = table_route_cache_loads(); + let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; + loads_with_invalidation += table_route_cache_loads() - loads_before; + + assert_eq!( + expected_join_rows(), + multiset(rows.clone()), + "round {round}: the nested plans returned an unexpected multiset of rows while the \ + unrelated table id {UNRELATED_TABLE_ID} was invalidated: {rows:?}" + ); + assert_eq!( + multiset(rows_cold.clone()), + multiset(rows.clone()), + "round {round}: the nested plans returned different rows while the unrelated table id \ + {UNRELATED_TABLE_ID} was invalidated, i.e. a retried route load returned another \ + route than the cold load" + ); + + // The retry loop of a route load that the unrelated invalidation interrupted must have + // loaded the route of the build table: the retry must not leave the cache without the + // route of the table that the nested plan asked for. + for datanode in datanodes(&probe_leaders) { + assert!( + table_route_cache(&cluster, datanode).contains_key(&build_table_id), + "round {round}: the datanode {datanode} must hold the route of {BUILD_TABLE} after \ + it executed the nested plan of its probe region" + ); + } + } + + stop.store(true, Ordering::Relaxed); + invalidation_task + .await + .expect("the task that invalidates the unrelated table id must stop"); + + // Best effort measurement, not an assertion: the unrelated invalidation retries a cold load + // only if it lands in the window between the version the load read and the version it compares + // afterwards, so the counts vary with the timing of the machine. The counter is the + // `CACHE_CONTAINER_CACHE_MISS` of the cache container, which every attempt of a version checked + // load increments; the cache name is shared by the datanodes and the frontend of the test + // process, but the frontend is idle while the nested plans of a round run. + let unrelated_invalidations = unrelated_invalidations.load(Ordering::Relaxed); + info!( + "the table route caches loaded a cold route of {BUILD_TABLE} {cold_loads} time(s) without \ + an unrelated invalidation and {loads_with_invalidation} time(s) while the unrelated table \ + id {UNRELATED_TABLE_ID} was invalidated {unrelated_invalidations} time(s); a count above 2 \ + per round means the shared version counter of the cache container made the cold loads \ + retry" + ); +} + /// A failing inner region makes the whole query fail: the datanode must not return the rows of the /// inner regions that succeed. #[tokio::test(flavor = "multi_thread")] @@ -665,6 +854,150 @@ fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer .region_server() } +/// Executes the nested plan of every probe region on the datanode that owns the region, consumes +/// every stream completely and returns the union of the rows of the results. +async fn run_nested_queries( + cluster: &GreptimeDbCluster, + probe_leaders: &BTreeMap, + plan: &[u8], + query_ctx: &QueryContextRef, +) -> Vec> { + let mut rows = Vec::new(); + for (probe_region, (probe_datanode, _)) in probe_leaders { + let stream = region_server(cluster, *probe_datanode) + .handle_remote_read( + QueryRequest { + header: Some(region_request_header(query_ctx)), + region_id: *probe_region, + plan: plan.to_vec(), + }, + query_ctx.clone(), + ) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region}: {e}" + ) + }); + let batches = RecordBatches::try_collect(stream) + .await + .unwrap_or_else(|e| { + panic!( + "the datanode {probe_datanode} must execute the nested plan of the region \ + {probe_region} completely: {e}" + ) + }); + let actual = batches + .pretty_print() + .expect("the result of the datanode must be printable"); + info!("join result of the nested plan of the region {probe_region}:\n{actual}"); + rows.extend(table_cells(&actual)); + } + + rows +} + +/// The id of the single table that the regions of `leaders` belong to. +fn table_id(leaders: &BTreeMap) -> TableId { + let table_ids = leaders + .keys() + .map(|region_id| RegionId::from_u64(*region_id).table_id()) + .collect::>(); + assert_eq!( + 1, + table_ids.len(), + "expected the regions to belong to a single table, actual region leaders: {leaders:?}" + ); + + *table_ids.first().expect("there is at least one region") +} + +/// Returns the table route cache of the datanode `datanode_id`: the cache that +/// `DatanodeRegionQueryHandler::select_target` reads to resolve the leader of a region. +fn table_route_cache(cluster: &GreptimeDbCluster, datanode_id: u64) -> TableRouteCacheRef { + cluster + .datanode_cache_registries + .get(&datanode_id) + .unwrap_or_else(|| panic!("expected the cache registry of the datanode {datanode_id}")) + .get::() + .unwrap_or_else(|| panic!("expected the table route cache of the datanode {datanode_id}")) +} + +/// Invalidates `table_ids` on every datanode, the way a datanode handles the invalidation +/// instruction of metasrv (see [`GreptimeDbCluster::datanode_cache_registries`]). +async fn invalidate_table_ids(cluster: &GreptimeDbCluster, table_ids: &[TableId]) { + let idents = table_ids + .iter() + .copied() + .map(CacheIdent::TableId) + .collect::>(); + for (datanode_id, registry) in &cluster.datanode_cache_registries { + registry + .invalidate(&Context::default(), &idents) + .await + .unwrap_or_else(|e| { + panic!("the datanode {datanode_id} must invalidate {idents:?}: {e}") + }); + } +} + +/// Invalidates [`UNRELATED_TABLE_ID`] on `registries` until `stop` is set, and counts the rounds. +async fn invalidate_unrelated_table_id( + registries: Vec>, + stop: Arc, + invalidations: Arc, +) { + while !stop.load(Ordering::Relaxed) { + for registry in ®istries { + registry + .invalidate( + &Context::default(), + &[CacheIdent::TableId(UNRELATED_TABLE_ID)], + ) + .await + .expect("the datanode must invalidate the unrelated table id"); + } + invalidations.fetch_add(1, Ordering::Relaxed); + tokio::time::sleep(UNRELATED_INVALIDATION_INTERVAL).await; + } +} + +/// Waits until no datanode holds the route of `table_id` any more, i.e. until the route is cold. +/// +/// The invalidation is awaited before this call, so the wait is a safety net for the internals of +/// the cache (e.g. its pending tasks) rather than a required step. +async fn assert_route_cold(cluster: &GreptimeDbCluster, table_id: TableId) { + for attempt in 0..100 { + let mut cached = Vec::new(); + for datanode_id in cluster.datanode_cache_registries.keys() { + if table_route_cache(cluster, *datanode_id).contains_key(&table_id) { + cached.push(*datanode_id); + } + } + if cached.is_empty() { + return; + } + assert!( + attempt < 99, + "the datanodes {cached:?} must drop the route of the table {table_id} after the \ + invalidation" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +/// The number of attempts that the table route caches of the process made to load a route: every +/// attempt of a version checked load (including its retries) increments this counter. +/// +/// The counter is labeled with the cache name only, so it adds up the table route caches of the +/// datanodes and of the frontend of the test process. +fn table_route_cache_loads() -> u64 { + common_meta::metrics::CACHE_CONTAINER_CACHE_MISS + .with_label_values(&[cache::TABLE_ROUTE_CACHE_NAME]) + .get() +} + /// Creates the probe and build tables of the tests and inserts the rows. async fn prepare_tables(frontend: &Arc) { execute_sql( From c069afbeebe19bc4e99539a4b5b717a15986912f Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 21 Sep 2026 20:23:27 +0800 Subject: [PATCH 07/37] docs(query): trim nested plan decoder commentary Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/query_engine/default_serializer.rs | 228 +++--------------- 1 file changed, 30 insertions(+), 198 deletions(-) diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 271961d6ebff..f41b7f545320 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -90,9 +90,7 @@ impl SerializerRegistry for DefaultSerializer { bytes: &[u8], ) -> Result> { if name == MergeScanLogicalPlan::name() { - // TODO(dennis): missing `session_state` to decode the logical plan in `MergeScanLogicalPlan`, - // so we only save the unoptimized logical plan for view currently. - // Plan decoding paths that need to decode a `MergeScan` use [`MergeScanAwareSerializer`]. + // `DefaultSerializer` has no session state; use `MergeScanAwareSerializer` to decode. Err(DataFusionError::Substrait(format!( "Unsupported plan node: {name}" ))) @@ -102,25 +100,11 @@ impl SerializerRegistry for DefaultSerializer { } } -/// Extended [`substrait::extension_serializer::ExtensionSerializer`] that supports both -/// serialization and deserialization of [`MergeScanLogicalPlan`]. -/// -/// [`DefaultSerializer`] is a unit struct that is shared by every encoding path, hence it cannot -/// decode a `MergeScan` node: the plan embedded in the `MergeScan` payload is decoded with a -/// [`SessionState`] which is only known at the decoding site. This registry keeps such a session -/// state and is installed (instead of `DefaultSerializer`) into the session state built by -/// [`DefaultPlanDecoder::decode`], so that plans containing a `MergeScan` can be decoded. +/// [`ExtensionSerializer`] with the session state required to decode [`MergeScanLogicalPlan`]. struct MergeScanAwareSerializer { - /// The session state used to decode the plans embedded in `MergeScan` payloads: the state of - /// the request being decoded, with the functions and the catalog list given to - /// [`DefaultPlanDecoder::decode`]. + /// Request state used to decode nested `MergeScan` payloads. session_state: SessionState, - /// The catalog of the query engine, which knows every table of the cluster. It resolves the - /// tables of a `MergeScan` payload, see [`Self::decode_payload`]. - /// - /// It is `None` when the session state doesn't carry the engine state (a session state built - /// by hand, e.g. in a unit test): the payload is then decoded with the catalog list of the - /// request, which is the behavior of every payload before this field existed. + /// Engine catalog for payload tables; absent for manually built session states. catalog_manager: Option, } @@ -135,7 +119,6 @@ impl fmt::Debug for MergeScanAwareSerializer { } impl MergeScanAwareSerializer { - /// Returns the query context of the session state, or an empty one if it doesn't carry any. fn query_ctx(&self) -> QueryContextRef { self.session_state .config() @@ -143,23 +126,9 @@ impl MergeScanAwareSerializer { .unwrap_or_else(QueryContext::arc) } - /// Returns the session state used to decode a `MergeScan` payload whose tables are resolved - /// with `catalog_list`. - /// - /// The state is rebuilt with this registry, so that a `MergeScan` nested in the payload (e.g. - /// the plan sent to a datanode that has to query another datanode) is decoded recursively. - /// - /// The default catalog and schema of the state are the ones of the request, because the table - /// references of a payload may be bare (a table name as it was written in the query): they are - /// resolved against that pair, which is how they refer to a `(catalog, schema, table)` of the - /// catalog. Decoding the payload with the catalog list of a region aware request used to make - /// that pair irrelevant (that list resolves every catalog and schema to the region); a catalog - /// that resolves the tables of the cluster needs it. + /// Rebuilds state for recursive payload decoding with request catalog/schema defaults. /// - /// The functions of the state are registered once more after it is built, see - /// [`register_greptime_functions`]: building a state may bind a function back to a built-in - /// DataFusion function of the same name or alias, so the payload has to be decoded with a state - /// that holds the GreptimeDB functions, like the state that decodes the enclosing plan. + /// Re-register Greptime functions after each build to preserve bindings over DataFusion aliases. fn payload_state(&self, catalog_list: Arc) -> Result { let query_ctx = self.query_ctx(); let mut config = self.session_state.config().clone(); @@ -182,28 +151,10 @@ impl MergeScanAwareSerializer { Ok(state) } - /// Decodes `payload`, the plan embedded in a `MergeScan`. + /// Decodes payload tables through the engine catalog, not the region-bound request catalog. /// - /// The tables of the payload are the tables that the merge scan reads, i.e. the tables of that - /// merge scan and not the region that the enclosing plan is decoded for: they are resolved by - /// `(catalog, schema, table)` name through the catalog of the query engine, which knows every - /// table of the cluster. - /// - /// The catalog list of the request is deliberately not used for the payload: on a datanode that - /// list is the region aware one (`NameAwareCatalogList`), which resolves *every* table name to - /// the region of the request. A nested `MergeScan` reads another table than the enclosing plan - /// (the build side of a nested join, e.g.), and binding it to the region of the request decodes - /// it with the schema of the table of that region (the columns don't even match when the two - /// tables don't share their column names). The tables of a payload are read by the datanodes - /// that own their regions: `MergeScanExec` sends the plan of the payload to them as is, and - /// they decode it against their own regions. - /// - /// When the session state carries a catalog manager (the query engine does), the payload is - /// decoded with it and a failure is returned as is: falling back to the catalog list of the - /// request would decode the payload with a *different* table binding semantics (on a datanode - /// it binds every table to the region of the request), masking a real catalog, metadata or - /// decoding error behind a wrong schema. The catalog list of the request is only used when no - /// catalog manager is available (e.g. a hand-built session state of a unit test). + /// Do not fall back after an engine-catalog failure: it can bind payload tables to the request + /// region with the wrong schema. Manually built states lack that catalog and use the request list. fn decode_payload(&self, payload: Vec) -> Result { if let Some(catalog_manager) = &self.catalog_manager { let engine_catalog = Arc::new( @@ -224,7 +175,6 @@ impl MergeScanAwareSerializer { impl SerializerRegistry for MergeScanAwareSerializer { fn serialize_logical_plan(&self, node: &dyn UserDefinedLogicalNode) -> Result> { - // Encoding doesn't need any session state, delegate to the default serializer. DefaultSerializer.serialize_logical_plan(node) } @@ -243,13 +193,7 @@ impl SerializerRegistry for MergeScanAwareSerializer { let input = self.decode_payload(merge_scan.input)?; - // `PbMergeScan` doesn't carry `partition_cols`, so the decoded node maps no partition - // column to its aliases (an empty `AliasMapping`). `MergeScanExec` uses it to declare the - // hash partitioning of its output and to keep a hash partitioned input of an upstream - // operator (see `MergeScanExec::try_with_new_distribution`), so an empty mapping only - // loses that optimization (an extra repartition may be inserted by the planner), it - // doesn't change the results. Carrying `partition_cols` in the payload requires a proto - // change and is left to a follow-up. + // `PbMergeScan` lacks `partition_cols`; decoded plans lose this optimization and may repartition. Ok(Arc::new(MergeScanLogicalPlan::new( input, merge_scan.is_placeholder, @@ -258,42 +202,11 @@ impl SerializerRegistry for MergeScanAwareSerializer { } } -/// Decodes `sub_plan`, the plan embedded in a `MergeScan` payload. -/// -/// This sync/async bridge exists because the two sides of it have incompatible shapes (a PoC -/// simplification): +/// Bridges synchronous serializer callbacks to asynchronous payload decoding. /// -/// * `SerializerRegistry::deserialize_logical_plan` is synchronous — it is called from the -/// synchronous `consume_extension_leaf` of the datafusion fork — while decoding a substrait plan -/// is asynchronous. -/// * `consume_extension_leaf` does not treat the input of a `MergeScan` as a subtree it recurses -/// into, so a nested `MergeScan` is only decoded by recursing from here, in our own registry. -/// Each recursion level therefore re-enters this synchronous bridge *from inside the future it -/// is currently driving*, which a plain `block_on` cannot do: the nested executor aborts with -/// `cannot execute LocalPool executor from within another executor: EnterError`. -/// -/// The future must thus be driven on a runtime other than the one the caller is already running -/// in. Driving it at all needs a runtime — not a bare polling loop — because resolving the tables -/// of the payload may read the metadata through the catalog (a remote read on a datanode), which -/// needs a reactor. -/// -/// * On a multi-threaded runtime (the datanode production path) `block_in_place` hands the other -/// tasks of the current worker thread over to another worker, then drives the future on the -/// current runtime; the remaining workers keep serving the I/O the future awaits. -/// * Otherwise (a current-thread runtime, e.g. the default of `#[tokio::test]`, or no runtime at -/// all) there is no other worker to hand the tasks over to, so driving the future here would -/// both starve the runtime and panic as soon as this function is re-entered recursively. The -/// future is moved to a fresh thread that builds its own current-thread runtime instead. That -/// thread is unrelated to whatever the caller runs in, so a recursive call simply spawns another -/// fresh thread; the caller thread only joins and never re-enters its own executor. -/// -/// Scope of that current-thread branch (best effort, not a general guarantee): the production -/// datanode runs on a multi-threaded runtime and uses the `block_in_place` branch above. The -/// current-thread branch exists for environments such as `#[tokio::test]`. It assumes the -/// catalog resolution the future awaits does *not* depend on tasks running on the caller's -/// runtime (the PoC uses in-memory / KV catalogs, for which that holds). `Send`-ness of the -/// future alone does not establish that independence, so this branch makes no safety claim for -/// arbitrary runtime setups; it also offers no cancellation of the synchronous `join`. +/// Nested `MergeScan` payloads recurse through this bridge, so plain nested `block_on` is invalid. +/// Multi-thread runtimes use `block_in_place`; current-thread or absent runtimes use a fresh thread. +/// The latter assumes catalog resolution does not require the caller runtime and has no cancellation. fn decode_sub_plan(sub_plan: Vec, session_state: SessionState) -> Result { let decode = async move { DFLogicalSubstraitConvertor @@ -336,25 +249,10 @@ impl DefaultPlanDecoder { } } -/// Registers every function of GreptimeDB into `session_state`, overwriting any built-in -/// function of the same name or alias. -/// -/// The sub strait decoder looks the functions of a plan up in the session state, and some of the -/// GreptimeDB functions collide with a built-in DataFusion function or one of its aliases: the -/// built-in `to_char` declares `date_format` as an alias, while GreptimeDB has a `date_format` -/// function of its own. +/// Re-registers Greptime functions after every state build. /// -/// That is why the functions have to be registered *after* the session state is built. Building a -/// state re-registers the functions of the state it was given — `SessionStateBuilder::build` -/// takes them from `new_from_existing`, which turns the function maps into vectors — in the -/// unspecified order in which a hash map yields them, and registering a function inserts its -/// aliases *before* its own name. A state that holds both the built-in `to_char` and the GreptimeDB -/// `date_format` therefore binds `date_format` back to the built-in `to_char` whenever it -/// re-registers `to_char` last. -/// -/// The query context must be passed to the functions to set the timezone. This function is -/// idempotent, and it must be applied to the state that actually decodes a plan, i.e. after *every* -/// `SessionStateBuilder::build`. +/// DataFusion aliases can rebind Greptime `date_format` to built-in `to_char`; registration must +/// happen after `SessionStateBuilder::build` on the state that decodes the plan. fn register_greptime_functions( session_state: &mut SessionState, query_ctx: &QueryContextRef, @@ -461,21 +359,12 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { catalog_list: Arc, optimize: bool, ) -> common_query::error::Result { - // The session_state already has the `DefaultSerialzier` as `SerializerRegistry`. It is - // replaced by `MergeScanAwareSerializer` below, once this state is fully built. let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) .with_catalog_list(catalog_list) .build(); register_greptime_functions(&mut session_state, &self.query_ctx)?; - // Install a registry that is also able to decode `MergeScan` nodes (the plans sent to a - // datanode may contain one), so that the sub-plans of a `MergeScan` see the same functions. - // - // The registry resolves the tables of a `MergeScan` payload with the catalog of the query - // engine, which it retrieves from the engine state the query engine leaves in the session - // config (`QueryEngine::engine_context`). A session state that carries no engine state has - // no catalog manager either, in which case the payloads are decoded with the catalog list - // of the request, see `MergeScanAwareSerializer::decode_payload`. + // Payloads need the engine catalog rather than the request's region-bound catalog. let catalog_manager = session_state .config() .get_extension::() @@ -486,10 +375,7 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { catalog_manager, })) .build(); - // Building the state above re-registered the functions of the previous state in hash order, - // which can bind a function back to a built-in DataFusion function of the same name or - // alias: register the GreptimeDB functions once more, so that the state that decodes the - // plan below holds the GreptimeDB bindings. + // Re-register after the build to avoid Greptime UDF alias collisions. register_greptime_functions(&mut session_state, &self.query_ctx)?; let logical_plan = DFLogicalSubstraitConvertor @@ -765,9 +651,7 @@ mod tests { #[tokio::test] async fn test_serializer_decode_merge_scan() { let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); - // The payload scans `numbers`: it has to be resolvable through the catalog of the query - // engine, which is what a datanode uses to decode the payload of a `MergeScan` (the - // catalog list of the request is not a fallback for the engine catalog). + // Payload tables must be registered in the engine catalog. catalog_manager .register_table_sync(RegisterTableRequest { catalog: DEFAULT_CATALOG_NAME.to_string(), @@ -788,8 +672,6 @@ mod tests { ); let engine = factory.query_engine(); - // A `MergeScan` embeds its input plan into its payload, so decoding a `MergeScan` has to - // decode that plan recursively. let input = LogicalPlanBuilder::scan( NUMBERS_TABLE_NAME, Arc::new(LogicalTableSource::new( @@ -810,8 +692,7 @@ mod tests { .engine_context(QueryContext::arc()) .new_plan_decoder() .unwrap(); - // The catalog list of the request is deliberately unrelated to the payload: the payload - // is resolved through the catalog of the query engine, not through this list. + // Must not bind payload tables through the request catalog. let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( mock_table_provider(1.into()), ))); @@ -835,19 +716,10 @@ mod tests { assert!(merge_scan.partition_cols().is_empty()); } - /// The plan embedded in a `MergeScan` is decoded with the catalog of the query engine, which - /// resolves its tables by `(catalog, schema, table)` name, and not with the catalog list of the - /// request. - /// - /// On a datanode the catalog list of the request is the region aware one - /// (`NameAwareCatalogList`), which resolves *every* table name to the region of the request: a - /// payload reading another table is then decoded with the schema of the table of that region, - /// which breaks the plan as soon as the two tables don't share their columns. This test covers - /// that binding with a `numbers` table whose column names differ from the request region's. + /// Payload table binding uses the engine catalog, not the request catalog. #[tokio::test] async fn test_serializer_decode_merge_scan_with_engine_catalog() { let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); - // The table of the payload: `numbers`, with a single `number` column. catalog_manager .register_table_sync(RegisterTableRequest { catalog: DEFAULT_CATALOG_NAME.to_string(), @@ -868,8 +740,6 @@ mod tests { ); let engine = factory.query_engine(); - // The payload scans `numbers` by name, so the catalog that decodes it has to resolve a - // bare table name too. let input = LogicalPlanBuilder::scan( NUMBERS_TABLE_NAME, Arc::new(LogicalTableSource::new( @@ -885,8 +755,7 @@ mod tests { .encode(&plan, DefaultSerializer) .unwrap(); - // The catalog list of the request resolves every table name to a provider whose schema is - // not the one of `numbers` (the region aware list of a datanode does the same). + // This request catalog deliberately has a different schema. let table_provider = Arc::new(mock_table_provider(1.into())); let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider)); @@ -910,20 +779,15 @@ mod tests { let LogicalPlan::TableScan(scan) = merge_scan.input() else { panic!("Expect a table scan, got: {}", merge_scan.input()); }; - // The scan is the `numbers` table of the catalog of the query engine, not the table of the - // request catalog list. assert_eq!(scan.table_name.table(), NUMBERS_TABLE_NAME); assert_eq!(scan.source.schema().fields().len(), 1); assert_eq!(scan.source.schema().field(0).name(), "number"); } - /// A `MergeScan` nested in the payload of another `MergeScan` is decoded recursively: this is - /// the shape of a datanode querying another datanode. + /// Nested `MergeScan` payloads decode recursively. #[tokio::test] async fn test_serializer_decode_nested_merge_scan() { let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); - // The payload of both `MergeScan` levels scans `numbers`, which has to be resolvable - // through the catalog of the query engine (the request catalog list is not a fallback). catalog_manager .register_table_sync(RegisterTableRequest { catalog: DEFAULT_CATALOG_NAME.to_string(), @@ -973,7 +837,6 @@ mod tests { .engine_context(QueryContext::arc()) .new_plan_decoder() .unwrap(); - // The catalog list of the request is deliberately unrelated to the payload. let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( mock_table_provider(1.into()), ))); @@ -984,12 +847,9 @@ mod tests { .unwrap(); assert_eq!(decoded.to_string(), plan.to_string()); - // Both levels are `MergeScan` nodes. assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); } - /// Returns the `date_format` implementation of GreptimeDB: the implementation a decoded plan - /// has to use. fn greptime_date_format_udf(query_ctx: &QueryContextRef) -> Arc { FUNCTION_REGISTRY .get_function("date_format") @@ -1001,16 +861,13 @@ mod tests { .into() } - /// The query context of the tests below: the timezone is not the default one, so that the - /// implementation bound to `date_format` is observable (only the GreptimeDB implementation - /// formats in the timezone of the query context). + /// Uses a non-default timezone to distinguish the Greptime implementation. fn date_format_query_ctx() -> QueryContextRef { let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()); Arc::new(query_ctx) } - /// A `date_format(timestamp, format)` expression that uses the GreptimeDB implementation. fn greptime_date_format_expr(query_ctx: &QueryContextRef, timestamp: Expr) -> Expr { Expr::ScalarFunction(ScalarFunction { func: greptime_date_format_udf(query_ctx), @@ -1018,9 +875,7 @@ mod tests { }) } - /// Returns the scalar functions of every expression of `plan`, including the expressions of its - /// inputs and the ones of the plan embedded in a `MergeScan` (the payload of a `MergeScan` is - /// not an input of the node, see `MergeScanLogicalPlan::inputs`). + /// Includes `MergeScan` payload expressions, which are not node inputs. fn scalar_functions_of(plan: &LogicalPlan) -> Vec> { use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; @@ -1051,12 +906,7 @@ mod tests { udfs } - /// Asserts that the only function of `plan` is the `date_format` of GreptimeDB. - /// - /// The built-in `to_char` of DataFusion declares `date_format` as an alias, so a plan where - /// `date_format` was bound to it holds a `to_char` function, which the substrait encoder - /// renders as `to_char(...) AS date_format(...)`: the binding itself is asserted here, and not - /// the way the expression is printed. + /// Asserts the Greptime `date_format` binding rather than DataFusion's `to_char` alias. fn assert_only_function_is_greptime_date_format( plan: &LogicalPlan, query_ctx: &QueryContextRef, @@ -1066,10 +916,6 @@ mod tests { let udf = &udfs[0]; assert_eq!(udf.name(), "date_format", "got: {plan}"); - // The built-in `to_char` of DataFusion declares `date_format` as an alias, so a plan - // where `date_format` was bound to it holds a `to_char` function (rendered - // `to_char(...) AS date_format(...)` by the substrait encoder). Its instance can never - // equal the implementation of GreptimeDB, which the assertion below checks. assert_eq!( udf.as_ref(), greptime_date_format_udf(query_ctx).as_ref(), @@ -1077,17 +923,10 @@ mod tests { ); } - /// The built-in `to_char` of DataFusion declares `date_format` as an alias, and - /// [`SessionStateBuilder`] re-registers the functions of a state in the (unspecified) order in - /// which a hash map yields them, so building the session state can bind `date_format` back to - /// the built-in `to_char` instead of the GreptimeDB implementation. Decoding a plan has to keep - /// the GreptimeDB binding, on the plain path (no `MergeScan` to decode) as well. + /// Rebuilding state must retain the Greptime `date_format` binding. #[tokio::test] async fn test_serializer_decode_keeps_greptime_date_format() { - // The functions of a session state are collected from a hash map, and the map order differs - // between the states (and between the decodings, which rebuild the state): decode the same - // plan several times, with a fresh engine each time, so that an order dependent binding is - // caught whatever the order of a single state is. + // Repeat to expose nondeterministic function-registration order. for _ in 0..8 { let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); let factory = QueryEngineFactory::new( @@ -1134,14 +973,10 @@ mod tests { } } - /// Same as [`test_serializer_decode_keeps_greptime_date_format`], for a `date_format` inside the - /// payload of a `MergeScan`: the payload is decoded with the session state built by - /// [`MergeScanAwareSerializer::payload_state`], which is rebuilt as well and thus has to keep - /// the GreptimeDB functions too. + /// Payload-state rebuilding must retain the Greptime `date_format` binding. #[tokio::test] async fn test_serializer_decode_merge_scan_payload_keeps_greptime_date_format() { let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); - // The payload scans `numbers`, which is resolved through the catalog of the query engine. catalog_manager .register_table_sync(RegisterTableRequest { catalog: DEFAULT_CATALOG_NAME.to_string(), @@ -1165,8 +1000,6 @@ mod tests { let engine = factory.query_engine(); let query_ctx = date_format_query_ctx(); - // `numbers` has a single `number` column of type `u32`: cast it to a timestamp, which - // is what `date_format` formats. let input = LogicalPlanBuilder::scan( NUMBERS_TABLE_NAME, Arc::new(LogicalTableSource::new( @@ -1195,7 +1028,6 @@ mod tests { .engine_context(query_ctx.clone()) .new_plan_decoder() .unwrap(); - // The catalog list of the request is deliberately unrelated to the payload. let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( mock_table_provider(1.into()), ))); From 13168e5cb7690694b7c2b0f6224e55ffb68b07b4 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Wed, 23 Sep 2026 14:38:36 +0800 Subject: [PATCH 08/37] refactor: remove redundant nested scan support plumbing Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cache/src/lib.rs | 175 ++++++-------- src/datanode/src/region_query.rs | 12 +- .../src/query_engine/default_serializer.rs | 3 +- .../nested_merge_scan_capability_test.rs | 213 +++--------------- 4 files changed, 97 insertions(+), 306 deletions(-) diff --git a/src/cache/src/lib.rs b/src/cache/src/lib.rs index 4fd9ae5fcc13..ff9106cbbcbe 100644 --- a/src/cache/src/lib.rs +++ b/src/cache/src/lib.rs @@ -18,16 +18,15 @@ use std::hash::Hash; use std::sync::Arc; use std::time::Duration; -use catalog::kvbackend::{TableCacheRef, new_table_cache}; +use catalog::kvbackend::new_table_cache; use common_meta::cache::{ CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder, - SchemaCacheRef, TableInfoCacheRef, TableNameCacheRef, TableRouteCacheRef, TableSchemaCacheRef, 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; use moka::future::{Cache, CacheBuilder}; -use partition::cache::{PartitionInfoCacheRef, new_partition_info_cache}; +use partition::cache::new_partition_info_cache; use snafu::OptionExt; use crate::error::Result; @@ -54,114 +53,6 @@ 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"; -/// The caches of a datanode. -/// -/// The `table_cache` is derived from the `table_info_cache` and the `table_name_cache`, and the -/// `partition_info_cache` is derived from the `table_route_cache`. -struct DatanodeCaches { - table_id_schema_cache: TableSchemaCacheRef, - schema_cache: SchemaCacheRef, - table_info_cache: TableInfoCacheRef, - table_name_cache: TableNameCacheRef, - table_route_cache: TableRouteCacheRef, - table_cache: TableCacheRef, - partition_info_cache: PartitionInfoCacheRef, -} - -impl DatanodeCaches { - /// Builds the caches of a datanode. - /// - /// The table/route caches are needed by the catalog manager and by the partition rule manager - /// of the datanode, which are used to plan the `MergeScan` nodes of the plans received from the - /// frontend (e.g. a nested merge scan that queries another datanode). - fn new(kv_backend: KvBackendRef) -> Self { - // Builds table id schema name cache that never expires. - 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(), - cache, - kv_backend.clone(), - )); - - // Builds schema cache - let cache = default_cache(); - let schema_cache = Arc::new(new_schema_cache( - SCHEMA_CACHE_NAME.to_string(), - cache, - kv_backend.clone(), - )); - - // Builds table info cache - let cache = default_cache(); - let table_info_cache = Arc::new(new_table_info_cache( - TABLE_INFO_CACHE_NAME.to_string(), - cache, - kv_backend.clone(), - )); - - // Builds table name cache - let cache = default_cache(); - let table_name_cache = Arc::new(new_table_name_cache( - TABLE_NAME_CACHE_NAME.to_string(), - cache, - kv_backend.clone(), - )); - - // Builds table cache - 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(), - )); - - // Builds table route cache - let cache = default_cache(); - let table_route_cache = Arc::new(new_table_route_cache( - TABLE_ROUTE_CACHE_NAME.to_string(), - cache, - kv_backend.clone(), - )); - - // Builds partition info cache - 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(), - )); - - Self { - table_id_schema_cache, - schema_cache, - table_info_cache, - table_name_cache, - table_route_cache, - table_cache, - partition_info_cache, - } - } - - /// Splits the caches into the layer of the caches that are read from the kv backend, and the - /// layer of the caches derived from them. - fn into_layers(self) -> (CacheRegistry, CacheRegistry) { - let base_registry = CacheRegistryBuilder::default() - .add_cache(self.table_id_schema_cache) - .add_cache(self.schema_cache) - .add_cache(self.table_info_cache) - .add_cache(self.table_name_cache) - .add_cache(self.table_route_cache) - .build(); - let derived_registry = CacheRegistryBuilder::default() - .add_cache(self.table_cache) - .add_cache(self.partition_info_cache) - .build(); - - (base_registry, derived_registry) - } -} - /// Builds the layered cache registry for datanode. /// /// The registry holds every cache a datanode needs: @@ -186,7 +77,67 @@ impl DatanodeCaches { /// derived cache until the entry expires. This is the only datanode cache registry: production code /// and test infrastructure build the same one. pub fn build_datanode_layered_cache_registry(kv_backend: KvBackendRef) -> LayeredCacheRegistry { - let (base_registry, derived_registry) = DatanodeCaches::new(kv_backend).into_layers(); + 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(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let schema_cache = Arc::new(new_schema_cache( + SCHEMA_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + 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) diff --git a/src/datanode/src/region_query.rs b/src/datanode/src/region_query.rs index d22a93457b9e..aa750b060416 100644 --- a/src/datanode/src/region_query.rs +++ b/src/datanode/src/region_query.rs @@ -37,14 +37,10 @@ use store_api::storage::RegionId; /// Serves the region queries issued by the datanode itself. /// -/// Known limitation: [`RegionQueryHandler::select_target`] resolves the leader of a region through -/// the table route cache, which uses `InitStrategy::VersionChecked`. The version counter of that -/// strategy belongs to the `CacheContainer` and not to a single table: any `TableId` invalidation -/// bumps the shared version, so a cold load of one table's route retries (and loads again) when an -/// *unrelated* table is invalidated concurrently. Clusters with frequent DDL or region movement can -/// therefore see amplified cold-load retries on the datanode; the tail latency of cold loads in -/// such clusters is worth verifying. The retries only add work, they keep the loaded route correct -/// (a stale route is never returned). +/// Known limitation: [`RegionQueryHandler::select_target`] resolves a region leader through the +/// version-checked table route cache, whose version is shared by the `CacheContainer`. An +/// unrelated `TableId` invalidation can therefore add route loads and cold-load latency; sustained +/// conflicts may exhaust the retry budget and fail the region query. pub struct DatanodeRegionQueryHandler { partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index f41b7f545320..211311c825a6 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -359,10 +359,9 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { catalog_list: Arc, optimize: bool, ) -> common_query::error::Result { - let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) + let session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) .with_catalog_list(catalog_list) .build(); - register_greptime_functions(&mut session_state, &self.query_ctx)?; // Payloads need the engine catalog rather than the request's region-bound catalog. let catalog_manager = session_state diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index f2e92d6428cd..4e4f36820ff0 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -55,19 +55,16 @@ //! executes: the inner region queries of the nested plans are admitted as execution stages of the //! outer queries, so they don't wait for the permit that their own query holds (see //! [`test_nested_merge_scan_capability_internal_stage_under_concurrency_limit`]), -//! * the inner `MergeScan` resolves the leader of the regions of the build table through the table -//! route cache of the datanode, whose version counter is shared by every table of the container: -//! a cold load of the route of the build table still returns the rows of the join while an -//! unrelated table id is invalidated in a loop (see +//! * a cold load of the route of the build table returns the rows of the join after an unrelated +//! table id is invalidated once (see //! [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]). use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; use api::v1::region::{QueryRequest, RegionRequestHeader}; -use common_meta::cache::{LayeredCacheRegistry, TableRouteCacheRef}; +use common_meta::cache::TableRouteCacheRef; use common_meta::cache_invalidator::{CacheInvalidator, Context}; use common_meta::instruction::CacheIdent; use common_query::Output; @@ -98,21 +95,10 @@ const BUILD_TABLE: &str = "nested_cap_build"; /// plan cannot be resolved. const MISSING_INNER_TABLE: &str = "nested_cap_missing_build"; -/// The table id that no table of the cluster uses. The test of the cold route load invalidates this -/// table id while it cold-loads the route of [`BUILD_TABLE`]: the version counter of the version -/// checked table route cache belongs to the cache container and not to a single table, so this -/// unrelated invalidation bumps the version that the cold load compares. +/// The table id that no table of the cluster uses. The cold-load test invalidates this unrelated +/// table once before it executes nested queries for [`BUILD_TABLE`]. const UNRELATED_TABLE_ID: TableId = 999_999; -/// The interval of the unrelated invalidation loop of -/// [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]. -/// -/// Short enough to hit an in-flight load of the table route cache, long enough to leave the retry -/// loop of the version checked load room to finish (a retry sleeps for at least `10ms`): the test -/// must fail because a route was stale or belonged to another table, not because the retry loop of -/// the cache gave up. -const UNRELATED_INVALIDATION_INTERVAL: Duration = Duration::from_millis(20); - /// `(a_id, probe_key, probe_v)` of the probe table: `a_id` is the primary key and the partition /// column, `probe_key` is the join key and `probe_v` is the value of the row. /// @@ -257,27 +243,8 @@ async fn test_nested_merge_scan_capability_join_across_datanodes() { assert_cross_datanode_region_query(&cluster, &expectation, &baseline_requests); } -/// The route of the build table is loaded cold while an *unrelated* table id is invalidated: the -/// version counter of the version checked table route cache belongs to the cache container and not -/// to a single table, so an invalidation of any `TableId` bumps the version that a cold load -/// compares, and a load that runs while an unrelated table is invalidated retries (retrying is what -/// keeps the loaded route correct; see the known limitation of `DatanodeRegionQueryHandler`). -/// -/// The invalidation of the unrelated table id is injected into the layered cache registry of every -/// datanode, which is exactly the local handling path of the invalidation instruction that a -/// datanode receives from metasrv; the broadcast to the datanode is the only part not exercised -/// here. -/// -/// The test pins the correctness of the retry path. The route of the build table is dropped on -/// every datanode (so the nested plans really load it from metasrv, the assertion below checks the -/// cache is cold), then every probe region executes the nested plan while a background task -/// invalidates the unrelated table id in a loop. The rows must be the rows of the join in every -/// round: a retried load must neither return the route of another table nor a stale route. -/// -/// The test also reports how often the table route caches loaded a route with and without the -/// unrelated invalidation. A version checked load counts every attempt, so the amplification of the -/// unrelated invalidation is visible in the counter; the test logs the numbers instead of asserting -/// them, because the unrelated invalidation only hits an in-flight load with some probability. +/// The route of the build table is cold before the nested queries run. An unrelated table id is +/// invalidated once before those queries, and the nested plan must still return the join rows. #[tokio::test(flavor = "multi_thread")] async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation() { common_telemetry::init_default_ut_logging(); @@ -309,13 +276,7 @@ async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation ); let build_table_id = table_id(&build_leaders); - info!( - "the regions of {BUILD_TABLE} belong to the table id {build_table_id}; the test \ - invalidates the unrelated table id {UNRELATED_TABLE_ID} while the datanodes load the \ - route of {BUILD_TABLE}" - ); - // The reference join of the frontend and the nested plan, as in the other capability tests. let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; info!("join result of the frontend:\n{reference}"); assert_eq!( @@ -327,99 +288,25 @@ async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation let query_ctx = query_ctx(); let plan = encode(&nested_plan(&frontend, &query_ctx).await); - // 1. The baseline: the route of the build table is dropped on every datanode and nothing - // invalidates a table while the nested plans run, so both datanodes load the route of the build - // table from metasrv exactly once. invalidate_table_ids(&cluster, &[build_table_id]).await; - assert_route_cold(&cluster, build_table_id).await; + assert_route_cold(&cluster, build_table_id); + invalidate_table_ids(&cluster, &[UNRELATED_TABLE_ID]).await; - let loads_before = table_route_cache_loads(); - let rows_cold = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; - let cold_loads = table_route_cache_loads() - loads_before; + let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; assert_eq!( expected_join_rows(), - multiset(rows_cold.clone()), + multiset(rows), "the nested plans returned an unexpected multiset of rows after the route of {BUILD_TABLE} \ - was dropped: {rows_cold:?}" - ); - // Every datanode owns a probe region, so every datanode resolved the leader of the regions of - // the build table with a cold route cache: the baseline really started without the route. - assert!( - cold_loads >= 2, - "expected the datanodes to load the dropped route of {BUILD_TABLE} from metasrv, but the \ - table route cache only loaded a route {cold_loads} time(s)" + was dropped and unrelated table id {UNRELATED_TABLE_ID} was invalidated" ); - // 2. The unrelated invalidation: a background task invalidates the table id of a table that the - // cluster does not have, the way the datanodes handle the invalidation instruction of a DDL - // that touched some other table. Every round drops the route of the build table again, so the - // nested plans load a cold route while the unrelated invalidation runs. - let stop = Arc::new(AtomicBool::new(false)); - let unrelated_invalidations = Arc::new(AtomicUsize::new(0)); - let invalidation_task = tokio::spawn(invalidate_unrelated_table_id( - cluster - .datanode_cache_registries - .values() - .cloned() - .collect::>>(), - stop.clone(), - unrelated_invalidations.clone(), - )); - - let mut loads_with_invalidation = 0; - for round in 0..3 { - invalidate_table_ids(&cluster, &[build_table_id]).await; - assert_route_cold(&cluster, build_table_id).await; - - let loads_before = table_route_cache_loads(); - let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; - loads_with_invalidation += table_route_cache_loads() - loads_before; - - assert_eq!( - expected_join_rows(), - multiset(rows.clone()), - "round {round}: the nested plans returned an unexpected multiset of rows while the \ - unrelated table id {UNRELATED_TABLE_ID} was invalidated: {rows:?}" - ); - assert_eq!( - multiset(rows_cold.clone()), - multiset(rows.clone()), - "round {round}: the nested plans returned different rows while the unrelated table id \ - {UNRELATED_TABLE_ID} was invalidated, i.e. a retried route load returned another \ - route than the cold load" + for datanode in datanodes(&probe_leaders) { + assert!( + table_route_cache(&cluster, datanode).contains_key(&build_table_id), + "the datanode {datanode} must hold the route of {BUILD_TABLE} after it executed the \ + nested plan of its probe region" ); - - // The retry loop of a route load that the unrelated invalidation interrupted must have - // loaded the route of the build table: the retry must not leave the cache without the - // route of the table that the nested plan asked for. - for datanode in datanodes(&probe_leaders) { - assert!( - table_route_cache(&cluster, datanode).contains_key(&build_table_id), - "round {round}: the datanode {datanode} must hold the route of {BUILD_TABLE} after \ - it executed the nested plan of its probe region" - ); - } } - - stop.store(true, Ordering::Relaxed); - invalidation_task - .await - .expect("the task that invalidates the unrelated table id must stop"); - - // Best effort measurement, not an assertion: the unrelated invalidation retries a cold load - // only if it lands in the window between the version the load read and the version it compares - // afterwards, so the counts vary with the timing of the machine. The counter is the - // `CACHE_CONTAINER_CACHE_MISS` of the cache container, which every attempt of a version checked - // load increments; the cache name is shared by the datanodes and the frontend of the test - // process, but the frontend is idle while the nested plans of a round run. - let unrelated_invalidations = unrelated_invalidations.load(Ordering::Relaxed); - info!( - "the table route caches loaded a cold route of {BUILD_TABLE} {cold_loads} time(s) without \ - an unrelated invalidation and {loads_with_invalidation} time(s) while the unrelated table \ - id {UNRELATED_TABLE_ID} was invalidated {unrelated_invalidations} time(s); a count above 2 \ - per round means the shared version counter of the cache container made the cold loads \ - retry" - ); } /// A failing inner region makes the whole query fail: the datanode must not return the rows of the @@ -942,60 +829,18 @@ async fn invalidate_table_ids(cluster: &GreptimeDbCluster, table_ids: &[TableId] } } -/// Invalidates [`UNRELATED_TABLE_ID`] on `registries` until `stop` is set, and counts the rounds. -async fn invalidate_unrelated_table_id( - registries: Vec>, - stop: Arc, - invalidations: Arc, -) { - while !stop.load(Ordering::Relaxed) { - for registry in ®istries { - registry - .invalidate( - &Context::default(), - &[CacheIdent::TableId(UNRELATED_TABLE_ID)], - ) - .await - .expect("the datanode must invalidate the unrelated table id"); - } - invalidations.fetch_add(1, Ordering::Relaxed); - tokio::time::sleep(UNRELATED_INVALIDATION_INTERVAL).await; - } -} - -/// Waits until no datanode holds the route of `table_id` any more, i.e. until the route is cold. -/// -/// The invalidation is awaited before this call, so the wait is a safety net for the internals of -/// the cache (e.g. its pending tasks) rather than a required step. -async fn assert_route_cold(cluster: &GreptimeDbCluster, table_id: TableId) { - for attempt in 0..100 { - let mut cached = Vec::new(); - for datanode_id in cluster.datanode_cache_registries.keys() { - if table_route_cache(cluster, *datanode_id).contains_key(&table_id) { - cached.push(*datanode_id); - } - } - if cached.is_empty() { - return; - } - assert!( - attempt < 99, - "the datanodes {cached:?} must drop the route of the table {table_id} after the \ - invalidation" - ); - tokio::time::sleep(Duration::from_millis(10)).await; - } -} - -/// The number of attempts that the table route caches of the process made to load a route: every -/// attempt of a version checked load (including its retries) increments this counter. -/// -/// The counter is labeled with the cache name only, so it adds up the table route caches of the -/// datanodes and of the frontend of the test process. -fn table_route_cache_loads() -> u64 { - common_meta::metrics::CACHE_CONTAINER_CACHE_MISS - .with_label_values(&[cache::TABLE_ROUTE_CACHE_NAME]) - .get() +/// Asserts that no datanode holds the route of `table_id`, i.e. the route is cold. +fn assert_route_cold(cluster: &GreptimeDbCluster, table_id: TableId) { + let cached = cluster + .datanode_cache_registries + .keys() + .filter(|datanode_id| table_route_cache(cluster, **datanode_id).contains_key(&table_id)) + .collect::>(); + assert!( + cached.is_empty(), + "the datanodes {cached:?} must drop the route of the table {table_id} after the \ + invalidation" + ); } /// Creates the probe and build tables of the tests and inserts the rows. From 9ab5e6217a854b1bdb5765c14947ef68267d8d21 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sun, 4 Oct 2026 16:29:26 +0800 Subject: [PATCH 09/37] fix(query): decode nested MergeScan payloads asynchronously Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/common/substrait/src/df_substrait.rs | 19 +- src/common/substrait/src/lib.rs | 1 + .../src/query_engine/default_serializer.rs | 880 ++++++++++++++++-- 3 files changed, 802 insertions(+), 98 deletions(-) diff --git a/src/common/substrait/src/df_substrait.rs b/src/common/substrait/src/df_substrait.rs index 08eeedd8316f..9570076ee7a7 100644 --- a/src/common/substrait/src/df_substrait.rs +++ b/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 { + 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, diff --git a/src/common/substrait/src/lib.rs b/src/common/substrait/src/lib.rs index 3eeb3c570e10..a6932451a62b 100644 --- a/src/common/substrait/src/lib.rs +++ b/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}; diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 211311c825a6..23b69abaaf2e 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::fmt; use std::sync::Arc; use bytes::Bytes; @@ -22,14 +21,14 @@ use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; use common_query::error::RegisterUdfSnafu; use common_query::logical_plan::SubstraitPlanDecoder; -use datafusion::catalog::CatalogProviderList; -use datafusion::common::DataFusionError; +use datafusion::catalog::{CatalogProviderList, TableProvider}; +use datafusion::common::{DFSchema, DataFusionError, TableReference, not_impl_err, substrait_err}; use datafusion::error::Result; use datafusion::execution::context::SessionState; use datafusion::execution::registry::SerializerRegistry; use datafusion::execution::{FunctionRegistry, SessionStateBuilder}; -use datafusion::logical_expr::LogicalPlan; -use datafusion_expr::UserDefinedLogicalNode; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion_expr::{Expr, UserDefinedLogicalNode}; use greptime_proto::substrait_extension::MergeScan as PbMergeScan; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, @@ -49,8 +48,13 @@ use promql::functions::{ use prost::Message; use session::context::{QueryContext, QueryContextRef}; use snafu::ResultExt; +use substrait::df_logical_plan::consumer::{DefaultSubstraitConsumer, SubstraitConsumer}; +use substrait::error::{DecodeDfPlanSnafu, DecodeRelSnafu}; use substrait::extension_serializer::ExtensionSerializer; -use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; +use substrait::substrait_proto_df::proto::{ + ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, Plan, Type, +}; +use substrait::{DFLogicalSubstraitConvertor, Extensions, SubstraitPlan}; use crate::dist_plan::MergeScanLogicalPlan; use crate::query_engine::QueryEngineState; @@ -90,7 +94,7 @@ impl SerializerRegistry for DefaultSerializer { bytes: &[u8], ) -> Result> { if name == MergeScanLogicalPlan::name() { - // `DefaultSerializer` has no session state; use `MergeScanAwareSerializer` to decode. + // This registry has no session state; the query engine decodes `MergeScan` payloads. Err(DataFusionError::Substrait(format!( "Unsupported plan node: {name}" ))) @@ -100,25 +104,16 @@ impl SerializerRegistry for DefaultSerializer { } } -/// [`ExtensionSerializer`] with the session state required to decode [`MergeScanLogicalPlan`]. -struct MergeScanAwareSerializer { - /// Request state used to decode nested `MergeScan` payloads. - session_state: SessionState, +/// Async `MergeScan` payload decoding over DataFusion's default Substrait consumer. +struct MergeScanSubstraitConsumer<'a> { + inner: DefaultSubstraitConsumer<'a>, + /// State of the level being decoded; payload states are rebuilt from it. + session_state: &'a SessionState, /// Engine catalog for payload tables; absent for manually built session states. catalog_manager: Option, } -impl fmt::Debug for MergeScanAwareSerializer { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // `dyn CatalogManager` is not `Debug`. - f.debug_struct("MergeScanAwareSerializer") - .field("session_state", &self.session_state) - .field("has_catalog_manager", &self.catalog_manager.is_some()) - .finish() - } -} - -impl MergeScanAwareSerializer { +impl MergeScanSubstraitConsumer<'_> { fn query_ctx(&self) -> QueryContextRef { self.session_state .config() @@ -126,9 +121,7 @@ impl MergeScanAwareSerializer { .unwrap_or_else(QueryContext::arc) } - /// Rebuilds state for recursive payload decoding with request catalog/schema defaults. - /// - /// Re-register Greptime functions after each build to preserve bindings over DataFusion aliases. + /// Rebuilds payload state with request defaults and re-registers Greptime function bindings. fn payload_state(&self, catalog_list: Arc) -> Result { let query_ctx = self.query_ctx(); let mut config = self.session_state.config().clone(); @@ -140,10 +133,6 @@ impl MergeScanAwareSerializer { let mut state = SessionStateBuilder::new_from_existing(self.session_state.clone()) .with_config(config) - .with_serializer_registry(Arc::new(Self { - session_state: self.session_state.clone(), - catalog_manager: self.catalog_manager.clone(), - })) .with_catalog_list(catalog_list) .build(); register_greptime_functions(&mut state, &query_ctx).map_err(DataFusionError::from)?; @@ -155,82 +144,156 @@ impl MergeScanAwareSerializer { /// /// Do not fall back after an engine-catalog failure: it can bind payload tables to the request /// region with the wrong schema. Manually built states lack that catalog and use the request list. - fn decode_payload(&self, payload: Vec) -> Result { - if let Some(catalog_manager) = &self.catalog_manager { + async fn decode_payload(&self, payload: Vec) -> Result { + let state = if let Some(catalog_manager) = &self.catalog_manager { let engine_catalog = Arc::new( catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( catalog_manager.clone(), self.query_ctx(), ), ); - return decode_sub_plan(payload, self.payload_state(engine_catalog)?); - } + self.payload_state(engine_catalog)? + } else { + self.payload_state(self.session_state.catalog_list().clone())? + }; - decode_sub_plan( - payload, - self.payload_state(self.session_state.catalog_list().clone())?, - ) + decode_plan(Bytes::from(payload), state) + .await + .map_err(|e| DataFusionError::External(Box::new(e))) } } -impl SerializerRegistry for MergeScanAwareSerializer { - fn serialize_logical_plan(&self, node: &dyn UserDefinedLogicalNode) -> Result> { - DefaultSerializer.serialize_logical_plan(node) +#[async_trait::async_trait] +impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> { + async fn resolve_table_ref( + &self, + table_ref: &TableReference, + ) -> Result>> { + self.inner.resolve_table_ref(table_ref).await } - fn deserialize_logical_plan( + fn get_extensions(&self) -> &Extensions { + self.inner.get_extensions() + } + + fn get_function_registry(&self) -> &impl FunctionRegistry { + self.inner.get_function_registry() + } + + fn push_outer_schema(&self, schema: Arc) { + self.inner.push_outer_schema(schema) + } + + fn pop_outer_schema(&self) { + self.inner.pop_outer_schema() + } + + fn get_outer_schema(&self, steps_out: usize) -> Option> { + self.inner.get_outer_schema(steps_out) + } + + fn push_lambda_parameters( &self, - name: &str, - bytes: &[u8], - ) -> Result> { - if name != MergeScanLogicalPlan::name() { - return DefaultSerializer.deserialize_logical_plan(name, bytes); + lambda_parameters: &[Type], + input_schema: &DFSchema, + ) -> Result> { + self.inner + .push_lambda_parameters(lambda_parameters, input_schema) + } + + fn pop_lambda_parameters(&self) { + self.inner.pop_lambda_parameters() + } + + fn lambda_variable(&self, steps_out: usize, field_idx: usize) -> Result { + self.inner.lambda_variable(steps_out, field_idx) + } + + async fn consume_extension_leaf(&self, rel: &ExtensionLeafRel) -> Result { + // Unknown and detail-less leaves keep the default consumer's errors. + let Some(detail) = rel.detail.as_ref() else { + return self.inner.consume_extension_leaf(rel).await; + }; + if detail.type_url != MergeScanLogicalPlan::name() { + return self.inner.consume_extension_leaf(rel).await; } - let merge_scan = PbMergeScan::decode(bytes).map_err(|e| { + let merge_scan = PbMergeScan::decode(detail.value.as_ref()).map_err(|e| { DataFusionError::Substrait(format!("Failed to decode the MergeScan plan node: {e}")) })?; - let input = self.decode_payload(merge_scan.input)?; + let input = self.decode_payload(merge_scan.input).await?; // `PbMergeScan` lacks `partition_cols`; decoded plans lose this optimization and may repartition. - Ok(Arc::new(MergeScanLogicalPlan::new( - input, - merge_scan.is_placeholder, - Default::default(), - ))) + Ok( + MergeScanLogicalPlan::new(input, merge_scan.is_placeholder, Default::default()) + .into_logical_plan(), + ) } -} -/// Bridges synchronous serializer callbacks to asynchronous payload decoding. -/// -/// Nested `MergeScan` payloads recurse through this bridge, so plain nested `block_on` is invalid. -/// Multi-thread runtimes use `block_in_place`; current-thread or absent runtimes use a fresh thread. -/// The latter assumes catalog resolution does not require the caller runtime and has no cancellation. -fn decode_sub_plan(sub_plan: Vec, session_state: SessionState) -> Result { - let decode = async move { - DFLogicalSubstraitConvertor - .decode(Bytes::from(sub_plan), session_state) - .await - .map_err(|e| DataFusionError::External(Box::new(e))) - }; + // Keep child dispatch on this consumer; delegating the parent bypasses async MergeScan decoding. + async fn consume_extension_single(&self, rel: &ExtensionSingleRel) -> Result { + let Some(detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionSingleRel"); + }; + let plan = self + .session_state + .serializer_registry() + .deserialize_logical_plan(&detail.type_url, &detail.value)?; + let Some(input_rel) = &rel.input else { + return substrait_err!( + "ExtensionSingleRel missing input rel, try using ExtensionLeafRel instead" + ); + }; + let input_plan = self.consume_rel(input_rel).await?; + let plan = plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?; + Ok(LogicalPlan::Extension(Extension { node: plan })) + } - match tokio::runtime::Handle::try_current() { - Ok(handle) if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread => { - tokio::task::block_in_place(|| handle.block_on(decode)) + // Multi-input extension children likewise need this consumer's dispatch. + async fn consume_extension_multi(&self, rel: &ExtensionMultiRel) -> Result { + let Some(detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionMultiRel"); + }; + let plan = self + .session_state + .serializer_registry() + .deserialize_logical_plan(&detail.type_url, &detail.value)?; + let mut inputs = Vec::with_capacity(rel.inputs.len()); + for input in &rel.inputs { + let input_plan = self.consume_rel(input).await?; + inputs.push(input_plan); } - _ => std::thread::spawn(move || { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .map_err(|e| DataFusionError::External(Box::new(e)))?; - runtime.block_on(decode) - }) - .join() - .unwrap_or_else(|panic| std::panic::resume_unwind(panic)), + let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; + Ok(LogicalPlan::Extension(Extension { node: plan })) } } +/// Decode with async MergeScan handling, preserving the default consumer's extension checks. +async fn decode_plan(bytes: Bytes, state: SessionState) -> substrait::error::Result { + let plan = Plan::decode(bytes).context(DecodeRelSnafu)?; + let extensions = Extensions::try_from(&plan.extensions).context(DecodeDfPlanSnafu)?; + if !extensions.type_variations.is_empty() { + return not_impl_err!("Type variation extensions are not supported") + .context(DecodeDfPlanSnafu); + } + + // Payloads need the engine catalog rather than the request's region-bound catalog. + let catalog_manager = state + .config() + .get_extension::() + .map(|engine_state| engine_state.catalog_manager().clone()); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, &state), + session_state: &state, + catalog_manager, + }; + + DFLogicalSubstraitConvertor + .decode_with_consumer(&plan, &state, &consumer) + .await +} + /// The datafusion `[LogicalPlan]` decoder. pub struct DefaultPlanDecoder { session_state: SessionState, @@ -359,26 +422,13 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { catalog_list: Arc, optimize: bool, ) -> common_query::error::Result { - let session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) + let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) .with_catalog_list(catalog_list) .build(); - - // Payloads need the engine catalog rather than the request's region-bound catalog. - let catalog_manager = session_state - .config() - .get_extension::() - .map(|engine_state| engine_state.catalog_manager().clone()); - let mut session_state = SessionStateBuilder::new_from_existing(session_state.clone()) - .with_serializer_registry(Arc::new(MergeScanAwareSerializer { - session_state, - catalog_manager, - })) - .build(); // Re-register after the build to avoid Greptime UDF alias collisions. register_greptime_functions(&mut session_state, &self.query_ctx)?; - let logical_plan = DFLogicalSubstraitConvertor - .decode(message, session_state) + let logical_plan = decode_plan(message, session_state) .await .map_err(BoxedError::new) .context(common_query::error::DecodePlanSnafu)?; @@ -395,21 +445,34 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { #[cfg(test)] mod tests { - use catalog::RegisterTableRequest; + use std::any::Any; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + use catalog::error::{QueryAccessDeniedSnafu, Result as CatalogResult}; + use catalog::{CatalogManager, RegisterTableRequest}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; + use common_function::aggrs::aggr_wrapper::StateWrapper; + use common_function::aggrs::aggr_wrapper::fix_order::FixStateUdafOrderingAnalyzer; + use common_query::logical_plan::SubstraitPlanDecoderRef; use common_query::native_histogram::native_histogram_value_type; use common_time::Timezone; use datafusion::catalog::TableProvider; use datafusion::datasource::MemTable; use datafusion::logical_expr::Extension; - use datafusion_expr::expr::{Cast, ScalarFunction}; - use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, ScalarUDF, col, lit}; + use datafusion::optimizer::AnalyzerRule; + use datafusion_expr::expr::{AggregateFunction, Cast, InSubquery, ScalarFunction, Sort}; + use datafusion_expr::{ + AggregateUDF, Expr, LogicalPlanBuilder, LogicalTableSource, ScalarUDF, Subquery, col, lit, + }; use datatypes::arrow::datatypes::{ DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, }; use datatypes::data_type::DataType; - use promql::extension_plan::RangeManipulate; + use futures::stream::BoxStream; + use promql::extension_plan::{RangeManipulate, SeriesNormalize, UnionDistinctOn}; use session::context::QueryContext; + use table::TableRef; + use table::metadata::{TableId, TableInfoRef}; use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable}; use super::*; @@ -1048,4 +1111,627 @@ mod tests { assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx); } } + + /// Registers a `numbers` table in a memory catalog manager. + fn numbers_catalog_manager() -> CatalogManagerRef { + let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); + catalog_manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + catalog_manager + } + + /// Scan plan over the `numbers` schema shared by the fixtures below. + fn numbers_scan() -> LogicalPlan { + LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .build() + .unwrap() + } + + fn merge_scan(input: LogicalPlan, is_placeholder: bool) -> LogicalPlan { + MergeScanLogicalPlan::new(input, is_placeholder, Default::default()).into_logical_plan() + } + + fn encode_plan(plan: &LogicalPlan) -> Bytes { + DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .unwrap() + } + + fn plan_decoder(engine: &dyn crate::QueryEngine) -> SubstraitPlanDecoderRef { + engine + .engine_context(QueryContext::arc()) + .new_plan_decoder() + .unwrap() + } + + /// Request catalog that resolves any table to a provider the payload must not use. + fn request_catalog_list() -> Arc { + Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))) + } + + /// How [`InterceptingCatalogManager`] resolves payload tables. + enum TableResolution { + /// Fails with the frontend's cross-catalog access error. + AccessDenied, + /// Resolves once another task on the caller runtime signals. + AwaitSignal { + notify: Arc, + entered: Arc, + }, + /// Never resolves, holding `live` until the caller drops the decode future. + Pending { + live: Arc, + entered: Arc, + }, + } + + /// Engine catalog manager that takes over `table` resolution while delegating the rest. + struct InterceptingCatalogManager { + inner: CatalogManagerRef, + resolution: TableResolution, + } + + #[async_trait::async_trait] + impl CatalogManager for InterceptingCatalogManager { + fn as_any(&self) -> &dyn Any { + self.inner.as_any() + } + + async fn catalog_names(&self) -> CatalogResult> { + self.inner.catalog_names().await + } + + async fn schema_names( + &self, + catalog: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + self.inner.schema_names(catalog, query_ctx).await + } + + async fn table_names( + &self, + catalog: &str, + schema: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + self.inner.table_names(catalog, schema, query_ctx).await + } + + async fn catalog_exists(&self, catalog: &str) -> CatalogResult { + self.inner.catalog_exists(catalog).await + } + + async fn schema_exists( + &self, + catalog: &str, + schema: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult { + self.inner.schema_exists(catalog, schema, query_ctx).await + } + + async fn table_exists( + &self, + catalog: &str, + schema: &str, + table: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult { + self.inner + .table_exists(catalog, schema, table, query_ctx) + .await + } + + async fn table( + &self, + catalog: &str, + schema: &str, + table_name: &str, + query_ctx: Option<&QueryContext>, + ) -> CatalogResult> { + match &self.resolution { + TableResolution::AccessDenied => QueryAccessDeniedSnafu { + catalog: catalog.to_string(), + schema: schema.to_string(), + } + .fail(), + TableResolution::AwaitSignal { notify, entered } => { + entered.store(true, Ordering::SeqCst); + notify.notified().await; + self.inner + .table(catalog, schema, table_name, query_ctx) + .await + } + TableResolution::Pending { live, entered } => { + entered.store(true, Ordering::SeqCst); + let _guard = LiveGuard::new(live); + std::future::pending::<()>().await; + unreachable!("pending resolution never completes") + } + } + } + + async fn table_info_by_id(&self, table_id: TableId) -> CatalogResult> { + self.inner.table_info_by_id(table_id).await + } + + async fn tables_by_ids( + &self, + catalog: &str, + schema: &str, + table_ids: &[TableId], + ) -> CatalogResult> { + self.inner.tables_by_ids(catalog, schema, table_ids).await + } + + fn tables<'a>( + &'a self, + catalog: &'a str, + schema: &'a str, + query_ctx: Option<&'a QueryContext>, + ) -> BoxStream<'a, CatalogResult> { + self.inner.tables(catalog, schema, query_ctx) + } + } + + /// Counts live payload resolutions, so a dropped decode leaves no work behind. + struct LiveGuard(Arc); + + impl LiveGuard { + fn new(live: &Arc) -> Self { + let live = live.clone(); + live.fetch_add(1, Ordering::SeqCst); + Self(live) + } + } + + impl Drop for LiveGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } + } + + fn engine_with_intercepting_catalog( + resolution: TableResolution, + ) -> Arc { + let catalog_manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager { + inner: numbers_catalog_manager(), + resolution, + }); + QueryEngineFactory::new( + catalog_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + } + + /// Decoding must not need a Tokio runtime: the public path is only `async`. + #[test] + fn test_serializer_decode_nested_merge_scan_without_tokio_runtime() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let plan = merge_scan(merge_scan(numbers_scan(), false), true); + let bytes = encode_plan(&plan); + + let decoded = futures::executor::block_on(plan_decoder(engine.as_ref()).decode( + bytes, + request_catalog_list(), + false, + )) + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + } + + /// Payload decoding must not block a multi-thread runtime either. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_serializer_decode_deeply_nested_merge_scan_without_blocking() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let plan = merge_scan( + merge_scan( + LogicalPlanBuilder::from(merge_scan(numbers_scan(), false)) + .filter(col("number").lt(lit(10u32))) + .unwrap() + .build() + .unwrap(), + false, + ), + false, + ); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 3); + } + + /// A payload lookup parked on a caller-runtime task must keep the decode async. + /// + /// The engine catalog resolves only once a task spawned on this runtime signals it, so a + /// blocking bridge would deadlock the test instead of decoding. + #[tokio::test(flavor = "current_thread")] + async fn test_serializer_decode_merge_scan_awaits_caller_runtime_task() { + let notify = Arc::new(tokio::sync::Notify::new()); + let entered = Arc::new(AtomicBool::new(false)); + let engine = engine_with_intercepting_catalog(TableResolution::AwaitSignal { + notify: notify.clone(), + entered: entered.clone(), + }); + let plan = merge_scan(numbers_scan(), false); + let bytes = encode_plan(&plan); + + let decoder = plan_decoder(engine.as_ref()); + let mut decode = Box::pin(decoder.decode(bytes, request_catalog_list(), false)); + // First poll reaches the payload table lookup, which waits for the signal below. + assert!(futures::poll!(decode.as_mut()).is_pending()); + assert!(entered.load(Ordering::SeqCst)); + + let signal = tokio::spawn(async move { + notify.notify_waiters(); + }); + tokio::task::yield_now().await; + let decoded = decode.await.unwrap(); + signal.await.unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + } + + /// Dropping a decode parked on the catalog cancels it; no detached job keeps running. + #[tokio::test(flavor = "current_thread")] + async fn test_serializer_decode_dropped_pending_payload_leaves_no_job() { + let live = Arc::new(AtomicUsize::new(0)); + let entered = Arc::new(AtomicBool::new(false)); + let engine = engine_with_intercepting_catalog(TableResolution::Pending { + live: live.clone(), + entered: entered.clone(), + }); + let bytes = encode_plan(&merge_scan(numbers_scan(), false)); + + let decoder = plan_decoder(engine.as_ref()); + { + let mut decode = Box::pin(decoder.decode(bytes, request_catalog_list(), false)); + assert!(futures::poll!(decode.as_mut()).is_pending()); + assert_eq!(live.load(Ordering::SeqCst), 1); + assert!(entered.load(Ordering::SeqCst)); + } + + assert_eq!( + live.load(Ordering::SeqCst), + 0, + "dropping the decode future must cancel its payload work" + ); + } + + /// A payload table missing from the engine catalog must not fall back to the request catalog. + #[tokio::test] + async fn test_serializer_decode_payload_missing_in_engine_catalog() { + // The engine catalog deliberately has no `numbers` table. + let catalog_manager: CatalogManagerRef = + catalog::memory::new_memory_catalog_manager().unwrap(); + let engine = QueryEngineFactory::new( + catalog_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let bytes = encode_plan(&merge_scan(numbers_scan(), false)); + + let err = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap_err(); + + let err = format!("{err:?}"); + assert!( + err.contains("Table not found: greptime.public.numbers"), + "unexpected error: {err}" + ); + } + + /// A payload table denied by the engine catalog must not fall back to the request catalog. + #[tokio::test] + async fn test_serializer_decode_payload_access_denied_in_engine_catalog() { + let engine = engine_with_intercepting_catalog(TableResolution::AccessDenied); + let bytes = encode_plan(&merge_scan(numbers_scan(), false)); + + let err = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap_err(); + + let err = format!("{err:?}"); + assert!( + err.contains("Illegal access to catalog"), + "unexpected error: {err}" + ); + } + + /// `MergeScan` decodes under an ordinary single-input extension parent. + #[tokio::test] + async fn test_serializer_decode_merge_scan_under_series_normalize() { + let outer = merge_scan(numbers_scan(), false); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 0, + "number".to_string(), + false, + Vec::new(), + outer.clone(), + )), + }); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a SeriesNormalize plan, got: {decoded}"); + }; + let normalize = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a SeriesNormalize plan node"); + assert_eq!(normalize.inputs()[0].to_string(), outer.to_string()); + } + + /// `MergeScan` decodes under an ordinary multi-input extension parent. + #[tokio::test] + async fn test_serializer_decode_merge_scan_under_union_distinct_on() { + let left = merge_scan(numbers_scan(), false); + let right = merge_scan(numbers_scan(), true); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new( + UnionDistinctOn::try_new(left.clone(), right.clone(), vec![], 0).unwrap(), + ), + }); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a UnionDistinctOn plan, got: {decoded}"); + }; + let union = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a UnionDistinctOn plan node"); + assert_eq!(union.inputs()[0].to_string(), left.to_string()); + assert_eq!(union.inputs()[1].to_string(), right.to_string()); + } + + /// `MergeScan` inside a subquery still decodes through the wrapper consumer. + /// + /// Correlated outer references cannot be covered here: the in-repo Substrait producer + /// rejects `Expr::OuterReferenceColumn`, so no supported helper can build such a plan. + #[tokio::test] + async fn test_serializer_decode_merge_scan_inside_subquery() { + let inner = merge_scan(numbers_scan(), false); + let payload = LogicalPlanBuilder::from(numbers_scan()) + .filter(Expr::InSubquery(InSubquery::new( + Box::new(col("number")), + Subquery { + subquery: Arc::new(inner.clone()), + outer_ref_columns: vec![], + spans: Default::default(), + }, + false, + ))) + .unwrap() + .build() + .unwrap(); + let plan = merge_scan(payload.clone(), false); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder( + QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine() + .as_ref(), + ) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + assert!( + decoded.to_string().contains("IN ()"), + "expect the decoded subquery filter, got: {decoded}" + ); + } + + /// Recursive payload decoding must restore ordered state aggregates, including their schema. + #[tokio::test] + async fn test_serializer_decode_payload_keeps_state_udaf_ordering() { + let last_value = (*datafusion::functions_aggregate::first_last::last_value_udaf()).clone(); + let state_udaf = AggregateUDF::new_from_impl(StateWrapper::new(last_value).unwrap()); + let aggr_expr = Expr::AggregateFunction(AggregateFunction::new_udf( + Arc::new(state_udaf), + vec![col("number")], + false, + None, + vec![Sort::new(col("number"), true, true)], + None, + )); + let payload = LogicalPlanBuilder::from(numbers_scan()) + .aggregate(Vec::::new(), vec![aggr_expr]) + .unwrap() + .build() + .unwrap(); + let bytes = encode_plan(&payload); + let encoded = Plan::decode(bytes.clone()).unwrap(); + let extensions = Extensions::try_from(&encoded.extensions).unwrap(); + let catalog_manager = numbers_catalog_manager(); + let engine = QueryEngineFactory::new( + catalog_manager.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let state = engine.engine_context(QueryContext::arc()).state().clone(); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, &state), + session_state: &state, + catalog_manager: Some(catalog_manager), + }; + let decoded = consumer.decode_payload(bytes.to_vec()).await.unwrap(); + + let fixed_payload = FixStateUdafOrderingAnalyzer {} + .analyze(payload.clone(), &Default::default()) + .unwrap(); + assert_ne!( + fixed_payload.schema().as_arrow(), + payload.schema().as_arrow() + ); + assert_eq!( + decoded.schema().as_arrow(), + fixed_payload.schema().as_arrow() + ); + let LogicalPlan::Aggregate(aggregate) = &decoded else { + panic!("Expected an aggregate payload, got: {decoded}"); + }; + let Expr::AggregateFunction(function) = &aggregate.aggr_expr[0] else { + panic!("Expected a state aggregate, got: {:?}", aggregate.aggr_expr); + }; + assert!(function.func.inner().is::()); + assert_eq!(function.params.order_by.len(), 1); + let ArrowDataType::Struct(fields) = decoded.schema().field(0).data_type() else { + panic!("Expected the ordered aggregate state type"); + }; + assert_eq!(fields.len(), 3); + assert_eq!(fields[1].name(), "number"); + assert!(fields[1].is_nullable()); + } + + /// Type variation extensions keep being rejected while decoding. + #[tokio::test] + async fn test_serializer_decode_rejects_type_variations() { + use substrait::substrait_proto_df::proto::extensions::SimpleExtensionDeclaration; + use substrait::substrait_proto_df::proto::extensions::simple_extension_declaration::{ + ExtensionTypeVariation, MappingType, + }; + + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let encoded = encode_plan(&merge_scan(numbers_scan(), false)); + let mut substrait_plan = Plan::decode(encoded).unwrap(); + substrait_plan.extensions.push(SimpleExtensionDeclaration { + mapping_type: Some(MappingType::ExtensionTypeVariation( + ExtensionTypeVariation { + extension_urn_reference: u32::MAX, + type_variation_anchor: 1, + name: "u!variation".to_string(), + }, + )), + }); + let bytes = Bytes::from(substrait_plan.encode_to_vec()); + + let err = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap_err(); + + let err = format!("{err:?}"); + assert!( + err.contains("Type variation extensions are not supported"), + "unexpected error: {err}" + ); + } } From 169b472716a6c21bd33d65da9b2d5d776e5471dd Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sun, 4 Oct 2026 16:42:39 +0800 Subject: [PATCH 10/37] test(query): avoid unnecessary extension name allocation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/query_engine/default_serializer.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 23b69abaaf2e..50283cecc9a7 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -1501,7 +1501,7 @@ mod tests { let plan = LogicalPlan::Extension(Extension { node: Arc::new(SeriesNormalize::new( 0, - "number".to_string(), + "number", false, Vec::new(), outer.clone(), From 2a46309a4f48d73fe038e4aa4631f718bf0ce818 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sun, 4 Oct 2026 17:12:21 +0800 Subject: [PATCH 11/37] fix(query): retain extension decoding for manual session states Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/query_engine/default_serializer.rs | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 50283cecc9a7..31b9e59be6f5 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -423,6 +423,7 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { optimize: bool, ) -> common_query::error::Result { let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) + .with_serializer_registry(Arc::new(DefaultSerializer)) .with_catalog_list(catalog_list) .build(); // Re-register after the build to avoid Greptime UDF alias collisions. @@ -1494,6 +1495,32 @@ mod tests { ); } + /// Manually built states need ordinary extension decoding without an engine catalog. + #[tokio::test] + async fn test_serializer_decode_manual_state_keeps_extension_registry() { + let state = SessionStateBuilder::new().with_default_features().build(); + let decoder = DefaultPlanDecoder::new(state, &QueryContext::arc()).unwrap(); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 0, + "number", + false, + Vec::new(), + merge_scan(numbers_scan(), false), + )), + }); + let request_catalog = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + MemTable::try_new(NumbersTable::schema().arrow_schema().clone(), vec![vec![]]).unwrap(), + ))); + let decoded = decoder + .decode(encode_plan(&plan), request_catalog, false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.schema().as_arrow(), plan.schema().as_arrow()); + } + /// `MergeScan` decodes under an ordinary single-input extension parent. #[tokio::test] async fn test_serializer_decode_merge_scan_under_series_normalize() { From 5137e10d4cbe58b3d3bc8e277563feec9c0aa2be Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 29 Sep 2026 14:26:57 +0800 Subject: [PATCH 12/37] feat(query): add opt-in nested broadcast join planner Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/datafusion.rs | 38 +- src/query/src/dist_plan.rs | 2 + src/query/src/dist_plan/analyzer.rs | 39 +- src/query/src/dist_plan/analyzer/test.rs | 1 + src/query/src/dist_plan/dist_join_planner.rs | 742 ++++++++++++++++++ src/query/src/dist_plan/planner.rs | 79 +- .../src/query_engine/default_serializer.rs | 9 +- src/query/src/query_engine/state.rs | 8 +- .../nested_merge_scan_capability_test.rs | 230 +++++- 9 files changed, 1116 insertions(+), 32 deletions(-) create mode 100644 src/query/src/dist_plan/dist_join_planner.rs diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 970f31098045..ea2a70c5824e 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -572,29 +572,37 @@ impl QueryEngine for DatafusionQueryEngine { Some(query_ctx.timezone().to_string()); // usually it's impossible to have both `set variable` set by sql client and - // hint in header by grpc client, so only need to deal with them separately + // hint in header by grpc client, so only need to deal with them separately. + // Start from the options already configured on the session config (e.g. the + // engine level `allow_query_fallback`) so that an insert below doesn't drop them. + let mut dist_planner_options = state + .config() + .options() + .extensions + .get::() + .cloned() + .unwrap_or_default(); + let mut has_dist_planner_options = false; if query_ctx.configuration_parameter().allow_query_fallback() { - state - .config_mut() - .options_mut() - .extensions - .insert(DistPlannerOptions { - allow_query_fallback: true, - }); + dist_planner_options.allow_query_fallback = true; + has_dist_planner_options = true; } else if let Some(fallback) = query_ctx.extension(QUERY_FALLBACK_HINT) { // also check the query context for fallback hint // if it is set, we will enable the fallback if fallback.to_lowercase().parse::().unwrap_or(false) { - state - .config_mut() - .options_mut() - .extensions - .insert(DistPlannerOptions { - allow_query_fallback: true, - }); + dist_planner_options.allow_query_fallback = true; + has_dist_planner_options = true; } } + if has_dist_planner_options { + state + .config_mut() + .options_mut() + .extensions + .insert(dist_planner_options); + } + state .config_mut() .options_mut() diff --git a/src/query/src/dist_plan.rs b/src/query/src/dist_plan.rs index 969f00e7c9cc..9b32cb8775f6 100644 --- a/src/query/src/dist_plan.rs +++ b/src/query/src/dist_plan.rs @@ -14,6 +14,7 @@ mod analyzer; mod commutativity; +mod dist_join_planner; mod dyn_filter_bridge; mod filter_id; mod merge_scan; @@ -25,6 +26,7 @@ mod remote_dyn_filter_receiver; mod remote_dyn_filter_registry; pub use analyzer::{DistPlannerAnalyzer, DistPlannerOptions}; +pub use dist_join_planner::DistJoinPlanner; pub use filter_id::{FilterFingerprint, FilterId, ParseFilterIdError, RemoteDynFilterProducerId}; pub use merge_scan::{MergeScanExec, MergeScanLogicalPlan}; pub(crate) use merge_sort::MergeSortExec; diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 2c1d066054ac..2877ee0f4f2f 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -68,9 +68,15 @@ pub(crate) use utils::AliasMapping; /// Placeholder for other physical partition columns that are not in logical table const OTHER_PHY_PART_COL_PLACEHOLDER: &str = "__OTHER_PHYSICAL_PART_COLS_PLACEHOLDER__"; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct DistPlannerOptions { pub allow_query_fallback: bool, + /// Manual opt-in for the nested broadcast join rewrite: name of the (small) build + /// side table, e.g. `device_limits`. `None` (the default) keeps it disabled. + /// + /// Not cost based in this slice: setting the option enables the rewrite for the + /// supported shape, see [`DistJoinPlanner`](crate::dist_plan::DistJoinPlanner). + pub nested_broadcast_join_build_table: Option, } impl ConfigExtension for DistPlannerOptions { @@ -91,17 +97,32 @@ impl ExtensionOptions for DistPlannerOptions { } fn set(&mut self, key: &str, value: &str) -> DfResult<()> { - Err(datafusion_common::DataFusionError::NotImplemented(format!( - "DistPlannerOptions does not support set key: {key} with value: {value}" - ))) + match key { + "nested_broadcast_join_build_table" => { + let value = value.trim().trim_matches(&['\'', '"'][..]); + self.nested_broadcast_join_build_table = + (!value.is_empty()).then(|| value.to_string()); + Ok(()) + } + _ => Err(datafusion_common::DataFusionError::NotImplemented(format!( + "DistPlannerOptions does not support set key: {key} with value: {value}" + ))), + } } fn entries(&self) -> Vec { - vec![datafusion::config::ConfigEntry { - key: "allow_query_fallback".to_string(), - value: Some(self.allow_query_fallback.to_string()), - description: "Allow query fallback to fallback plan rewriter", - }] + vec![ + datafusion::config::ConfigEntry { + key: "allow_query_fallback".to_string(), + value: Some(self.allow_query_fallback.to_string()), + description: "Allow query fallback to fallback plan rewriter", + }, + datafusion::config::ConfigEntry { + key: "nested_broadcast_join_build_table".to_string(), + value: self.nested_broadcast_join_build_table.clone(), + description: "Build side table of the nested broadcast join rewrite", + }, + ] } } diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 694f2bcde59b..d6547253348d 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -745,6 +745,7 @@ fn expand_sort_alias_conflict_limit() { let mut config = ConfigOptions::default(); config.extensions.insert(DistPlannerOptions { allow_query_fallback: true, + ..Default::default() }); let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); diff --git a/src/query/src/dist_plan/dist_join_planner.rs b/src/query/src/dist_plan/dist_join_planner.rs new file mode 100644 index 000000000000..af624503355a --- /dev/null +++ b/src/query/src/dist_plan/dist_join_planner.rs @@ -0,0 +1,742 @@ +// 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. + +//! Nested broadcast join rewrite for `INNER` joins of two distributed scans. +//! +//! Opt-in only: enabled when [`DistPlannerOptions::nested_broadcast_join_build_table`] +//! names the build side table. Runs after `DistPlannerAnalyzer`, which wraps the remote +//! scans in `MergeScan` and assigns the remote dynamic filter producer ids preserved +//! here. Not cost based: no statistics are consulted in this slice. + +use std::sync::Arc; + +use datafusion::error::Result as DfResult; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; +use datafusion_expr::{Join, JoinType, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_optimizer::analyzer::AnalyzerRule; +use datatypes::extension::json::is_json2_extension_type; +use table::table_name::TableName; + +use crate::dist_plan::analyzer::DistPlannerOptions; +use crate::dist_plan::merge_scan::MergeScanLogicalPlan; +use crate::dist_plan::planner::table_name_of; + +/// Nests the build side's `MergeScan` inside the probe side's one, so the join runs on +/// the datanodes holding the probe regions. +/// +/// Keeps the join node as is, so its schema, join conditions, residual filter and NULL +/// semantics are preserved; unmatched shapes are returned unchanged. +#[derive(Debug)] +pub struct DistJoinPlanner; + +impl AnalyzerRule for DistJoinPlanner { + fn name(&self) -> &str { + "DistJoinPlanner" + } + + fn analyze(&self, plan: LogicalPlan, config: &ConfigOptions) -> DfResult { + let Some(build_table) = config + .extensions + .get::() + .and_then(|options| options.nested_broadcast_join_build_table.as_deref()) + else { + return Ok(plan); + }; + + let mut rewriter = NestedBroadcastJoinRewriter { build_table }; + Ok(plan.rewrite(&mut rewriter)?.data) + } +} + +/// Rewriter nesting the build side `MergeScan` of a supported join inside the probe +/// side's `MergeScan`. +struct NestedBroadcastJoinRewriter<'a> { + /// Name of the build side table, either bare (`device_limits`) or fully qualified. + build_table: &'a str, +} + +impl NestedBroadcastJoinRewriter<'_> { + /// Returns the rewritten plan when `node` matches this slice's shape, otherwise + /// `None` to keep the plan unchanged. + /// + /// A rewrite changes where the join runs, so it must preserve the join output + /// schema, the row multiplicity, and the remote boundary identities; the guards + /// below only admit the shape where that holds. + fn try_rewrite_join(&self, node: &LogicalPlan) -> Option { + let LogicalPlan::Join(join) = node else { + return None; + }; + + // Scope restriction of this slice: only INNER equi-joins of two distributed + // scans are handled, everything else keeps the existing distributed plan. + if join.join_type != JoinType::Inner || join.on.is_empty() { + return None; + } + + let (probe_plan, probe_merge_scan) = strip_merge_scan(&join.left)?; + let (_, build_merge_scan) = strip_merge_scan(&join.right)?; + if probe_merge_scan.is_placeholder() || build_merge_scan.is_placeholder() { + return None; + } + + // Replacing the probe input must not change the join input schema. + if probe_plan.schema() != join.left.schema() { + return None; + } + + // A nested boundary is hidden from the remaining analyzer passes, so a boundary + // schema `JsonSchemaConcretizeRule` still has to fix up cannot be repaired + // anymore. + if has_json2_boundary(&probe_merge_scan) || has_json2_boundary(&build_merge_scan) { + return None; + } + + // The configured table must identify exactly one side, and the build side must be + // the right input: moving a left build side would change the join schema. + let left_is_build = table_name_of(probe_merge_scan.input()) + .is_some_and(|name| table_name_matches(&name, self.build_table)); + let right_is_build = table_name_of(build_merge_scan.input()) + .is_some_and(|name| table_name_matches(&name, self.build_table)); + if left_is_build || !right_is_build { + return None; + } + + // Keep the rewrite at two levels: the probe side must be local and the build side + // must not contain another distributed boundary. + if contains_merge_scan(&probe_plan) || contains_merge_scan(build_merge_scan.input()) { + return None; + } + + // Join the probe regions against the build plan kept above; the probe table's + // partition columns route the outer scan. + let new_join = Join { + left: Arc::new(probe_plan), + right: join.right.clone(), + ..join.clone() + }; + let mut merge_scan = MergeScanLogicalPlan::new( + LogicalPlan::Join(new_join), + false, + probe_merge_scan.partition_cols().clone(), + ); + // The outer boundary keeps the probe boundary's remote dynamic filter identity + // when it has one. A missing id only disables remote dynamic filter pushdown for + // that boundary (`MergeScanExec` fails open), and serialization drops the ids of + // the nested payload, so an absent id is not a rewrite blocker. + if let Some(producer_id) = probe_merge_scan.remote_dyn_filter_producer_id() { + merge_scan = merge_scan.with_remote_dyn_filter_producer_id(producer_id); + } + + Some(merge_scan.into_logical_plan()) + } +} + +impl TreeNodeRewriter for NestedBroadcastJoinRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DfResult> { + match self.try_rewrite_join(&node) { + Some(rewritten) => Ok(Transformed::yes(rewritten)), + None => Ok(Transformed::no(node)), + } + } +} + +/// Strips the outermost `MergeScan` of a join input, which may be hidden below +/// pass-through projections, and returns the input plan without the wrapper together +/// with the removed `MergeScan`. +/// +/// Returns `None` if the input is not wrapped in a `MergeScan`. +fn strip_merge_scan(plan: &LogicalPlan) -> Option<(LogicalPlan, MergeScanLogicalPlan)> { + match plan { + LogicalPlan::Extension(extension) => { + let merge_scan = extension + .node + .as_any() + .downcast_ref::()?; + Some((merge_scan.input().clone(), merge_scan.clone())) + } + LogicalPlan::Projection(projection) => { + let (input, merge_scan) = strip_merge_scan(&projection.input)?; + let plan = plan + .with_new_exprs(projection.expr.clone(), vec![input]) + .ok()?; + Some((plan, merge_scan)) + } + _ => None, + } +} + +/// Whether `plan` contains a visible `MergeScan` node. +fn contains_merge_scan(plan: &LogicalPlan) -> bool { + if let LogicalPlan::Extension(extension) = plan + && extension.node.as_any().is::() + { + return true; + } + + plan.inputs().into_iter().any(contains_merge_scan) +} + +/// Returns the local probe input that decides the regions of a supported nested broadcast +/// join, or `None` for any other plan. +/// +/// Region pruning clears the collected predicates at a node with several inputs (the +/// join), so the outer `MergeScan` of the rewritten plan would be routed by the join +/// itself and scan every region of the probe table. Routing by the returned side keeps +/// the probe predicate pruning, while the full payload is still dispatched to the +/// selected regions. +/// +/// Only the supported shape of [`DistJoinPlanner`] returns a side: an `INNER` equi-join +/// whose left input is a local plan and whose right input reaches its `MergeScan` +/// through projections. +pub(crate) fn nested_broadcast_join_probe_input(plan: &LogicalPlan) -> Option<&LogicalPlan> { + let LogicalPlan::Join(join) = plan else { + return None; + }; + if join.join_type != JoinType::Inner || join.on.is_empty() { + return None; + } + if contains_merge_scan(&join.left) { + return None; + } + let (_, build_merge_scan) = strip_merge_scan(&join.right)?; + if build_merge_scan.is_placeholder() { + return None; + } + + Some(&join.left) +} + +/// Whether the `MergeScan` boundary exposes a JSON2 column. +fn has_json2_boundary(merge_scan: &MergeScanLogicalPlan) -> bool { + merge_scan + .schema() + .fields() + .iter() + .any(|field| is_json2_extension_type(field)) +} + +/// Whether the configured build side table name refers to `name`. Accepts the bare +/// table name or the fully qualified `catalog.schema.table` form, ignoring case. +fn table_name_matches(name: &TableName, configured: &str) -> bool { + let configured = configured.trim().trim_matches(&['\'', '"'][..]); + !configured.is_empty() + && (name.table_name.eq_ignore_ascii_case(configured) + || name.to_string().eq_ignore_ascii_case(configured)) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType as _, + }; + use arrow_schema::{DataType, Field, Fields, Schema}; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use datafusion::datasource::DefaultTableSource; + use datafusion_common::{DFSchema, DFSchemaRef, JoinConstraint, NullEquality}; + use datafusion_expr::{Expr, LogicalPlanBuilder, build_join_schema, col, lit}; + use datatypes::data_type::ConcreteDataType; + use datatypes::extension::json::JsonExtensionType; + use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; + use pretty_assertions::assert_eq; + use table::TableRef; + use table::metadata::TableId; + use table::table::adapter::DfTableProviderAdapter; + use table::test_util::EmptyTable; + use table::test_util::table_info::test_table_info; + + use super::*; + use crate::dist_plan::DistPlannerAnalyzer; + + /// Two-column schema (`number`, `host`) shared by the test tables. + fn test_schema() -> SchemaRef { + let schema = SchemaBuilder::try_from_columns(vec![ + ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ]) + .unwrap() + .build() + .unwrap(); + Arc::new(schema) + } + + /// Base table named `name`. + fn test_table(table_id: TableId, name: &str) -> TableRef { + let info = test_table_info( + table_id, + name, + DEFAULT_SCHEMA_NAME, + DEFAULT_CATALOG_NAME, + test_schema(), + ); + EmptyTable::from_table_info(&info) + } + + /// Scan plan whose alias and base table name may differ. + fn table_scan(alias: &str, table_id: TableId, table_name: &str) -> LogicalPlan { + let source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(test_table(table_id, table_name)), + ))); + LogicalPlanBuilder::scan_with_filters(alias, source, None, vec![]) + .unwrap() + .build() + .unwrap() + } + + /// `t1 JOIN t2` on `number`, i.e. the plan the distributed planner sees. + fn join_plan(join_type: JoinType) -> LogicalPlan { + LogicalPlanBuilder::from(table_scan("t1", 1, "t1")) + .join_on( + table_scan("t2", 2, "t2"), + join_type, + vec![col("t1.number").eq(col("t2.number"))], + ) + .unwrap() + .build() + .unwrap() + } + + /// Builds an inner equi-join with the join keys in `on`, as the analyzer leaves them. + fn inner_join( + left: LogicalPlan, + right: LogicalPlan, + left_key: Expr, + right_key: Expr, + ) -> LogicalPlan { + let schema = build_join_schema(left.schema(), right.schema(), &JoinType::Inner).unwrap(); + LogicalPlan::Join(Join { + left: Arc::new(left), + right: Arc::new(right), + on: vec![(left_key, right_key)], + filter: None, + join_type: JoinType::Inner, + join_constraint: JoinConstraint::On, + schema: DFSchemaRef::new(schema), + null_equality: NullEquality::NullEqualsNothing, + null_aware: false, + }) + } + + /// Distributed plan of `t1 INNER JOIN t2`, as produced by `DistPlannerAnalyzer`. + fn distributed_join_plan() -> LogicalPlan { + DistPlannerAnalyzer {} + .analyze(join_plan(JoinType::Inner), &ConfigOptions::default()) + .unwrap() + } + + fn rewrite_config(build_table: &str) -> ConfigOptions { + let mut config = ConfigOptions::default(); + config.extensions.insert(DistPlannerOptions { + nested_broadcast_join_build_table: Some(build_table.to_string()), + ..Default::default() + }); + config + } + + fn rewrite(plan: LogicalPlan, build_table: &str) -> LogicalPlan { + DistJoinPlanner {} + .analyze(plan, &rewrite_config(build_table)) + .unwrap() + } + + /// All `MergeScan`s of `plan`, including the ones nested inside another `MergeScan`. + fn merge_scans(plan: &LogicalPlan) -> Vec { + let mut scans = Vec::new(); + collect_merge_scans(plan, &mut scans); + scans + } + + fn collect_merge_scans(plan: &LogicalPlan, scans: &mut Vec) { + if let LogicalPlan::Extension(extension) = plan + && let Some(merge_scan) = extension + .node + .as_any() + .downcast_ref::() + { + scans.push(merge_scan.clone()); + collect_merge_scans(merge_scan.input(), scans); + return; + } + + for input in plan.inputs() { + collect_merge_scans(input, scans); + } + } + + /// The join wrapped by the outer `MergeScan` of a rewritten plan. + fn outer_join(plan: &LogicalPlan) -> &Join { + let LogicalPlan::Extension(extension) = plan else { + panic!("expected the outer node to be a MergeScan, got: {plan}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .unwrap_or_else(|| panic!("expected the outer node to be a MergeScan, got: {plan}")); + let LogicalPlan::Join(join) = merge_scan.input() else { + panic!( + "expected the outer MergeScan to wrap a join, got: {}", + merge_scan.input() + ); + }; + + join + } + + /// Without the build table option the join stays above the two `MergeScan`s. + #[test] + fn nested_broadcast_join_rewrite_disabled_by_default() { + let plan = distributed_join_plan(); + let result = DistJoinPlanner {} + .analyze(plan.clone(), &ConfigOptions::default()) + .unwrap(); + + assert_eq!(plan.to_string(), result.to_string()); + assert!(matches!(result, LogicalPlan::Join(_))); + assert_eq!(2, merge_scans(&result).len()); + } + + /// With the build table option the build side `MergeScan` becomes the inner scan of + /// the outer `MergeScan`, which is routed by the probe table's regions. + /// + /// The rewrite must only move where the join runs: the join fields, the residual + /// filter, the NULL semantics and the output schema stay as the join had them. + #[test] + fn nested_broadcast_join_rewrite_nests_build_side_merge_scan() { + let plan = distributed_join_plan(); + let LogicalPlan::Join(original_join) = &plan else { + panic!("expected a join on top, got: {plan}"); + }; + let plan = LogicalPlan::Join(Join { + filter: Some(col("t1.number").gt(lit(1u32))), + null_equality: NullEquality::NullEqualsNull, + ..original_join.clone() + }); + let probe_scan = merge_scans(&plan).first().cloned().unwrap(); + + let result = rewrite(plan.clone(), "t2"); + + let join = outer_join(&result); + assert_eq!(JoinType::Inner, join.join_type); + assert_eq!(original_join.on, join.on); + assert_eq!( + Some(col("t1.number").gt(lit(1u32))), + join.filter, + "the residual filter must be preserved" + ); + assert_eq!(original_join.join_constraint, join.join_constraint); + assert_eq!(NullEquality::NullEqualsNull, join.null_equality); + assert_eq!(original_join.null_aware, join.null_aware); + assert!( + !contains_merge_scan(&join.left), + "probe side must be local, got: {}", + join.left + ); + assert!( + contains_merge_scan(&join.right), + "build side must keep its MergeScan, got: {}", + join.right + ); + + // The outer scan must be routed by the probe table's regions and keep the + // probe boundary's partition columns. + assert_eq!( + "t1", + table_name_of(&LogicalPlan::Join(join.clone())) + .unwrap() + .table_name + ); + let outer_scan = merge_scans(&result).first().cloned().unwrap(); + assert_eq!(probe_scan.partition_cols(), outer_scan.partition_cols()); + assert_eq!(plan.schema(), result.schema()); + } + + /// Both boundaries keep their remote dynamic filter producer ids. + #[test] + fn nested_broadcast_join_rewrite_preserves_remote_dyn_filter_producer_ids() { + let plan = distributed_join_plan(); + let scans = merge_scans(&plan); + let probe_id = scans[0].remote_dyn_filter_producer_id().unwrap(); + let build_id = scans[1].remote_dyn_filter_producer_id().unwrap(); + assert_ne!(probe_id, build_id); + + let result = rewrite(plan, "t2"); + + let scans = merge_scans(&result); + assert_eq!(2, scans.len()); + assert_eq!(Some(probe_id), scans[0].remote_dyn_filter_producer_id()); + assert_eq!(Some(build_id), scans[1].remote_dyn_filter_producer_id()); + } + + /// Boundaries without a remote dynamic filter producer id are still rewritten: a + /// missing id only disables remote dynamic filter pushdown for that boundary, and the + /// outer boundary then simply has none either. + #[test] + fn nested_broadcast_join_rewrite_without_producer_ids() { + let plan = distributed_join_plan_without_producer_ids(); + assert_eq!( + vec![None, None], + merge_scans(&plan) + .iter() + .map(|scan| scan.remote_dyn_filter_producer_id()) + .collect::>() + ); + + let result = rewrite(plan.clone(), "t2"); + + let scans = merge_scans(&result); + assert_eq!(2, scans.len()); + assert_eq!(None, scans[0].remote_dyn_filter_producer_id()); + assert_eq!(None, scans[1].remote_dyn_filter_producer_id()); + assert_eq!(plan.schema(), result.schema()); + } + + /// Distributed plan of `t1 INNER JOIN t2` whose boundaries carry no producer ids. + fn distributed_join_plan_without_producer_ids() -> LogicalPlan { + let plan = distributed_join_plan(); + let LogicalPlan::Join(join) = &plan else { + panic!("expected a join on top, got: {plan}"); + }; + let (probe_plan, probe_scan) = strip_merge_scan(&join.left).unwrap(); + let (build_plan, build_scan) = strip_merge_scan(&join.right).unwrap(); + let probe = + MergeScanLogicalPlan::new(probe_plan, false, probe_scan.partition_cols().clone()); + let build = + MergeScanLogicalPlan::new(build_plan, false, build_scan.partition_cols().clone()); + + inner_join( + probe.into_logical_plan(), + build.into_logical_plan(), + col("t1.number"), + col("t2.number"), + ) + } + + /// A `MergeScan` hidden below a pass-through projection is recognized on both sides. + #[test] + fn nested_broadcast_join_rewrite_handles_projection_wrappers() { + let projected = projected_join_plan(); + + let result = rewrite(projected.clone(), "t2"); + + let join = outer_join(&result); + assert!(matches!(*join.left, LogicalPlan::Projection(_))); + assert!(!contains_merge_scan(&join.left)); + assert!(matches!(join.right.as_ref(), LogicalPlan::Projection(_))); + assert!(contains_merge_scan(&join.right)); + assert_eq!(projected.schema(), result.schema()); + } + + /// Re-projects both join inputs of the distributed plan, so each side is a + /// pass-through projection above its `MergeScan`. + fn projected_join_plan() -> LogicalPlan { + let plan = distributed_join_plan(); + let LogicalPlan::Join(join) = &plan else { + panic!("expected a join on top, got: {plan}"); + }; + let left = LogicalPlanBuilder::from(join.left.as_ref().clone()) + .project(vec![col("t1.number"), col("t1.host")]) + .unwrap() + .build() + .unwrap(); + let right = LogicalPlanBuilder::from(join.right.as_ref().clone()) + .project(vec![col("t2.number"), col("t2.host")]) + .unwrap() + .build() + .unwrap(); + let schema = build_join_schema(left.schema(), right.schema(), &join.join_type).unwrap(); + + LogicalPlan::Join(Join { + left: Arc::new(left), + right: Arc::new(right), + schema: DFSchemaRef::new(schema), + ..join.clone() + }) + } + + /// A build table matching neither side keeps the plan unchanged. + #[test] + fn nested_broadcast_join_rewrite_ignores_unknown_build_table() { + let plan = distributed_join_plan(); + let result = rewrite(plan.clone(), "unknown_table"); + + assert_eq!(plan.to_string(), result.to_string()); + } + + /// Only `INNER` joins are rewritten. + #[test] + fn nested_broadcast_join_rewrite_ignores_non_inner_join() { + let plan = DistPlannerAnalyzer {} + .analyze(join_plan(JoinType::Left), &ConfigOptions::default()) + .unwrap(); + let result = rewrite(plan.clone(), "t2"); + + assert_eq!(plan.to_string(), result.to_string()); + } + + /// A build table on the left input is not rewritten: swapping the inputs would change + /// the join output schema, which the rewrite must preserve. + #[test] + fn nested_broadcast_join_rewrite_ignores_left_side_build() { + let plan = distributed_join_plan(); + let LogicalPlan::Join(join) = &plan else { + panic!("expected a join on top, got: {plan}"); + }; + let plan = LogicalPlan::Join(Join { + left: join.right.clone(), + right: join.left.clone(), + on: join + .on + .iter() + .map(|(left, right)| (right.clone(), left.clone())) + .collect(), + schema: DFSchemaRef::new( + build_join_schema(join.right.schema(), join.left.schema(), &join.join_type) + .unwrap(), + ), + ..join.clone() + }); + + let result = rewrite(plan.clone(), "t2"); + + assert_eq!(plan.to_string(), result.to_string()); + } + + /// Two sides resolving to the configured build table (e.g. a self join) are ambiguous. + #[test] + fn nested_broadcast_join_rewrite_ignores_ambiguous_build_side() { + let plan = DistPlannerAnalyzer {} + .analyze( + join_plan_with_matching_side_names(), + &ConfigOptions::default(), + ) + .unwrap(); + let result = rewrite(plan.clone(), "t1"); + + assert_eq!(plan.to_string(), result.to_string()); + } + + /// `l INNER JOIN r` where `l` and `r` read two base tables both named `t1`. + fn join_plan_with_matching_side_names() -> LogicalPlan { + LogicalPlanBuilder::from(table_scan("l", 1, "t1")) + .join_on( + table_scan("r", 2, "t1"), + JoinType::Inner, + vec![col("l.number").eq(col("r.number"))], + ) + .unwrap() + .build() + .unwrap() + } + + /// The rewrite stays at two levels: an already nested probe side is left alone. + #[test] + fn nested_broadcast_join_rewrite_ignores_more_than_two_levels() { + let nested = rewrite(distributed_join_plan(), "t2"); + let build = DistPlannerAnalyzer {} + .analyze(table_scan("t3", 3, "t3"), &ConfigOptions::default()) + .unwrap(); + let plan = inner_join(nested, build, col("t1.number"), col("t3.number")); + + let result = rewrite(plan.clone(), "t3"); + + assert_eq!(plan.to_string(), result.to_string()); + } + + /// The routing helper returns the probe side of the supported nested shape and `None` + /// for the join of two boundaries (the shape before the rewrite). + #[test] + fn nested_broadcast_join_probe_input_only_for_supported_shape() { + let plan = distributed_join_plan(); + let LogicalPlan::Join(join) = &plan else { + panic!("expected a join on top, got: {plan}"); + }; + + assert_eq!( + None, + nested_broadcast_join_probe_input(&plan).map(|plan| plan.to_string()) + ); + + let result = rewrite(plan.clone(), "t2"); + let LogicalPlan::Extension(extension) = &result else { + panic!("expected the rewritten plan to be a MergeScan, got: {result}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .unwrap(); + let (probe_plan, _) = strip_merge_scan(&join.left).unwrap(); + assert_eq!( + Some(probe_plan.to_string()), + nested_broadcast_join_probe_input(merge_scan.input()).map(|plan| plan.to_string()) + ); + } + + /// JSON2 boundaries are not rewritten: the nested build boundary is hidden from + /// `JsonSchemaConcretizeRule`. + #[test] + fn nested_broadcast_join_rewrite_ignores_json2_boundary() { + let plan = distributed_join_plan(); + let LogicalPlan::Join(join) = &plan else { + panic!("expected a join on top, got: {plan}"); + }; + let (probe_plan, probe_scan) = strip_merge_scan(&join.left).unwrap(); + let (build_plan, build_scan) = strip_merge_scan(&join.right).unwrap(); + let plan = inner_join( + MergeScanLogicalPlan::new(probe_plan, false, probe_scan.partition_cols().clone()) + .with_remote_dyn_filter_producer_id( + probe_scan.remote_dyn_filter_producer_id().unwrap(), + ) + .into_logical_plan(), + MergeScanLogicalPlan::new(build_plan, false, build_scan.partition_cols().clone()) + .with_output_schema(json2_boundary_schema()) + .with_remote_dyn_filter_producer_id( + build_scan.remote_dyn_filter_producer_id().unwrap(), + ) + .into_logical_plan(), + col("t1.number"), + col("t2.number"), + ); + + let result = rewrite(plan.clone(), "t2"); + + assert_eq!(plan.to_string(), result.to_string()); + } + + /// A one-column JSON2 boundary schema, as `JsonSchemaConcretizeRule` expects to fix up. + fn json2_boundary_schema() -> DFSchemaRef { + let field = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata( + HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + JsonExtensionType::NAME.to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + serde_json::json!({ "json_structure_settings": { "Structured": null } }) + .to_string(), + ), + ]), + ); + let schema = DFSchema::try_from(Schema::new(vec![field])).unwrap(); + + DFSchemaRef::new(schema) + } +} diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index b2a2ade9b4fa..35952815453c 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -46,6 +46,7 @@ use table::table::adapter::DfTableProviderAdapter; use table::table_name::TableName; use crate::dist_plan::PredicateExtractor; +use crate::dist_plan::dist_join_planner::nested_broadcast_join_probe_input; use crate::dist_plan::merge_scan::{MergeScanExec, MergeScanLogicalPlan}; use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan}; use crate::dist_plan::region_pruner::ConstraintPruner; @@ -195,12 +196,17 @@ impl ExtensionPlanner for DistExtensionPlanner { } let optimized_plan = input_plan; - let Some(table_name) = Self::extract_full_table_name(input_plan)? else { + // Region pruning clears the collected predicates at a node with several inputs, + // so the payload of a nested broadcast join would be routed by the join itself + // and read every probe region. The local probe input keeps the probe predicate + // pruning; the full payload is still dispatched to the selected regions. + let routing_plan = nested_broadcast_join_probe_input(input_plan).unwrap_or(input_plan); + let Some(table_name) = Self::extract_full_table_name(routing_plan)? else { // no relation found in input plan, going to execute them locally return fallback(optimized_plan).await; }; - let Ok(regions) = self.get_regions(&table_name, input_plan).await else { + let Ok(regions) = self.get_regions(&table_name, routing_plan).await else { // no peers found, going to execute them locally return fallback(optimized_plan).await; }; @@ -471,6 +477,14 @@ fn partition_column_types(table_info: &TableInfo) -> Vec<(String, ConcreteDataTy .collect() } +/// Returns the resolved name of the first base table found in `plan`. +pub(crate) fn table_name_of(plan: &LogicalPlan) -> Option { + let mut extractor = TableScanExtractor::default(); + let _ = plan.visit(&mut extractor).ok()?; + + extractor.table_name +} + /// Extract the scan name and captured table identity from a logical plan. #[derive(Default)] struct TableScanExtractor { @@ -555,7 +569,7 @@ mod tests { use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; use datafusion::datasource::DefaultTableSource; - use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col as df_col, lit}; + use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder, col as df_col, lit}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::value::Value; @@ -571,6 +585,8 @@ mod tests { use table::test_util::EmptyTable; use super::DistExtensionPlanner; + use super::nested_broadcast_join_probe_input; + use crate::dist_plan::merge_scan::MergeScanLogicalPlan; use crate::region_query::{RegionQueryHandler, RegionQueryTarget}; const LOGICAL_TABLE_ID: u32 = 1024; @@ -802,6 +818,63 @@ mod tests { ); } + /// Region pruning clears the collected predicates at a node with several inputs, so a + /// nested broadcast join must be routed by its local probe input: the probe predicate + /// prunes the probe regions, while routing by the whole payload would scan every + /// probe region. + #[tokio::test] + async fn nested_broadcast_join_routes_by_probe_input() { + let (planner, probe_plan, table_name) = + planner_and_plan(vec![0, 1], physical_partition_expressions()).await; + // The build side scans the same table under another qualifier, so the join sides + // have distinct field qualifiers. + let build_plan = + MergeScanLogicalPlan::new(build_side_plan(&probe_plan), false, Default::default()) + .into_logical_plan(); + let join_plan = LogicalPlanBuilder::from(probe_plan.clone()) + .join( + build_plan, + JoinType::Inner, + (vec!["host"], vec!["host"]), + None, + ) + .unwrap() + .build() + .unwrap(); + + // The whole payload clears the probe predicate at the join and reads all regions. + assert_all_logical_regions(planner.get_regions(&table_name, &join_plan).await.unwrap()); + + // The local probe input keeps the probe predicate pruning: one of three regions. + let routing_plan = nested_broadcast_join_probe_input(&join_plan) + .expect("the supported nested shape must route by its probe input"); + assert_eq!(probe_plan.to_string(), routing_plan.to_string()); + assert_eq!( + vec![RegionId::new(LOGICAL_TABLE_ID, 1)], + planner + .get_regions(&table_name, routing_plan) + .await + .unwrap() + ); + } + + /// Scans the table of `plan` under the qualifier `build` with the same probe predicate, + /// so a join of both sides has distinct field qualifiers. + fn build_side_plan(plan: &LogicalPlan) -> LogicalPlan { + let LogicalPlan::Filter(filter) = plan else { + panic!("expected the probe plan to be a filter, got: {plan}"); + }; + let LogicalPlan::TableScan(scan) = filter.input.as_ref() else { + panic!("expected the probe plan to scan a table, got: {plan}"); + }; + LogicalPlanBuilder::scan_with_filters("build", scan.source.clone(), None, vec![]) + .unwrap() + .filter(filter.predicate.clone()) + .unwrap() + .build() + .unwrap() + } + #[tokio::test] async fn missing_physical_partition_datatype_falls_back_to_all_logical_regions() { let (planner, plan, table_name) = diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 31b9e59be6f5..c9dbf3c22b6f 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -478,6 +478,7 @@ mod tests { use super::*; use crate::QueryEngineFactory; + use crate::dist_plan::RemoteDynFilterProducerId; use crate::dummy_catalog::DummyCatalogList; use crate::optimizer::test_util::mock_table_provider; use crate::options::QueryOptions; @@ -745,8 +746,9 @@ mod tests { .unwrap() .build() .unwrap(); - let plan = - MergeScanLogicalPlan::new(input.clone(), true, Default::default()).into_logical_plan(); + let plan = MergeScanLogicalPlan::new(input.clone(), true, Default::default()) + .with_remote_dyn_filter_producer_id(RemoteDynFilterProducerId::new(7)) + .into_logical_plan(); let bytes = DFLogicalSubstraitConvertor .encode(&plan, DefaultSerializer) @@ -777,6 +779,9 @@ mod tests { assert_eq!(merge_scan.input().to_string(), input.to_string()); // `PbMergeScan` doesn't carry `partition_cols`, so decoded nodes have an empty mapping. assert!(merge_scan.partition_cols().is_empty()); + // `PbMergeScan` doesn't carry the remote dynamic filter producer id either: a decoded + // boundary has none, and `MergeScanExec` fails open for it. + assert_eq!(None, merge_scan.remote_dyn_filter_producer_id()); } /// Payload table binding uses the engine catalog, not the request catalog. diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index aa3b9417b14e..45530a7b8c97 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -55,8 +55,8 @@ use table::table::adapter::DfTableProviderAdapter; use crate::QueryEngineContext; use crate::dist_plan::{ - DistExtensionPlanner, DistPlannerAnalyzer, DistPlannerOptions, DynFilterRegistryManager, - MergeSortExtensionPlanner, RemoteDynFilterReceiverExtensionPlanner, + DistExtensionPlanner, DistJoinPlanner, DistPlannerAnalyzer, DistPlannerOptions, + DynFilterRegistryManager, MergeSortExtensionPlanner, RemoteDynFilterReceiverExtensionPlanner, RemoteDynFilterRegistryLease, }; use crate::metrics::{QUERY_MEMORY_POOL_REJECTED_TOTAL, QUERY_MEMORY_POOL_USAGE_BYTES}; @@ -161,6 +161,7 @@ impl QueryEngineState { .extensions .insert(DistPlannerOptions { allow_query_fallback: true, + ..Default::default() }); } @@ -210,6 +211,9 @@ impl QueryEngineState { ); if with_dist_planner { analyzer.rules.push(Arc::new(DistPlannerAnalyzer)); + // Runs after `DistPlannerAnalyzer`, which wraps the remote scans in + // `MergeScan`: the nested broadcast join shape is only visible then. + analyzer.rules.push(Arc::new(DistJoinPlanner)); analyzer.rules.push(Arc::new(JsonSchemaConcretizeRule)); } analyzer.rules.push(Arc::new(FixStateUdafOrderingAnalyzer)); diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 4e4f36820ff0..28834ac40bb9 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -58,6 +58,11 @@ //! * a cold load of the route of the build table returns the rows of the join after an unrelated //! table id is invalidated once (see //! [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]). +//! +//! [`test_nested_broadcast_join_rewrite_executes_on_probe_regions`] is the exception to the +//! hand-built plans: it plans the join SQL with the distributed rules of the frontend, checks that +//! the outer `MergeScan` of that plan is planned for the probe region the probe predicate selects, +//! and executes the plan on the frontend (see that test). use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; @@ -70,11 +75,16 @@ use common_meta::instruction::CacheIdent; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::info; +use datafusion::common::config::ConfigOptions; +use datafusion::optimizer::AnalyzerRule; +use datafusion::physical_plan::ExecutionPlan; use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; use datanode::region_server::RegionServer; use frontend::instance::Instance; use query::datafusion::QUERY_PARALLELISM_HINT; -use query::dist_plan::MergeScanLogicalPlan; +use query::dist_plan::{ + DistJoinPlanner, DistPlannerAnalyzer, DistPlannerOptions, MergeScanExec, MergeScanLogicalPlan, +}; use query::parser::QueryLanguageParser; use query::query_engine::DefaultSerializer; use servers::query_handler::sql::SqlQueryHandler; @@ -311,6 +321,224 @@ async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation /// A failing inner region makes the whole query fail: the datanode must not return the rows of the /// inner regions that succeed. +/// The distributed rules produce the nested plan of the join SQL, the frontend plans the outer +/// `MergeScan` of that plan to the *selected* probe regions, and the datanodes execute it. +/// +/// Unlike the hand-built plans of the capability tests above, the plan here is the plan of a SQL +/// query: `DistPlannerAnalyzer` wraps the table scans of the join in `MergeScan`, and +/// `DistJoinPlanner` (opted in for the build table) nests the build side into the probe side's +/// scan. The probe predicate of the SQL selects one of the two probe regions, so the outer +/// `MergeScan` must be planned for that region only: routing by the whole join payload would read +/// both probe regions. The executed result must be the multiset of the equivalent join on the +/// frontend, which also covers the duplicated build keys, the NULL join keys, the residual join +/// filter and the projection of the query. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { + common_telemetry::init_default_ut_logging(); + + let cluster = + build_cluster("test_nested_broadcast_join_rewrite_executes_on_probe_regions").await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + // A NULL join key on both sides and a second probe row of the duplicated join key `10`. They + // are inserted here, so the shared fixture of the capability tests keeps its rows. + execute_sql( + &frontend, + &format!("INSERT INTO {PROBE_TABLE}(a_id, probe_key, probe_v, a_ts) VALUES (3, NULL, 500, 3000), (4, 10, 600, 4000)"), + ) + .await; + execute_sql( + &frontend, + &format!( + "INSERT INTO {BUILD_TABLE}(b_id, build_key, build_v, b_ts) VALUES (4, NULL, 6, 5000)" + ), + ) + .await; + + let sql = rewritten_join_sql(); + let query_ctx = query_ctx(); + + // The reference: the frontend executes the same SQL with its default distributed plan, i.e. + // both sides of the join are read with a `MergeScan` and the join runs on the frontend. + let reference = query_pretty(&frontend, &sql, query_ctx.clone()).await; + info!("reference result of the frontend:\n{reference}"); + assert_eq!( + expected_rewritten_join_rows(), + multiset(table_cells(&reference)), + "unexpected result of the join on the frontend:\n{reference}" + ); + + // The plan of the rules, not a hand-built plan. + let plan = rewritten_join_plan(&frontend, &query_ctx, &sql).await; + let outer = outer_merge_scan(&plan).unwrap_or_else(|| { + panic!("the rules must nest the join into the probe side's MergeScan, got: {plan}") + }); + assert!( + matches!(outer.input(), LogicalPlan::Join(_)), + "the outer MergeScan must carry the join, got: {}", + outer.input() + ); + + // Physical planning of the outer `MergeScan`: the probe predicate selects one of the two probe + // regions, so only that region is dispatched. + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + assert!( + probe_leaders.len() > 1, + "expected {PROBE_TABLE} to have several regions, actual region leaders: {probe_leaders:?}" + ); + let physical = create_physical_plan(&frontend, &plan, &query_ctx).await; + let merge_scan = find_merge_scan_exec(&physical) + .unwrap_or_else(|| panic!("expected the outer MergeScan to be planned, got: {physical:?}")); + let regions = merge_scan.regions().to_vec(); + info!("the outer MergeScan is planned for the regions {regions:?}"); + assert_eq!( + 1, + regions.len(), + "the probe predicate `a_id < 100` must select one of the {} probe regions, actual \ + selection: {regions:?}", + probe_leaders.len() + ); + assert!( + regions + .iter() + .all(|region| probe_leaders.contains_key(®ion.as_u64())), + "the outer MergeScan must select regions of {PROBE_TABLE}, actual selection: {regions:?}" + ); + + // Execution: the frontend serializes the payload of the outer `MergeScan` to the datanode that + // owns the selected region, and the nested `MergeScan` of that payload reads the build table. + let actual = output_batches( + frontend + .query_engine() + .execute(plan, query_ctx) + .await + .expect("the nested plan must be executable"), + ) + .await + .pretty_print() + .expect("the result of the nested plan must be printable"); + info!("result of the rule-produced nested plan:\n{actual}"); + assert_eq!( + expected_rewritten_join_rows(), + multiset(table_cells(&actual)), + "the nested plan returned an unexpected multiset of rows:\n{actual}" + ); + assert_eq!( + multiset(table_cells(&reference)), + multiset(table_cells(&actual)), + "the nested plan returned a different multiset of rows than the frontend:\n{actual}" + ); +} + +/// The SQL of [`test_nested_broadcast_join_rewrite_executes_on_probe_regions`], over the columns of +/// [`expected_rewritten_join_rows`] in order. +/// +/// The projection reorders the columns of the join, the join keys `10` of the build table are +/// duplicated, the residual join filter rejects the join partners whose values sum to at most +/// `105` (the probe row `100` sum up to `101` and `102`), and the probe predicate selects the +/// first probe region. +fn rewritten_join_sql() -> String { + format!( + "SELECT p.probe_v, p.a_id, b.build_v + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b + ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 + WHERE p.a_id < 100" + ) +} + +/// The rows of [`rewritten_join_sql`] as `(probe_v, a_id, build_v)`: the probe row `(4, 10, 600)` +/// joins the two build rows of the duplicated key `10` (`600 + 1` and `600 + 2`), the probe row +/// `(2, 20, 200)` joins `(3, 20, 3)`, the residual filter rejects the probe row `(1, 10, 100)` +/// (`100 + 1` and `100 + 2` are not greater than `105`), and the NULL join keys of both sides +/// match nothing. +fn expected_rewritten_join_rows() -> Vec> { + multiset(vec![ + vec!["200".to_string(), "2".to_string(), "3".to_string()], + vec!["600".to_string(), "4".to_string(), "1".to_string()], + vec!["600".to_string(), "4".to_string(), "2".to_string()], + ]) +} + +/// Plans `sql` on the frontend and applies the distributed rules of the frontend: the +/// [`DistPlannerAnalyzer`], which wraps the table scans of both sides of the join in `MergeScan`, +/// and the [`DistJoinPlanner`] opted in for the build table. +async fn rewritten_join_plan( + frontend: &Arc, + query_ctx: &QueryContextRef, + sql: &str, +) -> LogicalPlan { + let plan = plan_sql(frontend, sql, query_ctx).await; + info!("logical plan of the join SQL:\n{plan}"); + + let mut config = ConfigOptions::default(); + config.extensions.insert(DistPlannerOptions { + nested_broadcast_join_build_table: Some(BUILD_TABLE.to_string()), + ..Default::default() + }); + let plan = DistPlannerAnalyzer {} + .analyze(plan, &config) + .expect("the distributed plan of the join must be analyzable"); + let plan = DistJoinPlanner {} + .analyze(plan, &config) + .expect("the nested broadcast join rewrite must not fail"); + info!("plan of the distributed rules:\n{plan}"); + + plan +} + +/// Returns the outermost `MergeScan` of `plan`, descending through projections. +fn outer_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { + match plan { + LogicalPlan::Extension(extension) => extension + .node + .as_any() + .downcast_ref::(), + LogicalPlan::Projection(projection) => outer_merge_scan(&projection.input), + _ => None, + } +} + +/// Creates the physical plan of `plan` the way the query engine does: the session analyzer of the +/// frontend, the optimizer (skipped for a `MergeScan` root) and the physical planner with its +/// distributed extension planner. +async fn create_physical_plan( + frontend: &Arc, + plan: &LogicalPlan, + query_ctx: &QueryContextRef, +) -> Arc { + let ctx = frontend.query_engine().engine_context(query_ctx.clone()); + let state = ctx.state(); + let analyzed = state + .analyzer() + .execute_and_check(plan.clone(), state.config_options(), |_, _| {}) + .expect("the nested plan must pass the session analyzer"); + let optimized = if let LogicalPlan::Extension(extension) = &analyzed + && extension.node.name() == MergeScanLogicalPlan::name() + { + analyzed.clone() + } else { + state + .optimizer() + .optimize(analyzed, state, |_, _| {}) + .expect("the nested plan must be optimizable") + }; + state + .query_planner() + .create_physical_plan(&optimized, state) + .await + .expect("the nested plan must be physically plannable") +} + +/// Returns the `MergeScanExec` of `plan`, descending through the physical plan tree. +fn find_merge_scan_exec(plan: &Arc) -> Option<&MergeScanExec> { + if let Some(merge_scan) = plan.downcast_ref::() { + return Some(merge_scan); + } + + plan.children().into_iter().find_map(find_merge_scan_exec) +} + #[tokio::test(flavor = "multi_thread")] async fn test_nested_merge_scan_capability_inner_failure_fails_query() { common_telemetry::init_default_ut_logging(); From 76071571d530591594bd54b9a6f5e7e7372563fe Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 29 Sep 2026 15:53:30 +0800 Subject: [PATCH 13/37] feat(query): enable opt-in distributed join cost heuristic Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + src/operator/src/statement.rs | 5 +- src/operator/src/statement/set.rs | 75 +- src/query/Cargo.toml | 1 + src/query/src/datafusion.rs | 740 +++++++++++++++- src/query/src/dist_plan.rs | 6 + src/query/src/dist_plan/dist_join_planner.rs | 815 +++++++++++++++++- src/query/src/query_engine/context.rs | 9 + src/session/src/context.rs | 37 + .../nested_merge_scan_capability_test.rs | 371 ++++++-- 10 files changed, 1962 insertions(+), 98 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c16572c2fcc9..04cbd1c2a678 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11874,6 +11874,7 @@ dependencies = [ "common-macro", "common-meta", "common-plugins", + "common-procedure", "common-query", "common-recordbatch", "common-runtime", diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 17d29ab4259b..eb7e9df3647f 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -100,7 +100,7 @@ use crate::error::{ }; use crate::insert::InserterRef; use crate::statement::copy_database::{COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY}; -use crate::statement::set::set_allow_query_fallback; +use crate::statement::set::{set_allow_query_fallback, set_experimental_dist_join}; /// A configurator that customizes or enhances a [`StatementExecutor`]. #[async_trait::async_trait] @@ -555,6 +555,9 @@ impl StatementExecutor { // Allow query to fallback when failed to push down. "ALLOW_QUERY_FALLBACK" => set_allow_query_fallback(set_var.value, query_ctx)?, + // Opt into the cost heuristic of the nested broadcast join rewrite. + "EXPERIMENTAL_DIST_JOIN" => set_experimental_dist_join(set_var.value, query_ctx)?, + "CLIENT_ENCODING" => validate_client_encoding(set_var)?, "@@SESSION.MAX_EXECUTION_TIME" | "MAX_EXECUTION_TIME" => match query_ctx.channel() { Channel::Mysql => set_query_timeout(set_var.value, query_ctx)?, diff --git a/src/operator/src/statement/set.rs b/src/operator/src/statement/set.rs index 83924953e755..877c6c9ad66b 100644 --- a/src/operator/src/statement/set.rs +++ b/src/operator/src/statement/set.rs @@ -278,6 +278,33 @@ pub fn set_allow_query_fallback(exprs: Vec, ctx: QueryContextRef) -> Resul } } +/// Sets the session `experimental_dist_join` configuration parameter (disabled by default). +pub fn set_experimental_dist_join(exprs: Vec, ctx: QueryContextRef) -> Result<()> { + let Some((dist_join_expr, [])) = exprs.split_first() else { + return NotSupportedSnafu { + feat: "Set variable value must have one and only one value for experimental_dist_join", + } + .fail(); + }; + match dist_join_expr { + Expr::Value(ValueWithSpan { + value: Value::Boolean(enabled), + span: _, + }) => { + ctx.configuration_parameter() + .set_experimental_dist_join(*enabled); + Ok(()) + } + expr => NotSupportedSnafu { + feat: format!( + "Unsupported experimental dist join expr {} in set variable statement", + expr + ), + } + .fail(), + } +} + pub fn set_intervalstyle(exprs: Vec, ctx: QueryContextRef) -> Result<()> { let Some((var_value, [])) = exprs.split_first() else { return NotSupportedSnafu { @@ -415,7 +442,7 @@ mod test { use sql::ast::{Expr, Value}; use super::set_skip_wal; - use crate::statement::set::parse_pg_query_timeout_input; + use crate::statement::set::{parse_pg_query_timeout_input, set_experimental_dist_join}; #[test] fn test_set_skip_wal() { @@ -462,6 +489,52 @@ mod test { } } + #[test] + fn test_set_experimental_dist_join() { + let session = Session::new(None, Channel::Mysql, Default::default(), 0); + let other = Session::new(None, Channel::Mysql, Default::default(), 1); + let dist_join = |session: &Session| { + session + .new_query_context() + .configuration_parameter() + .experimental_dist_join() + }; + assert!(!dist_join(&session)); + + set_experimental_dist_join( + vec![Expr::Value(Value::Boolean(true).into())], + session.new_query_context(), + ) + .unwrap(); + // The value belongs to the session: later statements see it, other sessions do not. + assert!(dist_join(&session)); + assert!(!dist_join(&other)); + + set_experimental_dist_join( + vec![Expr::Value(Value::Boolean(false).into())], + session.new_query_context(), + ) + .unwrap(); + assert!(!dist_join(&session)); + + for values in [ + vec![], + vec![Expr::Value(Value::Number("1".to_string(), false).into())], + vec![Expr::Value( + Value::SingleQuotedString("true".to_string()).into(), + )], + // `true` first: if multiple values were accepted, the session would be + // mutated and the final assertion below would catch it. + vec![ + Expr::Value(Value::Boolean(true).into()), + Expr::Value(Value::Boolean(false).into()), + ], + ] { + assert!(set_experimental_dist_join(values, session.new_query_context()).is_err()); + } + assert!(!dist_join(&session)); + } + #[test] fn test_parse_pg_query_timeout_input() { assert!(parse_pg_query_timeout_input("").is_err()); diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index 1d95f7856d53..6f6896e0b1ef 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -84,6 +84,7 @@ arrow.workspace = true catalog = { workspace = true, features = ["testing"] } common-function.workspace = true common-macro.workspace = true +common-procedure.workspace = true common-query = { workspace = true, features = ["testing"] } fastrand = "2.0" moka = { workspace = true, features = ["future"] } diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index ea2a70c5824e..771e055dc21a 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -20,10 +20,11 @@ mod pg_oid_alias_expr_planner; mod planner; use std::any::Any; -use std::collections::HashMap; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; use async_trait::async_trait; +use catalog::kvbackend::KvBackendCatalogManager; use common_base::Plugins; use common_catalog::consts::is_readonly_table; use common_error::ext::BoxedError; @@ -56,7 +57,9 @@ use tracing::Span; use crate::analyze::DistAnalyzeExec; pub use crate::datafusion::planner::DfContextProviderAdapter; use crate::dist_plan::{ - DistPlannerOptions, MergeScanLogicalPlan, RemoteDynFilterReceiverInjectorRef, + DIST_JOIN_PLANNER_RULE_NAME, DistJoinStats, DistPlannerOptions, MergeScanLogicalPlan, + RemoteDynFilterReceiverInjectorRef, aggregate_region_stats, candidate_joins, + expected_region_ids, }; use crate::error::{ CatalogSnafu, CreateRecordBatchSnafu, MissingTableMutationHandlerSnafu, @@ -149,6 +152,121 @@ impl DatafusionQueryEngine { Self { state, plugins } } + /// Statistics of the candidate tables of `logical_plan`, or `None` when the nested broadcast + /// join rewrite has to keep the existing plan. + /// + /// Gated by the session opt-in and the registered distributed rules; the manual build table + /// option selects the build side by name and never consults statistics. + /// + /// The candidate tables are resolved to their physical routes before the statistics of the + /// query are fetched: the route must still describe the table the plan captured and at least + /// one candidate pair must have usable routes, otherwise this is `None` without any + /// `region_stats` call. The region sizes of the resolved tables then come from one + /// `region_stats` call of the information extension, so the result belongs to this query only + /// and never reaches the shared engine state. The rewrite is an optimization, so a table + /// whose statistics are unusable (missing, duplicate, zero-sized or overflowing) is left out + /// instead of failing the query. + async fn dist_join_stats( + &self, + ctx: &QueryEngineContext, + logical_plan: &LogicalPlan, + ) -> Option { + if !ctx + .query_ctx() + .configuration_parameter() + .experimental_dist_join() + { + return None; + } + + let state = ctx.state(); + if !state + .analyzer() + .rules + .iter() + .any(|rule| rule.name() == DIST_JOIN_PLANNER_RULE_NAME) + { + return None; + } + if state + .config() + .options() + .extensions + .get::() + .and_then(|options| options.nested_broadcast_join_build_table.as_ref()) + .is_some() + { + return None; + } + + let candidates = candidate_joins(logical_plan); + if candidates.is_empty() { + return None; + } + + let catalog_manager = self + .state + .catalog_manager() + .as_any() + .downcast_ref::()?; + let partition_manager = catalog_manager.partition_manager(); + + // Resolve the routes before fetching the statistics: a table whose route does not + // describe the id the plan captured is not priced, and without one candidate pair whose + // two sides both have usable routes the statistics cannot decide the rewrite. + let mut routes = BTreeMap::new(); + for table_id in candidates + .iter() + .flat_map(|candidate| [candidate.probe, candidate.build]) + .collect::>() + { + let Ok((physical_table_id, route)) = partition_manager + .find_physical_table_route_with_id(table_id) + .await + else { + continue; + }; + // A logical table resolves to the route of its physical table; pricing its regions + // under the captured id would describe another table. + if physical_table_id != table_id { + continue; + } + let Some(regions) = expected_region_ids(physical_table_id, &route.region_routes) else { + continue; + }; + + routes.insert(table_id, regions); + } + let has_usable_pair = candidates.iter().any(|candidate| { + routes.contains_key(&candidate.probe) && routes.contains_key(&candidate.build) + }); + if !has_usable_pair { + return None; + } + + let reports = match catalog_manager.information_extension().region_stats().await { + Ok(reports) => reports, + Err(err) => { + common_telemetry::debug!( + err = %err, + "Failed to fetch the region statistics of the nested broadcast join, keeping the existing plan" + ); + return None; + } + }; + + let mut stats = DistJoinStats::default(); + for (table_id, regions) in routes { + let Some(table_stats) = aggregate_region_stats(®ions, &reports) else { + continue; + }; + + stats.tables.insert(table_id, table_stats); + } + + (!stats.tables.is_empty()).then_some(stats) + } + #[tracing::instrument(skip_all)] async fn exec_query_plan( &self, @@ -389,10 +507,22 @@ impl DatafusionQueryEngine { }; let _timer = metrics::CREATE_PHYSICAL_ELAPSED.start_timer(); - let state = ctx.state(); common_telemetry::debug!("Create physical plan, input plan: {logical_plan}"); + // The nested broadcast join rewrite reads its statistics from a query-local config + // extension. This has to happen before the EXPLAIN branch below: the input of an EXPLAIN + // plan is analyzed by the physical planner of DataFusion, i.e. outside this function. + if let Some(stats) = self.dist_join_stats(ctx, logical_plan).await { + ctx.state_mut() + .config_mut() + .options_mut() + .extensions + .insert(stats); + } + + let state = ctx.state(); + // special handle EXPLAIN plan if matches!(logical_plan, DfLogicalPlan::Explain(_)) { return state @@ -745,20 +875,37 @@ impl QueryExecutor for DatafusionQueryEngine { #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::fmt; use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use api::v1::SemanticType; - use arrow::array::{ArrayRef, UInt64Array}; + use arrow::array::{ArrayRef, StringArray, UInt64Array}; use arrow_schema::SortOptions; use async_trait::async_trait; use catalog::RegisterTableRequest; + use catalog::information_schema::{ + DatanodeInspectRequest, InformationExtension, InformationExtensionRef, + }; + use catalog::kvbackend::KvBackendCatalogManagerBuilder; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; use common_error::ext::BoxedError; + use common_meta::cache::{ + CacheRegistryBuilder, LayeredCacheRegistryBuilder, new_table_route_cache, + }; + use common_meta::cluster::NodeInfo; + use common_meta::datanode::{RegionManifestInfo, RegionStat}; + use common_meta::key::flow::flow_state::FlowStat; + use common_meta::key::table_route::{TableRouteManager, TableRouteValue}; + use common_meta::kv_backend::TxnService; + use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::rpc::router::{Region, RegionRoute}; + use common_procedure::ProcedureInfo; use common_recordbatch::{ EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream, util, }; + use datafusion::datasource::DefaultTableSource; use datafusion::physical_plan::display::{DisplayAs, DisplayFormatType}; use datafusion::physical_plan::expressions::PhysicalSortExpr; use datafusion::physical_plan::joins::{HashJoinExec, JoinOn, PartitionMode}; @@ -766,19 +913,25 @@ mod tests { use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr}; use datafusion::prelude::{col, lit}; use datafusion_common::{JoinType, NullEquality, ScalarValue}; + use datafusion_expr::LogicalPlanBuilder; use datafusion_physical_expr::expressions::Column; use datatypes::prelude::ConcreteDataType; - use datatypes::schema::{ColumnSchema, SchemaRef}; + use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; use datatypes::vectors::{Helper, UInt32Vector, VectorRef}; + use partition::cache::new_partition_info_cache; use session::context::{QueryContext, QueryContextBuilder}; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; use store_api::region_engine::{ - PartitionRange, PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties, + PartitionRange, PrepareRequest, QueryScanContext, RegionRole, RegionScanner, + ScannerProperties, }; use store_api::storage::{RegionId, ScanRequest}; - use table::metadata::{TableInfoBuilder, TableMetaBuilder}; + use table::metadata::{TableId, TableInfoBuilder, TableMetaBuilder}; + use table::table::adapter::DfTableProviderAdapter; use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable}; use table::table::scan::RegionScanExec; + use table::test_util::EmptyTable; + use table::test_util::table_info::test_table_info; use super::*; use crate::options::QueryOptions; @@ -1701,4 +1854,577 @@ mod tests { ); } } + + /// The table ids of the two tables of the join of the automatic tests. + const PROBE_TABLE_ID: TableId = 1024; + const BUILD_TABLE_ID: TableId = 1025; + + /// The physical table id a logical candidate table resolves to in the route gating test. + const LOGICAL_PHYSICAL_TABLE_ID: TableId = 1026; + + /// An information extension with fixed region statistics and a counter of its `region_stats` + /// calls. + #[derive(Debug)] + struct TestInformationExtension { + reports: Vec, + calls: Arc, + } + + #[async_trait] + impl InformationExtension for TestInformationExtension { + type Error = catalog::error::Error; + + async fn nodes(&self) -> std::result::Result, Self::Error> { + Ok(vec![]) + } + + async fn procedures( + &self, + ) -> std::result::Result, Self::Error> { + Ok(vec![]) + } + + async fn region_stats(&self) -> std::result::Result, Self::Error> { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(self.reports.clone()) + } + + async fn flow_stats(&self) -> std::result::Result, Self::Error> { + Ok(None) + } + + async fn inspect_datanode( + &self, + _request: DatanodeInspectRequest, + ) -> std::result::Result { + Ok(common_recordbatch::RecordBatches::empty().as_stream()) + } + } + + /// A region report with `bytes` approximate disk bytes and `role`. + fn region_stat(region_id: RegionId, bytes: u64, role: RegionRole) -> RegionStat { + RegionStat { + id: region_id, + rcus: 0, + wcus: 0, + approximate_bytes: bytes, + engine: "mito".to_string(), + role, + num_rows: 0, + memtable_size: 0, + manifest_size: 0, + sst_size: 0, + sst_num: 0, + index_size: 0, + region_manifest: RegionManifestInfo::Mito { + manifest_version: 0, + flushed_entry_id: 0, + file_removed_cnt: 0, + }, + written_bytes: 0, + query_cpu_time: 0, + query_scanned_bytes: 0, + data_topic_latest_entry_id: 0, + metadata_topic_latest_entry_id: 0, + min_timestamp: None, + max_timestamp: None, + } + } + + /// The region statistics of a complete query: the probe table has two regions of 4_000 and + /// 6_000 bytes, the build table one region of 1_000 bytes. + fn dist_join_reports() -> Vec { + vec![ + region_stat(RegionId::new(PROBE_TABLE_ID, 1), 4_000, RegionRole::Leader), + region_stat(RegionId::new(PROBE_TABLE_ID, 2), 6_000, RegionRole::Leader), + region_stat(RegionId::new(BUILD_TABLE_ID, 1), 1_000, RegionRole::Leader), + ] + } + + /// Two-column schema (`number`, `host`) of the tables of the join of the tests. + fn dist_join_test_schema() -> SchemaRef { + let schema = SchemaBuilder::try_from_columns(vec![ + ColumnSchema::new("number", ConcreteDataType::uint32_datatype(), true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ]) + .unwrap() + .build() + .unwrap(); + Arc::new(schema) + } + + /// Base table `name` with `table_id`, backed by an empty table provider. + fn dist_join_test_table(table_id: TableId, name: &str) -> table::TableRef { + let info = test_table_info( + table_id, + name, + DEFAULT_SCHEMA_NAME, + DEFAULT_CATALOG_NAME, + dist_join_test_schema(), + ); + EmptyTable::from_table_info(&info) + } + + /// A scan of one of the tables of the join of the tests. + fn dist_join_test_scan(alias: &str, table_id: TableId, name: &str) -> LogicalPlan { + let source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(dist_join_test_table(table_id, name)), + ))); + LogicalPlanBuilder::scan_with_filters(alias, source, None, vec![]) + .unwrap() + .build() + .unwrap() + } + + /// `probe INNER JOIN build ON probe.number = build.number`: the shape `DistPlannerAnalyzer` + /// wraps in `MergeScan` and `DistJoinPlanner` can then nest. + fn dist_join_test_plan() -> LogicalPlan { + LogicalPlanBuilder::from(dist_join_test_scan("probe", PROBE_TABLE_ID, "probe")) + .join_on( + dist_join_test_scan("build", BUILD_TABLE_ID, "build"), + JoinType::Inner, + vec![col("probe.number").eq(col("build.number"))], + ) + .unwrap() + .build() + .unwrap() + } + + /// A query context with or without the session opt-in of the rewrite. + fn dist_join_query_ctx(opt_in: bool) -> QueryContextRef { + let query_ctx = QueryContextBuilder::default().build(); + query_ctx + .configuration_parameter() + .set_experimental_dist_join(opt_in); + Arc::new(query_ctx) + } + + /// The physical route of the table `table_id` with the given region numbers. + fn dist_join_physical_route( + table_id: TableId, + region_numbers: &[u32], + ) -> (TableId, TableRouteValue) { + let routes = region_numbers + .iter() + .map(|region_number| RegionRoute { + region: Region { + id: RegionId::new(table_id, *region_number), + ..Default::default() + }, + ..Default::default() + }) + .collect(); + (table_id, TableRouteValue::physical(routes)) + } + + /// The physical routes of the tables of the join of the tests: the probe table has the regions + /// 1 and 2, the build table the region 1. + fn dist_join_test_routes() -> Vec<(TableId, TableRouteValue)> { + vec![ + dist_join_physical_route(PROBE_TABLE_ID, &[1, 2]), + dist_join_physical_route(BUILD_TABLE_ID, &[1]), + ] + } + + /// A catalog manager whose KV backend holds the given table routes. Its region statistics are + /// the given reports. + async fn dist_join_catalog_manager( + routes: Vec<(TableId, TableRouteValue)>, + reports: Vec, + calls: Arc, + ) -> Arc { + let kv_backend = Arc::new(MemoryKvBackend::default()); + let table_route_manager = TableRouteManager::new(kv_backend.clone()); + for (table_id, table_route) in routes { + let (txn, _) = table_route_manager + .table_route_storage() + .build_create_txn(table_id, &table_route) + .unwrap(); + assert!(kv_backend.txn(txn).await.unwrap().succeeded); + } + + let table_route_cache = Arc::new(new_table_route_cache( + "test_table_route".to_string(), + moka::future::Cache::new(16), + kv_backend.clone(), + )); + let partition_info_cache = Arc::new(new_partition_info_cache( + "test_partition_info".to_string(), + moka::future::Cache::new(16), + table_route_cache.clone(), + )); + let cache_registry = LayeredCacheRegistryBuilder::default() + .add_cache_registry( + CacheRegistryBuilder::default() + .add_cache(table_route_cache) + .add_cache(partition_info_cache) + .build(), + ) + .build(); + let information_extension: InformationExtensionRef = + Arc::new(TestInformationExtension { reports, calls }); + + KvBackendCatalogManagerBuilder::new( + information_extension, + kv_backend, + Arc::new(cache_registry), + ) + .build() + } + + /// A query engine over the catalog manager of the tests, with the distributed rules when + /// `with_dist_planner` is set. + async fn dist_join_engine( + reports: Vec, + calls: Arc, + with_dist_planner: bool, + ) -> QueryEngineRef { + dist_join_engine_with_routes(dist_join_test_routes(), reports, calls, with_dist_planner) + .await + } + + /// A query engine over a catalog manager with the given table `routes`. + async fn dist_join_engine_with_routes( + routes: Vec<(TableId, TableRouteValue)>, + reports: Vec, + calls: Arc, + with_dist_planner: bool, + ) -> QueryEngineRef { + let catalog_manager = dist_join_catalog_manager(routes, reports, calls).await; + QueryEngineFactory::new_with_plugins( + catalog_manager.clone(), + Some(catalog_manager.partition_manager()), + None, + None, + None, + None, + with_dist_planner, + Plugins::new(), + QueryOptions::default(), + ) + .query_engine() + } + + /// The optimized logical plan of the EXPLAIN output of `plan`, as the query engine produces + /// it: `DatafusionQueryEngine::create_physical_plan` handles EXPLAIN, whose input plan is + /// analyzed and optimized (i.e. the distributed rules including `DistJoinPlanner` run on it) + /// before its stages are printed. + async fn explain_logical_plan( + engine: &DatafusionQueryEngine, + ctx: &mut QueryEngineContext, + plan: &LogicalPlan, + ) -> String { + let explain_plan = LogicalPlanBuilder::from(plan.clone()) + .explain(false, false) + .unwrap() + .build() + .unwrap(); + let physical = engine + .create_physical_plan(ctx, &explain_plan) + .await + .unwrap(); + let batches = datafusion::physical_plan::collect(physical, ctx.build_task_ctx()) + .await + .unwrap(); + + for batch in &batches { + let plan_types = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let plans = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + if plan_types.value(row) == "logical_plan" { + return plans.value(row).to_string(); + } + } + } + + panic!("the EXPLAIN output has no logical plan: {batches:?}"); + } + + /// Whether an EXPLAIN of a join shows the nested rewrite: the join sits inside a `MergeScan`, + /// so the enclosing boundary appears before the join in the plan text. The join of two + /// boundaries, i.e. the plan without the rewrite, shows the join first. + fn nests_join_in_merge_scan(plan: &str) -> bool { + match (plan.find("Join:"), plan.find("MergeScan")) { + (Some(join), Some(merge_scan)) => merge_scan < join, + _ => false, + } + } + + /// The automatic path of the statistics of the nested broadcast join rewrite: the engine + /// resolves the candidate tables to their physical routes, fetches the region statistics of + /// the query once, and the EXPLAIN of the same query then shows the nested rewrite. + #[tokio::test] + async fn test_dist_join_stats_from_physical_routes() { + common_telemetry::init_default_ut_logging(); + + let calls = Arc::new(AtomicUsize::new(0)); + let engine = dist_join_engine(dist_join_reports(), calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let mut ctx = query_engine.engine_context(dist_join_query_ctx(true)); + let plan = dist_join_test_plan(); + + let stats = query_engine.dist_join_stats(&ctx, &plan).await.unwrap(); + assert_eq!( + BTreeMap::from([ + ( + PROBE_TABLE_ID, + crate::dist_plan::DistJoinTableStats { + total_bytes: 10_000, + region_count: 2, + } + ), + ( + BUILD_TABLE_ID, + crate::dist_plan::DistJoinTableStats { + total_bytes: 1_000, + region_count: 1, + } + ), + ]), + stats.tables + ); + assert!(stats.favors_right_build(PROBE_TABLE_ID, BUILD_TABLE_ID)); + assert_eq!( + 1, + calls.load(Ordering::SeqCst), + "the statistics of one query must come from one region stats call" + ); + + // The same query through the EXPLAIN path of the engine. + let explained = explain_logical_plan(query_engine, &mut ctx, &plan).await; + assert!( + nests_join_in_merge_scan(&explained), + "the EXPLAIN must show the nested rewrite, got:\n{explained}" + ); + assert_eq!(2, calls.load(Ordering::SeqCst)); + } + + /// The gating of the automatic path: without the session opt-in, without the distributed + /// rules, with the manual build table option or without a candidate join the engine does not + /// fetch any statistics. + #[tokio::test] + async fn test_dist_join_stats_gating() { + common_telemetry::init_default_ut_logging(); + + // Without the session opt-in. + let calls = Arc::new(AtomicUsize::new(0)); + let engine = dist_join_engine(dist_join_reports(), calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let plan = dist_join_test_plan(); + let ctx = query_engine.engine_context(dist_join_query_ctx(false)); + assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); + assert_eq!(0, calls.load(Ordering::SeqCst)); + + // Without the distributed rules. + let calls = Arc::new(AtomicUsize::new(0)); + let engine = dist_join_engine(dist_join_reports(), calls.clone(), false).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let ctx = query_engine.engine_context(dist_join_query_ctx(true)); + assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); + assert_eq!(0, calls.load(Ordering::SeqCst)); + + // With the manual build table option, which selects the build side by name. + let calls = Arc::new(AtomicUsize::new(0)); + let engine = dist_join_engine(dist_join_reports(), calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let mut ctx = query_engine.engine_context(dist_join_query_ctx(true)); + ctx.state_mut() + .config_mut() + .options_mut() + .extensions + .insert(DistPlannerOptions { + nested_broadcast_join_build_table: Some("build".to_string()), + ..Default::default() + }); + assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); + assert_eq!(0, calls.load(Ordering::SeqCst)); + + // Without a candidate join. + let ctx = query_engine.engine_context(dist_join_query_ctx(true)); + let scan_plan = dist_join_test_scan("probe", PROBE_TABLE_ID, "probe"); + assert!( + query_engine + .dist_join_stats(&ctx, &scan_plan) + .await + .is_none() + ); + assert_eq!(0, calls.load(Ordering::SeqCst)); + + // The same plan runs through the EXPLAIN path without the opt-in: the join keeps its two + // boundaries. + let mut ctx = query_engine.engine_context(dist_join_query_ctx(false)); + let explained = explain_logical_plan(query_engine, &mut ctx, &plan).await; + assert!( + !nests_join_in_merge_scan(&explained), + "without the opt-in the join must keep its two MergeScan boundaries, got:\n{explained}" + ); + } + + /// An incomplete description of a candidate table (here a follower report instead of the + /// leader of the build table) leaves the table out of the statistics, so the EXPLAIN of the + /// same query keeps the existing plan. + #[tokio::test] + async fn test_dist_join_stats_missing_description_keeps_plan() { + common_telemetry::init_default_ut_logging(); + + let calls = Arc::new(AtomicUsize::new(0)); + let mut reports = dist_join_reports(); + reports[2].role = RegionRole::Follower; + let engine = dist_join_engine(reports, calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let mut ctx = query_engine.engine_context(dist_join_query_ctx(true)); + let plan = dist_join_test_plan(); + + let stats = query_engine.dist_join_stats(&ctx, &plan).await.unwrap(); + assert_eq!(1, stats.tables.len()); + assert!(!stats.favors_right_build(PROBE_TABLE_ID, BUILD_TABLE_ID)); + + let explained = explain_logical_plan(query_engine, &mut ctx, &plan).await; + assert!( + !nests_join_in_merge_scan(&explained), + "an incomplete description must keep the existing plan, got:\n{explained}" + ); + assert_eq!(2, calls.load(Ordering::SeqCst)); + } + + /// A candidate table that resolves to another physical table id is not priced under the id the + /// plan captured: the engine keeps the existing plan without fetching any statistics. + #[tokio::test] + async fn test_dist_join_stats_logical_route_keeps_plan() { + common_telemetry::init_default_ut_logging(); + + // The probe table is logical and resolves to the physical table LOGICAL_PHYSICAL_TABLE_ID, + // whose regions the reports describe; the build table is a usable physical table. Pricing + // the resolved regions under the captured probe id would make the pair usable, so the + // engine must reject the resolution and fetch nothing. + let calls = Arc::new(AtomicUsize::new(0)); + let routes = vec![ + ( + PROBE_TABLE_ID, + TableRouteValue::logical(LOGICAL_PHYSICAL_TABLE_ID), + ), + dist_join_physical_route(LOGICAL_PHYSICAL_TABLE_ID, &[1, 2]), + dist_join_physical_route(BUILD_TABLE_ID, &[1]), + ]; + let reports = vec![ + region_stat( + RegionId::new(LOGICAL_PHYSICAL_TABLE_ID, 1), + 4_000, + RegionRole::Leader, + ), + region_stat( + RegionId::new(LOGICAL_PHYSICAL_TABLE_ID, 2), + 6_000, + RegionRole::Leader, + ), + region_stat(RegionId::new(BUILD_TABLE_ID, 1), 1_000, RegionRole::Leader), + ]; + let engine = dist_join_engine_with_routes(routes, reports, calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let ctx = query_engine.engine_context(dist_join_query_ctx(true)); + let plan = dist_join_test_plan(); + + assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); + assert_eq!(0, calls.load(Ordering::SeqCst)); + + // The logical probe is the only candidate table with a route: the pair is still unusable. + let calls = Arc::new(AtomicUsize::new(0)); + let routes = vec![ + ( + PROBE_TABLE_ID, + TableRouteValue::logical(LOGICAL_PHYSICAL_TABLE_ID), + ), + dist_join_physical_route(LOGICAL_PHYSICAL_TABLE_ID, &[1, 2]), + ]; + let engine = dist_join_engine_with_routes( + routes, + vec![ + region_stat( + RegionId::new(LOGICAL_PHYSICAL_TABLE_ID, 1), + 4_000, + RegionRole::Leader, + ), + region_stat( + RegionId::new(LOGICAL_PHYSICAL_TABLE_ID, 2), + 6_000, + RegionRole::Leader, + ), + ], + calls.clone(), + true, + ) + .await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); + assert_eq!(0, calls.load(Ordering::SeqCst)); + } + + /// A candidate pair with a missing or empty build route is left alone: the engine keeps the + /// existing plan without fetching any statistics. + #[tokio::test] + async fn test_dist_join_stats_unusable_route_keeps_plan() { + common_telemetry::init_default_ut_logging(); + + // The build table has no route. + let calls = Arc::new(AtomicUsize::new(0)); + let routes = vec![dist_join_physical_route(PROBE_TABLE_ID, &[1, 2])]; + let engine = + dist_join_engine_with_routes(routes, dist_join_reports(), calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let ctx = query_engine.engine_context(dist_join_query_ctx(true)); + let plan = dist_join_test_plan(); + + assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); + assert_eq!(0, calls.load(Ordering::SeqCst)); + + // The build table route has no region. + let calls = Arc::new(AtomicUsize::new(0)); + let routes = vec![ + dist_join_physical_route(PROBE_TABLE_ID, &[1, 2]), + dist_join_physical_route(BUILD_TABLE_ID, &[]), + ]; + let engine = + dist_join_engine_with_routes(routes, dist_join_reports(), calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let ctx = query_engine.engine_context(dist_join_query_ctx(true)); + + assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); + assert_eq!(0, calls.load(Ordering::SeqCst)); + } } diff --git a/src/query/src/dist_plan.rs b/src/query/src/dist_plan.rs index 9b32cb8775f6..92bd687a64db 100644 --- a/src/query/src/dist_plan.rs +++ b/src/query/src/dist_plan.rs @@ -27,6 +27,12 @@ mod remote_dyn_filter_registry; pub use analyzer::{DistPlannerAnalyzer, DistPlannerOptions}; pub use dist_join_planner::DistJoinPlanner; +#[cfg(test)] +pub(crate) use dist_join_planner::DistJoinTableStats; +pub(crate) use dist_join_planner::{ + DIST_JOIN_PLANNER_RULE_NAME, DistJoinStats, aggregate_region_stats, candidate_joins, + expected_region_ids, +}; pub use filter_id::{FilterFingerprint, FilterId, ParseFilterIdError, RemoteDynFilterProducerId}; pub use merge_scan::{MergeScanExec, MergeScanLogicalPlan}; pub(crate) use merge_sort::MergeSortExec; diff --git a/src/query/src/dist_plan/dist_join_planner.rs b/src/query/src/dist_plan/dist_join_planner.rs index af624503355a..e83656854c46 100644 --- a/src/query/src/dist_plan/dist_join_planner.rs +++ b/src/query/src/dist_plan/dist_join_planner.rs @@ -14,25 +14,302 @@ //! Nested broadcast join rewrite for `INNER` joins of two distributed scans. //! -//! Opt-in only: enabled when [`DistPlannerOptions::nested_broadcast_join_build_table`] -//! names the build side table. Runs after `DistPlannerAnalyzer`, which wraps the remote -//! scans in `MergeScan` and assigns the remote dynamic filter producer ids preserved -//! here. Not cost based: no statistics are consulted in this slice. - +//! Opt-in only, in two ways. The manual selector is enabled when +//! [`DistPlannerOptions::nested_broadcast_join_build_table`] names the build side table. The +//! cost heuristic is enabled per query when the session sets `experimental_dist_join`: +//! [`DatafusionQueryEngine::create_physical_plan`](crate::datafusion::DatafusionQueryEngine) +//! then fetches the approximate disk sizes of the candidate tables into [`DistJoinStats`], a +//! query-local config extension, and the heuristic rewrites the join only when those statistics +//! favor the right build side. Runs after `DistPlannerAnalyzer`, which wraps the remote scans in +//! `MergeScan` and assigns the remote dynamic filter producer ids preserved here. + +use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; +use common_meta::datanode::RegionStat; +use common_meta::rpc::router::RegionRoute; +use datafusion::config::{ConfigExtension, ExtensionOptions}; +use datafusion::datasource::DefaultTableSource; use datafusion::error::Result as DfResult; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; -use datafusion_expr::{Join, JoinType, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_expr::utils::{can_hash, find_valid_equijoin_key_pair, split_conjunction}; +use datafusion_expr::{ + Expr, ExprSchemable, Join, JoinType, LogicalPlan, Operator, TableScan, + UserDefinedLogicalNodeCore, +}; use datafusion_optimizer::analyzer::AnalyzerRule; use datatypes::extension::json::is_json2_extension_type; +use store_api::region_engine::RegionRole; +use store_api::storage::RegionId; +use table::metadata::{TableId, TableType}; +use table::table::adapter::DfTableProviderAdapter; use table::table_name::TableName; use crate::dist_plan::analyzer::DistPlannerOptions; use crate::dist_plan::merge_scan::MergeScanLogicalPlan; use crate::dist_plan::planner::table_name_of; +/// Name of [`DistJoinPlanner`]. The query engine looks it up to check that this rewrite, the +/// only consumer of [`DistJoinStats`], runs for the query at all. +pub(crate) const DIST_JOIN_PLANNER_RULE_NAME: &str = "DistJoinPlanner"; + +/// Region count and total disk size of a table, as the cost heuristic compares them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct DistJoinTableStats { + /// Total approximate disk bytes of all regions of the table. + pub total_bytes: u64, + /// Number of regions of the table. + pub region_count: usize, +} + +/// Per-table statistics fetched for the cost heuristic of one query. +/// +/// A query-local `ConfigOptions` extension, like [`DistPlannerOptions`], built when +/// [`DatafusionQueryEngine::create_physical_plan`](crate::datafusion::DatafusionQueryEngine) +/// sees the session opt-in (`SET experimental_dist_join = true`) on a plausible candidate join. +#[derive(Debug, Clone, Default)] +pub(crate) struct DistJoinStats { + /// Statistics per table id of the candidate tables of the query. A table is only usable as + /// a probe or build side when it is present here with a known, non-zero size. + pub tables: BTreeMap, +} + +impl ConfigExtension for DistJoinStats { + const PREFIX: &'static str = "dist_join"; +} + +/// Only satisfies the [`ConfigExtension`] bound: the statistics are built by the query engine +/// and have no configurable entry. +impl ExtensionOptions for DistJoinStats { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn set(&mut self, _key: &str, value: &str) -> DfResult<()> { + Err(datafusion_common::DataFusionError::NotImplemented(format!( + "DistJoinStats cannot be set (key: {value})" + ))) + } + + fn entries(&self) -> Vec { + vec![] + } +} + +impl DistJoinStats { + /// Whether `build` should be the nested build side of a join whose probe side is `probe`. + /// + /// The probe side selects the regions the join runs on, and the build side is read from + /// every one of them. The heuristic compares the disk bytes of the build table per probe + /// region, `N_probe * B_build`, with the disk bytes of the probe table itself, `P_probe`, and + /// only takes the rewrite when the product is strictly smaller. Those sizes are the + /// approximate disk sizes the regions report, not a measured or predicted amount of network + /// traffic. + /// + /// The sides are never swapped, and a missing table, a table without regions or bytes, or an + /// overflowing product means "no". + pub fn favors_right_build(&self, probe: TableId, build: TableId) -> bool { + let (Some(probe), Some(build)) = (self.tables.get(&probe), self.tables.get(&build)) else { + return false; + }; + if probe.region_count == 0 || probe.total_bytes == 0 || build.total_bytes == 0 { + return false; + } + let Some(per_probe_region_bytes) = + (probe.region_count as u64).checked_mul(build.total_bytes) + else { + return false; + }; + + per_probe_region_bytes < probe.total_bytes + } +} + +/// The base tables of one plausible candidate join of the cost heuristic. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct CandidateJoin { + /// The base table of the left (probe) side. + pub probe: TableId, + /// The base table of the right (build) side. + pub build: TableId, +} + +/// Returns the candidate joins of `plan`, i.e. of its `EXPLAIN`/`EXPLAIN ANALYZE` input: the +/// `INNER` joins whose sides each resolve to a single physical base table and whose condition +/// compares the two sides. +/// +/// The plan is the one before the analyzers run: DataFusion's SQL planner leaves the whole `ON` +/// clause of `LogicalPlanBuilder::join_on` in `filter`, and `ExtractEquijoinPredicate` only +/// moves the equijoin into `on` later, so both shapes count. This only decides whose statistics +/// are worth fetching; the rewrite still requires the distributed `MergeScan` shape and usable +/// statistics for the two tables. +pub(crate) fn candidate_joins(plan: &LogicalPlan) -> Vec { + fn collect(plan: &LogicalPlan, candidates: &mut Vec) { + if let LogicalPlan::Join(join) = plan + && join.join_type == JoinType::Inner + && has_join_equality(join) + && let (Some(probe), Some(build)) = ( + side_base_table_id(&join.left), + side_base_table_id(&join.right), + ) + { + candidates.push(CandidateJoin { probe, build }); + } + + for input in plan.inputs() { + collect(input, candidates); + } + } + + let plan = match plan { + LogicalPlan::Explain(explain) => explain.plan.as_ref(), + LogicalPlan::Analyze(analyze) => analyze.input.as_ref(), + plan => plan, + }; + let mut candidates = Vec::new(); + collect(plan, &mut candidates); + candidates +} + +/// Whether the join compares its two inputs with a hashable equality, in `on` or still in +/// `filter`. +/// +/// Mirrors `ExtractEquijoinPredicate`, the rule that fills `on`: +/// [`find_valid_equijoin_key_pair`] rejects top-level equalities that do not compare the two +/// inputs (constants, same-side columns), and `can_hash` rejects operand types a hash join +/// cannot use. +fn has_join_equality(join: &Join) -> bool { + let is_equijoin_key = |left: &Expr, right: &Expr| { + let Ok(Some((left, right))) = + find_valid_equijoin_key_pair(left, right, join.left.schema(), join.right.schema()) + else { + return false; + }; + let (Ok(left_type), Ok(right_type)) = ( + left.get_type(join.left.schema()), + right.get_type(join.right.schema()), + ) else { + return false; + }; + + can_hash(&left_type) && can_hash(&right_type) + }; + + if join + .on + .iter() + .any(|(left, right)| is_equijoin_key(left, right)) + { + return true; + } + + join.filter.as_ref().is_some_and(|filter| { + split_conjunction(filter).into_iter().any(|expr| { + matches!( + expr, + Expr::BinaryExpr(binary) + if binary.op == Operator::Eq && is_equijoin_key(&binary.left, &binary.right) + ) + }) + }) +} + +/// The physical base table of a join side, which may be aliased, filtered or projected. +fn side_base_table_id(plan: &LogicalPlan) -> Option { + match plan { + LogicalPlan::TableScan(scan) => physical_base_table_id(scan), + LogicalPlan::Projection(projection) => side_base_table_id(&projection.input), + LogicalPlan::Filter(filter) => side_base_table_id(&filter.input), + LogicalPlan::SubqueryAlias(alias) => side_base_table_id(&alias.input), + _ => None, + } +} + +/// The table id of a scan of a physical base table, the only tables [`DistJoinPlanner`] +/// captures. +fn physical_base_table_id(scan: &TableScan) -> Option { + let source = scan.source.downcast_ref::()?; + let provider = source + .table_provider + .downcast_ref::()?; + let table = provider.table(); + + (table.table_type() == TableType::Base).then(|| table.table_info().table_id()) +} + +/// Totals the `approximate_bytes` of the `expected_regions` from `reports`. +/// +/// Returns `None`, i.e. no rewrite, unless every expected region has exactly one ordinary +/// leader report with a non-zero size: a missing or duplicate leader report, a zero size, or +/// an overflowing sum does not describe the table well enough to compare it with another one. +pub(crate) fn aggregate_region_stats( + expected_regions: &[RegionId], + reports: &[RegionStat], +) -> Option { + let expected = expected_regions.iter().copied().collect::>(); + if expected.is_empty() || expected.len() != expected_regions.len() { + return None; + } + + let mut sizes = BTreeMap::new(); + for report in reports { + if !expected.contains(&report.id) || report.role != RegionRole::Leader { + continue; + } + // A second leader report leaves the size of the region unknown. + if sizes.insert(report.id, report.approximate_bytes).is_some() { + return None; + } + } + if sizes.len() != expected.len() { + return None; + } + + let mut total_bytes = 0u64; + for bytes in sizes.values() { + if *bytes == 0 { + return None; + } + total_bytes = total_bytes.checked_add(*bytes)?; + } + + Some(DistJoinTableStats { + total_bytes, + region_count: expected.len(), + }) +} + +/// The expected region ids of the table with the physical table id `physical_table_id`, taken +/// from its physical table route. +/// +/// Returns `None` when the route lists no region or any region belongs to another table: the +/// heuristic compares the complete region list of one table with another table, so a route that +/// does not describe `physical_table_id` (e.g. the id a logical table was resolved to) is not +/// usable. +pub(crate) fn expected_region_ids( + physical_table_id: TableId, + routes: &[RegionRoute], +) -> Option> { + let mut regions = Vec::with_capacity(routes.len()); + for route in routes { + if route.region.id.table_id() != physical_table_id { + return None; + } + regions.push(route.region.id); + } + + (!regions.is_empty()).then_some(regions) +} + /// Nests the build side's `MergeScan` inside the probe side's one, so the join runs on /// the datanodes holding the probe regions. /// @@ -43,32 +320,77 @@ pub struct DistJoinPlanner; impl AnalyzerRule for DistJoinPlanner { fn name(&self) -> &str { - "DistJoinPlanner" + DIST_JOIN_PLANNER_RULE_NAME } fn analyze(&self, plan: LogicalPlan, config: &ConfigOptions) -> DfResult { - let Some(build_table) = config + let manual = config .extensions .get::() - .and_then(|options| options.nested_broadcast_join_build_table.as_deref()) - else { + .and_then(|options| options.nested_broadcast_join_build_table.as_deref()); + let selector = if let Some(build_table) = manual { + // The manual opt-in keeps its exact name matching and never consults statistics. + BuildSideSelector::Manual { build_table } + } else if let Some(stats) = config.extensions.get::() { + BuildSideSelector::Stats { stats } + } else { return Ok(plan); }; - let mut rewriter = NestedBroadcastJoinRewriter { build_table }; + let mut rewriter = NestedBroadcastJoinRewriter { selector }; Ok(plan.rewrite(&mut rewriter)?.data) } } -/// Rewriter nesting the build side `MergeScan` of a supported join inside the probe -/// side's `MergeScan`. +/// Rewriter nesting the build side `MergeScan` of a supported join inside the probe side's +/// `MergeScan`. struct NestedBroadcastJoinRewriter<'a> { - /// Name of the build side table, either bare (`device_limits`) or fully qualified. - build_table: &'a str, + /// How the build side of a supported join is selected. + selector: BuildSideSelector<'a>, +} + +/// How the rewriter finds the build side of a supported join. +enum BuildSideSelector<'a> { + /// Manual opt-in: the configured build table name has to be the right input. + Manual { + /// Name of the build side table, either bare (`device_limits`) or fully qualified. + build_table: &'a str, + }, + /// Cost heuristic: only the right input is considered, and only when the statistics of the + /// query show that broadcasting its table to the probe regions is cheaper. + Stats { + /// Per-table statistics fetched for this query, see [`DistJoinStats`]. + stats: &'a DistJoinStats, + }, +} + +impl BuildSideSelector<'_> { + /// Whether the build side of the supported join shape is its right input. + fn selects_right_build(&self, probe: &LogicalPlan, build: &LogicalPlan) -> bool { + match self { + Self::Manual { build_table } => { + let left_is_build = + table_name_of(probe).is_some_and(|name| table_name_matches(&name, build_table)); + let right_is_build = + table_name_of(build).is_some_and(|name| table_name_matches(&name, build_table)); + + !left_is_build && right_is_build + } + Self::Stats { stats } => { + let (Some(probe), Some(build)) = + (side_base_table_id(probe), side_base_table_id(build)) + else { + return false; + }; + + stats.favors_right_build(probe, build) + } + } + } } impl NestedBroadcastJoinRewriter<'_> { - /// Returns the rewritten plan when `node` matches this slice's shape, otherwise + /// Returns the rewritten plan when `node` matches the supported shape, otherwise /// `None` to keep the plan unchanged. /// /// A rewrite changes where the join runs, so it must preserve the join output @@ -79,8 +401,8 @@ impl NestedBroadcastJoinRewriter<'_> { return None; }; - // Scope restriction of this slice: only INNER equi-joins of two distributed - // scans are handled, everything else keeps the existing distributed plan. + // Only `INNER` equi-joins of two distributed scans are handled, everything else + // keeps the existing distributed plan. if join.join_type != JoinType::Inner || join.on.is_empty() { return None; } @@ -103,13 +425,14 @@ impl NestedBroadcastJoinRewriter<'_> { return None; } - // The configured table must identify exactly one side, and the build side must be - // the right input: moving a left build side would change the join schema. - let left_is_build = table_name_of(probe_merge_scan.input()) - .is_some_and(|name| table_name_matches(&name, self.build_table)); - let right_is_build = table_name_of(build_merge_scan.input()) - .is_some_and(|name| table_name_matches(&name, self.build_table)); - if left_is_build || !right_is_build { + // The selector picks the build side: the manual opt-in requires the configured table + // to be exactly the right input, the cost heuristic only ever considers the right + // input. The build side has to be the right input, because moving a left build side + // would change the join schema. + if !self + .selector + .selects_right_build(probe_merge_scan.input(), build_merge_scan.input()) + { return None; } @@ -226,7 +549,7 @@ fn has_json2_boundary(merge_scan: &MergeScanLogicalPlan) -> bool { .schema() .fields() .iter() - .any(|field| is_json2_extension_type(field)) + .any(is_json2_extension_type) } /// Whether the configured build side table name refers to `name`. Accepts the bare @@ -248,6 +571,8 @@ mod tests { }; use arrow_schema::{DataType, Field, Fields, Schema}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_meta::datanode::RegionManifestInfo; + use common_meta::rpc::router::Region; use datafusion::datasource::DefaultTableSource; use datafusion_common::{DFSchema, DFSchemaRef, JoinConstraint, NullEquality}; use datafusion_expr::{Expr, LogicalPlanBuilder, build_join_schema, col, lit}; @@ -720,6 +1045,444 @@ mod tests { assert_eq!(plan.to_string(), result.to_string()); } + /// A region report of the tests: one ordinary leader with `bytes` approximate disk bytes. + fn region_stat(region_id: RegionId, bytes: u64, role: RegionRole) -> RegionStat { + RegionStat { + id: region_id, + rcus: 0, + wcus: 0, + approximate_bytes: bytes, + engine: "mito".to_string(), + role, + num_rows: 0, + memtable_size: 0, + manifest_size: 0, + sst_size: 0, + sst_num: 0, + index_size: 0, + region_manifest: RegionManifestInfo::Mito { + manifest_version: 0, + flushed_entry_id: 0, + file_removed_cnt: 0, + }, + written_bytes: 0, + query_cpu_time: 0, + query_scanned_bytes: 0, + data_topic_latest_entry_id: 0, + metadata_topic_latest_entry_id: 0, + min_timestamp: None, + max_timestamp: None, + } + } + + /// Statistics of one table, as [`DistJoinStats`] reads them. + fn table_stats(total_bytes: u64, region_count: usize) -> DistJoinTableStats { + DistJoinTableStats { + total_bytes, + region_count, + } + } + + /// Statistics of the given tables. + fn stats(tables: &[(TableId, DistJoinTableStats)]) -> DistJoinStats { + DistJoinStats { + tables: tables.iter().copied().collect(), + } + } + + /// Config with the given statistics, as `create_physical_plan` inserts them. + fn stats_config(tables: &[(TableId, DistJoinTableStats)]) -> ConfigOptions { + let mut config = ConfigOptions::default(); + config.extensions.insert(stats(tables)); + config + } + + /// A region route of the tests. + fn region_route(region_id: RegionId) -> RegionRoute { + RegionRoute { + region: Region { + id: region_id, + ..Default::default() + }, + ..Default::default() + } + } + + /// The heuristic takes the rewrite only when the product of the build bytes per probe region + /// is strictly smaller than the probe bytes, and never with a missing, empty or overflowing + /// input. + #[test] + fn favors_right_build_compares_product_with_probe_bytes() { + let cases = [ + // 2 probe regions * 1_000 bytes of build = 2_000 < 10_000. + ( + Some(table_stats(10_000, 2)), + Some(table_stats(1_000, 1)), + true, + ), + // 10 * 1_000 == 10_000: not strictly smaller. + ( + Some(table_stats(10_000, 10)), + Some(table_stats(1_000, 1)), + false, + ), + // 10 * 1_001 > 10_000. + ( + Some(table_stats(10_000, 10)), + Some(table_stats(1_001, 1)), + false, + ), + // A table without bytes or without regions is not usable, even when the product + // would be smaller. + (Some(table_stats(10_000, 2)), Some(table_stats(0, 1)), false), + (Some(table_stats(10_000, 0)), Some(table_stats(1, 1)), false), + (Some(table_stats(0, 2)), Some(table_stats(1, 1)), false), + // 2 * u64::MAX overflows. + ( + Some(table_stats(u64::MAX, 2)), + Some(table_stats(u64::MAX, 1)), + false, + ), + ]; + + for (probe, build, expected) in cases { + let mut tables = BTreeMap::new(); + if let Some(probe) = probe { + tables.insert(1, probe); + } + if let Some(build) = build { + tables.insert(2, build); + } + + assert_eq!( + expected, + DistJoinStats { tables }.favors_right_build(1, 2), + "probe: {probe:?}, build: {build:?}" + ); + } + } + + /// A table without statistics selects neither side. + #[test] + fn favors_right_build_requires_both_tables() { + let stats = stats(&[(1, table_stats(10_000, 2))]); + + assert!(!stats.favors_right_build(1, 2)); + assert!(!stats.favors_right_build(2, 1)); + assert!(!stats.favors_right_build(3, 4)); + } + + /// The aggregation needs exactly one ordinary leader report per expected region, and ignores + /// the reports of other tables. + #[test] + fn aggregate_region_stats_requires_one_leader_per_region() { + let regions = vec![RegionId::new(1, 1), RegionId::new(1, 2)]; + let reports = vec![ + region_stat(RegionId::new(1, 1), 100, RegionRole::Leader), + region_stat(RegionId::new(1, 2), 200, RegionRole::Leader), + // Reports of other tables and of non-leaders are ignored. + region_stat(RegionId::new(2, 1), 1_000, RegionRole::Leader), + region_stat(RegionId::new(1, 3), 7, RegionRole::Follower), + ]; + + assert_eq!( + Some(table_stats(300, 2)), + aggregate_region_stats(®ions, &reports) + ); + } + + /// An incomplete or unusable description of the expected regions is not compared with another + /// table. + #[test] + fn aggregate_region_stats_rejects_incomplete_or_unusable_reports() { + let regions = vec![RegionId::new(1, 1), RegionId::new(1, 2)]; + let leader = |number: u32, bytes: u64| { + region_stat(RegionId::new(1, number), bytes, RegionRole::Leader) + }; + let cases = [ + ("missing report", regions.clone(), vec![leader(1, 100)]), + ( + "duplicate report", + regions.clone(), + vec![leader(1, 100), leader(1, 100), leader(2, 200)], + ), + ( + "follower instead of leader", + regions.clone(), + vec![ + region_stat(RegionId::new(1, 1), 100, RegionRole::Follower), + leader(2, 200), + ], + ), + ( + "staging leader instead of leader", + regions.clone(), + vec![ + region_stat(RegionId::new(1, 1), 100, RegionRole::StagingLeader), + leader(2, 200), + ], + ), + ( + "zero bytes", + regions.clone(), + vec![leader(1, 0), leader(2, 200)], + ), + ( + "overflowing sum", + regions.clone(), + vec![leader(1, u64::MAX), leader(2, 1)], + ), + ("no expected region", vec![], vec![leader(1, 100)]), + ( + "duplicate expected region", + vec![RegionId::new(1, 1), RegionId::new(1, 1)], + vec![leader(1, 100)], + ), + ]; + + for (name, regions, reports) in cases { + assert_eq!( + None, + aggregate_region_stats(®ions, &reports), + "case: {name}" + ); + } + } + + /// A physical table route provides the expected regions, and only when all of them belong to + /// the resolved physical table. + #[test] + fn expected_region_ids_of_physical_route() { + let routes = vec![ + region_route(RegionId::new(7, 1)), + region_route(RegionId::new(7, 2)), + ]; + + assert_eq!( + Some(vec![RegionId::new(7, 1), RegionId::new(7, 2)]), + expected_region_ids(7, &routes) + ); + // The route of another physical table, e.g. after the captured logical table id was + // resolved to a different physical table. + assert_eq!(None, expected_region_ids(8, &routes)); + assert_eq!(None, expected_region_ids(7, &[])); + assert_eq!( + None, + expected_region_ids( + 7, + &[ + region_route(RegionId::new(7, 1)), + region_route(RegionId::new(8, 1)), + ] + ) + ); + } + + /// The cost heuristic nests the build side only when the statistics favor it: the same + /// statistics in the other order never swap the sides, and missing statistics keep the plan. + #[test] + fn nested_broadcast_join_rewrite_with_stats() { + let plan = distributed_join_plan(); + + // The probe table `t1` (id 1) has two regions and 10_000 bytes, the build table `t2` + // (id 2) one region and 1_000 bytes: 2 * 1_000 < 10_000. + let favoring = stats_config(&[(1, table_stats(10_000, 2)), (2, table_stats(1_000, 1))]); + let result = DistJoinPlanner {}.analyze(plan.clone(), &favoring).unwrap(); + + let join = outer_join(&result); + assert_eq!(2, merge_scans(&result).len()); + assert!( + !contains_merge_scan(&join.left), + "probe side must become local, got: {}", + join.left + ); + assert!( + contains_merge_scan(&join.right), + "build side must keep its MergeScan, got: {}", + join.right + ); + + // The sides are never swapped: the same statistics in the other order do not favor the + // right build side. + let swapped = stats_config(&[(1, table_stats(1_000, 1)), (2, table_stats(10_000, 2))]); + let result = DistJoinPlanner {}.analyze(plan.clone(), &swapped).unwrap(); + assert_eq!(plan.to_string(), result.to_string()); + + // Statistics of the probe table only leave the build side unknown. + let partial = stats_config(&[(1, table_stats(10_000, 2))]); + let result = DistJoinPlanner {}.analyze(plan.clone(), &partial).unwrap(); + assert_eq!(plan.to_string(), result.to_string()); + } + + /// The manual build table keeps its exact name matching and is preferred over the statistics, + /// which its option never consults. + #[test] + fn manual_build_table_precedes_stats() { + let plan = distributed_join_plan(); + // These statistics do not favor the right build side. + let mut config = stats_config(&[(1, table_stats(1_000, 1)), (2, table_stats(10_000, 2))]); + config.extensions.insert(DistPlannerOptions { + nested_broadcast_join_build_table: Some("t2".to_string()), + ..Default::default() + }); + + let result = DistJoinPlanner {}.analyze(plan.clone(), &config).unwrap(); + + let join = outer_join(&result); + assert_eq!(2, merge_scans(&result).len()); + assert!(!contains_merge_scan(&join.left)); + assert!(contains_merge_scan(&join.right)); + } + + /// Both plan shapes the engine sees are covered: the equality in `on`, and the unanalyzed + /// `join_on` plan, whose whole `ON` clause is still in `filter`, including the input of an + /// EXPLAIN. + #[test] + fn candidate_joins_of_both_plan_shapes() { + // The equality in `on`. + let on_shape = inner_join( + table_scan("t1", 1, "t1"), + table_scan("t2", 2, "t2"), + col("t1.number"), + col("t2.number"), + ); + assert!( + matches!(&on_shape, LogicalPlan::Join(join) if !join.on.is_empty()), + "expected the equality in `on`, got: {on_shape}" + ); + assert_eq!( + vec![CandidateJoin { probe: 1, build: 2 }], + candidate_joins(&on_shape) + ); + + // The plan of `LogicalPlanBuilder::join_on`, as the SQL planner leaves it: the whole `ON` + // clause is the `filter`. + let filter_shape = join_plan(JoinType::Inner); + assert!( + matches!(&filter_shape, LogicalPlan::Join(join) if join.on.is_empty() && join.filter.is_some()), + "expected the unanalyzed join to carry its condition in `filter`, got: {filter_shape}" + ); + assert_eq!( + vec![CandidateJoin { probe: 1, build: 2 }], + candidate_joins(&filter_shape) + ); + + // The input of an EXPLAIN. + let explain = LogicalPlanBuilder::from(filter_shape.clone()) + .explain(false, false) + .unwrap() + .build() + .unwrap(); + assert_eq!(candidate_joins(&filter_shape), candidate_joins(&explain)); + + // A join without an equality is not a candidate, in either shape. + let cross = LogicalPlanBuilder::from(table_scan("t1", 1, "t1")) + .join_on( + table_scan("t2", 2, "t2"), + JoinType::Inner, + vec![col("t1.number").gt(col("t2.number"))], + ) + .unwrap() + .build() + .unwrap(); + assert!(candidate_joins(&cross).is_empty()); + + // Only `INNER` joins are candidates. + let left = LogicalPlanBuilder::from(table_scan("t1", 1, "t1")) + .join_on( + table_scan("t2", 2, "t2"), + JoinType::Left, + vec![col("t1.number").eq(col("t2.number"))], + ) + .unwrap() + .build() + .unwrap(); + assert!(candidate_joins(&left).is_empty()); + + // The join of two joins is no candidate itself (its sides are not base tables), while the + // candidates of its inputs are collected: the rewrite applies at every join node. + let other = inner_join( + table_scan("t3", 3, "t3"), + table_scan("t4", 4, "t4"), + col("t3.number"), + col("t4.number"), + ); + let nested = inner_join( + join_plan(JoinType::Inner), + other, + col("t1.number"), + col("t3.number"), + ); + assert_eq!( + vec![ + CandidateJoin { probe: 1, build: 2 }, + CandidateJoin { probe: 3, build: 4 } + ], + candidate_joins(&nested) + ); + } + + /// Candidate joins need a real cross-input equijoin: constants, same-side equalities and + /// non-equality conditions are no candidates, while a cross-input equality stays one next to + /// residual predicates. + #[test] + fn candidate_joins_require_cross_input_equijoin() { + let candidates_of = |on: Vec| { + let plan = LogicalPlanBuilder::from(table_scan("t1", 1, "t1")) + .join_on(table_scan("t2", 2, "t2"), JoinType::Inner, on) + .unwrap() + .build() + .unwrap(); + + candidate_joins(&plan) + }; + let cases = [ + ( + "cross-input equality", + vec![col("t1.number").eq(col("t2.number"))], + true, + ), + ( + "reversed cross-input equality", + vec![col("t2.number").eq(col("t1.number"))], + true, + ), + ( + "cross-input equality with residual predicate", + vec![ + col("t1.number").eq(col("t2.number")), + col("t1.number").gt(lit(1u32)), + ], + true, + ), + ("constants", vec![lit(1u32).eq(lit(1u32))], false), + ( + "same-side equality", + vec![col("t1.number").eq(col("t1.number"))], + false, + ), + ( + "column against constant", + vec![col("t1.number").eq(lit(1u32))], + false, + ), + ( + "cross-input non-equality", + vec![col("t1.number").gt(col("t2.number"))], + false, + ), + ]; + + for (name, on, expected) in cases { + let expected = if expected { + vec![CandidateJoin { probe: 1, build: 2 }] + } else { + Vec::new() + }; + assert_eq!(expected, candidates_of(on), "case: {name}"); + } + } + /// A one-column JSON2 boundary schema, as `JsonSchemaConcretizeRule` expects to fix up. fn json2_boundary_schema() -> DFSchemaRef { let field = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata( diff --git a/src/query/src/query_engine/context.rs b/src/query/src/query_engine/context.rs index 3386a38e238c..fee33bfdf5be 100644 --- a/src/query/src/query_engine/context.rs +++ b/src/query/src/query_engine/context.rs @@ -37,6 +37,15 @@ impl QueryEngineContext { &self.state } + /// Mutable access to the session state of this query. + /// + /// The state is built for a single query (see `DatafusionQueryEngine::engine_context`), so + /// changes are query local and do not reach the engine state. + #[inline] + pub(crate) fn state_mut(&mut self) -> &mut SessionState { + &mut self.state + } + #[inline] pub fn query_ctx(&self) -> QueryContextRef { self.query_ctx.clone() diff --git a/src/session/src/context.rs b/src/session/src/context.rs index c1448053f352..18f1207f92ff 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -705,6 +705,7 @@ pub struct ConfigurationVariables { pg_datestyle_format: ArcSwap<(PGDateTimeStyle, PGDateOrder)>, pg_intervalstyle_format: ArcSwap, allow_query_fallback: ArcSwap, + experimental_dist_join: ArcSwap, } impl Clone for ConfigurationVariables { @@ -714,6 +715,7 @@ impl Clone for ConfigurationVariables { pg_datestyle_format: ArcSwap::new(self.pg_datestyle_format.load().clone()), pg_intervalstyle_format: ArcSwap::new(self.pg_intervalstyle_format.load().clone()), allow_query_fallback: ArcSwap::new(self.allow_query_fallback.load().clone()), + experimental_dist_join: ArcSwap::new(self.experimental_dist_join.load().clone()), } } } @@ -754,6 +756,16 @@ impl ConfigurationVariables { pub fn set_allow_query_fallback(&self, allow: bool) { self.allow_query_fallback.swap(Arc::new(allow)); } + + /// Whether the session opted into the cost heuristic of the nested broadcast join + /// rewrite (`SET experimental_dist_join = true`). + pub fn experimental_dist_join(&self) -> bool { + **self.experimental_dist_join.load() + } + + pub fn set_experimental_dist_join(&self, enabled: bool) { + self.experimental_dist_join.swap(Arc::new(enabled)); + } } #[cfg(test)] @@ -827,6 +839,31 @@ mod test { assert!(fork.skip_wal()); } + /// The configuration parameters belong to the session, not to a single statement: a + /// statement that is planned or executed after `SET` sees the same value. + #[test] + fn test_experimental_dist_join_persists_across_statements() { + let session = Session::new(None, Channel::Mysql, Default::default(), 0); + assert!( + !session + .new_query_context() + .configuration_parameter() + .experimental_dist_join() + ); + + session + .new_query_context() + .configuration_parameter() + .set_experimental_dist_join(true); + + assert!( + session + .new_query_context() + .configuration_parameter() + .experimental_dist_join() + ); + } + #[test] fn test_write_admission_is_local_and_database_scoped() { let ctx = QueryContext::with_channel("greptime", "public", Channel::Otlp); diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 28834ac40bb9..3cf05e785747 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -319,8 +319,6 @@ async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation } } -/// A failing inner region makes the whole query fail: the datanode must not return the rows of the -/// inner regions that succeed. /// The distributed rules produce the nested plan of the join SQL, the frontend plans the outer /// `MergeScan` of that plan to the *selected* probe regions, and the datanodes execute it. /// @@ -356,7 +354,25 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { ) .await; - let sql = rewritten_join_sql(); + // The SQL under test: the projection reorders the columns of the join, the join keys `10` of + // the build table are duplicated, the residual join filter rejects the join partners whose + // values sum to at most `105` (the probe row `100` sums up to `101` and `102`), and the probe + // predicate selects the first probe region. + let sql = format!( + "SELECT p.probe_v, p.a_id, b.build_v + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b + ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 + WHERE p.a_id < 100" + ); + // The rows of the SQL as `(probe_v, a_id, build_v)`: the probe row `(4, 10, 600)` joins the two + // build rows of the duplicated key `10` (`600 + 1` and `600 + 2`), the probe row `(2, 20, 200)` + // joins `(3, 20, 3)`, the residual filter rejects the probe row `(1, 10, 100)` (`100 + 1` and + // `100 + 2` are not greater than `105`), and the NULL join keys of both sides match nothing. + let expected = multiset(vec![ + vec!["200".to_string(), "2".to_string(), "3".to_string()], + vec!["600".to_string(), "4".to_string(), "1".to_string()], + vec!["600".to_string(), "4".to_string(), "2".to_string()], + ]); let query_ctx = query_ctx(); // The reference: the frontend executes the same SQL with its default distributed plan, i.e. @@ -364,13 +380,28 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { let reference = query_pretty(&frontend, &sql, query_ctx.clone()).await; info!("reference result of the frontend:\n{reference}"); assert_eq!( - expected_rewritten_join_rows(), + expected, multiset(table_cells(&reference)), "unexpected result of the join on the frontend:\n{reference}" ); - // The plan of the rules, not a hand-built plan. - let plan = rewritten_join_plan(&frontend, &query_ctx, &sql).await; + // The plan of the rules, not a hand-built plan: `DistPlannerAnalyzer` wraps the table scans of + // the join in `MergeScan`, and `DistJoinPlanner` (opted in for the build table) nests the build + // side into the probe side's scan. + let plan = plan_sql(&frontend, &sql, &query_ctx).await; + info!("logical plan of the join SQL:\n{plan}"); + let mut config = ConfigOptions::default(); + config.extensions.insert(DistPlannerOptions { + nested_broadcast_join_build_table: Some(BUILD_TABLE.to_string()), + ..Default::default() + }); + let plan = DistPlannerAnalyzer {} + .analyze(plan, &config) + .expect("the distributed plan of the join must be analyzable"); + let plan = DistJoinPlanner {} + .analyze(plan, &config) + .expect("the nested broadcast join rewrite must not fail"); + info!("plan of the distributed rules:\n{plan}"); let outer = outer_merge_scan(&plan).unwrap_or_else(|| { panic!("the rules must nest the join into the probe side's MergeScan, got: {plan}") }); @@ -420,7 +451,7 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { .expect("the result of the nested plan must be printable"); info!("result of the rule-produced nested plan:\n{actual}"); assert_eq!( - expected_rewritten_join_rows(), + expected, multiset(table_cells(&actual)), "the nested plan returned an unexpected multiset of rows:\n{actual}" ); @@ -431,62 +462,6 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { ); } -/// The SQL of [`test_nested_broadcast_join_rewrite_executes_on_probe_regions`], over the columns of -/// [`expected_rewritten_join_rows`] in order. -/// -/// The projection reorders the columns of the join, the join keys `10` of the build table are -/// duplicated, the residual join filter rejects the join partners whose values sum to at most -/// `105` (the probe row `100` sum up to `101` and `102`), and the probe predicate selects the -/// first probe region. -fn rewritten_join_sql() -> String { - format!( - "SELECT p.probe_v, p.a_id, b.build_v - FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b - ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 - WHERE p.a_id < 100" - ) -} - -/// The rows of [`rewritten_join_sql`] as `(probe_v, a_id, build_v)`: the probe row `(4, 10, 600)` -/// joins the two build rows of the duplicated key `10` (`600 + 1` and `600 + 2`), the probe row -/// `(2, 20, 200)` joins `(3, 20, 3)`, the residual filter rejects the probe row `(1, 10, 100)` -/// (`100 + 1` and `100 + 2` are not greater than `105`), and the NULL join keys of both sides -/// match nothing. -fn expected_rewritten_join_rows() -> Vec> { - multiset(vec![ - vec!["200".to_string(), "2".to_string(), "3".to_string()], - vec!["600".to_string(), "4".to_string(), "1".to_string()], - vec!["600".to_string(), "4".to_string(), "2".to_string()], - ]) -} - -/// Plans `sql` on the frontend and applies the distributed rules of the frontend: the -/// [`DistPlannerAnalyzer`], which wraps the table scans of both sides of the join in `MergeScan`, -/// and the [`DistJoinPlanner`] opted in for the build table. -async fn rewritten_join_plan( - frontend: &Arc, - query_ctx: &QueryContextRef, - sql: &str, -) -> LogicalPlan { - let plan = plan_sql(frontend, sql, query_ctx).await; - info!("logical plan of the join SQL:\n{plan}"); - - let mut config = ConfigOptions::default(); - config.extensions.insert(DistPlannerOptions { - nested_broadcast_join_build_table: Some(BUILD_TABLE.to_string()), - ..Default::default() - }); - let plan = DistPlannerAnalyzer {} - .analyze(plan, &config) - .expect("the distributed plan of the join must be analyzable"); - let plan = DistJoinPlanner {} - .analyze(plan, &config) - .expect("the nested broadcast join rewrite must not fail"); - info!("plan of the distributed rules:\n{plan}"); - - plan -} - /// Returns the outermost `MergeScan` of `plan`, descending through projections. fn outer_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { match plan { @@ -539,6 +514,276 @@ fn find_merge_scan_exec(plan: &Arc) -> Option<&MergeScanExec> plan.children().into_iter().find_map(find_merge_scan_exec) } +/// The join keys of [`BUILD_ROWS`], so the extra probe rows of +/// [`test_experimental_dist_join_setting_activates_rewrite`] match every build key. +const BUILD_JOIN_KEYS: &[i32] = &[10, 20, 30, 40]; + +/// The number of extra rows of [`extra_probe_rows`]: enough for the cost heuristic of the rewrite +/// to favor the build side of the join by a wide margin, see +/// [`test_experimental_dist_join_setting_activates_rewrite`]. +const EXTRA_PROBE_ROWS: i32 = 4_000; + +/// `(a_id, probe_key, probe_v, a_ts)` of the extra probe rows: `a_id` alternates between the two +/// partitions of the probe table and `probe_key` cycles the build keys, so every build key is +/// matched. +fn extra_probe_rows() -> Vec<(i32, i32, i32, i64)> { + (0..EXTRA_PROBE_ROWS) + .map(|i| { + ( + if i % 2 == 0 { i } else { 1_000 + i }, + BUILD_JOIN_KEYS[i as usize % BUILD_JOIN_KEYS.len()], + i, + i as i64 + 10_000, + ) + }) + .collect() +} + +/// Inserts the extra probe rows of [`extra_probe_rows`] in batches. +async fn insert_extra_probe_rows(frontend: &Arc) { + for chunk in extra_probe_rows().chunks(500) { + let values = chunk + .iter() + .map(|(a_id, probe_key, probe_v, a_ts)| { + format!("({a_id}, {probe_key}, {probe_v}, {a_ts})") + }) + .collect::>() + .join(","); + execute_sql( + frontend, + &format!("INSERT INTO {PROBE_TABLE}(a_id, probe_key, probe_v, a_ts) VALUES {values}"), + ) + .await; + } +} + +/// Whether `plan` contains a hash join, the physical node of the join of the rewrite. +fn contains_hash_join(plan: &Arc) -> bool { + plan.name().contains("HashJoin") || plan.children().into_iter().any(contains_hash_join) +} + +/// The region statistics of `table`: `(leader reports, summed disk bytes, minimal disk bytes)`. +/// +/// A region that has not reported yet, reports zero bytes, or is not served by a leader makes the +/// cost heuristic of the rewrite skip the table, so a test that expects the rewrite has to wait +/// until every region of every table has an ordinary leader report. A missing or NULL aggregate +/// counts as `0`, i.e. as "not reported yet". +async fn region_statistics(frontend: &Arc, table: &str) -> (u64, u64, u64) { + let sql = format!( + "SELECT count(*) AS regions, sum(disk_size) AS bytes, min(disk_size) AS min_bytes + FROM information_schema.region_statistics + WHERE table_id IN (SELECT table_id FROM information_schema.tables + WHERE table_schema = 'public' AND table_name = '{table}') + AND region_role = 'Leader'" + ); + let pretty = query_pretty(frontend, &sql, query_ctx()).await; + let cells = table_cells(&pretty); + let row = cells.first().cloned().unwrap_or_default(); + let parse = |index: usize| { + row.get(index) + .and_then(|cell| cell.parse::().ok()) + .unwrap_or(0) + }; + + (parse(0), parse(1), parse(2)) +} + +/// Waits until both tables of the join have one ordinary leader report per region with a non-zero +/// size, and returns the summed disk bytes of the probe and of the build table. +/// +/// The frontend reads these statistics from the region statistics that the datanodes report to +/// metasrv, so they appear a few heartbeats after the tables are created. +async fn wait_for_region_statistics(frontend: &Arc) -> (u64, u64) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(120); + loop { + let mut stats = Vec::new(); + for table in [PROBE_TABLE, BUILD_TABLE] { + stats.push(region_statistics(frontend, table).await); + } + let probe = stats[0]; + let build = stats[1]; + info!("region statistics of {PROBE_TABLE}: {probe:?}, of {BUILD_TABLE}: {build:?}"); + if probe.0 == 2 && probe.2 > 0 && build.0 == 2 && build.2 > 0 { + return (probe.1, build.1); + } + assert!( + tokio::time::Instant::now() < deadline, + "the region statistics of {PROBE_TABLE} {probe:?} and {BUILD_TABLE} {build:?} must be \ + reported" + ); + tokio::time::sleep(Duration::from_secs(1)).await; + } +} + +/// Whether the logical plan text of an EXPLAIN of the join shows the nested rewrite: the outer +/// `MergeScan` of the rewritten plan wraps the join, so the boundary appears before the join, +/// while the plan without the rewrite starts at the join of the two boundaries. +fn nests_join_in_merge_scan(explained: &str) -> bool { + match (explained.find("Join:"), explained.find("MergeScan")) { + (Some(join), Some(merge_scan)) => merge_scan < join, + _ => false, + } +} + +/// `SET experimental_dist_join` is the session opt-in of the cost heuristic: with the setting, the +/// frontend nests the build side of a join into the probe side's `MergeScan` when the region +/// statistics favor it, and executes that plan; without it, the join of the two boundaries stays. +/// +/// The fixture makes the heuristic favor the build side by a wide margin: the probe table gets +/// [`EXTRA_PROBE_ROWS`] extra rows while the build table keeps the rows of [`BUILD_ROWS`], and the +/// test asserts the region statistics it reads really satisfy `regions * build_bytes < +/// probe_bytes`. The query of the test selects the rows of the fixture it can predict by hand, so +/// the parity of the execution with and without the opt-in is checked on the full multiset of the +/// projected rows, not on a row count: a rewritten plan that loses, duplicates or mixes up values +/// fails the comparison. +#[tokio::test(flavor = "multi_thread")] +async fn test_experimental_dist_join_setting_activates_rewrite() { + common_telemetry::init_default_ut_logging(); + + let cluster = build_cluster("test_experimental_dist_join_setting_activates_rewrite").await; + let frontend = cluster.fe_instance().clone(); + + prepare_tables(&frontend).await; + insert_extra_probe_rows(&frontend).await; + + // The join of the test: the projection reorders the columns, and the equality and the residual + // filter of the `ON` clause stay in the join. The predicate on the time index of the probe + // table is what keeps the result readable by hand: `10000::TIMESTAMP` is the raw `10000` of + // the (millisecond) time index, so it keeps the rows that [`prepare_tables`] inserted (raw + // timestamps `1000` to `4000`) and drops every row of [`insert_extra_probe_rows`] (raw + // timestamps `10000` and above). The fixture is not changed, and the heuristic reads the + // region statistics of the whole tables anyway. + let sql = format!( + "SELECT p.a_id, p.probe_v, b.b_id, b.build_v + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b + ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 + WHERE p.a_ts < 10000::TIMESTAMP" + ); + // The rows of the join as `(a_id, probe_v, b_id, build_v)`: the probe row `(2, 20, 200)` joins + // `(3, 20, 3)`, the probe row `(120, 30, 300)` joins `(120, 30, 4)`, the probe row + // `(121, 10, 400)` joins the two build rows of the duplicated key `10` (`400 + 1` and + // `400 + 2`), and the residual filter rejects the probe row `(1, 10, 100)` (`100 + 1` and + // `100 + 2` are not greater than `105`). + let expected = multiset(vec![ + vec![ + "2".to_string(), + "200".to_string(), + "3".to_string(), + "3".to_string(), + ], + vec![ + "120".to_string(), + "300".to_string(), + "120".to_string(), + "4".to_string(), + ], + vec![ + "121".to_string(), + "400".to_string(), + "1".to_string(), + "1".to_string(), + ], + vec![ + "121".to_string(), + "400".to_string(), + "2".to_string(), + "2".to_string(), + ], + ]); + let query_ctx = query_ctx(); + + let reference = run_sql(&frontend, &sql, query_ctx.clone()).await; + let plan = reference + .meta + .plan + .clone() + .expect("the output of the query must carry its physical plan"); + info!("the physical plan of the join without the opt-in:\n{plan:?}"); + assert!( + contains_hash_join(&plan), + "without the opt-in the join must run on the frontend, got:\n{plan:?}" + ); + let reference = reference.data.pretty_print().await; + info!("the join without the opt-in returns:\n{reference}"); + assert_eq!( + expected, + multiset(table_cells(&reference)), + "unexpected result of the join without the opt-in:\n{reference}" + ); + // The heuristic is off by default: the EXPLAIN of the same join keeps the join of the two + // boundaries. + let explained = query_pretty(&frontend, &format!("EXPLAIN {sql}"), query_ctx.clone()).await; + info!("EXPLAIN without the opt-in:\n{explained}"); + assert!( + !nests_join_in_merge_scan(&explained), + "without `SET experimental_dist_join = true` the join must keep its two MergeScan \ + boundaries:\n{explained}" + ); + + // The statistics the heuristic reads: every region of both tables has a leader report with a + // non-zero size, and broadcasting the build table to the probe regions is much cheaper than + // the probe table itself. + let (probe_bytes, build_bytes) = wait_for_region_statistics(&frontend).await; + assert!( + 2 * build_bytes < probe_bytes, + "the fixture must make the heuristic favor the build side, actual region statistics: \ + probe {probe_bytes} bytes in 2 regions, build {build_bytes} bytes in 2 regions" + ); + + // The session opt-in of the next statement. + run_sql( + &frontend, + "SET experimental_dist_join = true", + query_ctx.clone(), + ) + .await; + + let explained = query_pretty(&frontend, &format!("EXPLAIN {sql}"), query_ctx.clone()).await; + info!("EXPLAIN with the opt-in:\n{explained}"); + assert!( + nests_join_in_merge_scan(&explained), + "with `SET experimental_dist_join = true` the EXPLAIN must show the nested rewrite:\n\ + {explained}" + ); + + let rewritten = run_sql(&frontend, &sql, query_ctx.clone()).await; + let plan = rewritten + .meta + .plan + .clone() + .expect("the output of the query must carry its physical plan"); + info!("the physical plan of the join with the opt-in:\n{plan:?}"); + assert!( + !contains_hash_join(&plan) && find_merge_scan_exec(&plan).is_some(), + "the executed plan of the opted-in join must read the join through a MergeScan instead of \ + joining on the frontend, got:\n{plan:?}" + ); + let rewritten = rewritten.data.pretty_print().await; + info!("the join with the opt-in returns:\n{rewritten}"); + assert_eq!( + multiset(table_cells(&reference)), + multiset(table_cells(&rewritten)), + "the executed plan of the opted-in join must return the same multiset of rows as the \ + frontend without the opt-in:\nfrontend:\n{reference}\nrewritten:\n{rewritten}" + ); + + // Turning the setting off restores the existing plan for the same session. + run_sql( + &frontend, + "SET experimental_dist_join = false", + query_ctx.clone(), + ) + .await; + + let explained = query_pretty(&frontend, &format!("EXPLAIN {sql}"), query_ctx).await; + info!("EXPLAIN after turning the opt-in off:\n{explained}"); + assert!( + !nests_join_in_merge_scan(&explained), + "`SET experimental_dist_join = false` must restore the join of the two boundaries:\n\ + {explained}" + ); +} + #[tokio::test(flavor = "multi_thread")] async fn test_nested_merge_scan_capability_inner_failure_fails_query() { common_telemetry::init_default_ut_logging(); From c6912d425015fbd1ef6446f6fccabd75b95825e2 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Wed, 30 Sep 2026 13:30:46 +0800 Subject: [PATCH 14/37] test: add distributed sqlness case for experimental nested broadcast join Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../distributed/explain/dist_join.result | 263 ++++++++++++++++++ tests/cases/distributed/explain/dist_join.sql | 144 ++++++++++ 2 files changed, 407 insertions(+) create mode 100644 tests/cases/distributed/explain/dist_join.result create mode 100644 tests/cases/distributed/explain/dist_join.sql diff --git a/tests/cases/distributed/explain/dist_join.result b/tests/cases/distributed/explain/dist_join.result new file mode 100644 index 000000000000..1e829183c017 --- /dev/null +++ b/tests/cases/distributed/explain/dist_join.result @@ -0,0 +1,263 @@ +-- The distributed nested broadcast join rewrite of the query engine. +-- +-- `SET experimental_dist_join = true` opts the session into the rewrite: it nests the build +-- side's MergeScan inside the probe side's MergeScan, so that the join runs on the datanodes +-- that own the probe regions. The rewrite is additionally gated by the region statistics that +-- the datanodes report in their heartbeats: every region of both candidate tables has to have +-- a non-zero reported size, and broadcasting the build table to the probe regions has to be +-- cheaper than reading the probe table itself. +-- +-- The tables below are a partitioned fact table (the probe side, three regions) that is much +-- larger than the small build table (two regions), so that the cost heuristic favors the +-- rewrite as soon as the statistics are available. The join keys are duplicated on both sides, +-- a NULL key on both sides matches nothing, one key of each side has no partner, the residual +-- predicate rejects the pairs whose values sum to at most 105, and the probe predicate below +-- both prunes the third probe region and excludes a probe row that would otherwise match. +-- +-- `dist_join_gen` is a tiny numbers table, only used to pad the probe table below. +CREATE TABLE dist_join_gen (n INT, ts TIMESTAMP TIME INDEX); + +Affected Rows: 0 + +INSERT INTO dist_join_gen VALUES + (0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6), (7, 7), (8, 8), (9, 9); + +Affected Rows: 10 + +CREATE TABLE dist_join_probe ( + p_id INT, + p_key INT, + p_val INT, + p_pad STRING, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (p_id) +) PARTITION ON COLUMNS (p_id) ( + p_id < 100, + p_id >= 100 AND p_id < 200, + p_id >= 200 +); + +Affected Rows: 0 + +CREATE TABLE dist_join_build ( + b_id INT, + b_key INT, + b_val INT, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (b_id) +) PARTITION ON COLUMNS (b_id) ( + b_id < 100, + b_id >= 100 +); + +Affected Rows: 0 + +-- The rows of the join: the probe rows of key 10 join both rows of the duplicated build key +-- (601, 602 and 701, 702), the probe row of key 10 with value 100 is rejected by the residual +-- predicate (101, 102), key 20 joins across the two build partitions, key 30 has no build +-- partner, key 99 has no probe partner, the NULL keys match nothing, and the probe row 301 in +-- the third region matches too but is excluded by the probe predicate of the query below. +INSERT INTO dist_join_probe (p_id, p_key, p_val, p_pad, ts) VALUES + (1, 10, 100, '', 1001), + (2, 10, 600, '', 1002), + (3, NULL, 700, '', 1003), + (4, 20, 200, '', 1004), + (5, 30, 30, '', 1005), + (6, 10, 700, '', 1006), + (301, 10, 800, '', 1007); + +Affected Rows: 7 + +INSERT INTO dist_join_build (b_id, b_key, b_val, ts) VALUES + (1, 10, 1, 2001), + (2, 10, 2, 2002), + (101, 20, 3, 2003), + (102, 99, 4, 2004), + (103, NULL, 5, 2005); + +Affected Rows: 5 + +-- Padding of the third probe region: it makes the probe table larger than the whole build +-- table, the comparison of the cost heuristic. These rows join with no build row. +INSERT INTO dist_join_probe (p_id, p_key, p_val, p_pad, ts) +SELECT n1.n + n2.n * 10 + n3.n * 100 + 1000, + -1, + 0, + 'padding-padding-padding-padding-padding', + 1 +FROM dist_join_gen n1, dist_join_gen n2, dist_join_gen n3; + +Affected Rows: 1000 + +-- The region statistics are reported in the datanodes' heartbeats (every second in this +-- environment, and the metasrv flushes them on every heartbeat), so a heartbeat has to pass +-- after the inserts before the rewrite can see the sizes of the regions above. +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS SLEEP 3s +SET experimental_dist_join = false; + +affected_rows: 0 + +-- Without the session opt-in the join keeps its plain distributed plan: both sides are read +-- with their own MergeScan and the join runs on the frontend. +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val +FROM dist_join_probe p INNER JOIN dist_join_build b +ON p.p_key = b.b_key AND p.p_val + b.b_val > 105 +WHERE p.p_id < 200 +ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Sort: p.p_val ASC NULLS LAST, p.p_id ASC NULLS LAST, b.b_val ASC NULLS LAST, b.b_id ASC NULLS LAST_| +|_|_Inner Join: p.p_key = b.b_key Filter: p.p_val + b.b_val > Int32(105)_| +|_|_Projection: p.p_id, p.p_key, p.p_val_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| SubqueryAlias: p_| +|_|_Filter: dist_join_probe.p_key IS NOT NULL AND dist_join_probe.p_id < Int32(200)_| +|_|_TableScan: dist_join_probe, partial_filters=[dist_join_probe.p_key IS NOT NULL, dist_join_probe.p_id < Int32(200)]_| +|_| ]]_| +|_|_Projection: b.b_id, b.b_key, b.b_val_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| SubqueryAlias: b_| +|_|_Filter: dist_join_build.b_key IS NOT NULL_| +|_|_TableScan: dist_join_build, partial_filters=[dist_join_build.b_key IS NOT NULL]_| +|_| ]]_| +| physical_plan | SortPreservingMergeExec: [p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST]_| +|_|_SortExec: expr=[p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST], preserve_partitioning=[true] | +|_|_REDACTED +|_|_RepartitionExec: partitioning=REDACTED +|_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val]_| +|_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED +|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- With the opt-in the build side is nested into the probe side's MergeScan. +-- SQLNESS PROTOCOL MYSQL +SET experimental_dist_join = true; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val +FROM dist_join_probe p INNER JOIN dist_join_build b +ON p.p_key = b.b_key AND p.p_val + b.b_val > 105 +WHERE p.p_id < 200 +ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Sort: p.p_val ASC NULLS LAST, p.p_id ASC NULLS LAST, b.b_val ASC NULLS LAST, b.b_id ASC NULLS LAST_| +|_|_Projection: p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| Inner Join: p.p_key = b.b_key Filter: p.p_val + b.b_val > Int32(105)_| +|_|_Projection: p.p_id, p.p_key, p.p_val, p.p_pad, p.ts_| +|_|_SubqueryAlias: p_| +|_|_Filter: dist_join_probe.p_key IS NOT NULL AND dist_join_probe.p_id < Int32(200)_| +|_|_TableScan: dist_join_probe, partial_filters=[dist_join_probe.p_key IS NOT NULL, dist_join_probe.p_id < Int32(200)]_| +|_|_Projection: b.b_id, b.b_key, b.b_val, b.ts_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| SubqueryAlias: b_| +|_|_Filter: dist_join_build.b_key IS NOT NULL_| +|_|_TableScan: dist_join_build, partial_filters=[dist_join_build.b_key IS NOT NULL]_| +|_| ]]_| +|_| ]]_| +| physical_plan | SortPreservingMergeExec: [p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST]_| +|_|_SortExec: expr=[p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST], preserve_partitioning=[true] | +|_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val, b_id@5 as b_id, b_key@6 as b_key, b_val@7 as b_val]_| +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +-- The rewritten plan returns the same rows as the plain one. +-- SQLNESS PROTOCOL MYSQL +SELECT p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val +FROM dist_join_probe p INNER JOIN dist_join_build b +ON p.p_key = b.b_key AND p.p_val + b.b_val > 105 +WHERE p.p_id < 200 +ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; + ++------+-------+-------+------+-------+-------+ +| p_id | p_key | p_val | b_id | b_key | b_val | ++------+-------+-------+------+-------+-------+ +| 4 | 20 | 200 | 101 | 20 | 3 | +| 2 | 10 | 600 | 1 | 10 | 1 | +| 2 | 10 | 600 | 2 | 10 | 2 | +| 6 | 10 | 700 | 1 | 10 | 1 | +| 6 | 10 | 700 | 2 | 10 | 2 | ++------+-------+-------+------+-------+-------+ + +-- Unsetting the opt-in reverts the plan. +-- SQLNESS PROTOCOL MYSQL +SET experimental_dist_join = false; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val +FROM dist_join_probe p INNER JOIN dist_join_build b +ON p.p_key = b.b_key AND p.p_val + b.b_val > 105 +WHERE p.p_id < 200 +ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; + ++-+-+ +| plan_type_| plan_| ++-+-+ +| logical_plan_| Sort: p.p_val ASC NULLS LAST, p.p_id ASC NULLS LAST, b.b_val ASC NULLS LAST, b.b_id ASC NULLS LAST_| +|_|_Inner Join: p.p_key = b.b_key Filter: p.p_val + b.b_val > Int32(105)_| +|_|_Projection: p.p_id, p.p_key, p.p_val_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| SubqueryAlias: p_| +|_|_Filter: dist_join_probe.p_key IS NOT NULL AND dist_join_probe.p_id < Int32(200)_| +|_|_TableScan: dist_join_probe, partial_filters=[dist_join_probe.p_key IS NOT NULL, dist_join_probe.p_id < Int32(200)]_| +|_| ]]_| +|_|_Projection: b.b_id, b.b_key, b.b_val_| +|_|_MergeScan [is_placeholder=false, remote_input=[_| +|_| SubqueryAlias: b_| +|_|_Filter: dist_join_build.b_key IS NOT NULL_| +|_|_TableScan: dist_join_build, partial_filters=[dist_join_build.b_key IS NOT NULL]_| +|_| ]]_| +| physical_plan | SortPreservingMergeExec: [p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST]_| +|_|_SortExec: expr=[p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST], preserve_partitioning=[true] | +|_|_REDACTED +|_|_RepartitionExec: partitioning=REDACTED +|_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val]_| +|_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED +|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| +|_|_MergeScanExec: REDACTED +|_|_| ++-+-+ + +DROP TABLE dist_join_probe; + +Affected Rows: 0 + +DROP TABLE dist_join_build; + +Affected Rows: 0 + +DROP TABLE dist_join_gen; + +Affected Rows: 0 + diff --git a/tests/cases/distributed/explain/dist_join.sql b/tests/cases/distributed/explain/dist_join.sql new file mode 100644 index 000000000000..83a5ee8e0ba3 --- /dev/null +++ b/tests/cases/distributed/explain/dist_join.sql @@ -0,0 +1,144 @@ +-- The distributed nested broadcast join rewrite of the query engine. +-- +-- `SET experimental_dist_join = true` opts the session into the rewrite: it nests the build +-- side's MergeScan inside the probe side's MergeScan, so that the join runs on the datanodes +-- that own the probe regions. The rewrite is additionally gated by the region statistics that +-- the datanodes report in their heartbeats: every region of both candidate tables has to have +-- a non-zero reported size, and broadcasting the build table to the probe regions has to be +-- cheaper than reading the probe table itself. +-- +-- The tables below are a partitioned fact table (the probe side, three regions) that is much +-- larger than the small build table (two regions), so that the cost heuristic favors the +-- rewrite as soon as the statistics are available. The join keys are duplicated on both sides, +-- a NULL key on both sides matches nothing, one key of each side has no partner, the residual +-- predicate rejects the pairs whose values sum to at most 105, and the probe predicate below +-- both prunes the third probe region and excludes a probe row that would otherwise match. +-- +-- `dist_join_gen` is a tiny numbers table, only used to pad the probe table below. + +CREATE TABLE dist_join_gen (n INT, ts TIMESTAMP TIME INDEX); + +INSERT INTO dist_join_gen VALUES + (0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6), (7, 7), (8, 8), (9, 9); + +CREATE TABLE dist_join_probe ( + p_id INT, + p_key INT, + p_val INT, + p_pad STRING, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (p_id) +) PARTITION ON COLUMNS (p_id) ( + p_id < 100, + p_id >= 100 AND p_id < 200, + p_id >= 200 +); + +CREATE TABLE dist_join_build ( + b_id INT, + b_key INT, + b_val INT, + ts TIMESTAMP TIME INDEX, + PRIMARY KEY (b_id) +) PARTITION ON COLUMNS (b_id) ( + b_id < 100, + b_id >= 100 +); + +-- The rows of the join: the probe rows of key 10 join both rows of the duplicated build key +-- (601, 602 and 701, 702), the probe row of key 10 with value 100 is rejected by the residual +-- predicate (101, 102), key 20 joins across the two build partitions, key 30 has no build +-- partner, key 99 has no probe partner, the NULL keys match nothing, and the probe row 301 in +-- the third region matches too but is excluded by the probe predicate of the query below. +INSERT INTO dist_join_probe (p_id, p_key, p_val, p_pad, ts) VALUES + (1, 10, 100, '', 1001), + (2, 10, 600, '', 1002), + (3, NULL, 700, '', 1003), + (4, 20, 200, '', 1004), + (5, 30, 30, '', 1005), + (6, 10, 700, '', 1006), + (301, 10, 800, '', 1007); + +INSERT INTO dist_join_build (b_id, b_key, b_val, ts) VALUES + (1, 10, 1, 2001), + (2, 10, 2, 2002), + (101, 20, 3, 2003), + (102, 99, 4, 2004), + (103, NULL, 5, 2005); + +-- Padding of the third probe region: it makes the probe table larger than the whole build +-- table, the comparison of the cost heuristic. These rows join with no build row. +INSERT INTO dist_join_probe (p_id, p_key, p_val, p_pad, ts) +SELECT n1.n + n2.n * 10 + n3.n * 100 + 1000, + -1, + 0, + 'padding-padding-padding-padding-padding', + 1 +FROM dist_join_gen n1, dist_join_gen n2, dist_join_gen n3; + +-- The region statistics are reported in the datanodes' heartbeats (every second in this +-- environment, and the metasrv flushes them on every heartbeat), so a heartbeat has to pass +-- after the inserts before the rewrite can see the sizes of the regions above. +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS SLEEP 3s +SET experimental_dist_join = false; + +-- Without the session opt-in the join keeps its plain distributed plan: both sides are read +-- with their own MergeScan and the join runs on the frontend. +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val +FROM dist_join_probe p INNER JOIN dist_join_build b +ON p.p_key = b.b_key AND p.p_val + b.b_val > 105 +WHERE p.p_id < 200 +ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; + +-- With the opt-in the build side is nested into the probe side's MergeScan. +-- SQLNESS PROTOCOL MYSQL +SET experimental_dist_join = true; + +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val +FROM dist_join_probe p INNER JOIN dist_join_build b +ON p.p_key = b.b_key AND p.p_val + b.b_val > 105 +WHERE p.p_id < 200 +ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; + +-- The rewritten plan returns the same rows as the plain one. +-- SQLNESS PROTOCOL MYSQL +SELECT p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val +FROM dist_join_probe p INNER JOIN dist_join_build b +ON p.p_key = b.b_key AND p.p_val + b.b_val > 105 +WHERE p.p_id < 200 +ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; + +-- Unsetting the opt-in reverts the plan. +-- SQLNESS PROTOCOL MYSQL +SET experimental_dist_join = false; + +-- SQLNESS PROTOCOL MYSQL +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +EXPLAIN SELECT p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val +FROM dist_join_probe p INNER JOIN dist_join_build b +ON p.p_key = b.b_key AND p.p_val + b.b_val > 105 +WHERE p.p_id < 200 +ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; + +DROP TABLE dist_join_probe; + +DROP TABLE dist_join_build; + +DROP TABLE dist_join_gen; From 678cf968fc866b291b83fa1cb662239ecc84340f Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Wed, 30 Sep 2026 13:30:54 +0800 Subject: [PATCH 15/37] test(perf): add distributed broadcast join regression case with harness support Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cmd/src/bin/query_perf_fixture/case.rs | 97 ++++++++- .../src/bin/query_perf_fixture/direct_sst.rs | 40 ++-- .../bin/query_regression_runner/measure.rs | 195 +++++++++++++++++- .../src/bin/query_regression_runner/model.rs | 92 +++++++++ .../src/bin/query_regression_runner/sql.rs | 25 +++ tests/perf/README.md | 21 ++ tests/perf/fixture-format.md | 37 ++++ .../sql_join_dist_broadcast/case.toml | 127 ++++++++++++ 8 files changed, 602 insertions(+), 32 deletions(-) create mode 100644 tests/perf/query_cases/sql_join_dist_broadcast/case.toml diff --git a/src/cmd/src/bin/query_perf_fixture/case.rs b/src/cmd/src/bin/query_perf_fixture/case.rs index 271c3928e352..5a110e4c6e39 100644 --- a/src/cmd/src/bin/query_perf_fixture/case.rs +++ b/src/cmd/src/bin/query_perf_fixture/case.rs @@ -496,6 +496,77 @@ candidate_setup_sql = ["CREATE TABLE candidate_table", "ALTER TABLE candidate_ta ] ); } + + #[test] + fn per_table_layout_overrides_only_the_table_volume() { + let case: CaseFile = toml::from_str( + r#" +[scenario] +kind = "direct_readable_sst" + +[scenario.layout] +regions = 1 +sst_count = 2 +rows_per_sst = 8 +row_group_size = 4 +series_count = 4 +start_unix_nanos = 1704067200000000000 +step_nanos = 1000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "round_robin" + +[[scenario.tables]] +database = "public" +name = "fact" +engine = "mito" +sst_count = 5 +rows_per_sst = 100 +primary_key = ["host"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +[[scenario.tables]] +database = "public" +name = "dim" +engine = "mito" +primary_key = ["host"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" +"#, + ) + .unwrap(); + let scenario = case.scenario.direct_readable_sst(); + + let fact = scenario.layout.for_table(&scenario.tables[0]); + assert_eq!(fact.sst_count, 5); + assert_eq!(fact.rows_per_sst, 100); + // Everything else stays shared. + assert_eq!(fact.row_group_size, 4); + assert_eq!(fact.series_count.get(), 4); + assert_eq!(fact.series_layout, "round_robin"); + + let dim = scenario.layout.for_table(&scenario.tables[1]); + assert_eq!(dim.sst_count, scenario.layout.sst_count); + assert_eq!(dim.rows_per_sst, scenario.layout.rows_per_sst); + } } impl Scenario { @@ -527,6 +598,14 @@ pub(super) struct TableConfig { pub(super) primary_key: Vec, pub(super) time_index: String, pub(super) columns: Vec, + /// Per-table override of the shared [`LayoutConfig::sst_count`]. Cases that mix tables of + /// very different volume, e.g. a fact table joined with a small dimension table, size each + /// table on its own; everything else of the shared layout still applies. + #[serde(default)] + pub(super) sst_count: Option, + /// Per-table override of the shared [`LayoutConfig::rows_per_sst`]. + #[serde(default)] + pub(super) rows_per_sst: Option, } #[derive(Debug, Deserialize, Serialize)] @@ -550,7 +629,7 @@ pub(super) enum Distribution { DeterministicWave { min: f64, max: f64 }, } -#[derive(Debug, Deserialize, Serialize)] +#[derive(Clone, Debug, Deserialize, Serialize)] pub(super) struct LayoutConfig { pub(super) regions: usize, pub(super) sst_count: usize, @@ -562,3 +641,19 @@ pub(super) struct LayoutConfig { pub(super) time_range_layout: String, pub(super) series_layout: String, } + +impl LayoutConfig { + /// The layout of `table`: the shared layout with the per-table size overrides of + /// [`TableConfig::sst_count`] and [`TableConfig::rows_per_sst`] applied. + /// + /// The remaining fields stay shared, so the tables of a case keep one deterministic time + /// layout and one deterministic series layout; only the volume of a table is decided by the + /// table itself. + pub(super) fn for_table(&self, table: &TableConfig) -> LayoutConfig { + LayoutConfig { + sst_count: table.sst_count.unwrap_or(self.sst_count), + rows_per_sst: table.rows_per_sst.unwrap_or(self.rows_per_sst), + ..self.clone() + } + } +} diff --git a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs index 27a6b4109cde..67f78d73089b 100644 --- a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs +++ b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs @@ -370,19 +370,6 @@ pub(super) async fn run_direct_sst(args: DirectArgs) { if scenario.layout.time_range_layout != "non_overlapping_per_sst" { panic!("MVP supports time_range_layout=non_overlapping_per_sst"); } - if scenario.layout.series_layout == "timestamp_major" - && !scenario - .layout - .rows_per_sst - .is_multiple_of(scenario.layout.series_count.get()) - { - panic!( - "series_layout=timestamp_major requires rows_per_sst to be divisible by series_count" - ); - } - if scenario.layout.sst_count > 1000 && !args.allow_large { - panic!("sst_count exceeds 1000; pass --allow-large"); - } let seed = scenario.seed.unwrap_or(0); let table_index = match args.table.as_deref() { Some(name) => scenario @@ -401,6 +388,19 @@ pub(super) async fn run_direct_sst(args: DirectArgs) { if table.engine != "mito" { panic!("MVP supports engine=mito"); } + let layout = scenario.layout.for_table(table); + if layout.series_layout == "timestamp_major" + && !layout + .rows_per_sst + .is_multiple_of(layout.series_count.get()) + { + panic!( + "series_layout=timestamp_major requires rows_per_sst to be divisible by series_count" + ); + } + if layout.sst_count > 1000 && !args.allow_large { + panic!("sst_count exceeds 1000; pass --allow-large"); + } let region_id = RegionId::from(args.region_id); let table_dir = args .table_dir @@ -426,8 +426,8 @@ pub(super) async fn run_direct_sst(args: DirectArgs) { case.scenario.kind(), table.database, table.name, - scenario.layout.sst_count, - scenario.layout.rows_per_sst + layout.sst_count, + layout.rows_per_sst ); println!( "out_dir={} table_dir={} region_dir={}", @@ -446,11 +446,11 @@ pub(super) async fn run_direct_sst(args: DirectArgs) { let ostorage = ObjectStore::new(FsBuilder::default().root(&obj_store_dir.to_string_lossy())) .expect("failed to create filesystem object store for fixture output"); let metadata: RegionMetadataRef = Arc::new(build_region_metadata(table, region_id)); - let mut files = HashMap::with_capacity(scenario.layout.sst_count); + let mut files = HashMap::with_capacity(layout.sst_count); let mut next_file_index = 1; - for i in 0..scenario.layout.sst_count { + for i in 0..layout.sst_count { let sequence = 1000 + i as u64; - let batch = generate_record_batch(table, &metadata, &scenario.layout, i, sequence); + let batch = generate_record_batch(table, &metadata, &layout, i, sequence); let source = FlatSource::new_iter(batch.schema(), Box::new(vec![batch].into_iter().map(Ok))); let mut metrics = Metrics::new(WriteType::Flush); @@ -467,7 +467,7 @@ pub(super) async fn run_direct_sst(args: DirectArgs) { .await; let opts = WriteOptions { write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, - row_group_size: scenario.layout.row_group_size, + row_group_size: layout.row_group_size, max_file_size: None, float_field_encoding: Default::default(), }; @@ -533,7 +533,7 @@ pub(super) async fn run_direct_sst(args: DirectArgs) { writeln!(jsonl, "{}", serde_json::to_string(&serde_json::json!({ "file_id": file_id.to_string(), "region_id": meta.region_id.as_u64(), "object_path": mito2::sst::location::sst_file_path(&table_dir, RegionFileId::new(meta.region_id, *file_id), PathType::Bare), "time_range_start": meta.time_range.0.value(), "time_range_end": meta.time_range.1.value(), "num_rows": meta.num_rows, "num_row_groups": meta.num_row_groups, "file_size": meta.file_size, "num_series": meta.num_series, "sequence": meta.sequence.map(|s| s.get()) })).expect("failed to serialize fixture SST metadata JSONL entry")).expect("failed to write fixture SST metadata JSONL entry"); } jsonl.flush().expect("failed to flush fixture files.jsonl"); - fs::write(out_dir.join("summary.json"), serde_json::to_vec_pretty(&serde_json::json!({ "case": case_name, "seed": seed, "table_index": table_index, "table": table.name, "database": table.database, "region_id": region_id.as_u64(), "table_dir": table_dir, "region_dir": region_dir, "sst_format": format!("{format:?}"), "sst_count": scenario.layout.sst_count, "rows_per_sst": scenario.layout.rows_per_sst, "row_group_size": scenario.layout.row_group_size, "total_rows": scenario.layout.sst_count * scenario.layout.rows_per_sst, "checkpoint_path": checkpoint_path, "files_jsonl_path": files_jsonl_path, "readback_validated": false, "metadata_source": "synthetic" })).expect("failed to serialize fixture summary")).expect("failed to write fixture summary.json"); + fs::write(out_dir.join("summary.json"), serde_json::to_vec_pretty(&serde_json::json!({ "case": case_name, "seed": seed, "table_index": table_index, "table": table.name, "database": table.database, "region_id": region_id.as_u64(), "table_dir": table_dir, "region_dir": region_dir, "sst_format": format!("{format:?}"), "sst_count": layout.sst_count, "rows_per_sst": layout.rows_per_sst, "row_group_size": layout.row_group_size, "total_rows": layout.sst_count * layout.rows_per_sst, "checkpoint_path": checkpoint_path, "files_jsonl_path": files_jsonl_path, "readback_validated": false, "metadata_source": "synthetic" })).expect("failed to serialize fixture summary")).expect("failed to write fixture summary.json"); println!("Done. wrote {} SST file entries", manifest.files.len()); } diff --git a/src/cmd/src/bin/query_regression_runner/measure.rs b/src/cmd/src/bin/query_regression_runner/measure.rs index fd80813f40ed..54b1198d2c82 100644 --- a/src/cmd/src/bin/query_regression_runner/measure.rs +++ b/src/cmd/src/bin/query_regression_runner/measure.rs @@ -21,9 +21,24 @@ use serde_json::{Map, Value, json}; use crate::query_regression_runner::model::{Measurement, Query, QueryResult, Scenario, Table}; use crate::query_regression_runner::plan::{load_plan, normalize_scenario}; -use crate::query_regression_runner::sql::{http_post_prom_range_query, http_post_sql, sql_ident}; +use crate::query_regression_runner::sql::{ + extract_rows, http_post_multi_statement_sql, http_post_prom_range_query, http_post_sql, + row_u64, sql_ident, sql_string, +}; use crate::query_regression_runner::{MeasureArgs, Result}; +/// How long the measurement of a direct-SST case waits for the region statistics of its tables +/// to settle before it measures anything. +/// +/// Datanodes report the role and the size of their regions on every heartbeat (`interval` of the +/// heartbeat options, three seconds by default) and a region is reported as a leader with a +/// non-zero size only after the metasrv granted its lease and the datanode applied the reported +/// role. A rewrite that prices tables by those statistics, like the nested broadcast join, sees +/// nothing before that, so a case would measure the un-rewritten plan without noticing. +const REGION_STATS_TIMEOUT: Duration = Duration::from_secs(60); +/// Poll interval of [`wait_for_region_statistics`]. +const REGION_STATS_POLL_INTERVAL: Duration = Duration::from_secs(1); + pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { if !args.http_timeout.is_finite() || args.http_timeout < 0.0 { return Err("--http-timeout must be a non-negative finite number".into()); @@ -49,17 +64,28 @@ pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { let client = Client::builder() .timeout(Duration::from_secs_f64(args.http_timeout)) .build()?; - let base = run_target(args.base_http_port, &tables, &configured_queries, &client).await; - let candidate = run_target( + let direct_fixture = matches!(scenario_kind(&plan), Some("direct_readable_sst")); + let base = run_target( + args.base_http_port, + &tables, + &configured_queries, + &client, + false, + direct_fixture, + ) + .await; + let candidate_target = run_target( args.candidate_http_port, &tables, &configured_queries, &client, + true, + direct_fixture, ) .await; - let thresholds = enforce_thresholds(&configured_queries, &base, &candidate)?; + let thresholds = enforce_thresholds(&configured_queries, &base, &candidate_target)?; let status = if base.status == "failed" - || candidate.status == "failed" + || candidate_target.status == "failed" || thresholds .iter() .any(|threshold| threshold["status"] == "failed") @@ -75,7 +101,7 @@ pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { "queries": configured_queries, "query_mode": "endpoint", "http_timeout": args.http_timeout, - "targets": [target_report("base", args.base_http_port, base), target_report("candidate", args.candidate_http_port, candidate)], + "targets": [target_report("base", args.base_http_port, base), target_report("candidate", args.candidate_http_port, candidate_target)], "thresholds": thresholds, "status": status, }); @@ -109,7 +135,7 @@ fn target_report(name: &str, http_port: u16, result: QueryResult) -> Value { /// Routes a configured query to the right endpoint: `prom_http` queries hit /// the Prometheus HTTP range API (which exercises the Prometheus JSON response /// builder), everything else goes through `/v1/sql` (including `TQL ANALYZE`). -async fn post_query(client: &Client, port: u16, query: &Query, db: &str) -> Value { +async fn post_query(client: &Client, port: u16, query: &Query, db: &str, candidate: bool) -> Value { if query.kind.as_deref() == Some("prom_http") { http_post_prom_range_query( client, @@ -122,8 +148,118 @@ async fn post_query(client: &Client, port: u16, query: &Query, db: &str) -> Valu ) .await } else { - http_post_sql(client, port, &query.query, db).await + let sql = query.request_sql(candidate); + if sql == query.query { + http_post_sql(client, port, &sql, db).await + } else { + // Statements of one request share the session context of that request, which is how + // a session setting such as `SET experimental_dist_join = true` reaches the query of + // the same request. + http_post_multi_statement_sql(client, port, &sql, db).await + } + } +} + +/// The scenario kind of the normalized plan, if it carries one. +fn scenario_kind(plan: &Value) -> Option<&str> { + plan.get("scenario")?.get("kind")?.as_str() +} + +/// The region statistics of `table`: `(regions, ready_regions)`, where a ready region has an +/// ordinary leader report with a non-zero disk size. +/// +/// This is the shape a rewrite that prices tables by region statistics needs, e.g. the cost +/// heuristic of the nested broadcast join: a region that has not reported, reports zero bytes, or +/// is not served by a leader makes the statistics of its table unusable. +async fn region_statistics(client: &Client, port: u16, db: &str, table: &str) -> Value { + let sql = format!( + "SELECT count(*) AS regions, \ + sum(CASE WHEN region_role = 'Leader' AND disk_size > 0 THEN 1 ELSE 0 END) \ + AS ready_regions \ + FROM information_schema.region_statistics \ + WHERE table_id = (SELECT table_id FROM information_schema.tables \ + WHERE table_schema = {} AND table_name = {})", + sql_string(db), + sql_string(table) + ); + http_post_sql(client, port, &sql, db).await +} + +/// Waits until every table of a direct-SST case has one leader report with a non-zero size per +/// region, and returns the samples of the last round plus an error sample when the statistics did +/// not settle in time. +/// +/// The measurement of a case that depends on those statistics would otherwise start before the +/// datanode reported them and silently measure the plan the case is not about. +async fn wait_for_region_statistics( + client: &Client, + port: u16, + tables: &[Table], +) -> (Vec, Option) { + let deadline = tokio::time::Instant::now() + REGION_STATS_TIMEOUT; + loop { + let mut samples = Vec::with_capacity(tables.len()); + let mut pending = Vec::new(); + for table in tables { + let sample = region_statistics(client, port, &table.database, &table.name).await; + match sample.get("response").and_then(parse_region_statistics) { + Some((regions, ready)) if regions > 0 && ready == regions => {} + Some((regions, ready)) => pending.push(json!({ + "table": table.name, + "regions": regions, + "ready_regions": ready, + })), + None => { + return ( + samples, + Some(json!({ + "sql": "information_schema.region_statistics", + "phase": "region_statistics", + "table": table.name, + "error": "region statistics query failed", + "response": sample.get("response"), + })), + ); + } + } + samples.push(with_table(sample, table)); + } + if pending.is_empty() { + return (samples, None); + } + if tokio::time::Instant::now() >= deadline { + return ( + samples, + Some(json!({ + "sql": "information_schema.region_statistics", + "phase": "region_statistics", + "error": format!( + "region statistics did not settle within {}s", + REGION_STATS_TIMEOUT.as_secs() + ), + "pending": pending, + })), + ); + } + tokio::time::sleep(REGION_STATS_POLL_INTERVAL).await; + } +} + +fn with_table(mut sample: Value, table: &Table) -> Value { + if let Some(object) = sample.as_object_mut() { + object.insert("table".to_string(), Value::String(table.name.clone())); } + sample +} + +/// The `(regions, ready_regions)` row of a region statistics query. +fn parse_region_statistics(body: &Value) -> Option<(u64, u64)> { + let rows = extract_rows(body); + let row = rows.first()?; + Some(( + row_u64(row, 0, "regions").ok()?, + row_u64(row, 1, "ready_regions").ok()?, + )) } async fn run_target( @@ -131,6 +267,8 @@ async fn run_target( tables: &[Table], configured_queries: &[Query], client: &Client, + candidate: bool, + wait_for_region_stats: bool, ) -> QueryResult { let mut queries = configured_queries.to_vec(); if queries.is_empty() { @@ -141,6 +279,7 @@ async fn run_target( start: None, end: None, step: None, + candidate_session_sql: Vec::new(), warmup: 0, iterations: 1, thresholds: Map::new(), @@ -150,6 +289,13 @@ async fn run_target( let db = &tables[0].database; let mut validation = Vec::new(); let mut validation_errors = Vec::new(); + if wait_for_region_stats { + let (stats, error) = wait_for_region_statistics(client, port, tables).await; + if let Some(error) = error { + validation_errors.push(error); + } + validation.extend(stats); + } for table in tables { let sql = format!("SHOW CREATE TABLE {}", sql_ident(&table.name)); let sample = http_post_sql(client, port, &sql, &table.database).await; @@ -170,7 +316,7 @@ async fn run_target( } validation.push(sample); } - let first = post_query(client, port, &queries[0], db).await; + let first = post_query(client, port, &queries[0], db, candidate).await; if !first["ok"].as_bool().unwrap_or(false) { validation_errors.push(json!({ "sql": queries[0].query, @@ -183,7 +329,7 @@ async fn run_target( let mut measurements = Vec::with_capacity(queries.len()); for query in &queries { for _ in 0..query.warmup { - let warmup = post_query(client, port, query, db).await; + let warmup = post_query(client, port, query, db, candidate).await; if !warmup["ok"].as_bool().unwrap_or(false) { validation_errors.push(json!({ "sql": query.query, @@ -196,7 +342,7 @@ async fn run_target( let mut samples = Vec::with_capacity(query.iterations); let mut good_latencies = Vec::with_capacity(query.iterations); for _ in 0..query.iterations { - let mut sample = post_query(client, port, query, db).await; + let mut sample = post_query(client, port, query, db, candidate).await; let execution_time = sample .get("response") .and_then(|response| { @@ -412,6 +558,32 @@ mod tests { ); } + #[test] + fn parses_region_statistics_counts() { + let body = json!({ + "output": [{ + "records": { + "schema": {"column_schemas": [ + {"name": "regions", "data_type": "UInt64"}, + {"name": "ready_regions", "data_type": "UInt64"} + ]}, + "rows": [[1, 1]], + "total_rows": 1 + } + }] + }); + assert_eq!(Some((1, 1)), parse_region_statistics(&body)); + assert_eq!( + Some((2, 0)), + parse_region_statistics(&json!({"data": [[2, 0]]})) + ); + // A round without a row describes no regions and never settles. + assert_eq!( + None, + parse_region_statistics(&json!({"output": [{"records": {"rows": []}}]})) + ); + } + #[test] fn prom_http_does_not_extract_execution_time_from_response() { let response = json!({ @@ -439,6 +611,7 @@ mod tests { start: None, end: None, step: None, + candidate_session_sql: Vec::new(), warmup: 0, iterations: 1, thresholds: Map::from_iter([ diff --git a/src/cmd/src/bin/query_regression_runner/model.rs b/src/cmd/src/bin/query_regression_runner/model.rs index 5fafdb5635b2..86a1c7cce119 100644 --- a/src/cmd/src/bin/query_regression_runner/model.rs +++ b/src/cmd/src/bin/query_regression_runner/model.rs @@ -220,6 +220,16 @@ pub(super) struct Query { pub(super) end: Option, #[serde(default)] pub(super) step: Option, + /// Statements the candidate target executes in the same request before `query`. + /// + /// `/v1/sql` gives every request its own session context, so a statement that configures the + /// session, e.g. `SET experimental_dist_join = true`, only affects the statements of the + /// request it is sent with. The runner therefore sends these statements and `query` as one + /// multi-statement SQL string for every request of this query (validation, warmup and + /// measurement). Only the candidate target runs them, so a base build that does not know the + /// setting yet still measures the same query without it. + #[serde(default)] + pub(super) candidate_session_sql: Vec, #[serde(default)] pub(super) warmup: usize, #[serde(default = "one")] @@ -228,6 +238,27 @@ pub(super) struct Query { pub(super) thresholds: Map, } +impl Query { + /// The SQL text of one request of this query for the target: the candidate session + /// statements followed by `query`, or just `query`. + /// + /// Statements are joined with `;` separators, so both a prefixed and a bare statement stay + /// valid SQL. + pub(super) fn request_sql(&self, candidate: bool) -> String { + if candidate && !self.candidate_session_sql.is_empty() { + let prefix = self + .candidate_session_sql + .iter() + .map(|statement| statement.trim().trim_end_matches(';')) + .collect::>() + .join(";\n"); + format!("{prefix};\n{}", self.query) + } else { + self.query.clone() + } + } +} + const fn one() -> usize { 1 } @@ -250,3 +281,64 @@ pub(super) struct Measurement { pub(super) latency_ms_p95: Option, pub(super) status: String, } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn query(candidate_session_sql: Vec<&str>) -> Query { + Query { + name: Some("q".to_string()), + kind: Some("sql".to_string()), + query: "SELECT 1".to_string(), + start: None, + end: None, + step: None, + candidate_session_sql: candidate_session_sql + .into_iter() + .map(ToString::to_string) + .collect(), + warmup: 0, + iterations: 1, + thresholds: Map::new(), + } + } + + #[test] + fn request_sql_prefixes_only_the_candidate() { + let query = query(vec!["SET experimental_dist_join = true"]); + assert_eq!( + query.request_sql(true), + "SET experimental_dist_join = true;\nSELECT 1" + ); + // The base target keeps the plain statement. + assert_eq!(query.request_sql(false), "SELECT 1"); + } + + #[test] + fn request_sql_without_a_prefix_is_the_plain_statement() { + let query = query(vec![]); + assert_eq!(query.request_sql(true), "SELECT 1"); + assert_eq!(query.request_sql(false), "SELECT 1"); + } + + #[test] + fn request_sql_normalizes_statement_separators() { + let query = query(vec![ + "SET experimental_dist_join = true;", + " SET statement_timeout = '60s' ", + ]); + assert_eq!( + query.request_sql(true), + "SET experimental_dist_join = true;\nSET statement_timeout = '60s';\nSELECT 1" + ); + } + + #[test] + fn query_deserializes_without_session_sql() { + let query: Query = serde_json::from_value(json!({"query": "SELECT 1"})).unwrap(); + assert!(query.candidate_session_sql.is_empty()); + } +} diff --git a/src/cmd/src/bin/query_regression_runner/sql.rs b/src/cmd/src/bin/query_regression_runner/sql.rs index b33c0a1c22b1..2307b03b697e 100644 --- a/src/cmd/src/bin/query_regression_runner/sql.rs +++ b/src/cmd/src/bin/query_regression_runner/sql.rs @@ -130,6 +130,31 @@ pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &st sample } +/// Posts `sql` to `/v1/sql` without an explicit response format. +/// +/// A request that carries several statements is answered with the default `greptimedb_v1` +/// format, whose response holds one output per statement; the `json` format used for single +/// statements rejects multi-statement results. The response of a failing statement is an error +/// response, so a failed prefix statement still fails the sample. +pub(super) async fn http_post_multi_statement_sql( + client: &Client, + port: u16, + sql: &str, + db: &str, +) -> Value { + let mut sample = post_form( + client, + format!("http://127.0.0.1:{port}/v1/sql"), + &[("sql", sql), ("db", db)], + ) + .await; + sample + .as_object_mut() + .expect("HTTP samples are objects") + .insert("sql".to_string(), Value::String(sql.to_string())); + sample +} + /// Posts a Prometheus HTTP API range query (`/v1/prometheus/api/v1/query_range`) /// and measures the full request-to-body latency. This exercises the Prometheus /// JSON response building path (`PrometheusJsonResponse::record_batches_to_data`) diff --git a/tests/perf/README.md b/tests/perf/README.md index 0361f598a3ba..c9089e317f88 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -348,6 +348,20 @@ seed = 12345 The outer driver currently supports `direct_readable_sst`, `prom_remote_write_then_query`, and `otlp_trace_load`. +A direct-SST case sizes its tables with the shared `[scenario.layout]` and can +override `sst_count` and `rows_per_sst` per table, so a case can join a large fact +table with a small dimension table. Before measuring, the runner waits until +every table of the case has one leader region report with a non-zero disk size in +`information_schema.region_statistics`: the datanode reports those statistics on +each heartbeat, and a rewrite that prices tables by them (such as the nested +broadcast join) would otherwise measure the un-rewritten plan. A query can list +`candidate_session_sql` statements that only the candidate target runs; the +runner sends them in the same `/v1/sql` request as the query itself (validation, +warmup and measurement), because every request gets its own session context and a +`SET` statement only affects the statements of its own request. A base build that +does not know a setting yet therefore keeps measuring the query without it, which +is what makes `SET experimental_dist_join = true` a candidate-only opt-in. + ## Metrics Primary gates should compare query work rather than plan text: @@ -565,3 +579,10 @@ Additional SQL optimizer cases: `LIMIT`. - `sql_join_filter_order`: two direct-SST tables joined on a shared tag with time filters, aggregate ordering, and `LIMIT`. +- `sql_join_dist_broadcast`: a 32-SST fact table joined with a 16-row dimension + table, comparing the join on the frontend with the nested broadcast join + rewrite of the experimental `experimental_dist_join` opt-in. Only the + candidate target sends the opt-in, through `candidate_session_sql`. + +These cases are not part of the routine default case group; run them by path, for +example `--cases tests/perf/query_cases/sql_join_dist_broadcast/case.toml`. diff --git a/tests/perf/fixture-format.md b/tests/perf/fixture-format.md index c7d306d852a8..9de4d75084a4 100644 --- a/tests/perf/fixture-format.md +++ b/tests/perf/fixture-format.md @@ -58,6 +58,25 @@ warmup = 0 iterations = 1 ``` +A table can override the volume of the shared layout with its own `sst_count` +and `rows_per_sst`; the remaining layout fields stay shared. This keeps mixed +cases, such as a join of a large fact table with a small dimension table, in one +deterministic layout: + +```toml +[[scenario.tables]] +database = "public" +name = "fact" +sst_count = 32 +rows_per_sst = 8192 + +[[scenario.tables]] +database = "public" +name = "dim" +sst_count = 1 +rows_per_sst = 16 +``` + ### Query kinds `kind = "prom_http"` runs a Prometheus range query by POSTing form fields to @@ -84,6 +103,24 @@ encoded primary key ascending, timestamp ascending, and sequence descending, as required by the Mito Parquet writer; sorting preserves every generated column and row, including deterministic-wave values. +A query can list `candidate_session_sql` statements that the candidate target +sends before the query itself, for example: + +```toml +[[scenario.queries]] +name = "join_dist_broadcast" +kind = "sql" +query = "SELECT f.host, avg(f.value) FROM fact f JOIN dim d ON f.host = d.host GROUP BY f.host" +candidate_session_sql = ["SET experimental_dist_join = true"] +``` + +Every `/v1/sql` request is answered in its own session context, so a session +setting only affects the statements of the request that carries it: the runner +sends the statements and the query as one multi-statement SQL string for the +validation, warmup, and measurement requests of that query. Only the candidate +target runs them, so a base build that does not know the setting measures the +same query with the default behavior. + `[scenario]` is required. Other scenario variants are intentionally unsupported for now, but `scenario.kind` leaves room for future `write_then_query` and `cache_warm_query` configuration. diff --git a/tests/perf/query_cases/sql_join_dist_broadcast/case.toml b/tests/perf/query_cases/sql_join_dist_broadcast/case.toml new file mode 100644 index 000000000000..1b09993f9070 --- /dev/null +++ b/tests/perf/query_cases/sql_join_dist_broadcast/case.toml @@ -0,0 +1,127 @@ +# SQL distributed broadcast-join case over direct SSTs. +# +# The case measures the nested broadcast join rewrite of the experimental +# `experimental_dist_join` opt-in: the join of a large probe (fact) table with a +# small build (dimension) table. The candidate target opts into the rewrite for +# its queries with `candidate_session_sql`, the base target runs the same query +# without the opt-in (and without a build that knows the setting yet), so the +# report compares "join on the frontend" with "join next to the probe data". +# +# The heuristic compares `probe_regions * build_bytes` with `probe_bytes`, so a +# single-region fact table with a much smaller dimension table already favors the +# build side by a wide margin. A single region per table is enough for the +# heuristic and keeps the case on the supported one-region-per-table direct-SST +# fixture path; the layout stays `regions = 1`, and the fact table is made large +# with rows instead of regions. +# +# Both tables also have to be reported as leaders with a non-zero disk size +# before the rewrite can fire: `query_regression_runner measure` waits for that +# in `information_schema.region_statistics` before it measures a direct-SST case. + +[case] +name = "sql_join_dist_broadcast" +description = "Distributed broadcast join (nested MergeScan rewrite) over a large fact table and a small dimension table" + +[scenario] +kind = "direct_readable_sst" +seed = 1004 + +# The probe side of the join: the large fact table that the join runs on. +[[scenario.tables]] +database = "public" +name = "sql_join_dist_fact" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["host", "instance"] +time_index = "ts" +sst_count = 32 +rows_per_sst = 8192 + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 16, prefix = "host" } + +[[scenario.tables.columns]] +name = "instance" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 64, prefix = "instance" } + +[[scenario.tables.columns]] +name = "value" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 1.0, max = 500.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +# The build side of the join: one row per host, so the join matches once per fact +# row and the build side stays small enough to be broadcast to the probe region. +[[scenario.tables]] +database = "public" +name = "sql_join_dist_dim" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["host", "zone"] +time_index = "ts" +sst_count = 1 +rows_per_sst = 16 + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 16, prefix = "host" } + +[[scenario.tables.columns]] +name = "zone" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 4, prefix = "zone" } + +[[scenario.tables.columns]] +name = "weight" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.1, max = 10.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +[scenario.layout] +regions = 1 +# Shared fallbacks; the tables above size themselves with their own sst_count and +# rows_per_sst. 32 x 8192 fact rows and one 16-row dimension SST keep the fact +# table orders of magnitude larger than the dimension table. +sst_count = 32 +rows_per_sst = 8192 +row_group_size = 1024 +series_count = 64 +start_unix_nanos = 1704067200000000000 +step_nanos = 1000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "round_robin" + +[[scenario.queries]] +name = "join_dist_broadcast_aggregate_top" +kind = "sql" +query = "SELECT f.host, d.zone, avg(f.value * d.weight) AS score FROM sql_join_dist_fact f JOIN sql_join_dist_dim d ON f.host = d.host WHERE f.ts >= '2024-01-01 00:00:00+0000'::TIMESTAMP AND f.ts < '2024-01-05 00:00:00+0000'::TIMESTAMP AND d.ts >= '2024-01-01 00:00:00+0000'::TIMESTAMP AND d.ts < '2024-01-01 01:00:00+0000'::TIMESTAMP AND d.zone IN ('zone1', 'zone2') GROUP BY f.host, d.zone ORDER BY score DESC LIMIT 10" +# Every request of `/v1/sql` gets its own session context, so the opt-in has to +# travel with the query it should affect. Only the candidate target runs these +# statements; the base target measures the same query with the default rewrite +# behavior. +candidate_session_sql = ["SET experimental_dist_join = true"] +warmup = 1 +iterations = 3 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 50 From 10add67458b101f7a5a5465dad7aa348d476fa2b Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sun, 4 Oct 2026 15:25:03 +0800 Subject: [PATCH 16/37] fix(query): skip impossible broadcast self-join candidates Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/datafusion.rs | 163 ++++++++++++++++--- src/query/src/dist_plan/dist_join_planner.rs | 58 ++++++- 2 files changed, 200 insertions(+), 21 deletions(-) diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 771e055dc21a..31ec8398b11f 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -1862,6 +1862,9 @@ mod tests { /// The physical table id a logical candidate table resolves to in the route gating test. const LOGICAL_PHYSICAL_TABLE_ID: TableId = 1026; + /// The table id of the table an aliased self-join joins with itself. + const SELF_JOIN_TABLE_ID: TableId = 1027; + /// An information extension with fixed region statistics and a counter of its `region_stats` /// calls. #[derive(Debug)] @@ -1990,6 +1993,42 @@ mod tests { .unwrap() } + /// The shape of [`dist_join_test_plan`] with the given tables under the aliases `a` and `b`: + /// an aliased self-join passes the same table id twice. + fn dist_join_aliased_test_plan(left: (TableId, &str), right: (TableId, &str)) -> LogicalPlan { + LogicalPlanBuilder::from(dist_join_test_scan("a", left.0, left.1)) + .join_on( + dist_join_test_scan("b", right.0, right.1), + JoinType::Inner, + vec![col("a.number").eq(col("b.number"))], + ) + .unwrap() + .build() + .unwrap() + } + + /// The statistics the engine derives from [`dist_join_reports`] and [`dist_join_test_routes`]: + /// the probe table with 10_000 bytes in two regions and the build table with 1_000 bytes in + /// one. + fn dist_join_test_stats() -> BTreeMap { + BTreeMap::from([ + ( + PROBE_TABLE_ID, + crate::dist_plan::DistJoinTableStats { + total_bytes: 10_000, + region_count: 2, + }, + ), + ( + BUILD_TABLE_ID, + crate::dist_plan::DistJoinTableStats { + total_bytes: 1_000, + region_count: 1, + }, + ), + ]) + } + /// A query context with or without the session opt-in of the rewrite. fn dist_join_query_ctx(opt_in: bool) -> QueryContextRef { let query_ctx = QueryContextBuilder::default().build(); @@ -2175,25 +2214,7 @@ mod tests { let plan = dist_join_test_plan(); let stats = query_engine.dist_join_stats(&ctx, &plan).await.unwrap(); - assert_eq!( - BTreeMap::from([ - ( - PROBE_TABLE_ID, - crate::dist_plan::DistJoinTableStats { - total_bytes: 10_000, - region_count: 2, - } - ), - ( - BUILD_TABLE_ID, - crate::dist_plan::DistJoinTableStats { - total_bytes: 1_000, - region_count: 1, - } - ), - ]), - stats.tables - ); + assert_eq!(dist_join_test_stats(), stats.tables); assert!(stats.favors_right_build(PROBE_TABLE_ID, BUILD_TABLE_ID)); assert_eq!( 1, @@ -2427,4 +2448,108 @@ mod tests { assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); assert_eq!(0, calls.load(Ordering::SeqCst)); } + + /// An aliased self-join prices one table against itself, which can never satisfy + /// `N * B < B`: the engine resolves no route and fetches no statistics for it, and the + /// EXPLAIN of the same query keeps the existing plan. + #[tokio::test] + async fn test_dist_join_stats_self_join_keeps_plan() { + common_telemetry::init_default_ut_logging(); + + let calls = Arc::new(AtomicUsize::new(0)); + let engine = dist_join_engine(dist_join_reports(), calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let mut ctx = query_engine.engine_context(dist_join_query_ctx(true)); + let plan = + dist_join_aliased_test_plan((PROBE_TABLE_ID, "probe"), (PROBE_TABLE_ID, "probe")); + + assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); + assert_eq!( + 0, + calls.load(Ordering::SeqCst), + "a self-join must not fetch any statistics" + ); + + let explained = explain_logical_plan(query_engine, &mut ctx, &plan).await; + assert!( + !nests_join_in_merge_scan(&explained), + "the self-join must keep its two MergeScan boundaries, got:\n{explained}" + ); + assert_eq!(0, calls.load(Ordering::SeqCst)); + } + + /// The aliases of the two sides do not change the candidate: the same shape over two tables is + /// priced, the heuristic favors the build side, and the EXPLAIN shows the nested rewrite. + #[tokio::test] + async fn test_dist_join_stats_aliased_tables_stay_candidates() { + common_telemetry::init_default_ut_logging(); + + let calls = Arc::new(AtomicUsize::new(0)); + let engine = dist_join_engine(dist_join_reports(), calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let mut ctx = query_engine.engine_context(dist_join_query_ctx(true)); + let plan = + dist_join_aliased_test_plan((PROBE_TABLE_ID, "probe"), (BUILD_TABLE_ID, "build")); + + let stats = query_engine.dist_join_stats(&ctx, &plan).await.unwrap(); + assert_eq!(dist_join_test_stats(), stats.tables); + assert!(stats.favors_right_build(PROBE_TABLE_ID, BUILD_TABLE_ID)); + assert_eq!(1, calls.load(Ordering::SeqCst)); + + let explained = explain_logical_plan(query_engine, &mut ctx, &plan).await; + assert!( + nests_join_in_merge_scan(&explained), + "the EXPLAIN must show the nested rewrite, got:\n{explained}" + ); + assert_eq!(2, calls.load(Ordering::SeqCst)); + } + + /// A rejected self-join does not stop the traversal of the plan: the legitimate join next to + /// it is still priced, the self-join table gets no statistics, and the query still comes from + /// one `region_stats` call. + #[tokio::test] + async fn test_dist_join_stats_self_join_next_to_candidate() { + common_telemetry::init_default_ut_logging(); + + let calls = Arc::new(AtomicUsize::new(0)); + let mut routes = dist_join_test_routes(); + routes.push(dist_join_physical_route(SELF_JOIN_TABLE_ID, &[1])); + let mut reports = dist_join_reports(); + reports.push(region_stat( + RegionId::new(SELF_JOIN_TABLE_ID, 1), + 2_000, + RegionRole::Leader, + )); + let engine = dist_join_engine_with_routes(routes, reports, calls.clone(), true).await; + let query_engine = engine + .as_any() + .downcast_ref::() + .unwrap(); + let ctx = query_engine.engine_context(dist_join_query_ctx(true)); + let self_join = + dist_join_aliased_test_plan((SELF_JOIN_TABLE_ID, "logs"), (SELF_JOIN_TABLE_ID, "logs")); + let plan = LogicalPlanBuilder::from(self_join) + .join_on( + dist_join_test_plan(), + JoinType::Inner, + vec![col("a.number").eq(col("probe.number"))], + ) + .unwrap() + .build() + .unwrap(); + + let stats = query_engine.dist_join_stats(&ctx, &plan).await.unwrap(); + assert_eq!( + dist_join_test_stats(), + stats.tables, + "the rejected self-join must not be priced" + ); + assert_eq!(1, calls.load(Ordering::SeqCst)); + } } diff --git a/src/query/src/dist_plan/dist_join_planner.rs b/src/query/src/dist_plan/dist_join_planner.rs index e83656854c46..25732a9ff0fa 100644 --- a/src/query/src/dist_plan/dist_join_planner.rs +++ b/src/query/src/dist_plan/dist_join_planner.rs @@ -144,8 +144,11 @@ pub(crate) struct CandidateJoin { } /// Returns the candidate joins of `plan`, i.e. of its `EXPLAIN`/`EXPLAIN ANALYZE` input: the -/// `INNER` joins whose sides each resolve to a single physical base table and whose condition -/// compares the two sides. +/// `INNER` joins whose condition compares two sides that each resolve to a single physical base +/// table, with two different tables on the two sides. +/// +/// A join of one table with itself, e.g. under two aliases, can never satisfy `N * B < B`, so it +/// is left out and its statistics are never fetched. /// /// The plan is the one before the analyzers run: DataFusion's SQL planner leaves the whole `ON` /// clause of `LogicalPlanBuilder::join_on` in `filter`, and `ExtractEquijoinPredicate` only @@ -161,6 +164,7 @@ pub(crate) fn candidate_joins(plan: &LogicalPlan) -> Vec { side_base_table_id(&join.left), side_base_table_id(&join.right), ) + && probe != build { candidates.push(CandidateJoin { probe, build }); } @@ -1483,6 +1487,56 @@ mod tests { } } + /// The two sides of a candidate join must be two different physical tables: the aliased + /// self-join `logs AS l JOIN logs AS r` can never satisfy `N * B < B`, so it is no candidate, + /// while the same shape over two tables stays one. Rejecting a self-join keeps the candidates + /// of its inputs. + #[test] + fn candidate_joins_reject_same_physical_table() { + let self_join = inner_join( + table_scan("l", 7, "logs"), + table_scan("r", 7, "logs"), + col("l.number"), + col("r.number"), + ); + let aliased_tables = inner_join( + table_scan("l", 7, "logs"), + table_scan("r", 8, "archive"), + col("l.number"), + col("r.number"), + ); + let legitimate = inner_join( + table_scan("t1", 1, "t1"), + table_scan("t2", 2, "t2"), + col("t1.number"), + col("t2.number"), + ); + let mixed = inner_join( + self_join.clone(), + legitimate, + col("l.number"), + col("t1.number"), + ); + + let cases = [ + ("aliased self-join", self_join, vec![]), + ( + "aliased different tables", + aliased_tables, + vec![CandidateJoin { probe: 7, build: 8 }], + ), + ( + "self-join next to a legitimate join", + mixed, + vec![CandidateJoin { probe: 1, build: 2 }], + ), + ]; + + for (name, plan, expected) in cases { + assert_eq!(expected, candidate_joins(&plan), "case: {name}"); + } + } + /// A one-column JSON2 boundary schema, as `JsonSchemaConcretizeRule` expects to fix up. fn json2_boundary_schema() -> DFSchemaRef { let field = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata( From 53dc09375a429c0a3ddf64df3a720f5b09d143e6 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sun, 4 Oct 2026 15:49:10 +0800 Subject: [PATCH 17/37] test(perf): require remote join execution before measuring Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../bin/query_regression_runner/measure.rs | 438 +++++++++++++++++- .../src/bin/query_regression_runner/model.rs | 38 ++ tests/perf/README.md | 25 +- tests/perf/fixture-format.md | 22 + .../sql_join_dist_broadcast/case.toml | 14 +- 5 files changed, 519 insertions(+), 18 deletions(-) diff --git a/src/cmd/src/bin/query_regression_runner/measure.rs b/src/cmd/src/bin/query_regression_runner/measure.rs index 54b1198d2c82..463ec3acf63e 100644 --- a/src/cmd/src/bin/query_regression_runner/measure.rs +++ b/src/cmd/src/bin/query_regression_runner/measure.rs @@ -23,12 +23,12 @@ use crate::query_regression_runner::model::{Measurement, Query, QueryResult, Sce use crate::query_regression_runner::plan::{load_plan, normalize_scenario}; use crate::query_regression_runner::sql::{ extract_rows, http_post_multi_statement_sql, http_post_prom_range_query, http_post_sql, - row_u64, sql_ident, sql_string, + row_u64, row_value, sql_ident, sql_string, }; use crate::query_regression_runner::{MeasureArgs, Result}; -/// How long the measurement of a direct-SST case waits for the region statistics of its tables -/// to settle before it measures anything. +/// How long a candidate run waits for the region statistics of the case tables when a query opts +/// in with `require_region_statistics`. /// /// Datanodes report the role and the size of their regions on every heartbeat (`interval` of the /// heartbeat options, three seconds by default) and a region is reported as a leader with a @@ -38,6 +38,8 @@ use crate::query_regression_runner::{MeasureArgs, Result}; const REGION_STATS_TIMEOUT: Duration = Duration::from_secs(60); /// Poll interval of [`wait_for_region_statistics`]. const REGION_STATS_POLL_INTERVAL: Duration = Duration::from_secs(1); +/// Validation phase of the separate executed-plan preflight of a candidate query. +const EXECUTED_PLAN_PREFLIGHT: &str = "executed_plan_preflight"; pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { if !args.http_timeout.is_finite() || args.http_timeout < 0.0 { @@ -65,13 +67,16 @@ pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { .timeout(Duration::from_secs_f64(args.http_timeout)) .build()?; let direct_fixture = matches!(scenario_kind(&plan), Some("direct_readable_sst")); + // Only the candidate target waits for region statistics, and only when a query of the case + // opts in: the base target does not need the reports, and a case without such a query keeps + // measuring immediately. let base = run_target( args.base_http_port, &tables, &configured_queries, &client, false, - direct_fixture, + waits_for_region_statistics(false, direct_fixture, &configured_queries), ) .await; let candidate_target = run_target( @@ -80,7 +85,7 @@ pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { &configured_queries, &client, true, - direct_fixture, + waits_for_region_statistics(true, direct_fixture, &configured_queries), ) .await; let thresholds = enforce_thresholds(&configured_queries, &base, &candidate_target)?; @@ -165,6 +170,15 @@ fn scenario_kind(plan: &Value) -> Option<&str> { plan.get("scenario")?.get("kind")?.as_str() } +/// Whether a run of `queries` on a target waits for the region statistics of the case tables +/// before it measures anything. +/// +/// Only candidate direct-SST cases with an opted-in query wait; base targets and candidate +/// cases without an opted-in query do not wait. +fn waits_for_region_statistics(candidate: bool, direct_fixture: bool, queries: &[Query]) -> bool { + candidate && direct_fixture && queries.iter().any(|query| query.require_region_statistics) +} + /// The region statistics of `table`: `(regions, ready_regions)`, where a ready region has an /// ordinary leader report with a non-zero disk size. /// @@ -185,11 +199,11 @@ async fn region_statistics(client: &Client, port: u16, db: &str, table: &str) -> http_post_sql(client, port, &sql, db).await } -/// Waits until every table of a direct-SST case has one leader report with a non-zero size per -/// region, and returns the samples of the last round plus an error sample when the statistics did -/// not settle in time. +/// Waits until every table of the case has one leader report with a non-zero size per region, and +/// returns the samples of the last round plus an error sample when the statistics did not settle +/// in time. /// -/// The measurement of a case that depends on those statistics would otherwise start before the +/// The measurement of a query that depends on those statistics would otherwise start before the /// datanode reported them and silently measure the plan the case is not about. async fn wait_for_region_statistics( client: &Client, @@ -262,6 +276,121 @@ fn parse_region_statistics(body: &Value) -> Option<(u64, u64)> { )) } +/// The `executed_plan_preflight` validation error of `query`, with the HTTP response of the +/// separate preflight request when one was made. +fn preflight_failure(query: &Query, error: &str, sample: Option<&Value>) -> Value { + json!({ + "query": query.name, + "phase": EXECUTED_PLAN_PREFLIGHT, + "error": error, + "response": sample.and_then(|sample| sample.get("response")), + }) +} + +/// A target result that failed before its measurement loops: the validation samples so far and +/// no measurement at all. +fn failed_before_measurement(validation: Vec, validation_errors: Vec) -> QueryResult { + QueryResult { + validation, + validation_errors, + measurements: Vec::new(), + status: "failed".to_string(), + } +} + +/// The `EXPLAIN ANALYZE` clone of a candidate query that must run `candidate_remote_operator` +/// remotely, or the reason it cannot be checked. `Ok(None)` for a query without the field; the +/// operator is matched as a token, so an empty pattern is rejected because it would match every +/// plan. +fn candidate_preflight(query: &Query) -> std::result::Result, &'static str> { + let Some(operator) = query.candidate_remote_operator.as_deref().map(str::trim) else { + return Ok(None); + }; + if !matches!(query.kind.as_deref(), None | Some("sql")) { + return Err("candidate_remote_operator needs a plain sql query"); + } + if operator.is_empty() { + return Err("candidate_remote_operator must not be empty"); + } + let mut explain = query.clone(); + explain.query = format!("EXPLAIN ANALYZE {}", query.query); + Ok(Some((explain, operator))) +} + +/// Tags a preflight sample with its phase and returns the error of an executed TEXT DistAnalyze +/// output that does not keep `operator` out of stage 0 and in a later stage. Anything +/// unrecognized fails closed. +fn executed_plan_preflight(sample: &mut Value, operator: &str) -> Option { + sample + .as_object_mut() + .expect("HTTP samples are objects") + .insert( + "phase".to_string(), + Value::String(EXECUTED_PLAN_PREFLIGHT.to_string()), + ); + if !sample["ok"].as_bool().unwrap_or(false) { + return Some("EXPLAIN ANALYZE request failed".to_string()); + } + let mut stages = Vec::new(); + for row in &extract_rows(&sample["response"]) { + match plan_row(row) { + Ok(Some((stage, plan))) => stages.push((stage, plan_has_operator(plan, operator))), + Ok(None) => {} + Err(error) => return Some(error), + } + } + if !stages.iter().any(|(stage, _)| *stage == 0) { + return Some("executed plan has no stage 0".to_string()); + } + if stages + .iter() + .any(|(stage, matched)| *stage == 0 && *matched) + { + return Some(format!("stage 0 still runs {operator} on the frontend")); + } + if !stages.iter().any(|(stage, matched)| *stage > 0 && *matched) { + return Some(format!("no remote stage runs {operator}")); + } + None +} + +/// The `(stage, plan)` of one executed TEXT DistAnalyze row, or `Ok(None)` for the explicit +/// null/null `Total rows:` trailer row. A row with a missing, non-null, or unparseable stage or +/// node, or without non-empty plan text, is malformed. +fn plan_row(row: &Value) -> std::result::Result, String> { + // Only a row whose `stage` and `node` are explicitly null is the trailer; a row with missing + // fields or values that merely fail to parse must not pass as one. + if row_value(row, 0, "stage") == Some(&Value::Null) + && row_value(row, 1, "node") == Some(&Value::Null) + && row_value(row, 2, "plan") + .and_then(Value::as_str) + .is_some_and(|plan| plan.trim_start().starts_with("Total rows:")) + { + return Ok(None); + } + match ( + row_u64(row, 0, "stage"), + row_u64(row, 1, "node"), + row_value(row, 2, "plan").and_then(Value::as_str), + ) { + // An empty plan can never hold an operator, so it would fake a clean stage 0 or a remote + // stage and must not pass as a plan row. + (Ok(stage), Ok(_), Some(plan)) if !plan.trim().is_empty() => Ok(Some((stage, plan))), + _ => Err(format!("malformed executed plan row {row}")), + } +} + +/// Whether any line of `plan` starts with exactly `operator` as its own token, e.g. +/// `HashJoinExec:` or a bare `HashJoinExec`. Mentioning the name inside a parameter or a +/// predicate does not match. +fn plan_has_operator(plan: &str, operator: &str) -> bool { + plan.lines().any(|line| { + line.trim().strip_prefix(operator).is_some_and(|rest| { + rest.is_empty() || rest.starts_with(':') || rest.starts_with(char::is_whitespace) + }) + }) +} + async fn run_target( port: u16, tables: &[Table], @@ -280,6 +409,8 @@ async fn run_target( end: None, step: None, candidate_session_sql: Vec::new(), + require_region_statistics: false, + candidate_remote_operator: None, warmup: 0, iterations: 1, thresholds: Map::new(), @@ -291,10 +422,12 @@ async fn run_target( let mut validation_errors = Vec::new(); if wait_for_region_stats { let (stats, error) = wait_for_region_statistics(client, port, tables).await; + validation.extend(stats); if let Some(error) = error { + // Measuring now would silently time the plan the case is not about. validation_errors.push(error); + return failed_before_measurement(validation, validation_errors); } - validation.extend(stats); } for table in tables { let sql = format!("SHOW CREATE TABLE {}", sql_ident(&table.name)); @@ -326,6 +459,31 @@ async fn run_target( } validation.push(first); + // The candidate target only measures a query that expects a remote operator once its + // executed plan was checked: the preflight is one separate `EXPLAIN ANALYZE` request with + // the same session prefix, kept in `validation`, and any failure reports the query without a + // timed sample. It proves this execution ran the operator remotely, not the timed requests. + if candidate { + for query in &queries { + let (explain, operator) = match candidate_preflight(query) { + Ok(Some(preflight)) => preflight, + Ok(None) => continue, + Err(reason) => { + validation_errors.push(preflight_failure(query, reason, None)); + return failed_before_measurement(validation, validation_errors); + } + }; + let mut sample = post_query(client, port, &explain, db, true).await; + let error = executed_plan_preflight(&mut sample, operator) + .map(|error| preflight_failure(query, &error, Some(&sample))); + validation.push(sample); + if let Some(error) = error { + validation_errors.push(error); + return failed_before_measurement(validation, validation_errors); + } + } + } + let mut measurements = Vec::with_capacity(queries.len()); for query in &queries { for _ in 0..query.warmup { @@ -602,6 +760,38 @@ mod tests { ); } + #[test] + fn region_statistics_wait_is_candidate_only_and_query_opt_in() { + let plain: Query = serde_json::from_value(json!({"query": "SELECT 1"})).unwrap(); + let opted_in: Query = serde_json::from_value(json!({ + "query": "SELECT 1", + "require_region_statistics": true, + })) + .unwrap(); + // The default query of a direct-SST case never waits. + assert!(!waits_for_region_statistics(true, true, &[plain.clone()])); + assert!(!waits_for_region_statistics(true, true, &[])); + // An opted-in query waits on the candidate target only. + assert!(waits_for_region_statistics( + true, + true, + std::slice::from_ref(&opted_in) + )); + assert!(!waits_for_region_statistics( + false, + true, + std::slice::from_ref(&opted_in) + )); + // No query can make a run of another scenario kind wait. + assert!(!waits_for_region_statistics( + true, + false, + std::slice::from_ref(&opted_in) + )); + // One opted-in query is enough when a case mixes queries. + assert!(waits_for_region_statistics(true, true, &[plain, opted_in])); + } + #[test] fn threshold_rejects_unknown_keys_and_zero_base() { let query = Query { @@ -612,6 +802,8 @@ mod tests { end: None, step: None, candidate_session_sql: Vec::new(), + require_region_statistics: false, + candidate_remote_operator: None, warmup: 0, iterations: 1, thresholds: Map::from_iter([ @@ -647,4 +839,230 @@ mod tests { ); assert_eq!(results[1]["reason"], "unsupported threshold"); } + + #[test] + fn candidate_preflight_explains_only_the_candidate_and_rejects_bad_config() { + let configured = |kind: Option<&str>, operator: Option<&str>, prefix: &[&str]| { + serde_json::from_value::(json!({ + "query": "SELECT 1", + "kind": kind, + "candidate_session_sql": prefix, + "candidate_remote_operator": operator, + })) + .unwrap() + }; + // An unconfigured query is never checked. + assert!(matches!( + candidate_preflight(&configured(None, None, &[])), + Ok(None) + )); + // The preflight carries the candidate session prefix, and the query itself is untouched. + let query = configured( + Some("sql"), + Some(" HashJoinExec "), + &["SET experimental_dist_join = true"], + ); + let (explain, operator) = candidate_preflight(&query).unwrap().unwrap(); + assert_eq!("HashJoinExec", operator); + assert_eq!( + "SET experimental_dist_join = true;\nEXPLAIN ANALYZE SELECT 1", + explain.request_sql(true) + ); + assert_eq!("EXPLAIN ANALYZE SELECT 1", explain.request_sql(false)); + assert_eq!("SELECT 1", query.query); + // An operator without a kind defaults to a plain SQL query. + assert!(candidate_preflight(&configured(None, Some("HashJoinExec"), &[])).is_ok()); + assert_eq!( + Err("candidate_remote_operator must not be empty"), + candidate_preflight(&configured(Some("sql"), Some(" "), &[])).map(|_| ()) + ); + assert_eq!( + Err("candidate_remote_operator needs a plain sql query"), + candidate_preflight(&configured(Some("prom_http"), Some("HashJoinExec"), &[])) + .map(|_| ()) + ); + } + + #[test] + fn executed_plan_preflight_fails_closed() { + let native = |rows: Value| { + json!({ + "ok": true, + "response": {"output": [{"records": { + "schema": {"column_schemas": [ + {"name": "stage", "data_type": "UInt32"}, + {"name": "node", "data_type": "UInt32"}, + {"name": "plan", "data_type": "Utf8"} + ]}, + "rows": rows, + "total_rows": 3 + }}]} + }) + }; + let cases = [ + ( + "a remote join passes", + native(json!([ + [0, 0, "SortExec: fetch=10\n MergeScanExec: peer_id=1"], + [1, 0, "HashJoinExec: mode=CollectLeft"], + [null, null, "Total rows: 10"], + ])), + None, + ), + ( + "object rows pass as well", + json!({"ok": true, "response": {"data": [ + {"stage": 0, "node": 0, "plan": "MergeScanExec: peer_id=1"}, + {"stage": 1, "node": 3, "plan": " HashJoinExec: mode=CollectLeft"}, + ]}}), + None, + ), + ( + "a failed request fails closed", + json!({"ok": false, "error": "HTTP 500"}), + Some("EXPLAIN ANALYZE request failed"), + ), + ( + "a join in stage 0 fails closed", + native(json!([ + [ + 0, + 0, + "HashJoinExec: mode=Partitioned\n MergeScanExec: peer_id=1" + ], + [1, 0, "RegionScanExec: region_id=1"], + [null, null, "Total rows: 10"], + ])), + Some("stage 0 still runs HashJoinExec"), + ), + ( + "a plan without a remote stage fails closed", + native(json!([ + [0, 0, "MergeScanExec: peer_id=1"], + [null, null, "Total rows: 10"], + ])), + Some("no remote stage runs HashJoinExec"), + ), + ( + "a plan without stage 0 fails closed", + native(json!([[1, 0, "HashJoinExec: mode=CollectLeft"]])), + Some("executed plan has no stage 0"), + ), + ( + "a malformed stage fails closed", + native(json!([ + [0, 0, "MergeScanExec: peer_id=1"], + ["stage one", 0, "HashJoinExec: mode=CollectLeft"], + ])), + Some("malformed executed plan row"), + ), + ( + "a malformed node fails closed", + native(json!([ + [0, 0, "MergeScanExec: peer_id=1"], + [1, "node one", "HashJoinExec: mode=CollectLeft"], + ])), + Some("malformed executed plan row"), + ), + ( + "a half-empty stage row fails closed", + native(json!([ + [0, 0, "MergeScanExec: peer_id=1"], + [1, null, "HashJoinExec: mode=CollectLeft"], + ])), + Some("malformed executed plan row"), + ), + ( + "an unrecognized row without a stage fails closed", + native(json!([ + [0, 0, "MergeScanExec: peer_id=1"], + [null, null, "some totals"], + ])), + Some("malformed executed plan row"), + ), + ( + "a plan row without plan text fails closed", + native(json!([[0, 0, null]])), + Some("malformed executed plan row"), + ), + ( + "a non-null bad stage and node do not pass as trailer", + native(json!([ + [0, 0, "MergeScanExec: peer_id=1"], + [1, 0, "HashJoinExec: mode=CollectLeft"], + ["stage one", "node one", "Total rows: 10"], + ])), + Some("malformed executed plan row"), + ), + ( + "missing stage and node fields do not pass as trailer", + json!({"ok": true, "response": {"data": [ + {"stage": 0, "node": 0, "plan": "MergeScanExec: peer_id=1"}, + {"stage": 1, "node": 2, "plan": "HashJoinExec: mode=CollectLeft"}, + {"plan": "Total rows: 10"}, + ]}}), + Some("malformed executed plan row"), + ), + ( + "an empty frontend plan does not pass as a clean stage 0", + native(json!([ + [0, 0, ""], + [1, 0, "HashJoinExec: mode=CollectLeft"], + ])), + Some("malformed executed plan row"), + ), + ( + "an empty remote plan does not pass as a remote stage", + native(json!([[0, 0, "MergeScanExec: peer_id=1"], [1, 0, " "],])), + Some("malformed executed plan row"), + ), + ( + "the name inside a parameter does not match", + native(json!([ + [0, 0, "FilterExec: note = 'HashJoinExec: not an operator'"], + [1, 0, "RegionScanExec: region_id=1"], + [null, null, "Total rows: 1"], + ])), + Some("no remote stage runs HashJoinExec"), + ), + ( + "a longer name does not match", + native(json!([ + [0, 0, "MergeScanExec: peer_id=1"], + [1, 0, "HashJoinExecExtra: mode=CollectLeft"], + [null, null, "Total rows: 10"], + ])), + Some("no remote stage runs HashJoinExec"), + ), + ]; + for (name, mut sample, expected) in cases { + let error = executed_plan_preflight(&mut sample, "HashJoinExec"); + assert_eq!( + Some("executed_plan_preflight"), + sample["phase"].as_str(), + "{name}" + ); + match expected { + Some(expected) => { + let error = error.unwrap_or_default(); + assert!(error.contains(expected), "{name}: unexpected {error:?}"); + } + None => assert!(error.is_none(), "{name}: unexpected {error:?}"), + } + } + } + + #[test] + fn failure_before_measurement_reports_no_measurement() { + let result = failed_before_measurement( + vec![json!({"sql": "EXPLAIN ANALYZE SELECT 1", "phase": "executed_plan_preflight"})], + vec![ + json!({"phase": "executed_plan_preflight", "error": "no remote stage runs HashJoinExec"}), + ], + ); + assert_eq!("failed", result.status); + assert!(result.measurements.is_empty()); + assert_eq!(1, result.validation.len()); + assert_eq!(1, result.validation_errors.len()); + } } diff --git a/src/cmd/src/bin/query_regression_runner/model.rs b/src/cmd/src/bin/query_regression_runner/model.rs index 86a1c7cce119..a435023a45fe 100644 --- a/src/cmd/src/bin/query_regression_runner/model.rs +++ b/src/cmd/src/bin/query_regression_runner/model.rs @@ -230,6 +230,15 @@ pub(super) struct Query { /// setting yet still measures the same query without it. #[serde(default)] pub(super) candidate_session_sql: Vec, + /// Opts the candidate direct-SST target into waiting for its tables' leader/size reports + /// before any measurements. Defaults to false; the base target never waits. + #[serde(default)] + pub(super) require_region_statistics: bool, + /// An operator that a separate candidate `EXPLAIN ANALYZE` must execute remotely, not in + /// stage 0. Failure aborts all candidate measurements. Only SQL queries are supported; + /// the preflight uses the same session prefix but proves only its own execution. + #[serde(default)] + pub(super) candidate_remote_operator: Option, #[serde(default)] pub(super) warmup: usize, #[serde(default = "one")] @@ -300,6 +309,8 @@ mod tests { .into_iter() .map(ToString::to_string) .collect(), + require_region_statistics: false, + candidate_remote_operator: None, warmup: 0, iterations: 1, thresholds: Map::new(), @@ -341,4 +352,31 @@ mod tests { let query: Query = serde_json::from_value(json!({"query": "SELECT 1"})).unwrap(); assert!(query.candidate_session_sql.is_empty()); } + + #[test] + fn query_defaults_to_not_requiring_region_statistics() { + let query: Query = serde_json::from_value(json!({"query": "SELECT 1"})).unwrap(); + assert!(!query.require_region_statistics); + let opted_in: Query = serde_json::from_value(json!({ + "query": "SELECT 1", + "require_region_statistics": true, + })) + .unwrap(); + assert!(opted_in.require_region_statistics); + } + + #[test] + fn query_defaults_to_no_candidate_remote_operator() { + let query: Query = serde_json::from_value(json!({"query": "SELECT 1"})).unwrap(); + assert!(query.candidate_remote_operator.is_none()); + let configured: Query = serde_json::from_value(json!({ + "query": "SELECT 1", + "candidate_remote_operator": "HashJoinExec", + })) + .unwrap(); + assert_eq!( + Some("HashJoinExec"), + configured.candidate_remote_operator.as_deref() + ); + } } diff --git a/tests/perf/README.md b/tests/perf/README.md index c9089e317f88..06172b66efae 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -350,11 +350,12 @@ The outer driver currently supports `direct_readable_sst`, A direct-SST case sizes its tables with the shared `[scenario.layout]` and can override `sst_count` and `rows_per_sst` per table, so a case can join a large fact -table with a small dimension table. Before measuring, the runner waits until -every table of the case has one leader region report with a non-zero disk size in -`information_schema.region_statistics`: the datanode reports those statistics on -each heartbeat, and a rewrite that prices tables by them (such as the nested -broadcast join) would otherwise measure the un-rewritten plan. A query can list +table with a small dimension table. Setting `require_region_statistics = true` +on any query makes the candidate direct-SST target wait for every case table's +leader reports with non-zero disk size in `information_schema.region_statistics`. +This untimed wait precedes all candidate measurements; timeout or request failure +aborts that target with no timed samples. The base target and candidate cases +without an opted-in query do not wait. A query can list `candidate_session_sql` statements that only the candidate target runs; the runner sends them in the same `/v1/sql` request as the query itself (validation, warmup and measurement), because every request gets its own session context and a @@ -362,6 +363,13 @@ warmup and measurement), because every request gets its own session context and does not know a setting yet therefore keeps measuring the query without it, which is what makes `SET experimental_dist_join = true` a candidate-only opt-in. +A SQL query can set `candidate_remote_operator`, e.g. `HashJoinExec`, to require +one separate untimed `EXPLAIN ANALYZE` with the same session prefix. It must show +the operator in a positive stage and absent from stage 0; failed requests or +malformed stage/node/plan rows abort all candidate measurements. This proves only +the preflight execution, not later timed SELECT plans. The base target and +candidate cases without the field are unchanged. + ## Metrics Primary gates should compare query work rather than plan text: @@ -582,7 +590,12 @@ Additional SQL optimizer cases: - `sql_join_dist_broadcast`: a 32-SST fact table joined with a 16-row dimension table, comparing the join on the frontend with the nested broadcast join rewrite of the experimental `experimental_dist_join` opt-in. Only the - candidate target sends the opt-in, through `candidate_session_sql`. + candidate target sends the opt-in, through `candidate_session_sql`, and the + candidate query sets `require_region_statistics = true` so it waits for the + region statistics the rewrite needs before it is measured. It also sets + `candidate_remote_operator = "HashJoinExec"`, so a candidate run in which the + join stayed on the frontend fails without a timed sample instead of only + reporting the un-rewritten latency. These cases are not part of the routine default case group; run them by path, for example `--cases tests/perf/query_cases/sql_join_dist_broadcast/case.toml`. diff --git a/tests/perf/fixture-format.md b/tests/perf/fixture-format.md index 9de4d75084a4..ecf7a267dadb 100644 --- a/tests/perf/fixture-format.md +++ b/tests/perf/fixture-format.md @@ -121,6 +121,28 @@ validation, warmup, and measurement requests of that query. Only the candidate target runs them, so a base build that does not know the setting measures the same query with the default behavior. +For a statistics-dependent rewrite, the query may also set: + +```toml +require_region_statistics = true +candidate_remote_operator = "HashJoinExec" +``` + +`require_region_statistics` defaults to `false`. If any query opts in, the +candidate direct-SST target waits for every case table's leader reports with +non-zero disk size in `information_schema.region_statistics` before any timed +samples. Timeout or request failure aborts all measurements on that target. +Base targets and candidate cases without the flag do not wait. + +`candidate_remote_operator` defaults to unset and supports SQL queries only. +Before any measurements, the candidate runs a separate `EXPLAIN ANALYZE` with +the query's session prefix. Its TEXT DistAnalyze rows must have valid `stage`, +`node`, and non-empty `plan` values (apart from the null/null `Total rows:` +trailer), with the exact operator token absent from stage 0 and present in a +positive stage. Failure aborts all candidate measurements. The complete preflight +sample is retained in validation; it proves only that execution, not the plans +of later timed requests. The base target does not run this preflight. + `[scenario]` is required. Other scenario variants are intentionally unsupported for now, but `scenario.kind` leaves room for future `write_then_query` and `cache_warm_query` configuration. diff --git a/tests/perf/query_cases/sql_join_dist_broadcast/case.toml b/tests/perf/query_cases/sql_join_dist_broadcast/case.toml index 1b09993f9070..d000d6324f48 100644 --- a/tests/perf/query_cases/sql_join_dist_broadcast/case.toml +++ b/tests/perf/query_cases/sql_join_dist_broadcast/case.toml @@ -15,8 +15,11 @@ # with rows instead of regions. # # Both tables also have to be reported as leaders with a non-zero disk size -# before the rewrite can fire: `query_regression_runner measure` waits for that -# in `information_schema.region_statistics` before it measures a direct-SST case. +# before the rewrite can fire. `require_region_statistics = true` opts this +# query into waiting for that in `information_schema.region_statistics` before +# the candidate target measures it; no other query or target needs the wait. +# The wait only makes the statistics available in time, so the query also pins +# the rewrite with `candidate_remote_operator` below. [case] name = "sql_join_dist_broadcast" @@ -120,6 +123,13 @@ query = "SELECT f.host, d.zone, avg(f.value * d.weight) AS score FROM sql_join_d # statements; the base target measures the same query with the default rewrite # behavior. candidate_session_sql = ["SET experimental_dist_join = true"] +# Before it measures, the candidate target runs one separate EXPLAIN ANALYZE +# request of this query and fails it without a timed sample unless the executed +# plan keeps HashJoinExec out of the frontend stage and runs it in a later one. +candidate_remote_operator = "HashJoinExec" +# Wait for both tables' leader/size reports before any candidate measurements. +# A failed wait or preflight aborts the candidate target without timed samples. +require_region_statistics = true warmup = 1 iterations = 3 From 694489fa209f75b0aed48cb32bbfcd568538a2c6 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Sun, 4 Oct 2026 16:01:05 +0800 Subject: [PATCH 18/37] test(perf): avoid cloning query in readiness test Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cmd/src/bin/query_regression_runner/measure.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/cmd/src/bin/query_regression_runner/measure.rs b/src/cmd/src/bin/query_regression_runner/measure.rs index 463ec3acf63e..a217d7a64497 100644 --- a/src/cmd/src/bin/query_regression_runner/measure.rs +++ b/src/cmd/src/bin/query_regression_runner/measure.rs @@ -769,7 +769,11 @@ mod tests { })) .unwrap(); // The default query of a direct-SST case never waits. - assert!(!waits_for_region_statistics(true, true, &[plain.clone()])); + assert!(!waits_for_region_statistics( + true, + true, + std::slice::from_ref(&plain) + )); assert!(!waits_for_region_statistics(true, true, &[])); // An opted-in query waits on the candidate target only. assert!(waits_for_region_statistics( From 80980f71170cae553fd39dfed1cb35f1a72b0c4a Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:27:41 +0800 Subject: [PATCH 19/37] refactor(query): remove manual distributed join selector Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/datafusion.rs | 36 +--- src/query/src/dist_plan/analyzer.rs | 37 +--- src/query/src/dist_plan/analyzer/test.rs | 1 - src/query/src/dist_plan/dist_join_planner.rs | 189 ++++-------------- src/query/src/dist_plan/planner.rs | 1 + src/query/src/query_engine/state.rs | 1 - .../nested_merge_scan_capability_test.rs | 174 +++++++--------- 7 files changed, 117 insertions(+), 322 deletions(-) diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 31ec8398b11f..0a8cf80358ef 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -155,8 +155,7 @@ impl DatafusionQueryEngine { /// Statistics of the candidate tables of `logical_plan`, or `None` when the nested broadcast /// join rewrite has to keep the existing plan. /// - /// Gated by the session opt-in and the registered distributed rules; the manual build table - /// option selects the build side by name and never consults statistics. + /// Gated by the session opt-in and the registered distributed rules. /// /// The candidate tables are resolved to their physical routes before the statistics of the /// query are fetched: the route must still describe the table the plan captured and at least @@ -188,17 +187,6 @@ impl DatafusionQueryEngine { { return None; } - if state - .config() - .options() - .extensions - .get::() - .and_then(|options| options.nested_broadcast_join_build_table.as_ref()) - .is_some() - { - return None; - } - let candidates = candidate_joins(logical_plan); if candidates.is_empty() { return None; @@ -2232,8 +2220,7 @@ mod tests { } /// The gating of the automatic path: without the session opt-in, without the distributed - /// rules, with the manual build table option or without a candidate join the engine does not - /// fetch any statistics. + /// rules or without a candidate join the engine does not fetch any statistics. #[tokio::test] async fn test_dist_join_stats_gating() { common_telemetry::init_default_ut_logging(); @@ -2261,25 +2248,6 @@ mod tests { assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); assert_eq!(0, calls.load(Ordering::SeqCst)); - // With the manual build table option, which selects the build side by name. - let calls = Arc::new(AtomicUsize::new(0)); - let engine = dist_join_engine(dist_join_reports(), calls.clone(), true).await; - let query_engine = engine - .as_any() - .downcast_ref::() - .unwrap(); - let mut ctx = query_engine.engine_context(dist_join_query_ctx(true)); - ctx.state_mut() - .config_mut() - .options_mut() - .extensions - .insert(DistPlannerOptions { - nested_broadcast_join_build_table: Some("build".to_string()), - ..Default::default() - }); - assert!(query_engine.dist_join_stats(&ctx, &plan).await.is_none()); - assert_eq!(0, calls.load(Ordering::SeqCst)); - // Without a candidate join. let ctx = query_engine.engine_context(dist_join_query_ctx(true)); let scan_plan = dist_join_test_scan("probe", PROBE_TABLE_ID, "probe"); diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 2877ee0f4f2f..7728fc668f47 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -71,12 +71,6 @@ const OTHER_PHY_PART_COL_PLACEHOLDER: &str = "__OTHER_PHYSICAL_PART_COLS_PLACEHO #[derive(Debug, Clone, Default)] pub struct DistPlannerOptions { pub allow_query_fallback: bool, - /// Manual opt-in for the nested broadcast join rewrite: name of the (small) build - /// side table, e.g. `device_limits`. `None` (the default) keeps it disabled. - /// - /// Not cost based in this slice: setting the option enables the rewrite for the - /// supported shape, see [`DistJoinPlanner`](crate::dist_plan::DistJoinPlanner). - pub nested_broadcast_join_build_table: Option, } impl ConfigExtension for DistPlannerOptions { @@ -97,32 +91,17 @@ impl ExtensionOptions for DistPlannerOptions { } fn set(&mut self, key: &str, value: &str) -> DfResult<()> { - match key { - "nested_broadcast_join_build_table" => { - let value = value.trim().trim_matches(&['\'', '"'][..]); - self.nested_broadcast_join_build_table = - (!value.is_empty()).then(|| value.to_string()); - Ok(()) - } - _ => Err(datafusion_common::DataFusionError::NotImplemented(format!( - "DistPlannerOptions does not support set key: {key} with value: {value}" - ))), - } + Err(datafusion_common::DataFusionError::NotImplemented(format!( + "DistPlannerOptions does not support set key: {key} with value: {value}" + ))) } fn entries(&self) -> Vec { - vec![ - datafusion::config::ConfigEntry { - key: "allow_query_fallback".to_string(), - value: Some(self.allow_query_fallback.to_string()), - description: "Allow query fallback to fallback plan rewriter", - }, - datafusion::config::ConfigEntry { - key: "nested_broadcast_join_build_table".to_string(), - value: self.nested_broadcast_join_build_table.clone(), - description: "Build side table of the nested broadcast join rewrite", - }, - ] + vec![datafusion::config::ConfigEntry { + key: "allow_query_fallback".to_string(), + value: Some(self.allow_query_fallback.to_string()), + description: "Allow query fallback to fallback plan rewriter", + }] } } diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index d6547253348d..694f2bcde59b 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -745,7 +745,6 @@ fn expand_sort_alias_conflict_limit() { let mut config = ConfigOptions::default(); config.extensions.insert(DistPlannerOptions { allow_query_fallback: true, - ..Default::default() }); let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap(); diff --git a/src/query/src/dist_plan/dist_join_planner.rs b/src/query/src/dist_plan/dist_join_planner.rs index 25732a9ff0fa..a60894727bfe 100644 --- a/src/query/src/dist_plan/dist_join_planner.rs +++ b/src/query/src/dist_plan/dist_join_planner.rs @@ -14,9 +14,7 @@ //! Nested broadcast join rewrite for `INNER` joins of two distributed scans. //! -//! Opt-in only, in two ways. The manual selector is enabled when -//! [`DistPlannerOptions::nested_broadcast_join_build_table`] names the build side table. The -//! cost heuristic is enabled per query when the session sets `experimental_dist_join`: +//! The cost heuristic is enabled per query when the session sets `experimental_dist_join`: //! [`DatafusionQueryEngine::create_physical_plan`](crate::datafusion::DatafusionQueryEngine) //! then fetches the approximate disk sizes of the candidate tables into [`DistJoinStats`], a //! query-local config extension, and the heuristic rewrites the join only when those statistics @@ -44,11 +42,8 @@ use store_api::region_engine::RegionRole; use store_api::storage::RegionId; use table::metadata::{TableId, TableType}; use table::table::adapter::DfTableProviderAdapter; -use table::table_name::TableName; -use crate::dist_plan::analyzer::DistPlannerOptions; use crate::dist_plan::merge_scan::MergeScanLogicalPlan; -use crate::dist_plan::planner::table_name_of; /// Name of [`DistJoinPlanner`]. The query engine looks it up to check that this rewrite, the /// only consumer of [`DistJoinStats`], runs for the query at all. @@ -65,7 +60,7 @@ pub(crate) struct DistJoinTableStats { /// Per-table statistics fetched for the cost heuristic of one query. /// -/// A query-local `ConfigOptions` extension, like [`DistPlannerOptions`], built when +/// A query-local `ConfigOptions` extension, built when /// [`DatafusionQueryEngine::create_physical_plan`](crate::datafusion::DatafusionQueryEngine) /// sees the session opt-in (`SET experimental_dist_join = true`) on a plausible candidate join. #[derive(Debug, Clone, Default)] @@ -328,20 +323,11 @@ impl AnalyzerRule for DistJoinPlanner { } fn analyze(&self, plan: LogicalPlan, config: &ConfigOptions) -> DfResult { - let manual = config - .extensions - .get::() - .and_then(|options| options.nested_broadcast_join_build_table.as_deref()); - let selector = if let Some(build_table) = manual { - // The manual opt-in keeps its exact name matching and never consults statistics. - BuildSideSelector::Manual { build_table } - } else if let Some(stats) = config.extensions.get::() { - BuildSideSelector::Stats { stats } - } else { + let Some(stats) = config.extensions.get::() else { return Ok(plan); }; - let mut rewriter = NestedBroadcastJoinRewriter { selector }; + let mut rewriter = NestedBroadcastJoinRewriter { stats }; Ok(plan.rewrite(&mut rewriter)?.data) } } @@ -349,48 +335,8 @@ impl AnalyzerRule for DistJoinPlanner { /// Rewriter nesting the build side `MergeScan` of a supported join inside the probe side's /// `MergeScan`. struct NestedBroadcastJoinRewriter<'a> { - /// How the build side of a supported join is selected. - selector: BuildSideSelector<'a>, -} - -/// How the rewriter finds the build side of a supported join. -enum BuildSideSelector<'a> { - /// Manual opt-in: the configured build table name has to be the right input. - Manual { - /// Name of the build side table, either bare (`device_limits`) or fully qualified. - build_table: &'a str, - }, - /// Cost heuristic: only the right input is considered, and only when the statistics of the - /// query show that broadcasting its table to the probe regions is cheaper. - Stats { - /// Per-table statistics fetched for this query, see [`DistJoinStats`]. - stats: &'a DistJoinStats, - }, -} - -impl BuildSideSelector<'_> { - /// Whether the build side of the supported join shape is its right input. - fn selects_right_build(&self, probe: &LogicalPlan, build: &LogicalPlan) -> bool { - match self { - Self::Manual { build_table } => { - let left_is_build = - table_name_of(probe).is_some_and(|name| table_name_matches(&name, build_table)); - let right_is_build = - table_name_of(build).is_some_and(|name| table_name_matches(&name, build_table)); - - !left_is_build && right_is_build - } - Self::Stats { stats } => { - let (Some(probe), Some(build)) = - (side_base_table_id(probe), side_base_table_id(build)) - else { - return false; - }; - - stats.favors_right_build(probe, build) - } - } - } + /// Per-table statistics fetched for this query, see [`DistJoinStats`]. + stats: &'a DistJoinStats, } impl NestedBroadcastJoinRewriter<'_> { @@ -429,14 +375,13 @@ impl NestedBroadcastJoinRewriter<'_> { return None; } - // The selector picks the build side: the manual opt-in requires the configured table - // to be exactly the right input, the cost heuristic only ever considers the right - // input. The build side has to be the right input, because moving a left build side + // The right input is the only candidate build side, because moving a left build side // would change the join schema. - if !self - .selector - .selects_right_build(probe_merge_scan.input(), build_merge_scan.input()) - { + let (probe_table, build_table) = ( + side_base_table_id(probe_merge_scan.input())?, + side_base_table_id(build_merge_scan.input())?, + ); + if !self.stats.favors_right_build(probe_table, build_table) { return None; } @@ -556,15 +501,6 @@ fn has_json2_boundary(merge_scan: &MergeScanLogicalPlan) -> bool { .any(is_json2_extension_type) } -/// Whether the configured build side table name refers to `name`. Accepts the bare -/// table name or the fully qualified `catalog.schema.table` form, ignoring case. -fn table_name_matches(name: &TableName, configured: &str) -> bool { - let configured = configured.trim().trim_matches(&['\'', '"'][..]); - !configured.is_empty() - && (name.table_name.eq_ignore_ascii_case(configured) - || name.to_string().eq_ignore_ascii_case(configured)) -} - #[cfg(test)] mod tests { use std::collections::HashMap; @@ -592,6 +528,7 @@ mod tests { use super::*; use crate::dist_plan::DistPlannerAnalyzer; + use crate::dist_plan::planner::table_name_of; /// Two-column schema (`number`, `host`) shared by the test tables. fn test_schema() -> SchemaRef { @@ -669,18 +606,20 @@ mod tests { .unwrap() } - fn rewrite_config(build_table: &str) -> ConfigOptions { - let mut config = ConfigOptions::default(); - config.extensions.insert(DistPlannerOptions { - nested_broadcast_join_build_table: Some(build_table.to_string()), - ..Default::default() - }); - config + fn rewrite_config(tables: &[(TableId, DistJoinTableStats)]) -> ConfigOptions { + stats_config(tables) } - fn rewrite(plan: LogicalPlan, build_table: &str) -> LogicalPlan { + fn rewrite(plan: LogicalPlan) -> LogicalPlan { DistJoinPlanner {} - .analyze(plan, &rewrite_config(build_table)) + .analyze( + plan, + &rewrite_config(&[ + (1, table_stats(10_000, 2)), + (2, table_stats(1_000, 1)), + (3, table_stats(500, 1)), + ]), + ) .unwrap() } @@ -741,8 +680,8 @@ mod tests { assert_eq!(2, merge_scans(&result).len()); } - /// With the build table option the build side `MergeScan` becomes the inner scan of - /// the outer `MergeScan`, which is routed by the probe table's regions. + /// When statistics favor the right build side, its `MergeScan` becomes the inner scan + /// of the outer `MergeScan`, which is routed by the probe table's regions. /// /// The rewrite must only move where the join runs: the join fields, the residual /// filter, the NULL semantics and the output schema stay as the join had them. @@ -759,7 +698,7 @@ mod tests { }); let probe_scan = merge_scans(&plan).first().cloned().unwrap(); - let result = rewrite(plan.clone(), "t2"); + let result = rewrite(plan.clone()); let join = outer_join(&result); assert_eq!(JoinType::Inner, join.join_type); @@ -805,7 +744,7 @@ mod tests { let build_id = scans[1].remote_dyn_filter_producer_id().unwrap(); assert_ne!(probe_id, build_id); - let result = rewrite(plan, "t2"); + let result = rewrite(plan); let scans = merge_scans(&result); assert_eq!(2, scans.len()); @@ -827,7 +766,7 @@ mod tests { .collect::>() ); - let result = rewrite(plan.clone(), "t2"); + let result = rewrite(plan.clone()); let scans = merge_scans(&result); assert_eq!(2, scans.len()); @@ -862,7 +801,7 @@ mod tests { fn nested_broadcast_join_rewrite_handles_projection_wrappers() { let projected = projected_join_plan(); - let result = rewrite(projected.clone(), "t2"); + let result = rewrite(projected.clone()); let join = outer_join(&result); assert!(matches!(*join.left, LogicalPlan::Projection(_))); @@ -899,22 +838,13 @@ mod tests { }) } - /// A build table matching neither side keeps the plan unchanged. - #[test] - fn nested_broadcast_join_rewrite_ignores_unknown_build_table() { - let plan = distributed_join_plan(); - let result = rewrite(plan.clone(), "unknown_table"); - - assert_eq!(plan.to_string(), result.to_string()); - } - /// Only `INNER` joins are rewritten. #[test] fn nested_broadcast_join_rewrite_ignores_non_inner_join() { let plan = DistPlannerAnalyzer {} .analyze(join_plan(JoinType::Left), &ConfigOptions::default()) .unwrap(); - let result = rewrite(plan.clone(), "t2"); + let result = rewrite(plan.clone()); assert_eq!(plan.to_string(), result.to_string()); } @@ -942,48 +872,21 @@ mod tests { ..join.clone() }); - let result = rewrite(plan.clone(), "t2"); + let result = rewrite(plan.clone()); assert_eq!(plan.to_string(), result.to_string()); } - /// Two sides resolving to the configured build table (e.g. a self join) are ambiguous. - #[test] - fn nested_broadcast_join_rewrite_ignores_ambiguous_build_side() { - let plan = DistPlannerAnalyzer {} - .analyze( - join_plan_with_matching_side_names(), - &ConfigOptions::default(), - ) - .unwrap(); - let result = rewrite(plan.clone(), "t1"); - - assert_eq!(plan.to_string(), result.to_string()); - } - - /// `l INNER JOIN r` where `l` and `r` read two base tables both named `t1`. - fn join_plan_with_matching_side_names() -> LogicalPlan { - LogicalPlanBuilder::from(table_scan("l", 1, "t1")) - .join_on( - table_scan("r", 2, "t1"), - JoinType::Inner, - vec![col("l.number").eq(col("r.number"))], - ) - .unwrap() - .build() - .unwrap() - } - /// The rewrite stays at two levels: an already nested probe side is left alone. #[test] fn nested_broadcast_join_rewrite_ignores_more_than_two_levels() { - let nested = rewrite(distributed_join_plan(), "t2"); + let nested = rewrite(distributed_join_plan()); let build = DistPlannerAnalyzer {} .analyze(table_scan("t3", 3, "t3"), &ConfigOptions::default()) .unwrap(); let plan = inner_join(nested, build, col("t1.number"), col("t3.number")); - let result = rewrite(plan.clone(), "t3"); + let result = rewrite(plan.clone()); assert_eq!(plan.to_string(), result.to_string()); } @@ -1002,7 +905,7 @@ mod tests { nested_broadcast_join_probe_input(&plan).map(|plan| plan.to_string()) ); - let result = rewrite(plan.clone(), "t2"); + let result = rewrite(plan.clone()); let LogicalPlan::Extension(extension) = &result else { panic!("expected the rewritten plan to be a MergeScan, got: {result}"); }; @@ -1044,7 +947,7 @@ mod tests { col("t2.number"), ); - let result = rewrite(plan.clone(), "t2"); + let result = rewrite(plan.clone()); assert_eq!(plan.to_string(), result.to_string()); } @@ -1318,26 +1221,6 @@ mod tests { assert_eq!(plan.to_string(), result.to_string()); } - /// The manual build table keeps its exact name matching and is preferred over the statistics, - /// which its option never consults. - #[test] - fn manual_build_table_precedes_stats() { - let plan = distributed_join_plan(); - // These statistics do not favor the right build side. - let mut config = stats_config(&[(1, table_stats(1_000, 1)), (2, table_stats(10_000, 2))]); - config.extensions.insert(DistPlannerOptions { - nested_broadcast_join_build_table: Some("t2".to_string()), - ..Default::default() - }); - - let result = DistJoinPlanner {}.analyze(plan.clone(), &config).unwrap(); - - let join = outer_join(&result); - assert_eq!(2, merge_scans(&result).len()); - assert!(!contains_merge_scan(&join.left)); - assert!(contains_merge_scan(&join.right)); - } - /// Both plan shapes the engine sees are covered: the equality in `on`, and the unanalyzed /// `join_on` plan, whose whole `ON` clause is still in `filter`, including the input of an /// EXPLAIN. diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index 35952815453c..617a43e2f367 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -478,6 +478,7 @@ fn partition_column_types(table_info: &TableInfo) -> Vec<(String, ConcreteDataTy } /// Returns the resolved name of the first base table found in `plan`. +#[cfg(test)] pub(crate) fn table_name_of(plan: &LogicalPlan) -> Option { let mut extractor = TableScanExtractor::default(); let _ = plan.visit(&mut extractor).ok()?; diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index 45530a7b8c97..c720229fb59e 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -161,7 +161,6 @@ impl QueryEngineState { .extensions .insert(DistPlannerOptions { allow_query_fallback: true, - ..Default::default() }); } diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 3cf05e785747..1ef7739bab40 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -59,10 +59,9 @@ //! table id is invalidated once (see //! [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]). //! -//! [`test_nested_broadcast_join_rewrite_executes_on_probe_regions`] is the exception to the -//! hand-built plans: it plans the join SQL with the distributed rules of the frontend, checks that -//! the outer `MergeScan` of that plan is planned for the probe region the probe predicate selects, -//! and executes the plan on the frontend (see that test). +//! [`test_nested_broadcast_join_rewrite_executes_on_probe_regions`] exercises the production +//! selector through SQL `SET` and the regular query entrypoint, checks that the outer `MergeScan` +//! is planned for the probe region selected by the predicate, and executes the SQL on the frontend. use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; @@ -75,16 +74,12 @@ use common_meta::instruction::CacheIdent; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::info; -use datafusion::common::config::ConfigOptions; -use datafusion::optimizer::AnalyzerRule; use datafusion::physical_plan::ExecutionPlan; use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; use datanode::region_server::RegionServer; use frontend::instance::Instance; use query::datafusion::QUERY_PARALLELISM_HINT; -use query::dist_plan::{ - DistJoinPlanner, DistPlannerAnalyzer, DistPlannerOptions, MergeScanExec, MergeScanLogicalPlan, -}; +use query::dist_plan::{MergeScanExec, MergeScanLogicalPlan}; use query::parser::QueryLanguageParser; use query::query_engine::DefaultSerializer; use servers::query_handler::sql::SqlQueryHandler; @@ -319,17 +314,11 @@ async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation } } -/// The distributed rules produce the nested plan of the join SQL, the frontend plans the outer -/// `MergeScan` of that plan to the *selected* probe regions, and the datanodes execute it. -/// -/// Unlike the hand-built plans of the capability tests above, the plan here is the plan of a SQL -/// query: `DistPlannerAnalyzer` wraps the table scans of the join in `MergeScan`, and -/// `DistJoinPlanner` (opted in for the build table) nests the build side into the probe side's -/// scan. The probe predicate of the SQL selects one of the two probe regions, so the outer -/// `MergeScan` must be planned for that region only: routing by the whole join payload would read -/// both probe regions. The executed result must be the multiset of the equivalent join on the -/// frontend, which also covers the duplicated build keys, the NULL join keys, the residual join -/// filter and the projection of the query. +/// The SQL setting and region statistics select the nested plan through the production frontend +/// path. The probe predicate selects one of the two probe regions, so the outer `MergeScan` must be +/// planned for that region only: routing by the whole join payload would read both probe regions. +/// The executed result must preserve the duplicated build keys, NULL join keys, residual filter, +/// and projection of the query. #[tokio::test(flavor = "multi_thread")] async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { common_telemetry::init_default_ut_logging(); @@ -353,6 +342,7 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { ), ) .await; + insert_extra_probe_rows(&frontend).await; // The SQL under test: the projection reorders the columns of the join, the join keys `10` of // the build table are duplicated, the residual join filter rejects the join partners whose @@ -362,7 +352,7 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { "SELECT p.probe_v, p.a_id, b.build_v FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 - WHERE p.a_id < 100" + WHERE p.a_id < 100 AND p.a_ts < 10000::TIMESTAMP" ); // The rows of the SQL as `(probe_v, a_id, build_v)`: the probe row `(4, 10, 600)` joins the two // build rows of the duplicated key `10` (`600 + 1` and `600 + 2`), the probe row `(2, 20, 200)` @@ -385,40 +375,36 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { "unexpected result of the join on the frontend:\n{reference}" ); - // The plan of the rules, not a hand-built plan: `DistPlannerAnalyzer` wraps the table scans of - // the join in `MergeScan`, and `DistJoinPlanner` (opted in for the build table) nests the build - // side into the probe side's scan. - let plan = plan_sql(&frontend, &sql, &query_ctx).await; - info!("logical plan of the join SQL:\n{plan}"); - let mut config = ConfigOptions::default(); - config.extensions.insert(DistPlannerOptions { - nested_broadcast_join_build_table: Some(BUILD_TABLE.to_string()), - ..Default::default() - }); - let plan = DistPlannerAnalyzer {} - .analyze(plan, &config) - .expect("the distributed plan of the join must be analyzable"); - let plan = DistJoinPlanner {} - .analyze(plan, &config) - .expect("the nested broadcast join rewrite must not fail"); - info!("plan of the distributed rules:\n{plan}"); - let outer = outer_merge_scan(&plan).unwrap_or_else(|| { - panic!("the rules must nest the join into the probe side's MergeScan, got: {plan}") - }); + // Use the actual session setting and SQL entrypoint. The region statistics above drive the + // production selector; no analyzer options or plans are injected by this test. + let (probe_bytes, build_bytes) = wait_for_region_statistics(&frontend).await; assert!( - matches!(outer.input(), LogicalPlan::Join(_)), - "the outer MergeScan must carry the join, got: {}", - outer.input() + 2 * build_bytes < probe_bytes, + "the fixture must make the heuristic favor the build side, actual region statistics: \ + probe {probe_bytes} bytes in 2 regions, build {build_bytes} bytes in 2 regions" ); + run_sql( + &frontend, + "SET experimental_dist_join = true", + query_ctx.clone(), + ) + .await; + + // The time predicate excludes the padding rows inserted to make the probe side larger. + let actual = run_sql(&frontend, &sql, query_ctx.clone()).await; + let physical = actual + .meta + .plan + .clone() + .expect("the output of the query must carry its physical plan"); + info!("physical plan of the SQL-selected nested join:\\n{physical:?}"); - // Physical planning of the outer `MergeScan`: the probe predicate selects one of the two probe - // regions, so only that region is dispatched. + // The probe predicate selects one of the two probe regions, so only that region is dispatched. let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; assert!( probe_leaders.len() > 1, "expected {PROBE_TABLE} to have several regions, actual region leaders: {probe_leaders:?}" ); - let physical = create_physical_plan(&frontend, &plan, &query_ctx).await; let merge_scan = find_merge_scan_exec(&physical) .unwrap_or_else(|| panic!("expected the outer MergeScan to be planned, got: {physical:?}")); let regions = merge_scan.regions().to_vec(); @@ -437,19 +423,8 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { "the outer MergeScan must select regions of {PROBE_TABLE}, actual selection: {regions:?}" ); - // Execution: the frontend serializes the payload of the outer `MergeScan` to the datanode that - // owns the selected region, and the nested `MergeScan` of that payload reads the build table. - let actual = output_batches( - frontend - .query_engine() - .execute(plan, query_ctx) - .await - .expect("the nested plan must be executable"), - ) - .await - .pretty_print() - .expect("the result of the nested plan must be printable"); - info!("result of the rule-produced nested plan:\n{actual}"); + let actual = actual.data.pretty_print().await; + info!("result of the SQL-selected nested plan:\n{actual}"); assert_eq!( expected, multiset(table_cells(&actual)), @@ -462,49 +437,6 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { ); } -/// Returns the outermost `MergeScan` of `plan`, descending through projections. -fn outer_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { - match plan { - LogicalPlan::Extension(extension) => extension - .node - .as_any() - .downcast_ref::(), - LogicalPlan::Projection(projection) => outer_merge_scan(&projection.input), - _ => None, - } -} - -/// Creates the physical plan of `plan` the way the query engine does: the session analyzer of the -/// frontend, the optimizer (skipped for a `MergeScan` root) and the physical planner with its -/// distributed extension planner. -async fn create_physical_plan( - frontend: &Arc, - plan: &LogicalPlan, - query_ctx: &QueryContextRef, -) -> Arc { - let ctx = frontend.query_engine().engine_context(query_ctx.clone()); - let state = ctx.state(); - let analyzed = state - .analyzer() - .execute_and_check(plan.clone(), state.config_options(), |_, _| {}) - .expect("the nested plan must pass the session analyzer"); - let optimized = if let LogicalPlan::Extension(extension) = &analyzed - && extension.node.name() == MergeScanLogicalPlan::name() - { - analyzed.clone() - } else { - state - .optimizer() - .optimize(analyzed, state, |_, _| {}) - .expect("the nested plan must be optimizable") - }; - state - .query_planner() - .create_physical_plan(&optimized, state) - .await - .expect("the nested plan must be physically plannable") -} - /// Returns the `MergeScanExec` of `plan`, descending through the physical plan tree. fn find_merge_scan_exec(plan: &Arc) -> Option<&MergeScanExec> { if let Some(merge_scan) = plan.downcast_ref::() { @@ -730,13 +662,47 @@ async fn test_experimental_dist_join_setting_activates_rewrite() { probe {probe_bytes} bytes in 2 regions, build {build_bytes} bytes in 2 regions" ); - // The session opt-in of the next statement. + // The setting is only an opt-in: with the small build table on the left and the padded probe + // table on the right, it must leave the unfavorable join placement on the frontend. run_sql( &frontend, "SET experimental_dist_join = true", query_ctx.clone(), ) .await; + let reversed_sql = format!( + "SELECT p.a_id, p.probe_v, b.b_id, b.build_v + FROM {BUILD_TABLE} b JOIN {PROBE_TABLE} p + ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 + WHERE p.a_ts < 10000::TIMESTAMP" + ); + let reversed_explained = query_pretty( + &frontend, + &format!("EXPLAIN {reversed_sql}"), + query_ctx.clone(), + ) + .await; + assert!( + !nests_join_in_merge_scan(&reversed_explained), + "the opt-in must leave the unfavorable small-left/big-right placement on the frontend:\n\ + {reversed_explained}" + ); + let reversed = run_sql(&frontend, &reversed_sql, query_ctx.clone()).await; + let reversed_plan = reversed + .meta + .plan + .clone() + .expect("the output of the reversed query must carry its physical plan"); + assert!( + contains_hash_join(&reversed_plan), + "the unfavorable small-left/big-right placement must keep its frontend hash join, got:\n\ + {reversed_plan:?}" + ); + assert_eq!( + expected, + multiset(table_cells(&reversed.data.pretty_print().await)), + "the reversed join must return the same expected rows" + ); let explained = query_pretty(&frontend, &format!("EXPLAIN {sql}"), query_ctx.clone()).await; info!("EXPLAIN with the opt-in:\n{explained}"); From 288d073152f8dcb97f9a5b9ccfa7063e892a046e Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:52:04 +0800 Subject: [PATCH 20/37] refactor(datanode): keep nested scan query admission unchanged Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/catalog/src/information_extension.rs | 1 - src/client/src/region.rs | 163 +------------- src/common/query/src/request.rs | 58 ----- src/datanode/src/region_query.rs | 11 +- src/datanode/src/region_server.rs | 159 +------------- src/query/src/dist_plan/merge_scan.rs | 1 - src/session/src/hints.rs | 30 +-- src/standalone/src/information_extension.rs | 1 - tests-integration/src/cluster.rs | 25 --- .../nested_merge_scan_capability_test.rs | 203 ------------------ 10 files changed, 12 insertions(+), 640 deletions(-) diff --git a/src/catalog/src/information_extension.rs b/src/catalog/src/information_extension.rs index b9b0baa0a585..9d55f9f7a1ee 100644 --- a/src/catalog/src/information_extension.rs +++ b/src/catalog/src/information_extension.rs @@ -82,7 +82,6 @@ impl DistributedInformationExtension { plan, region_id: RegionId::default(), header: None, - internal: false, }) .await .context(crate::error::HandleQuerySnafu) diff --git a/src/client/src/region.rs b/src/client/src/region.rs index e357e18a4209..fb2e4d10ca16 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -30,9 +30,7 @@ use common_error::status_code::StatusCode; use common_grpc::flight::{FlightDecoder, FlightMessage}; use common_meta::error::{self as meta_error, Result as MetaResult}; use common_meta::node_manager::Datanode; -use common_query::request::{ - QUERY_INTERNAL_STAGE_EXTENSION_KEY, QueryRequest, mark_query_internal_stage, -}; +use common_query::request::QueryRequest; use common_recordbatch::error::ExternalSnafu; use common_recordbatch::{RecordBatch, RecordBatchStreamWrapper, SendableRecordBatchStream}; use common_telemetry::error; @@ -87,7 +85,7 @@ impl Datanode for RegionRequester { .context(meta_error::ExternalSnafu)? .to_vec(); let request = api::v1::region::QueryRequest { - header: encode_region_request_header(request.header, request.internal), + header: request.header, region_id: request.region_id.as_u64(), plan, }; @@ -102,50 +100,6 @@ impl Datanode for RegionRequester { } } -/// Encodes the header of a region query, carrying [`QueryRequest::internal`] to the peer datanode. -/// -/// The wire format of a region query (`api::v1::region::QueryRequest`) has no field for the marker, -/// so an internal request (an execution stage of the plan the sender received) is marked with the -/// [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] extension of the query context of the header, which the -/// peer reads back in `RegionServer::handle_remote_read_inner`. Without it the peer can't tell an -/// inner stage of another query from an independent query and takes a concurrency permit for it, -/// which makes the outer query wait for the inner stage while the inner stage waits for the permit -/// held by the outer query. -/// -/// A header-less request keeps no header unless it is internal. An internal request that has no -/// header is encoded with a header that carries the marker only, so the peer builds the query -/// context of the request from that header, i.e. without a remote query id (the peer generates a -/// fresh query context, with an id, only for a request that has no header at all). That only -/// disables the remote dynamic filter cleanup of the request. -fn encode_region_request_header( - header: Option, - internal: bool, -) -> Option { - let mut header = match header { - Some(header) => header, - // A request that has no header is an independent query: it has no marker to carry. - None if !internal => return None, - None => RegionRequestHeader::default(), - }; - if internal { - // Creates the query context if the caller didn't set one: the peer honors the marker only - // in the query context of the header. - mark_query_internal_stage( - &mut header - .query_context - .get_or_insert_with(api::v1::QueryContext::default) - .extensions, - ); - } else if let Some(query_context) = header.query_context.as_mut() { - // [`QueryRequest::internal`] owns the marker: a header that carries it while the request - // claims to be an independent query must not make the peer skip its concurrency limiter. - query_context - .extensions - .remove(QUERY_INTERNAL_STAGE_EXTENSION_KEY); - } - Some(header) -} - impl RegionRequester { pub fn new(client: Client, send_compression: bool, accept_compression: bool) -> Self { Self { @@ -884,116 +838,3 @@ mod test { assert!(recordbatches.next().await.is_none()); } } - -#[cfg(test)] -mod internal_stage_test { - use std::collections::HashMap; - - use common_query::request::{QUERY_INTERNAL_STAGE_EXTENSION_VALUE, is_internal_stage}; - - use super::*; - - /// Reads the marker of a header the way the receiver does: `RegionServer::do_get` builds the - /// query context of the request from the header, `session::context::QueryContext` copies the - /// extensions of that context verbatim. - fn extensions_of(header: &RegionRequestHeader) -> HashMap { - header - .query_context - .as_ref() - .map(|ctx| ctx.extensions.clone()) - .unwrap_or_default() - } - - #[test] - fn test_encode_region_request_header_marks_internal_request() { - // An internal request carries the marker even when its caller set no header at all. - let header = encode_region_request_header(None, true).unwrap(); - - assert!(is_internal_stage(&extensions_of(&header))); - } - - #[test] - fn test_encode_region_request_header_keeps_context_of_internal_request() { - let header = encode_region_request_header( - Some(RegionRequestHeader { - dbname: "greptime".to_string(), - query_context: Some(api::v1::QueryContext { - current_schema: "public".to_string(), - ..Default::default() - }), - ..Default::default() - }), - true, - ) - .unwrap(); - - assert_eq!(header.dbname, "greptime"); - assert_eq!( - header - .query_context - .as_ref() - .map(|ctx| ctx.current_schema.as_str()), - Some("public") - ); - assert!(is_internal_stage(&extensions_of(&header))); - } - - #[test] - fn test_encode_region_request_header_does_not_mark_independent_request() { - // An independent request with no header keeps no header at all. - assert!(encode_region_request_header(None, false).is_none()); - - // A marker that a caller left in the header of an independent request is dropped: the - // marker is owned by `QueryRequest::internal`. - let header = encode_region_request_header( - Some(RegionRequestHeader { - query_context: Some(api::v1::QueryContext { - extensions: HashMap::from([( - QUERY_INTERNAL_STAGE_EXTENSION_KEY.to_string(), - QUERY_INTERNAL_STAGE_EXTENSION_VALUE.to_string(), - )]), - ..Default::default() - }), - ..Default::default() - }), - false, - ) - .unwrap(); - - assert!(!is_internal_stage(&extensions_of(&header))); - } - - #[test] - fn test_internal_marker_survives_the_wire_format() { - let request = api::v1::region::QueryRequest { - header: encode_region_request_header( - Some(RegionRequestHeader { - query_context: Some(api::v1::QueryContext { - current_schema: "public".to_string(), - extensions: HashMap::from([( - "remote_query_id".to_string(), - "1".to_string(), - )]), - ..Default::default() - }), - ..Default::default() - }), - true, - ), - region_id: 1, - plan: Vec::new(), - }; - let bytes = request.encode_to_vec(); - - // The receiver decodes the ticket and builds the query context of the request from the - // header of the decoded request. - let decoded = api::v1::region::QueryRequest::decode(bytes.as_slice()).unwrap(); - let extensions = extensions_of(decoded.header.as_ref().unwrap()); - - assert!(is_internal_stage(&extensions)); - assert_eq!( - extensions.get("remote_query_id").map(String::as_str), - Some("1") - ); - } -} diff --git a/src/common/query/src/request.rs b/src/common/query/src/request.rs index 223500b21d8f..86c6758aa255 100644 --- a/src/common/query/src/request.rs +++ b/src/common/query/src/request.rs @@ -16,7 +16,6 @@ mod base64_serde; mod initial_remote_dyn_filter_reg; use std::cell::Cell; -use std::collections::HashMap; use std::sync::Arc; use api::v1::region::RegionRequestHeader; @@ -340,24 +339,6 @@ impl DynFilterUpdate { } } -/// The extension key of the marker that a region query is an execution stage of another query, -/// instead of an independent query. -/// -/// The wire format of a region query (`api::v1::region::QueryRequest`) has no field for -/// [`QueryRequest::internal`], so the marker travels to the peer datanode as an extension of the -/// query context of the [`RegionRequestHeader`] of the request (the only map-shaped part of the -/// header that survives the round trip). -/// -/// The key is reserved (see `session::hints::RESERVED_EXTENSION_KEYS`), so the hint path -/// (`x-greptime-hints`) can't set it. The marker is trusted on the region query port of a datanode, -/// like the other extensions of the header and the plan itself. -pub const QUERY_INTERNAL_STAGE_EXTENSION_KEY: &str = "query.internal_stage"; - -/// The value of [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] that marks a region query as an execution -/// stage of another query. Any other value, including an absent one, leaves the query subject to -/// the concurrency limiter. -pub const QUERY_INTERNAL_STAGE_EXTENSION_VALUE: &str = "true"; - /// The query request to be handled by the RegionServer (Datanode). #[derive(Clone, Debug)] pub struct QueryRequest { @@ -369,45 +350,6 @@ pub struct QueryRequest { /// The form of the query: a logical plan. pub plan: LogicalPlan, - - /// Whether this request is an execution stage of another query, instead of an independent - /// query. - /// - /// A datanode executes the `MergeScan` nodes of the plan it received by querying the regions - /// of those nodes from the datanodes hosting them. Such an inner region query is a stage of - /// the outer query: the outer query waits for it while holding its own concurrency permit, so - /// the inner query must not acquire a permit of its own, otherwise the outer query waits for - /// the inner stage while the inner stage waits for the permit held by the outer query. - /// - /// `false` means an ordinary, independent query, which is subject to the concurrency limiter. - /// - /// A request that crosses the wire carries the marker as the - /// [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] extension of the query context of its header (the - /// sender sets it, the receiver reads it back into this field), so the marker also survives a - /// dispatch to a peer datanode. - pub internal: bool, -} - -/// Marks the query context of a region query as an execution stage of another query: sets the -/// [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] extension to [`QUERY_INTERNAL_STAGE_EXTENSION_VALUE`]. -/// -/// The sender side of a region query uses it to carry [`QueryRequest::internal`] to a peer -/// datanode; [`is_internal_stage`] reads the marker back on the receiver side. -pub fn mark_query_internal_stage(extensions: &mut HashMap) { - extensions.insert( - QUERY_INTERNAL_STAGE_EXTENSION_KEY.to_string(), - QUERY_INTERNAL_STAGE_EXTENSION_VALUE.to_string(), - ); -} - -/// Returns whether `extensions` marks a region query as an execution stage of another query, i.e. -/// whether the sender of the query set [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] to -/// [`QUERY_INTERNAL_STAGE_EXTENSION_VALUE`]. -pub fn is_internal_stage(extensions: &HashMap) -> bool { - extensions - .get(QUERY_INTERNAL_STAGE_EXTENSION_KEY) - .map(String::as_str) - == Some(QUERY_INTERNAL_STAGE_EXTENSION_VALUE) } #[cfg(test)] diff --git a/src/datanode/src/region_query.rs b/src/datanode/src/region_query.rs index aa750b060416..0327395d1647 100644 --- a/src/datanode/src/region_query.rs +++ b/src/datanode/src/region_query.rs @@ -82,17 +82,8 @@ impl RegionQueryHandler for DatanodeRegionQueryHandler { async fn do_get( &self, target: &RegionQueryTarget, - mut request: QueryRequest, + request: QueryRequest, ) -> QueryResult { - // A datanode serves the region queries of the `MergeScan` nodes of the plan it received. - // Such a region query is an execution stage of that plan, not an independent query, so it - // must leave the concurrency permits to the query that dispatched the plan. The marker is - // carried to the peer with the request: the client encodes it into the query context of the - // header (`common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_KEY`), where the peer reads - // it back and admits the stage without a permit of its own (see - // `RegionServer::handle_remote_read_inner`). - request.internal = true; - self.node_manager .datanode(target.peer()) .await diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index d8ef5e123189..7cebd73f6a68 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -74,7 +74,6 @@ use servers::grpc::region_server::RegionServerHandler; use session::context::{ FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef, }; -use session::hints::{QUERY_INTERNAL_STAGE_EXTENSION_KEY, QUERY_INTERNAL_STAGE_EXTENSION_VALUE}; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metric_engine_consts::{ FILE_ENGINE_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, @@ -291,13 +290,10 @@ impl RegionServer { request: api::v1::region::QueryRequest, query_ctx: QueryContextRef, ) -> Result { - // A peer datanode that executes the `MergeScan` nodes of the plan it received queries the - // regions of those nodes from this datanode and marks the region query as an execution - // stage of its query (see [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`]). Such a stage must not - // take a permit of its own: the outer query holds one until its stream is fully consumed - // and waits for the stage to return. - let internal = is_internal_stage_query(&query_ctx); - let permit = acquire_query_permit(&self.inner.parallelism, internal).await?; + let permit = match &self.inner.parallelism { + Some(p) => Some(p.acquire().await?), + None => None, + }; let region_id = RegionId::from_u64(request.region_id); let catalog_list = Arc::new(NameAwareCatalogList::new( @@ -331,10 +327,6 @@ impl RegionServer { header: request.header, region_id, plan, - // The stage marker is passed on: the plan of this read is executed as part of - // the stage this remote read was admitted for, so nothing on that path may - // take a permit of its own either. - internal, }, query_ctx.clone(), ) @@ -355,10 +347,10 @@ impl RegionServer { } async fn handle_read_inner(&self, request: QueryRequest) -> Result { - // The datanode executes the `MergeScan` nodes of the plan it received by querying regions - // from the datanodes hosting them. Those inner region queries are marked as internal - // stages of the query that dispatched the plan and must not take a permit of their own. - let permit = acquire_query_permit(&self.inner.parallelism, request.internal).await?; + let permit = match &self.inner.parallelism { + Some(p) => Some(p.acquire().await?), + None => None, + }; let ctx = request.header.as_ref().map(|h| h.into()); let query_ctx = Arc::new(ctx.unwrap_or_else(|| QueryContextBuilder::default().build())); @@ -1092,39 +1084,6 @@ impl RegionServerParallelism { } } -/// Returns whether the query context of a region query received over the wire marks the query as -/// an execution stage of another query, instead of an independent query. -/// -/// The sender of the query sets the [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] extension of the query -/// context to [`QUERY_INTERNAL_STAGE_EXTENSION_VALUE`] when the query is an execution stage of the -/// plan it received (see [`QueryRequest::internal`]), i.e. when the query is dispatched by a -/// datanode for the `MergeScan` nodes of that plan. -fn is_internal_stage_query(query_ctx: &QueryContextRef) -> bool { - query_ctx.extension(QUERY_INTERNAL_STAGE_EXTENSION_KEY) - == Some(QUERY_INTERNAL_STAGE_EXTENSION_VALUE) -} - -/// Acquires a concurrency permit for a request that is an independent query, if the limiter is -/// enabled. -/// -/// A request that is an execution stage of another query (see [`QueryRequest::internal`]) doesn't -/// acquire a permit: the query that dispatched it already holds one, holds it until its stream is -/// fully consumed and waits for the stage to return. Letting the stage acquire a permit of its own -/// makes the outer query wait for the inner stage while the inner stage waits for the permit held -/// by the outer query. -async fn acquire_query_permit( - parallelism: &Option, - internal: bool, -) -> Result> { - if internal { - return Ok(None); - } - match parallelism { - Some(p) => Ok(Some(p.acquire().await?)), - None => Ok(None), - } -} - /// Wraps a record batch stream and holds a concurrency permit until the stream is /// fully consumed (dropped), so `max_concurrent_queries` bounds the number of /// in-flight read streams, not just query planning. @@ -2825,108 +2784,6 @@ mod tests { assert!(forth_query.is_ok()); } - #[tokio::test] - async fn test_internal_query_skips_concurrency_permit() { - let parallelism = - Some(RegionServerParallelism::from_opts(1, Duration::from_millis(1)).unwrap()); - - // The query that dispatched the nested plan holds the only permit of the server for as - // long as its stream is alive. - let outer_query = acquire_query_permit(¶llelism, false).await.unwrap(); - assert!(outer_query.is_some()); - - // An inner stage of that query must neither wait for nor consume a permit: it runs while - // the outer query waits for it to return. - let inner_stage = acquire_query_permit(¶llelism, true).await.unwrap(); - assert!(inner_stage.is_none()); - - // An independent query is still subject to the limiter. - assert!(acquire_query_permit(¶llelism, false).await.is_err()); - - // The inner stage reserved nothing: the permit goes back to the limiter with the outer - // query's stream. - drop(outer_query); - assert!( - acquire_query_permit(¶llelism, false) - .await - .unwrap() - .is_some() - ); - - // The limiter stays disabled when `max_concurrent_queries` is 0. - assert!(acquire_query_permit(&None, false).await.unwrap().is_none()); - } - - #[test] - fn test_internal_stage_extension_keys_match_between_crates() { - // `session` can't depend on `common-query`: the sender writes the marker with the constant - // of `common-query`, the receiver reads it with the constant of `session`. - assert_eq!( - QUERY_INTERNAL_STAGE_EXTENSION_KEY, - common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_KEY - ); - assert_eq!( - QUERY_INTERNAL_STAGE_EXTENSION_VALUE, - common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_VALUE - ); - } - - /// Asserts that a region query received over the wire is admitted without a permit only when the - /// query context of its header marks it as an execution stage of another query. - #[tokio::test] - async fn test_internal_stage_query_skips_concurrency_permit() { - let parallelism = - Some(RegionServerParallelism::from_opts(1, Duration::from_millis(1)).unwrap()); - - // A datanode that fans a nested `MergeScan` out to its peer marks the region query it sends - // as an execution stage of its own query. - let mut extensions = std::collections::HashMap::new(); - common_query::request::mark_query_internal_stage(&mut extensions); - let internal_ctx: QueryContextRef = Arc::new( - QueryContextBuilder::default() - .extensions(extensions) - .build(), - ); - assert!(is_internal_stage_query(&internal_ctx)); - - // The query of the peer itself holds the only permit of the server. - let outer_query = acquire_query_permit(¶llelism, false).await.unwrap(); - assert!(outer_query.is_some()); - - // The stage is admitted while the outer query waits for it to return. - let stage_permit = - acquire_query_permit(¶llelism, is_internal_stage_query(&internal_ctx)) - .await - .unwrap(); - assert!(stage_permit.is_none()); - - // A query with no marker, a marker with another value or an empty query context is an - // independent query and is subject to the limiter. - for query_ctx in [ - QueryContext::arc(), - Arc::new( - QueryContextBuilder::default() - .set_extension( - QUERY_INTERNAL_STAGE_EXTENSION_KEY.to_string(), - "false".to_string(), - ) - .build(), - ), - Arc::new( - QueryContextBuilder::default() - .set_extension("query.internal_stage.other".to_string(), "true".to_string()) - .build(), - ), - ] { - assert!(!is_internal_stage_query(&query_ctx)); - assert!( - acquire_query_permit(¶llelism, is_internal_stage_query(&query_ctx)) - .await - .is_err() - ); - } - } - fn mock_region_metadata(region_id: RegionId) -> RegionMetadata { let mut metadata_builder = RegionMetadataBuilder::new(region_id); metadata_builder.push_column_metadata(ColumnMetadata { diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 7993e2b9e1b5..ee4e9574280c 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -728,7 +728,6 @@ impl MergeScanExec { }), region_id, plan: plan.clone(), - internal: false, }; if explain_verbose { common_telemetry::info!( diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index b97d1cf94cda..971426f4cc03 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -27,36 +27,11 @@ 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"; -/// The extension key of the marker that a query is an execution stage of another query, instead of -/// an independent query. -/// -/// A datanode executes the `MergeScan` nodes of the plan it received by querying the regions of -/// those nodes from the datanodes hosting them. Such an inner region query is a stage of the outer -/// query: the outer query waits for it while holding its own concurrency permit, so the inner query -/// must not acquire a permit of its own, otherwise the outer query waits for the inner stage while -/// the inner stage waits for the permit held by the outer query. The marker travels to the peer -/// datanode as an extension of the query context of the `RegionRequestHeader` of the request. -/// -/// The key is reserved (see [`RESERVED_EXTENSION_KEYS`]), so the hint path (`x-greptime-hints`) -/// can't set it: a SQL/PromQL client of the frontend can't make its own query bypass the -/// concurrency limiter of a datanode. The marker is trusted on the region query port of a datanode, -/// like the other extensions of the header and the plan itself. -/// -/// Keep in sync with `common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_KEY` -/// (`session` can't depend on `common-query`). -pub const QUERY_INTERNAL_STAGE_EXTENSION_KEY: &str = "query.internal_stage"; -/// The value of [`QUERY_INTERNAL_STAGE_EXTENSION_KEY`] that marks an execution stage of another -/// query. Any other value, including an absent one, leaves the query subject to the concurrency -/// limiter. -/// -/// Keep in sync with `common_query::request::QUERY_INTERNAL_STAGE_EXTENSION_VALUE`. -pub const QUERY_INTERNAL_STAGE_EXTENSION_VALUE: &str = "true"; -pub const RESERVED_EXTENSION_KEYS: [&str; 5] = [ +pub const RESERVED_EXTENSION_KEYS: [&str; 4] = [ REMOTE_QUERY_ID_EXTENSION_KEY, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, LIVE_ANALYZE_METRICS_EXTENSION_KEY, - QUERY_INTERNAL_STAGE_EXTENSION_KEY, ]; /// Deprecated, use `HINTS_KEY` instead. @@ -90,9 +65,6 @@ mod tests { assert!(is_reserved_extension_key( LIVE_ANALYZE_METRICS_EXTENSION_KEY )); - assert!(is_reserved_extension_key( - QUERY_INTERNAL_STAGE_EXTENSION_KEY - )); assert!(!is_reserved_extension_key(READ_PREFERENCE_HINT)); } } diff --git a/src/standalone/src/information_extension.rs b/src/standalone/src/information_extension.rs index b465d1b1270d..8d52eb12ccc1 100644 --- a/src/standalone/src/information_extension.rs +++ b/src/standalone/src/information_extension.rs @@ -169,7 +169,6 @@ impl InformationExtension for StandaloneInformationExtension { .context(catalog::error::DatafusionSnafu)?, region_id: RegionId::default(), header: None, - internal: false, }; self.region_server diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 965ee867e1e3..a01aed240afc 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -205,11 +205,6 @@ impl GreptimeDbCluster { } } -/// Adjusts the options of a datanode of the cluster before the datanode is created, see -/// [`GreptimeDbClusterBuilder::with_datanode_options_override`]. The first argument is the id of -/// the datanode, so a test can vary the options per datanode. -type DatanodeOptionsOverride = Box; - pub struct GreptimeDbClusterBuilder { cluster_name: String, kv_backend: KvBackendRef, @@ -226,7 +221,6 @@ pub struct GreptimeDbClusterBuilder { local_file_access: LocalFileAccess, event_recorder_options: EventRecorderOptions, real_datanode_grpc_addr: bool, - datanode_options_override: Option, } impl GreptimeDbClusterBuilder { @@ -264,7 +258,6 @@ impl GreptimeDbClusterBuilder { local_file_access: LocalFileAccess::default(), event_recorder_options: EventRecorderOptions::default(), real_datanode_grpc_addr: false, - datanode_options_override: None, } } @@ -358,20 +351,6 @@ impl GreptimeDbClusterBuilder { self } - /// Adjusts the options of every datanode before the datanode is created, e.g. to set the - /// concurrency limit of the region server of the datanode (`max_concurrent_queries`). - /// - /// The hook runs after the id and the gRPC address of the datanode have been set, and it - /// receives the id of the datanode, so a test can vary the options per datanode. - #[must_use] - pub fn with_datanode_options_override( - mut self, - override_options: impl Fn(u64, &mut DatanodeOptions) + Send + Sync + 'static, - ) -> Self { - self.datanode_options_override = Some(Box::new(override_options)); - self - } - pub async fn build_with( &self, datanode_options: Vec, @@ -516,10 +495,6 @@ impl GreptimeDbClusterBuilder { opts.grpc.bind_addr = addr.clone(); opts.grpc.server_addr = addr; } - if let Some(override_options) = &self.datanode_options_override { - override_options(datanode_id, &mut opts); - } - options.push(opts); } (options, guards) diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 4e4f36820ff0..bf0986391008 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -51,17 +51,12 @@ //! * a nested plan whose inner relation does not exist, and a nested plan with an inner region that //! is not served any more, both make the whole query fail: the datanode never returns the rows of //! the inner regions that happen to succeed, -//! * with the concurrency limit of every datanode set to a single query, every probe region still -//! executes: the inner region queries of the nested plans are admitted as execution stages of the -//! outer queries, so they don't wait for the permit that their own query holds (see -//! [`test_nested_merge_scan_capability_internal_stage_under_concurrency_limit`]), //! * a cold load of the route of the build table returns the rows of the join after an unrelated //! table id is invalidated once (see //! [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]). use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; -use std::time::Duration; use api::v1::region::{QueryRequest, RegionRequestHeader}; use common_meta::cache::TableRouteCacheRef; @@ -439,180 +434,6 @@ async fn test_nested_merge_scan_capability_inner_failure_fails_query() { ); } -/// The nested plans execute under the concurrency limit of their datanodes: a datanode admits a -/// single query at a time, the outer query of a probe region holds the only permit of its datanode -/// until its stream is fully consumed, and the nested plan then queries the regions of the build -/// table from the datanode itself and from the other datanode. -/// -/// Every inner region query has to be admitted as an execution stage of the outer query, without a -/// permit of its own. The marker of the stage travels from the region query handler of the datanode -/// (`request.internal = true`), through the client that encodes it into the header of the request, -/// to the region server of the peer that reads it back and skips its limiter. If any one of them is -/// missing, the inner region query of the *same* datanode waits for the permit held by the outer -/// query and fails after `concurrent_query_limiter_timeout`, and the inner region query of the -/// *other* datanode waits for the outer query of that datanode, which waits for its own nested -/// plan: the consumption of the outer queries below is bounded by a timeout, so both cases fail the -/// test instead of hanging it. -#[tokio::test(flavor = "multi_thread")] -async fn test_nested_merge_scan_capability_internal_stage_under_concurrency_limit() { - common_telemetry::init_default_ut_logging(); - - let cluster = build_cluster_with_query_limit( - "test_nested_merge_scan_capability_internal_stage_under_concurrency_limit", - ) - .await; - let frontend = cluster.fe_instance().clone(); - - prepare_tables(&frontend).await; - - let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; - let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; - info!("region leaders of {PROBE_TABLE}: {probe_leaders:?}"); - info!("region leaders of {BUILD_TABLE}: {build_leaders:?}"); - // Every table is partitioned into two regions. Both tables span both datanodes, so a nested - // plan of a probe region always has to read a region of the build table from the other - // datanode, and the two probe regions sit on different datanodes: asking both datanodes for - // their probe region exhausts the single permit of every datanode of the cluster. - assert_eq!( - 2, - datanodes(&probe_leaders).len(), - "expected the regions of {PROBE_TABLE} to be spread over both datanodes, actual region \ - leaders: {probe_leaders:?}" - ); - assert_eq!( - 2, - datanodes(&build_leaders).len(), - "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual region \ - leaders: {build_leaders:?}" - ); - - // The reference: the frontend executes the same join with its default distributed plan. - let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; - info!("join result of the frontend:\n{reference}"); - assert_eq!( - expected_join_rows(), - multiset(table_cells(&reference)), - "unexpected result of the join on the frontend:\n{reference}" - ); - - let query_ctx = query_ctx(); - let nested_plan = nested_plan(&frontend, &query_ctx).await; - let plan = encode(&nested_plan); - - let expectation = cross_datanode_expectation(&probe_leaders, &build_leaders); - let baseline_requests = expectation - .remote_datanodes - .keys() - .map(|datanode| (*datanode, rpc_requests(&cluster, *datanode))) - .collect::>(); - - // Ask every datanode for its probe region before consuming anything: the region server takes - // the permit of the datanode when it admits the request and holds it until the stream of the - // request is consumed, so no datanode has a free permit left while the nested plans run. - let mut outer_queries = Vec::with_capacity(probe_leaders.len()); - for (probe_region, (probe_datanode, _)) in &probe_leaders { - let stream = region_server(&cluster, *probe_datanode) - .handle_remote_read( - QueryRequest { - header: Some(region_request_header(&query_ctx)), - region_id: *probe_region, - plan: plan.clone(), - }, - query_ctx.clone(), - ) - .await - .unwrap_or_else(|e| { - panic!( - "the datanode {probe_datanode} must admit the nested plan of the region \ - {probe_region} while it serves a single query at a time: {e}" - ) - }); - outer_queries.push((*probe_region, *probe_datanode, stream)); - } - - // Consume the outer queries of both datanodes at the same time, so that every nested plan - // queries the regions of the build table while the only permit of every datanode is held by the - // outer query of that datanode. An inner region query that misses the stage marker waits for - // that permit up to `concurrent_query_limiter_timeout` and fails, or the two outer queries wait - // for each other: the timeout turns the second case into a failure instead of a hang. - let collected = tokio::time::timeout( - Duration::from_secs(30), - futures::future::join_all(outer_queries.into_iter().map( - |(probe_region, probe_datanode, stream)| async move { - ( - probe_region, - probe_datanode, - RecordBatches::try_collect(stream).await, - ) - }, - )), - ) - .await - .unwrap_or_else(|_| { - panic!( - "the nested plans of the probe regions did not return within the timeout while every \ - datanode served a single query at a time: an inner region query is waiting for the \ - concurrency permit held by the outer query that dispatched it, i.e. the stage marker \ - of the inner region query was lost" - ) - }); - - let mut rows = Vec::new(); - for (probe_region, probe_datanode, result) in collected { - let batches = result.unwrap_or_else(|e| { - panic!( - "the datanode {probe_datanode} must execute the nested plan of the region \ - {probe_region} completely while it serves a single query at a time: {e}" - ) - }); - let actual = batches - .pretty_print() - .expect("the result of the datanode must be printable"); - info!("join result of the nested plan of the region {probe_region}:\n{actual}"); - rows.extend(table_cells(&actual)); - } - - // The nested plan of a probe region reads that region of the probe table, so the union of the - // results of all probe regions is the result of the whole join. Comparing the rows as a - // multiset catches both a lost row and a row that was counted twice. - assert_eq!( - expected_join_rows(), - multiset(rows.clone()), - "the nested plans of the datanodes returned an unexpected multiset of rows under the \ - concurrency limit: {rows:?}" - ); - assert_eq!( - multiset(table_cells(&reference)), - multiset(rows), - "the nested plans of the datanodes returned a different multiset of rows than the \ - frontend:\nfrontend:\n{reference}" - ); - - assert_cross_datanode_region_query(&cluster, &expectation, &baseline_requests); - - // The nested plan really reached the other datanode over the network: the region query of a - // `MergeScan` is a Flight `DoGet` (see `RegionRequester::handle_query`), and the remote - // requests of a datanode are counted at the address it registered, so a `DoGet` recorded there - // was sent by another datanode of the cluster. - for datanode in expectation.remote_datanodes.keys() { - let stats = cluster - .datanode_rpc_stats(*datanode) - .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {datanode}")); - info!( - "the datanode {datanode} served {} remote gRPC requests, paths: {:?}", - stats.requests(), - stats.paths() - ); - assert!( - stats.paths().iter().any(|path| path.contains("DoGet")), - "expected the nested plans of {PROBE_TABLE} to query the datanode {datanode} for the \ - regions of {BUILD_TABLE} it owns with a Flight `DoGet`, actual remote gRPC paths: \ - {:?}", - stats.paths() - ); - } -} - async fn build_cluster(test_name: &str) -> GreptimeDbCluster { GreptimeDbClusterBuilder::new(test_name) .await @@ -626,30 +447,6 @@ async fn build_cluster(test_name: &str) -> GreptimeDbCluster { .await } -/// Same as [`build_cluster`], but every datanode admits a single query at a time. -/// -/// The outer region query of a probe region holds the only permit of its datanode until its stream -/// is fully consumed, so a region query that the nested plan of that query issues to the same -/// datanode can only succeed if it is admitted as an execution stage of the query instead of -/// taking a permit of its own (see -/// [`test_nested_merge_scan_capability_internal_stage_under_concurrency_limit`]). -async fn build_cluster_with_query_limit(test_name: &str) -> GreptimeDbCluster { - GreptimeDbClusterBuilder::new(test_name) - .await - .with_real_datanode_grpc_addr(true) - .with_datanodes(2) - .with_datanode_options_override(|_, opts| { - opts.max_concurrent_queries = 1; - // Wider than the default: the test must fail because the stage marker of an inner - // region query is lost, not because a slow machine made an admitted query wait for a - // permit for too long. The timeout around the consumption of the outer queries is much - // larger than this. - opts.concurrent_query_limiter_timeout = Duration::from_secs(1); - }) - .build(false) - .await -} - /// The `DoGet` requests that the nested plans of the probe regions have to send to the datanodes /// that own the regions of the build table. fn cross_datanode_expectation( From 236a3dd0dcd1eb65d3fa1b327978819a8bafc460 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 12:29:36 +0800 Subject: [PATCH 21/37] fix(query): build nested broadcast joins from the small input Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datanode/src/datanode.rs | 42 +- src/meta-client/src/client.rs | 7 + src/query/src/dist_plan/dist_join_planner.rs | 360 +++++++++++++----- src/query/src/dist_plan/merge_scan.rs | 117 +++++- src/query/src/dist_plan/planner.rs | 290 +++++++++++++- .../nested_merge_scan_capability_test.rs | 261 ++++++++++--- 6 files changed, 896 insertions(+), 181 deletions(-) diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 4fef618b90dd..2fa638935d0f 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -20,7 +20,8 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use catalog::CatalogManagerRef; -use catalog::information_schema::NoopInformationExtension; +use catalog::information_extension::DistributedInformationExtension; +use catalog::information_schema::{InformationExtensionRef, NoopInformationExtension}; use catalog::kvbackend::KvBackendCatalogManagerBuilder; use client::client_manager::NodeClients; use common_base::Plugins; @@ -362,6 +363,7 @@ impl DatanodeBuilder { region_event_listener, file_ref_manager, cache_registry.clone(), + meta_client.as_ref(), ) .await?; @@ -503,6 +505,7 @@ impl DatanodeBuilder { event_listener: RegionServerEventListenerRef, file_ref_manager: FileReferenceManagerRef, cache_registry: Arc, + meta_client: Option<&MetaClientRef>, ) -> Result { let opts: &DatanodeOptions = &self.opts; @@ -518,8 +521,29 @@ impl DatanodeBuilder { // `with_dist_planner`. // - `with_dist_planner` stays `false` on purpose: the datanode only executes the plans it // receives, it must not rewrite its own plans into distributed ones. + // Clients to the other datanodes. The same pool backs distributed information + // queries and MergeScan dispatch. A region query is a streaming request, so it must not + // be cut off by the request timeout, same as the frontend to datanode client. + 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_clients = Arc::new(NodeClients::new(channel_config)); + let node_manager: NodeManagerRef = node_clients.clone(); + + let information_extension: InformationExtensionRef = match meta_client { + Some(meta_client) => Arc::new(DistributedInformationExtension::new( + meta_client.clone(), + node_clients.clone(), + )), + None => Arc::new(NoopInformationExtension), + }; let catalog_manager: CatalogManagerRef = KvBackendCatalogManagerBuilder::new( - Arc::new(NoopInformationExtension), + information_extension, self.kv_backend.clone(), cache_registry.clone(), ) @@ -535,20 +559,6 @@ impl DatanodeBuilder { partition_info_cache, )); - // Clients to the other datanodes. They are only used when a plan contains a `MergeScan` - // node, i.e. they stay idle for the plans that the datanode executes locally. - // - // A region query is a streaming request, so it must not be cut off by the request timeout: - // same as the frontend to datanode client. - 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 = DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager); diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index e0932af5a694..dbf733a70c97 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -130,6 +130,7 @@ impl MetaClientBuilder { Self::new(member_id, Role::Datanode) .enable_store() .enable_heartbeat() + .enable_access_cluster_info() } /// Returns the role of Flownode's default options. @@ -1014,6 +1015,12 @@ mod tests { client } + #[test] + fn test_datanode_default_options_enable_cluster_info_access() { + let builder = MetaClientBuilder::datanode_default_options(0); + assert!(builder.enable_access_cluster_info); + } + #[tokio::test] async fn test_meta_client_builder() { let urls = &["127.0.0.1:3001", "127.0.0.1:3002"]; diff --git a/src/query/src/dist_plan/dist_join_planner.rs b/src/query/src/dist_plan/dist_join_planner.rs index a60894727bfe..a04c21454c62 100644 --- a/src/query/src/dist_plan/dist_join_planner.rs +++ b/src/query/src/dist_plan/dist_join_planner.rs @@ -33,7 +33,7 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_expr::utils::{can_hash, find_valid_equijoin_key_pair, split_conjunction}; use datafusion_expr::{ - Expr, ExprSchemable, Join, JoinType, LogicalPlan, Operator, TableScan, + Expr, ExprSchemable, Join, JoinType, LogicalPlan, Operator, Projection, TableScan, UserDefinedLogicalNodeCore, }; use datafusion_optimizer::analyzer::AnalyzerRule; @@ -110,8 +110,7 @@ impl DistJoinStats { /// approximate disk sizes the regions report, not a measured or predicted amount of network /// traffic. /// - /// The sides are never swapped, and a missing table, a table without regions or bytes, or an - /// overflowing product means "no". + /// A missing table, a table without regions or bytes, or an overflowing product means "no". pub fn favors_right_build(&self, probe: TableId, build: TableId) -> bool { let (Some(probe), Some(build)) = (self.tables.get(&probe), self.tables.get(&build)) else { return false; @@ -132,9 +131,9 @@ impl DistJoinStats { /// The base tables of one plausible candidate join of the cost heuristic. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct CandidateJoin { - /// The base table of the left (probe) side. + /// The base table of the left side, used as the candidate region-driving side. pub probe: TableId, - /// The base table of the right (build) side. + /// The base table of the right side, considered as the candidate small build side. pub build: TableId, } @@ -309,11 +308,13 @@ pub(crate) fn expected_region_ids( (!regions.is_empty()).then_some(regions) } -/// Nests the build side's `MergeScan` inside the probe side's one, so the join runs on -/// the datanodes holding the probe regions. +/// Nests the small side's `MergeScan` inside the large side's boundary, so each datanode +/// hashes the small side and probes with its local large-side rows. The outer boundary is +/// routed by the large table's regions. /// -/// Keeps the join node as is, so its schema, join conditions, residual filter and NULL -/// semantics are preserved; unmatched shapes are returned unchanged. +/// The join inputs are reversed for physical hash-build/probe roles, then a projection restores +/// the original output schema and column order. Join conditions, residual filter and NULL +/// semantics are preserved; unaddressable schemas and unmatched shapes stay unchanged. #[derive(Debug)] pub struct DistJoinPlanner; @@ -332,8 +333,7 @@ impl AnalyzerRule for DistJoinPlanner { } } -/// Rewriter nesting the build side `MergeScan` of a supported join inside the probe side's -/// `MergeScan`. +/// Rewriter nesting the small-side `MergeScan` inside the large-side boundary of a supported join. struct NestedBroadcastJoinRewriter<'a> { /// Per-table statistics fetched for this query, see [`DistJoinStats`]. stats: &'a DistJoinStats, @@ -357,57 +357,68 @@ impl NestedBroadcastJoinRewriter<'_> { return None; } - let (probe_plan, probe_merge_scan) = strip_merge_scan(&join.left)?; - let (_, build_merge_scan) = strip_merge_scan(&join.right)?; - if probe_merge_scan.is_placeholder() || build_merge_scan.is_placeholder() { + let (large_plan, large_merge_scan) = strip_merge_scan(&join.left)?; + let (_, small_merge_scan) = strip_merge_scan(&join.right)?; + if large_merge_scan.is_placeholder() || small_merge_scan.is_placeholder() { return None; } - // Replacing the probe input must not change the join input schema. - if probe_plan.schema() != join.left.schema() { + // Replacing the large input must not change its schema. The small input remains under + // its boundary, nested on each large-side region. + if large_plan.schema() != join.left.schema() { return None; } // A nested boundary is hidden from the remaining analyzer passes, so a boundary // schema `JsonSchemaConcretizeRule` still has to fix up cannot be repaired // anymore. - if has_json2_boundary(&probe_merge_scan) || has_json2_boundary(&build_merge_scan) { + if has_json2_boundary(&large_merge_scan) || has_json2_boundary(&small_merge_scan) { return None; } - // The right input is the only candidate build side, because moving a left build side - // would change the join schema. - let (probe_table, build_table) = ( - side_base_table_id(probe_merge_scan.input())?, - side_base_table_id(build_merge_scan.input())?, + let (large_table, small_table) = ( + side_base_table_id(large_merge_scan.input())?, + side_base_table_id(small_merge_scan.input())?, ); - if !self.stats.favors_right_build(probe_table, build_table) { + if !self.stats.favors_right_build(large_table, small_table) { return None; } - // Keep the rewrite at two levels: the probe side must be local and the build side + // Keep the rewrite at two levels: the large side must be local and the small side // must not contain another distributed boundary. - if contains_merge_scan(&probe_plan) || contains_merge_scan(build_merge_scan.input()) { + if contains_merge_scan(&large_plan) || contains_merge_scan(small_merge_scan.input()) { return None; } - // Join the probe regions against the build plan kept above; the probe table's - // partition columns route the outer scan. - let new_join = Join { - left: Arc::new(probe_plan), - right: join.right.clone(), - ..join.clone() - }; + // Swapping the logical inputs gives DataFusion's hash join its intended small build + // side and large probe side. Refuse schemas whose original output columns cannot be + // mapped uniquely to the swapped output; the restoration projection then preserves + // SQL-visible names, qualifiers, metadata, and order exactly. + let swapped_join = Join::try_new( + join.right.clone(), + Arc::new(large_plan), + join.on + .iter() + .map(|(left, right)| (right.clone(), left.clone())) + .collect(), + join.filter.clone(), + join.join_type, + join.join_constraint, + join.null_equality, + join.null_aware, + ) + .ok()?; + let restored_plan = restore_join_schema(swapped_join, join.schema.clone())?; let mut merge_scan = MergeScanLogicalPlan::new( - LogicalPlan::Join(new_join), + restored_plan, false, - probe_merge_scan.partition_cols().clone(), + large_merge_scan.partition_cols().clone(), ); - // The outer boundary keeps the probe boundary's remote dynamic filter identity + // The outer boundary keeps the large boundary's remote dynamic filter identity // when it has one. A missing id only disables remote dynamic filter pushdown for // that boundary (`MergeScanExec` fails open), and serialization drops the ids of // the nested payload, so an absent id is not a rewrite blocker. - if let Some(producer_id) = probe_merge_scan.remote_dyn_filter_producer_id() { + if let Some(producer_id) = large_merge_scan.remote_dyn_filter_producer_id() { merge_scan = merge_scan.with_remote_dyn_filter_producer_id(producer_id); } @@ -451,6 +462,41 @@ fn strip_merge_scan(plan: &LogicalPlan) -> Option<(LogicalPlan, MergeScanLogical } } +/// Wraps a swapped join in a projection restoring the original schema, only when each +/// original output column uniquely addresses its corresponding swapped input position. +fn restore_join_schema(join: Join, schema: datafusion_common::DFSchemaRef) -> Option { + let left_len = join.right.schema().fields().len(); + let right_len = join.left.schema().fields().len(); + if schema.fields().len() != left_len + right_len { + return None; + } + + let columns = schema.columns(); + for (index, column) in columns.iter().enumerate() { + let intended_index = if index < left_len { + right_len + index + } else { + index - left_len + }; + if join + .schema + .columns() + .iter() + .filter(|candidate| *candidate == column) + .count() + != 1 + || join.schema.maybe_index_of_column(column) != Some(intended_index) + { + return None; + } + } + + Some(LogicalPlan::Projection(Projection::new_from_schema( + Arc::new(LogicalPlan::Join(join)), + schema, + ))) +} + /// Whether `plan` contains a visible `MergeScan` node. fn contains_merge_scan(plan: &LogicalPlan) -> bool { if let LogicalPlan::Extension(extension) = plan @@ -462,34 +508,35 @@ fn contains_merge_scan(plan: &LogicalPlan) -> bool { plan.inputs().into_iter().any(contains_merge_scan) } -/// Returns the local probe input that decides the regions of a supported nested broadcast +/// Returns the local large-side input that decides regions for a supported nested broadcast /// join, or `None` for any other plan. /// -/// Region pruning clears the collected predicates at a node with several inputs (the -/// join), so the outer `MergeScan` of the rewritten plan would be routed by the join -/// itself and scan every region of the probe table. Routing by the returned side keeps -/// the probe predicate pruning, while the full payload is still dispatched to the -/// selected regions. +/// Region pruning clears collected predicates at a multi-input join, so route by the local +/// large-side input rather than the join itself. This preserves large-table pruning while +/// dispatching the full nested payload to selected regions. /// -/// Only the supported shape of [`DistJoinPlanner`] returns a side: an `INNER` equi-join -/// whose left input is a local plan and whose right input reaches its `MergeScan` -/// through projections. +/// Only the supported shape returns a side: an `INNER` equi-join (possibly behind schema +/// restoration projections) with a small-side `MergeScan` and a local large-side input. pub(crate) fn nested_broadcast_join_probe_input(plan: &LogicalPlan) -> Option<&LogicalPlan> { + let mut plan = plan; + while let LogicalPlan::Projection(projection) = plan { + plan = &projection.input; + } let LogicalPlan::Join(join) = plan else { return None; }; if join.join_type != JoinType::Inner || join.on.is_empty() { return None; } - if contains_merge_scan(&join.left) { + if contains_merge_scan(&join.right) { return None; } - let (_, build_merge_scan) = strip_merge_scan(&join.right)?; - if build_merge_scan.is_placeholder() { + let (_, small_merge_scan) = strip_merge_scan(&join.left)?; + if small_merge_scan.is_placeholder() { return None; } - Some(&join.left) + Some(&join.right) } /// Whether the `MergeScan` boundary exposes a JSON2 column. @@ -514,7 +561,7 @@ mod tests { use common_meta::datanode::RegionManifestInfo; use common_meta::rpc::router::Region; use datafusion::datasource::DefaultTableSource; - use datafusion_common::{DFSchema, DFSchemaRef, JoinConstraint, NullEquality}; + use datafusion_common::{DFSchema, DFSchemaRef, JoinConstraint, NullEquality, TableReference}; use datafusion_expr::{Expr, LogicalPlanBuilder, build_join_schema, col, lit}; use datatypes::data_type::ConcreteDataType; use datatypes::extension::json::JsonExtensionType; @@ -611,15 +658,22 @@ mod tests { } fn rewrite(plan: LogicalPlan) -> LogicalPlan { + rewrite_with_stats( + plan, + &[ + (1, table_stats(10_000, 2)), + (2, table_stats(1_000, 1)), + (3, table_stats(500, 1)), + ], + ) + } + + fn rewrite_with_stats( + plan: LogicalPlan, + tables: &[(TableId, DistJoinTableStats)], + ) -> LogicalPlan { DistJoinPlanner {} - .analyze( - plan, - &rewrite_config(&[ - (1, table_stats(10_000, 2)), - (2, table_stats(1_000, 1)), - (3, table_stats(500, 1)), - ]), - ) + .analyze(plan, &rewrite_config(tables)) .unwrap() } @@ -657,7 +711,11 @@ mod tests { .as_any() .downcast_ref::() .unwrap_or_else(|| panic!("expected the outer node to be a MergeScan, got: {plan}")); - let LogicalPlan::Join(join) = merge_scan.input() else { + let mut input = merge_scan.input(); + while let LogicalPlan::Projection(projection) = input { + input = &projection.input; + } + let LogicalPlan::Join(join) = input else { panic!( "expected the outer MergeScan to wrap a join, got: {}", merge_scan.input() @@ -680,8 +738,8 @@ mod tests { assert_eq!(2, merge_scans(&result).len()); } - /// When statistics favor the right build side, its `MergeScan` becomes the inner scan - /// of the outer `MergeScan`, which is routed by the probe table's regions. + /// When statistics favor the small right side as build, its `MergeScan` becomes the + /// inner scan, while the large left side probes locally on regions selected by its routes. /// /// The rewrite must only move where the join runs: the join fields, the residual /// filter, the NULL semantics and the output schema stay as the join had them. @@ -696,13 +754,20 @@ mod tests { null_equality: NullEquality::NullEqualsNull, ..original_join.clone() }); - let probe_scan = merge_scans(&plan).first().cloned().unwrap(); + let large_scan = merge_scans(&plan).first().cloned().unwrap(); let result = rewrite(plan.clone()); let join = outer_join(&result); assert_eq!(JoinType::Inner, join.join_type); - assert_eq!(original_join.on, join.on); + assert_eq!( + original_join + .on + .iter() + .map(|(l, r)| (r.clone(), l.clone())) + .collect::>(), + join.on + ); assert_eq!( Some(col("t1.number").gt(lit(1u32))), join.filter, @@ -712,26 +777,25 @@ mod tests { assert_eq!(NullEquality::NullEqualsNull, join.null_equality); assert_eq!(original_join.null_aware, join.null_aware); assert!( - !contains_merge_scan(&join.left), - "probe side must be local, got: {}", + contains_merge_scan(&join.left), + "small hash-build side must keep its MergeScan, got: {}", join.left ); assert!( - contains_merge_scan(&join.right), - "build side must keep its MergeScan, got: {}", + !contains_merge_scan(&join.right), + "large probe side must be local, got: {}", join.right ); - // The outer scan must be routed by the probe table's regions and keep the - // probe boundary's partition columns. + // The outer scan must be routed by the large table's regions and keep the + // large boundary's partition columns. assert_eq!( "t1", - table_name_of(&LogicalPlan::Join(join.clone())) - .unwrap() - .table_name + table_name_of(&join.right).unwrap().table_name, + "the large local probe input controls outer region routing" ); let outer_scan = merge_scans(&result).first().cloned().unwrap(); - assert_eq!(probe_scan.partition_cols(), outer_scan.partition_cols()); + assert_eq!(large_scan.partition_cols(), outer_scan.partition_cols()); assert_eq!(plan.schema(), result.schema()); } @@ -750,6 +814,15 @@ mod tests { assert_eq!(2, scans.len()); assert_eq!(Some(probe_id), scans[0].remote_dyn_filter_producer_id()); assert_eq!(Some(build_id), scans[1].remote_dyn_filter_producer_id()); + let join = outer_join(&result); + assert!(contains_merge_scan(&join.left)); + assert!(!contains_merge_scan(&join.right)); + assert_eq!( + "t1", + table_name_of(nested_broadcast_join_probe_input(scans[0].input()).unwrap()) + .unwrap() + .table_name + ); } /// Boundaries without a remote dynamic filter producer id are still rewritten: a @@ -772,6 +845,15 @@ mod tests { assert_eq!(2, scans.len()); assert_eq!(None, scans[0].remote_dyn_filter_producer_id()); assert_eq!(None, scans[1].remote_dyn_filter_producer_id()); + let join = outer_join(&result); + assert!(contains_merge_scan(&join.left)); + assert!(!contains_merge_scan(&join.right)); + assert_eq!( + "t1", + table_name_of(nested_broadcast_join_probe_input(scans[0].input()).unwrap()) + .unwrap() + .table_name + ); assert_eq!(plan.schema(), result.schema()); } @@ -804,10 +886,10 @@ mod tests { let result = rewrite(projected.clone()); let join = outer_join(&result); - assert!(matches!(*join.left, LogicalPlan::Projection(_))); - assert!(!contains_merge_scan(&join.left)); + assert!(matches!(join.left.as_ref(), LogicalPlan::Projection(_))); + assert!(contains_merge_scan(&join.left)); assert!(matches!(join.right.as_ref(), LogicalPlan::Projection(_))); - assert!(contains_merge_scan(&join.right)); + assert!(!contains_merge_scan(&join.right)); assert_eq!(projected.schema(), result.schema()); } @@ -838,6 +920,89 @@ mod tests { }) } + /// Original output columns with the same basename but distinct qualifiers remain + /// addressable in the reversed join and preserve their source positions. + #[test] + fn restore_join_schema_supports_same_basename_under_distinct_qualifiers() { + let left = table_scan("large", 1, "large"); + let right = table_scan("small", 2, "small"); + let original_schema = DFSchemaRef::new( + build_join_schema(left.schema(), right.schema(), &JoinType::Inner).unwrap(), + ); + let swapped = LogicalPlan::Join( + Join::try_new( + Arc::new(right), + Arc::new(left), + vec![(col("small.number"), col("large.number"))], + None, + JoinType::Inner, + JoinConstraint::On, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + + let LogicalPlan::Join(swapped_join) = swapped else { + unreachable!("the schema restoration test constructs a join"); + }; + let restored = restore_join_schema(swapped_join, original_schema.clone()).unwrap(); + + assert_eq!(&original_schema, restored.schema()); + let LogicalPlan::Projection(projection) = restored else { + panic!("expected output-order restoration projection"); + }; + assert_eq!( + original_schema.columns(), + projection + .expr + .iter() + .map(|expr| { + match expr { + Expr::Column(column) => column.clone(), + other => panic!("expected column projection, got {other}"), + } + }) + .collect::>() + ); + } + + /// Bare and partial references with the same table name are distinct for DFSchema's + /// exact-name uniqueness check, but `resolved_eq` treats them as the same qualifier. + /// The first matching swapped position is wrong for one original field, so restoration + /// must decline. + #[test] + fn restore_join_schema_rejects_resolved_qualifier_ambiguity() { + let left = table_scan("same", 1, "large"); + let source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(test_table(2, "small")), + ))); + let right = LogicalPlanBuilder::scan( + TableReference::partial(DEFAULT_SCHEMA_NAME, "same"), + source, + None, + ) + .unwrap() + .build() + .unwrap(); + let original_schema = DFSchemaRef::new( + build_join_schema(left.schema(), right.schema(), &JoinType::Inner).unwrap(), + ); + let swapped = Join::try_new( + Arc::new(right), + Arc::new(left), + vec![], + None, + JoinType::Inner, + JoinConstraint::On, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(); + + assert!(restore_join_schema(swapped, original_schema).is_none()); + } + /// Only `INNER` joins are rewritten. #[test] fn nested_broadcast_join_rewrite_ignores_non_inner_join() { @@ -849,10 +1014,10 @@ mod tests { assert_eq!(plan.to_string(), result.to_string()); } - /// A build table on the left input is not rewritten: swapping the inputs would change - /// the join output schema, which the rewrite must preserve. + /// The cost heuristic currently selects only an original right-side build candidate. + /// A plan whose original right table is larger is therefore left unchanged. #[test] - fn nested_broadcast_join_rewrite_ignores_left_side_build() { + fn nested_broadcast_join_does_not_rewrite_unfavorable_right_candidate() { let plan = distributed_join_plan(); let LogicalPlan::Join(join) = &plan else { panic!("expected a join on top, got: {plan}"); @@ -891,8 +1056,8 @@ mod tests { assert_eq!(plan.to_string(), result.to_string()); } - /// The routing helper returns the probe side of the supported nested shape and `None` - /// for the join of two boundaries (the shape before the rewrite). + /// The routing helper returns the large local side of the supported nested shape and + /// `None` for the join of two boundaries (the shape before the rewrite). #[test] fn nested_broadcast_join_probe_input_only_for_supported_shape() { let plan = distributed_join_plan(); @@ -914,11 +1079,25 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - let (probe_plan, _) = strip_merge_scan(&join.left).unwrap(); + let large_plan = nested_broadcast_join_probe_input(merge_scan.input()).unwrap(); + let (expected_large, _) = strip_merge_scan(&join.left).unwrap(); assert_eq!( - Some(probe_plan.to_string()), + Some(expected_large.to_string()), nested_broadcast_join_probe_input(merge_scan.input()).map(|plan| plan.to_string()) ); + assert_eq!(large_plan.schema(), expected_large.schema()); + + let mut projected = merge_scan.input().clone(); + for _ in 0..4 { + let schema = projected.schema().clone(); + projected = + LogicalPlan::Projection(Projection::new_from_schema(Arc::new(projected), schema)); + } + assert_eq!( + Some(large_plan.to_string()), + nested_broadcast_join_probe_input(&projected).map(ToString::to_string), + "routing should look through four ordinary projection wrappers" + ); } /// JSON2 boundaries are not rewritten: the nested build boundary is hidden from @@ -1185,8 +1364,8 @@ mod tests { ); } - /// The cost heuristic nests the build side only when the statistics favor it: the same - /// statistics in the other order never swap the sides, and missing statistics keep the plan. + /// The cost heuristic nests the candidate build side only when its statistics favor it: the same + /// statistics in the other order do not favor this candidate order, and missing statistics keep the plan. #[test] fn nested_broadcast_join_rewrite_with_stats() { let plan = distributed_join_plan(); @@ -1199,18 +1378,17 @@ mod tests { let join = outer_join(&result); assert_eq!(2, merge_scans(&result).len()); assert!( - !contains_merge_scan(&join.left), - "probe side must become local, got: {}", + contains_merge_scan(&join.left), + "small hash-build side must keep its MergeScan, got: {}", join.left ); assert!( - contains_merge_scan(&join.right), - "build side must keep its MergeScan, got: {}", + !contains_merge_scan(&join.right), + "large probe side must become local, got: {}", join.right ); - // The sides are never swapped: the same statistics in the other order do not favor the - // right build side. + // The reversed candidate sizes do not favor the right build side for this input order. let swapped = stats_config(&[(1, table_stats(1_000, 1)), (2, table_stats(10_000, 2))]); let result = DistJoinPlanner {}.analyze(plan.clone(), &swapped).unwrap(); assert_eq!(plan.to_string(), result.to_string()); diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 7993e2b9e1b5..ea69dbb55d39 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -491,6 +491,7 @@ pub struct MergeScanExec { target_partition: usize, partition_cols: AliasMapping, enable_per_region_metrics: bool, + row_estimate: Option, } impl std::fmt::Debug for MergeScanExec { @@ -592,24 +593,39 @@ impl MergeScanExec { target_partition, partition_cols, enable_per_region_metrics, + row_estimate: None, }) } - /// Conservative row-count upper bound for all selected regions. + /// Attaches an ordinary-scan row estimate to this merge scan. + pub fn with_row_estimate(mut self, row_estimate: usize) -> Self { + self.row_estimate = Some(row_estimate); + self + } + + /// Row estimate for the whole merge scan, or a conservative plan-derived upper bound. /// - /// This is not an expected cardinality: DataFusion receives it as an - /// inexact estimate only because `Statistics` has no upper-bound precision. + /// The explicit estimate is inexact, not a bound. The fallback plan-derived upper bound is + /// also represented as inexact because `Statistics` has no upper-bound precision. fn estimated_num_rows(&self) -> Precision { if self.regions.is_empty() { return Precision::Inexact(0); } let Some(rows_per_region) = remote_plan_row_bound(&self.plan) else { - return Precision::Absent; + return self + .row_estimate + .map_or(Precision::Absent, |estimate| Precision::Inexact(estimate)); + }; + let Some(row_bound) = rows_per_region.checked_mul(self.regions.len()) else { + return self + .row_estimate + .map_or(Precision::Absent, |estimate| Precision::Inexact(estimate)); }; - rows_per_region - .checked_mul(self.regions.len()) - .map_or(Precision::Absent, Precision::Inexact) + Precision::Inexact( + self.row_estimate + .map_or(row_bound, |estimate| estimate.min(row_bound)), + ) } /// Number of partitions populated by the region striping in [`Self::to_stream`]. @@ -982,6 +998,7 @@ impl MergeScanExec { target_partition: self.target_partition, partition_cols: self.partition_cols.clone(), enable_per_region_metrics: self.enable_per_region_metrics, + row_estimate: self.row_estimate, }) } @@ -1665,6 +1682,88 @@ mod tests { } } + #[test] + fn merge_scan_reports_attached_row_estimate_and_preserves_plan_bound() { + let regions = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)]; + let uncapped = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .build() + .unwrap(); + let exec = merge_scan_exec_with_plan(regions, uncapped, 10).with_row_estimate(123); + let stats = merge_scan_statistics(&exec); + assert_eq!(stats.num_rows, Precision::Inexact(123)); + assert_eq!(stats.total_byte_size, Precision::Absent); + assert!(stats.column_statistics.iter().all(|column| { + column.null_count == Precision::Absent + && column.max_value == Precision::Absent + && column.min_value == Precision::Absent + && column.distinct_count == Precision::Absent + })); + assert_eq!( + exec.statistics_from_inputs(&[], &StatisticsArgs::new().with_partition(Some(0)),) + .unwrap() + .num_rows, + Precision::Absent + ); + let limited = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(50)) + .unwrap() + .build() + .unwrap(); + let exec = merge_scan_exec_with_plan(vec![RegionId::new(1024, 1)], limited, 10) + .with_row_estimate(100); + assert_eq!( + merge_scan_statistics(&exec).num_rows, + Precision::Inexact(50) + ); + + let two_region_limit = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(50)) + .unwrap() + .build() + .unwrap(); + let exec = merge_scan_exec_with_plan( + vec![RegionId::new(1024, 1), RegionId::new(1024, 2)], + two_region_limit, + 10, + ) + .with_row_estimate(123); + assert_eq!( + merge_scan_statistics(&exec).num_rows, + Precision::Inexact(100) + ); + + let overflow_limit = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(i64::MAX as usize)) + .unwrap() + .build() + .unwrap(); + let overflow_cap = i64::MAX as usize; + assert_eq!(Some(overflow_cap), remote_plan_row_bound(&overflow_limit)); + assert_eq!(None, overflow_cap.checked_mul(3)); + let exec = merge_scan_exec_with_plan( + vec![ + RegionId::new(1024, 1), + RegionId::new(1024, 2), + RegionId::new(1024, 3), + ], + overflow_limit, + 10, + ) + .with_row_estimate(123); + assert_eq!( + merge_scan_statistics(&exec).num_rows, + Precision::Inexact(123) + ); + } + #[test] fn merge_scan_reports_only_deterministic_plan_bounds() { use datafusion::functions_aggregate::expr_fn::count; @@ -3330,7 +3429,8 @@ mod tests { Some(remote_dyn_filter_producer_id), false, ) - .unwrap(); + .unwrap() + .with_row_estimate(17); assert_eq!( exec.remote_dyn_filter_producer_id(), @@ -3354,6 +3454,7 @@ mod tests { Some(remote_dyn_filter_producer_id), "try_with_new_distribution must preserve remote dynamic filter producer id" ); + assert_eq!(cloned.row_estimate, Some(17)); } #[test] diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index 617a43e2f367..005ccb246f80 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -20,7 +20,9 @@ use ahash::HashMap; use arrow_schema::SortOptions; use async_trait::async_trait; use catalog::CatalogManagerRef; +use catalog::kvbackend::KvBackendCatalogManager; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; +use common_meta::datanode::RegionStat; use common_telemetry::debug; use datafusion::catalog::Session; use datafusion::common::Result; @@ -38,6 +40,7 @@ use partition::expr::PartitionExpr; use partition::manager::{PartitionRuleManagerRef, create_partitions_from_region_routes}; use session::context::QueryContext; use snafu::{OptionExt, ResultExt}; +use store_api::region_engine::RegionRole; use store_api::storage::RegionId; use table::TableRef; use table::metadata::TableInfo; @@ -46,7 +49,7 @@ use table::table::adapter::DfTableProviderAdapter; use table::table_name::TableName; use crate::dist_plan::PredicateExtractor; -use crate::dist_plan::dist_join_planner::nested_broadcast_join_probe_input; +use crate::dist_plan::dist_join_planner::{expected_region_ids, nested_broadcast_join_probe_input}; use crate::dist_plan::merge_scan::{MergeScanExec, MergeScanLogicalPlan}; use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan}; use crate::dist_plan::region_pruner::ConstraintPruner; @@ -196,10 +199,9 @@ impl ExtensionPlanner for DistExtensionPlanner { } let optimized_plan = input_plan; - // Region pruning clears the collected predicates at a node with several inputs, - // so the payload of a nested broadcast join would be routed by the join itself - // and read every probe region. The local probe input keeps the probe predicate - // pruning; the full payload is still dispatched to the selected regions. + // Region pruning clears collected predicates at a multi-input join, so routing the + // nested payload by the join would read every large-side region. The local large-side + // input preserves its predicate pruning; the full payload is still dispatched there. let routing_plan = nested_broadcast_join_probe_input(input_plan).unwrap_or(input_plan); let Some(table_name) = Self::extract_full_table_name(routing_plan)? else { // no relation found in input plan, going to execute them locally @@ -225,6 +227,7 @@ impl ExtensionPlanner for DistExtensionPlanner { .config() .get_extension() .unwrap_or_else(QueryContext::arc); + let row_estimate = self.ordinary_scan_row_estimate(input_plan, ®ions).await; let merge_scan_plan = MergeScanExec::new( session_state, table_name, @@ -238,11 +241,107 @@ impl ExtensionPlanner for DistExtensionPlanner { merge_scan.remote_dyn_filter_producer_id(), self.enable_per_region_metrics, )?; + let merge_scan_plan = match row_estimate { + Some(estimate) => merge_scan_plan.with_row_estimate(estimate), + None => merge_scan_plan, + }; Ok(Some(Arc::new(merge_scan_plan) as _)) } } +/// Sums positive ordinary-leader row counts for exactly the selected regions. +fn aggregate_region_rows(selected_regions: &[RegionId], reports: &[RegionStat]) -> Option { + let selected = selected_regions + .iter() + .copied() + .collect::>(); + if selected_regions.is_empty() || selected.len() != selected_regions.len() { + return None; + } + + let mut rows = 0u64; + for region in selected_regions { + let mut leaders = reports + .iter() + .filter(|report| report.id == *region && report.role == RegionRole::Leader); + let report = leaders.next()?; + if leaders.next().is_some() || report.num_rows == 0 { + return None; + } + rows = rows.checked_add(report.num_rows)?; + } + usize::try_from(rows).ok() +} + impl DistExtensionPlanner { + /// Returns a row estimate only for a simple scan of the captured physical base table. + async fn ordinary_scan_row_estimate( + &self, + plan: &LogicalPlan, + selected_regions: &[RegionId], + ) -> Option { + fn captured_base_table_id(plan: &LogicalPlan) -> Option { + match plan { + LogicalPlan::TableScan(scan) => { + let source = scan.source.downcast_ref::()?; + let provider = source + .table_provider + .downcast_ref::()?; + (provider.table().table_type() == TableType::Base) + .then(|| provider.table().table_info().table_id()) + } + LogicalPlan::Filter(filter) => captured_base_table_id(&filter.input), + LogicalPlan::Projection(projection) => captured_base_table_id(&projection.input), + LogicalPlan::SubqueryAlias(alias) => captured_base_table_id(&alias.input), + _ => None, + } + } + + let table_id = captured_base_table_id(plan)?; + if selected_regions.is_empty() { + return None; + } + let selected = selected_regions + .iter() + .copied() + .collect::>(); + if selected.len() != selected_regions.len() + || selected.iter().any(|region| region.table_id() != table_id) + { + return None; + } + let catalog_manager = self + .catalog_manager + .as_any() + .downcast_ref::()?; + let (physical_table_id, route) = self + .partition_rule_manager + .find_physical_table_route_with_id(table_id) + .await + .ok()?; + if physical_table_id != table_id { + return None; + } + let routed_region_ids = expected_region_ids(physical_table_id, &route.region_routes)?; + let routed_regions = routed_region_ids + .iter() + .copied() + .collect::>(); + if routed_regions.is_empty() + || routed_regions.len() != routed_region_ids.len() + || !selected.is_subset(&routed_regions) + { + return None; + } + + let reports = catalog_manager + .information_extension() + .region_stats() + .await + .ok()?; + aggregate_region_rows(selected_regions, &reports) + } + /// Extract fully resolved table name from logical plan fn extract_full_table_name(plan: &LogicalPlan) -> Result> { let mut extractor = TableScanExtractor::default(); @@ -563,6 +662,7 @@ mod tests { use catalog::{CatalogManagerRef, RegisterTableRequest}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_meta::cache::new_table_route_cache; + use common_meta::datanode::RegionStat; use common_meta::key::TableMetadataManager; use common_meta::key::table_route::TableRouteValue; use common_meta::kv_backend::memory::MemoryKvBackend; @@ -570,7 +670,9 @@ mod tests { use common_query::request::QueryRequest; use common_recordbatch::SendableRecordBatchStream; use datafusion::datasource::DefaultTableSource; - use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder, col as df_col, lit}; + use datafusion_expr::{ + JoinType, LogicalPlan, LogicalPlanBuilder, Projection, col as df_col, lit, + }; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::value::Value; @@ -579,14 +681,15 @@ mod tests { use partition::expr::{PartitionExpr, col as partition_col}; use partition::manager::PartitionRuleManager; use session::ReadPreference; + use store_api::region_engine::RegionRole; use store_api::storage::RegionId; use table::metadata::{TableInfo, TableInfoBuilder, TableMeta, TableType}; use table::table::adapter::DfTableProviderAdapter; use table::table_name::TableName; use table::test_util::EmptyTable; - use super::DistExtensionPlanner; use super::nested_broadcast_join_probe_input; + use super::{DistExtensionPlanner, aggregate_region_rows}; use crate::dist_plan::merge_scan::MergeScanLogicalPlan; use crate::region_query::{RegionQueryHandler, RegionQueryTarget}; @@ -820,21 +923,20 @@ mod tests { } /// Region pruning clears the collected predicates at a node with several inputs, so a - /// nested broadcast join must be routed by its local probe input: the probe predicate - /// prunes the probe regions, while routing by the whole payload would scan every - /// probe region. + /// nested broadcast join must be routed by its local large-side input: its predicate + /// prunes the large-table regions, while routing by the whole payload scans all regions. #[tokio::test] - async fn nested_broadcast_join_routes_by_probe_input() { - let (planner, probe_plan, table_name) = + async fn nested_broadcast_join_routes_by_large_input() { + let (planner, large_plan, table_name) = planner_and_plan(vec![0, 1], physical_partition_expressions()).await; - // The build side scans the same table under another qualifier, so the join sides - // have distinct field qualifiers. + // The small build side can use the same table under another qualifier; only the + // large-side predicate must decide which regions are selected. let build_plan = - MergeScanLogicalPlan::new(build_side_plan(&probe_plan), false, Default::default()) + MergeScanLogicalPlan::new(build_side_plan(&large_plan), false, Default::default()) .into_logical_plan(); - let join_plan = LogicalPlanBuilder::from(probe_plan.clone()) + let join_plan = LogicalPlanBuilder::from(build_plan) .join( - build_plan, + large_plan.clone(), JoinType::Inner, (vec!["host"], vec!["host"]), None, @@ -842,14 +944,21 @@ mod tests { .unwrap() .build() .unwrap(); + assert!( + planner + .ordinary_scan_row_estimate(&join_plan, &[RegionId::new(LOGICAL_TABLE_ID, 1)]) + .await + .is_none(), + "a join payload must not inherit a base table's row count" + ); - // The whole payload clears the probe predicate at the join and reads all regions. + // The whole payload clears the large-side predicate at the join and reads all regions. assert_all_logical_regions(planner.get_regions(&table_name, &join_plan).await.unwrap()); - // The local probe input keeps the probe predicate pruning: one of three regions. + // The local large input keeps its predicate pruning: one of three regions. let routing_plan = nested_broadcast_join_probe_input(&join_plan) - .expect("the supported nested shape must route by its probe input"); - assert_eq!(probe_plan.to_string(), routing_plan.to_string()); + .expect("the supported nested shape must route by its large-side input"); + assert_eq!(large_plan.to_string(), routing_plan.to_string()); assert_eq!( vec![RegionId::new(LOGICAL_TABLE_ID, 1)], planner @@ -857,6 +966,28 @@ mod tests { .await .unwrap() ); + + let mut full_payload = join_plan.clone(); + for _ in 0..4 { + let schema = full_payload.schema().clone(); + full_payload = LogicalPlan::Projection(Projection::new_from_schema( + Arc::new(full_payload), + schema, + )); + } + assert!( + planner + .ordinary_scan_row_estimate(&full_payload, &[RegionId::new(LOGICAL_TABLE_ID, 1)],) + .await + .is_none(), + "a join under projection wrappers is not an ordinary scan" + ); + let large_input = nested_broadcast_join_probe_input(&full_payload) + .expect("projection-wrapped join should retain its large-side routing input"); + assert_eq!( + vec![RegionId::new(LOGICAL_TABLE_ID, 1)], + planner.get_regions(&table_name, large_input).await.unwrap() + ); } /// Scans the table of `plan` under the qualifier `build` with the same probe predicate, @@ -893,6 +1024,123 @@ mod tests { assert_all_logical_regions(planner.get_regions(&table_name, &plan).await.unwrap()); } + fn region_stat(id: RegionId, role: RegionRole, num_rows: u64) -> RegionStat { + RegionStat { + id, + rcus: 0, + wcus: 0, + approximate_bytes: 1, + engine: "mito".to_string(), + role, + num_rows, + memtable_size: 0, + manifest_size: 0, + sst_size: 1, + sst_num: 1, + index_size: 0, + region_manifest: common_meta::datanode::RegionManifestInfo::Mito { + manifest_version: 0, + flushed_entry_id: 0, + file_removed_cnt: 0, + }, + written_bytes: 0, + query_cpu_time: 0, + query_scanned_bytes: 0, + data_topic_latest_entry_id: 0, + metadata_topic_latest_entry_id: 0, + min_timestamp: None, + max_timestamp: None, + } + } + + #[test] + fn aggregate_region_rows_requires_one_positive_ordinary_leader_per_selected_region() { + let first = RegionId::new(LOGICAL_TABLE_ID, 1); + let second = RegionId::new(LOGICAL_TABLE_ID, 2); + let unrelated = RegionId::new(LOGICAL_TABLE_ID, 3); + + assert_eq!( + Some(30), + aggregate_region_rows( + &[first, second], + &[ + region_stat(first, RegionRole::Leader, 10), + region_stat(second, RegionRole::Leader, 20), + region_stat(unrelated, RegionRole::Leader, 999), + ], + ) + ); + assert_eq!( + Some(10), + aggregate_region_rows( + &[first], + &[ + region_stat(first, RegionRole::Leader, 10), + region_stat(second, RegionRole::Leader, 20), + ], + ), + "pruned regions must not contribute to the selected scan estimate" + ); + assert_eq!( + Some(10), + aggregate_region_rows( + &[first], + &[ + region_stat(first, RegionRole::Follower, 100), + region_stat(first, RegionRole::Leader, 10), + ], + ) + ); + + let invalid_cases = [ + ( + vec![first, second], + vec![region_stat(first, RegionRole::Leader, 10)], + ), + (vec![first], vec![region_stat(first, RegionRole::Leader, 0)]), + ( + vec![first], + vec![ + region_stat(first, RegionRole::Leader, 10), + region_stat(first, RegionRole::Leader, 0), + ], + ), + ( + vec![first], + vec![region_stat(first, RegionRole::Follower, 10)], + ), + ( + vec![first], + vec![region_stat(first, RegionRole::DowngradingLeader, 10)], + ), + ( + vec![first, second], + vec![ + region_stat(first, RegionRole::Leader, u64::MAX), + region_stat(second, RegionRole::Leader, 1), + ], + ), + ]; + for (selected, reports) in invalid_cases { + assert_eq!(None, aggregate_region_rows(&selected, &reports)); + } + assert_eq!(None, aggregate_region_rows(&[], &[])); + assert_eq!( + None, + aggregate_region_rows( + &[first, first], + &[region_stat(first, RegionRole::Leader, 10)] + ) + ); + assert_eq!( + None, + aggregate_region_rows( + &[first], + &[region_stat(RegionId::new(9999, 1), RegionRole::Leader, 10)] + ) + ); + } + fn assert_all_logical_regions(mut regions: Vec) { regions.sort_unstable(); assert_eq!( diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 1ef7739bab40..050e90e2885e 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -60,8 +60,8 @@ //! [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]). //! //! [`test_nested_broadcast_join_rewrite_executes_on_probe_regions`] exercises the production -//! selector through SQL `SET` and the regular query entrypoint, checks that the outer `MergeScan` -//! is planned for the probe region selected by the predicate, and executes the SQL on the frontend. +//! selector through SQL `SET` and the regular query entrypoint: an unfiltered case proves BIG local +//! probing and SMALL remote building, while a filtered case checks BIG-region pruning. use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; @@ -315,10 +315,9 @@ async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation } /// The SQL setting and region statistics select the nested plan through the production frontend -/// path. The probe predicate selects one of the two probe regions, so the outer `MergeScan` must be -/// planned for that region only: routing by the whole join payload would read both probe regions. -/// The executed result must preserve the duplicated build keys, NULL join keys, residual filter, -/// and projection of the query. +/// path. An unfiltered BIG-table case proves the local probe and remote build roles, complete +/// output, and schema; a separate filtered query verifies BIG-region pruning. Both exercise the +/// executed SQL result with duplicate keys, NULL keys, and a residual join filter. #[tokio::test(flavor = "multi_thread")] async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { common_telemetry::init_default_ut_logging(); @@ -344,40 +343,74 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { .await; insert_extra_probe_rows(&frontend).await; - // The SQL under test: the projection reorders the columns of the join, the join keys `10` of - // the build table are duplicated, the residual join filter rejects the join partners whose - // values sum to at most `105` (the probe row `100` sums up to `101` and `102`), and the probe - // predicate selects the first probe region. + // Leave the BIG probe table unfiltered so all padding rows reach the local probe input. The + // projection order, duplicate keys, NULLs, and residual condition remain part of the check. let sql = format!( "SELECT p.probe_v, p.a_id, b.build_v FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b - ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 - WHERE p.a_id < 100 AND p.a_ts < 10000::TIMESTAMP" + ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105" ); - // The rows of the SQL as `(probe_v, a_id, build_v)`: the probe row `(4, 10, 600)` joins the two - // build rows of the duplicated key `10` (`600 + 1` and `600 + 2`), the probe row `(2, 20, 200)` - // joins `(3, 20, 3)`, the residual filter rejects the probe row `(1, 10, 100)` (`100 + 1` and - // `100 + 2` are not greater than `105`), and the NULL join keys of both sides match nothing. - let expected = multiset(vec![ - vec!["200".to_string(), "2".to_string(), "3".to_string()], - vec!["600".to_string(), "4".to_string(), "1".to_string()], - vec!["600".to_string(), "4".to_string(), "2".to_string()], - ]); + let expected = expected_full_broadcast_rows(); let query_ctx = query_ctx(); // The reference: the frontend executes the same SQL with its default distributed plan, i.e. // both sides of the join are read with a `MergeScan` and the join runs on the frontend. - let reference = query_pretty(&frontend, &sql, query_ctx.clone()).await; - info!("reference result of the frontend:\n{reference}"); + let reference = output_batches(run_sql(&frontend, &sql, query_ctx.clone()).await).await; + assert_eq!( + vec!["probe_v", "a_id", "build_v"], + reference + .schema() + .column_schemas() + .iter() + .map(|column| column.name.as_str()) + .collect::>(), + "the unfiltered query must preserve its explicit projection order" + ); + let reference_pretty = reference.pretty_print().unwrap(); + info!("reference result of the frontend:\n{reference_pretty}"); + let star_sql = format!( + "SELECT * + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b + ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 + WHERE p.a_id < 100 AND p.a_ts < 10000::TIMESTAMP" + ); + let star_reference = + output_batches(run_sql(&frontend, &star_sql, query_ctx.clone()).await).await; + let expected_star_columns = vec![ + "a_id", + "probe_key", + "probe_v", + "a_ts", + "b_id", + "build_key", + "build_v", + "b_ts", + ]; + assert_eq!( + expected_star_columns, + star_reference + .schema() + .column_schemas() + .iter() + .map(|column| column.name.as_str()) + .collect::>(), + "SELECT * must retain BIG columns followed by SMALL columns" + ); + assert_eq!( + 3, + table_cells(&star_reference.pretty_print().unwrap()).len(), + "the filtered SELECT * baseline must contain all three rows" + ); assert_eq!( expected, - multiset(table_cells(&reference)), - "unexpected result of the join on the frontend:\n{reference}" + multiset(table_cells(&reference_pretty)), + "unexpected result of the join on the frontend:\n{reference_pretty}" ); // Use the actual session setting and SQL entrypoint. The region statistics above drive the // production selector; no analyzer options or plans are injected by this test. - let (probe_bytes, build_bytes) = wait_for_region_statistics(&frontend).await; + let (probe_bytes, build_bytes) = + wait_for_region_statistics(&frontend, BUILD_ROWS.len() as u64 + 1).await; assert!( 2 * build_bytes < probe_bytes, "the fixture must make the heuristic favor the build side, actual region statistics: \ @@ -390,7 +423,6 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { ) .await; - // The time predicate excludes the padding rows inserted to make the probe side larger. let actual = run_sql(&frontend, &sql, query_ctx.clone()).await; let physical = actual .meta @@ -399,7 +431,7 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { .expect("the output of the query must carry its physical plan"); info!("physical plan of the SQL-selected nested join:\\n{physical:?}"); - // The probe predicate selects one of the two probe regions, so only that region is dispatched. + // The unfiltered BIG probe scan dispatches both of its regions. let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; assert!( probe_leaders.len() > 1, @@ -410,10 +442,9 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { let regions = merge_scan.regions().to_vec(); info!("the outer MergeScan is planned for the regions {regions:?}"); assert_eq!( - 1, + probe_leaders.len(), regions.len(), - "the probe predicate `a_id < 100` must select one of the {} probe regions, actual \ - selection: {regions:?}", + "the unfiltered BIG table scan must dispatch all {} regions, actual selection: {regions:?}", probe_leaders.len() ); assert!( @@ -423,18 +454,147 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { "the outer MergeScan must select regions of {PROBE_TABLE}, actual selection: {regions:?}" ); - let actual = actual.data.pretty_print().await; - info!("result of the SQL-selected nested plan:\n{actual}"); + let actual = output_batches(actual).await; + assert_eq!( + reference.schema(), + actual.schema(), + "broadcast execution must preserve full schema and types" + ); + let actual_pretty = actual.pretty_print().unwrap(); + info!("result of the SQL-selected nested plan:\n{actual_pretty}"); + let stages = merge_scan.sub_stage_metrics(); + assert_eq!( + regions.len(), + stages.len(), + "each dispatched outer region must report its stage metrics" + ); + let hash_joins = stages + .iter() + .flat_map(|stage| &stage.plan_metrics) + .filter(|metric| metric.plan_name == "HashJoinExec") + .collect::>(); + assert_eq!( + regions.len(), + hash_joins.len(), + "every dispatched outer region must report its actual DN HashJoinExec: {stages:?}" + ); + // The ordinary equijoin pre-scan optimizer filters NULL join keys before the DN HashJoin. + // The six-row SMALL table (including its NULL-key row) is still covered by result/schema + // checks and heartbeat readiness; the actual hash build therefore contains BUILD_ROWS only. + let expected_build_rows = BUILD_ROWS.len(); + assert!( + hash_joins.iter().all(|metric| { + metric + .metrics + .iter() + .any(|(name, value)| *name == "build_input_rows" && *value == expected_build_rows) + && metric + .metrics + .iter() + .any(|(name, value)| *name == "input_rows" && *value > expected_build_rows) + }), + "each DN HashJoinExec must build {expected_build_rows} non-NULL remote SMALL rows and consume more BIG probe rows: {hash_joins:?}" + ); + assert_eq!( expected, - multiset(table_cells(&actual)), - "the nested plan returned an unexpected multiset of rows:\n{actual}" + multiset(table_cells(&actual_pretty)), + "the nested plan returned an unexpected multiset of rows:\n{actual_pretty}" ); assert_eq!( - multiset(table_cells(&reference)), - multiset(table_cells(&actual)), - "the nested plan returned a different multiset of rows than the frontend:\n{actual}" + multiset(table_cells(&reference_pretty)), + multiset(table_cells(&actual_pretty)), + "the nested plan returned a different multiset of rows than the frontend:\n{actual_pretty}" ); + + // Keep the original selective query as a separate routing regression. Its filtered local + // side can be smaller than the remote table, so this asserts pruning/results only, not build + // orientation. + let pruned_sql = format!( + "SELECT p.probe_v, p.a_id, b.build_v + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b + ON p.probe_key = b.build_key AND p.probe_v + b.build_v > 105 + WHERE p.a_id < 100 AND p.a_ts < 10000::TIMESTAMP" + ); + let pruned_expected = multiset(vec![ + vec!["200".to_string(), "2".to_string(), "3".to_string()], + vec!["600".to_string(), "4".to_string(), "1".to_string()], + vec!["600".to_string(), "4".to_string(), "2".to_string()], + ]); + let pruned = run_sql(&frontend, &pruned_sql, query_ctx.clone()).await; + let pruned_plan = pruned + .meta + .plan + .clone() + .expect("the pruned query output must carry its physical plan"); + let pruned_merge_scan = find_merge_scan_exec(&pruned_plan) + .unwrap_or_else(|| panic!("expected pruned query outer MergeScan: {pruned_plan:?}")); + let pruned_regions = pruned_merge_scan.regions(); + assert_eq!( + 1, + pruned_regions.len(), + "the predicate must prune to one BIG region" + ); + assert!( + pruned_regions + .iter() + .all(|region| probe_leaders.contains_key(®ion.as_u64())), + "the pruned region must belong to BIG table {PROBE_TABLE}: {pruned_regions:?}" + ); + let pruned_batches = output_batches(pruned).await; + assert_eq!( + pruned_expected, + multiset(table_cells(&pruned_batches.pretty_print().unwrap())), + "pruned query must preserve the expected complete result" + ); + + // SELECT * must also restore the original BIG-then-SMALL schema after rewriting. + let star_actual = output_batches(run_sql(&frontend, &star_sql, query_ctx.clone()).await).await; + assert_eq!( + star_reference.schema(), + star_actual.schema(), + "rewritten SELECT * must preserve all column types and order" + ); + assert_eq!( + expected_star_columns, + star_actual + .schema() + .column_schemas() + .iter() + .map(|column| column.name.as_str()) + .collect::>(), + "rewritten SELECT * must retain BIG columns followed by SMALL columns" + ); + assert_eq!( + multiset(table_cells(&star_reference.pretty_print().unwrap())), + multiset(table_cells(&star_actual.pretty_print().unwrap())), + "rewritten SELECT * must preserve all three baseline rows" + ); +} + +/// Complete expected output for the unfiltered broadcast SQL, in its projected column order. +fn expected_full_broadcast_rows() -> Vec> { + let mut rows = Vec::new(); + let mut probes = PROBE_ROWS.to_vec(); + // SQL equality never matches the explicit NULL-key probe row. + probes.extend( + extra_probe_rows() + .into_iter() + .map(|(a_id, key, value, _)| (a_id, key, value)), + ); + probes.push((4, 10, 600)); + for (a_id, probe_key, probe_v) in probes { + for (_, build_key, build_v) in BUILD_ROWS.iter().copied() { + if probe_key == build_key && probe_v + build_v > 105 { + rows.push(vec![ + probe_v.to_string(), + a_id.to_string(), + build_v.to_string(), + ]); + } + } + } + multiset(rows) } /// Returns the `MergeScanExec` of `plan`, descending through the physical plan tree. @@ -494,15 +654,16 @@ fn contains_hash_join(plan: &Arc) -> bool { plan.name().contains("HashJoin") || plan.children().into_iter().any(contains_hash_join) } -/// The region statistics of `table`: `(leader reports, summed disk bytes, minimal disk bytes)`. +/// The region statistics of `table`: `(leader reports, summed disk bytes, minimum disk bytes, summed rows)`. /// /// A region that has not reported yet, reports zero bytes, or is not served by a leader makes the /// cost heuristic of the rewrite skip the table, so a test that expects the rewrite has to wait /// until every region of every table has an ordinary leader report. A missing or NULL aggregate /// counts as `0`, i.e. as "not reported yet". -async fn region_statistics(frontend: &Arc, table: &str) -> (u64, u64, u64) { +async fn region_statistics(frontend: &Arc, table: &str) -> (u64, u64, u64, u64) { let sql = format!( - "SELECT count(*) AS regions, sum(disk_size) AS bytes, min(disk_size) AS min_bytes + "SELECT count(*) AS regions, sum(disk_size) AS bytes, min(disk_size) AS min_bytes, + sum(region_rows) AS rows FROM information_schema.region_statistics WHERE table_id IN (SELECT table_id FROM information_schema.tables WHERE table_schema = 'public' AND table_name = '{table}') @@ -517,7 +678,7 @@ async fn region_statistics(frontend: &Arc, table: &str) -> (u64, u64, .unwrap_or(0) }; - (parse(0), parse(1), parse(2)) + (parse(0), parse(1), parse(2), parse(3)) } /// Waits until both tables of the join have one ordinary leader report per region with a non-zero @@ -525,7 +686,10 @@ async fn region_statistics(frontend: &Arc, table: &str) -> (u64, u64, /// /// The frontend reads these statistics from the region statistics that the datanodes report to /// metasrv, so they appear a few heartbeats after the tables are created. -async fn wait_for_region_statistics(frontend: &Arc) -> (u64, u64) { +async fn wait_for_region_statistics( + frontend: &Arc, + expected_build_rows: u64, +) -> (u64, u64) { let deadline = tokio::time::Instant::now() + Duration::from_secs(120); loop { let mut stats = Vec::new(); @@ -535,7 +699,13 @@ async fn wait_for_region_statistics(frontend: &Arc) -> (u64, u64) { let probe = stats[0]; let build = stats[1]; info!("region statistics of {PROBE_TABLE}: {probe:?}, of {BUILD_TABLE}: {build:?}"); - if probe.0 == 2 && probe.2 > 0 && build.0 == 2 && build.2 > 0 { + if probe.0 == 2 + && probe.2 > 0 + && probe.3 > 0 + && build.0 == 2 + && build.2 > 0 + && build.3 == expected_build_rows + { return (probe.1, build.1); } assert!( @@ -655,7 +825,8 @@ async fn test_experimental_dist_join_setting_activates_rewrite() { // The statistics the heuristic reads: every region of both tables has a leader report with a // non-zero size, and broadcasting the build table to the probe regions is much cheaper than // the probe table itself. - let (probe_bytes, build_bytes) = wait_for_region_statistics(&frontend).await; + let (probe_bytes, build_bytes) = + wait_for_region_statistics(&frontend, BUILD_ROWS.len() as u64).await; assert!( 2 * build_bytes < probe_bytes, "the fixture must make the heuristic favor the build side, actual region statistics: \ From 42ee0eeeebd2e33aaca0d40f193b8bb6b3ace1ab Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 12:30:56 +0800 Subject: [PATCH 22/37] style(query): simplify row estimate construction Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/merge_scan.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index ea69dbb55d39..8f84198820fc 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -615,12 +615,12 @@ impl MergeScanExec { let Some(rows_per_region) = remote_plan_row_bound(&self.plan) else { return self .row_estimate - .map_or(Precision::Absent, |estimate| Precision::Inexact(estimate)); + .map_or(Precision::Absent, Precision::Inexact); }; let Some(row_bound) = rows_per_region.checked_mul(self.regions.len()) else { return self .row_estimate - .map_or(Precision::Absent, |estimate| Precision::Inexact(estimate)); + .map_or(Precision::Absent, Precision::Inexact); }; Precision::Inexact( self.row_estimate From 5719b845abe6581ca8dc8d0228452862e8adedba Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:11:11 +0800 Subject: [PATCH 23/37] test(query): refresh small-build distributed join plans Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../distributed/explain/dist_join.result | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/tests/cases/distributed/explain/dist_join.result b/tests/cases/distributed/explain/dist_join.result index 1e829183c017..31ea97a4959e 100644 --- a/tests/cases/distributed/explain/dist_join.result +++ b/tests/cases/distributed/explain/dist_join.result @@ -130,13 +130,14 @@ ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; |_|_TableScan: dist_join_build, partial_filters=[dist_join_build.b_key IS NOT NULL]_| |_| ]]_| | physical_plan | SortPreservingMergeExec: [p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST]_| -|_|_SortExec: expr=[p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST], preserve_partitioning=[true] | +|_|_ProjectionExec: expr=[p_id@3 as p_id, p_key@4 as p_key, p_val@5 as p_val, b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| +|_|_SortExec: expr=[p_val@5 ASC NULLS LAST, p_id@3 ASC NULLS LAST, b_val@2 ASC NULLS LAST, b_id@0 ASC NULLS LAST], preserve_partitioning=[true] | |_|_REDACTED -|_|_RepartitionExec: partitioning=REDACTED -|_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val]_| +|_|_CoalescePartitionsExec_| +|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| |_|_MergeScanExec: REDACTED |_|_RepartitionExec: partitioning=REDACTED -|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| +|_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -165,17 +166,18 @@ ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; | logical_plan_| Sort: p.p_val ASC NULLS LAST, p.p_id ASC NULLS LAST, b.b_val ASC NULLS LAST, b.b_id ASC NULLS LAST_| |_|_Projection: p.p_id, p.p_key, p.p_val, b.b_id, b.b_key, b.b_val_| |_|_MergeScan [is_placeholder=false, remote_input=[_| -|_| Inner Join: p.p_key = b.b_key Filter: p.p_val + b.b_val > Int32(105)_| -|_|_Projection: p.p_id, p.p_key, p.p_val, p.p_pad, p.ts_| -|_|_SubqueryAlias: p_| -|_|_Filter: dist_join_probe.p_key IS NOT NULL AND dist_join_probe.p_id < Int32(200)_| -|_|_TableScan: dist_join_probe, partial_filters=[dist_join_probe.p_key IS NOT NULL, dist_join_probe.p_id < Int32(200)]_| +|_| Projection: p.p_id, p.p_key, p.p_val, p.p_pad, p.ts, b.b_id, b.b_key, b.b_val, b.ts_| +|_|_Inner Join: b.b_key = p.p_key Filter: p.p_val + b.b_val > Int32(105)_| |_|_Projection: b.b_id, b.b_key, b.b_val, b.ts_| |_|_MergeScan [is_placeholder=false, remote_input=[_| |_| SubqueryAlias: b_| |_|_Filter: dist_join_build.b_key IS NOT NULL_| |_|_TableScan: dist_join_build, partial_filters=[dist_join_build.b_key IS NOT NULL]_| |_| ]]_| +|_|_Projection: p.p_id, p.p_key, p.p_val, p.p_pad, p.ts_| +|_|_SubqueryAlias: p_| +|_|_Filter: dist_join_probe.p_key IS NOT NULL AND dist_join_probe.p_id < Int32(200)_| +|_|_TableScan: dist_join_probe, partial_filters=[dist_join_probe.p_key IS NOT NULL, dist_join_probe.p_id < Int32(200)]_| |_| ]]_| | physical_plan | SortPreservingMergeExec: [p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST]_| |_|_SortExec: expr=[p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST], preserve_partitioning=[true] | @@ -238,13 +240,14 @@ ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; |_|_TableScan: dist_join_build, partial_filters=[dist_join_build.b_key IS NOT NULL]_| |_| ]]_| | physical_plan | SortPreservingMergeExec: [p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST]_| -|_|_SortExec: expr=[p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST], preserve_partitioning=[true] | +|_|_ProjectionExec: expr=[p_id@3 as p_id, p_key@4 as p_key, p_val@5 as p_val, b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| +|_|_SortExec: expr=[p_val@5 ASC NULLS LAST, p_id@3 ASC NULLS LAST, b_val@2 ASC NULLS LAST, b_id@0 ASC NULLS LAST], preserve_partitioning=[true] | |_|_REDACTED -|_|_RepartitionExec: partitioning=REDACTED -|_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val]_| +|_|_CoalescePartitionsExec_| +|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| |_|_MergeScanExec: REDACTED |_|_RepartitionExec: partitioning=REDACTED -|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| +|_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val]_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ From dd675c226ced80d3017caf633f1ae4f4eddea1b5 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:23:02 +0800 Subject: [PATCH 24/37] fix(datanode): honor query handler factories and trim nested scan fixtures Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cache/src/lib.rs | 49 +- src/datanode/src/datanode.rs | 87 ++-- src/datanode/src/region_query.rs | 24 +- src/datanode/src/region_server.rs | 14 +- tests-integration/src/cluster.rs | 68 +-- .../nested_merge_scan_capability_test.rs | 470 ++---------------- 6 files changed, 141 insertions(+), 571 deletions(-) diff --git a/src/cache/src/lib.rs b/src/cache/src/lib.rs index ff9106cbbcbe..58ad4c14e5fc 100644 --- a/src/cache/src/lib.rs +++ b/src/cache/src/lib.rs @@ -55,19 +55,9 @@ pub const PARTITION_INFO_CACHE_NAME: &str = "partition_info_cache"; /// Builds the layered cache registry for datanode. /// -/// The registry holds every cache a datanode needs: -/// - Schema cache. -/// - Table id to schema name cache. -/// - Table info cache. -/// - Table name cache. -/// - Table cache. -/// - Table route cache. -/// - Partition info cache. -/// -/// The first layer holds the caches the datanode reads from the kv backend: the schema cache, the -/// table id to schema name cache, the table info cache, the table name cache and the table route -/// cache. The second layer holds the caches derived from them: the table cache (from the table info -/// and the table name caches) and the partition info cache (from the table route cache). +/// Layers separate caches derived from kv-backed metadata. The first layer holds the schema, table +/// id-to-schema-name, table info, table name and table route caches; the second holds the table cache +/// (derived from table info and name) and partition info cache (derived from table routes). /// /// [LayeredCacheRegistry] invalidates a layer only after the previous layer finished invalidating /// (see [LayeredCacheRegistryBuilder::add_cache_registry]), so the derived caches are invalidated @@ -263,34 +253,23 @@ mod tests { SchemaCacheRef, TableInfoCacheRef, TableNameCacheRef, TableRouteCacheRef, TableSchemaCacheRef, }; - use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use partition::cache::PartitionInfoCacheRef; use super::*; - fn memory_backend() -> KvBackendRef { - Arc::new(MemoryKvBackend::::new()) - } - - /// Asserts that the registry holds every cache the catalog manager and the partition rule - /// manager of a datanode look up. - macro_rules! assert_datanode_caches { - ($registry:expr) => { - assert!($registry.get::().is_some()); - assert!($registry.get::().is_some()); - assert!($registry.get::().is_some()); - assert!($registry.get::().is_some()); - assert!($registry.get::().is_some()); - assert!($registry.get::().is_some()); - assert!($registry.get::().is_some()); - }; - } - #[test] fn test_datanode_layered_cache_registry_holds_all_caches() { - let registry = build_datanode_layered_cache_registry(memory_backend()); - - assert_datanode_caches!(registry); + let registry = build_datanode_layered_cache_registry(Arc::new(MemoryKvBackend::< + common_meta::error::Error, + >::new())); + + assert!(registry.get::().is_some()); + assert!(registry.get::().is_some()); + assert!(registry.get::().is_some()); + assert!(registry.get::().is_some()); + assert!(registry.get::().is_some()); + assert!(registry.get::().is_some()); + assert!(registry.get::().is_some()); } } diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index 4fef618b90dd..e6f82602a4eb 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -19,7 +19,6 @@ use std::path::Path; use std::sync::Arc; use std::time::{Duration, Instant}; -use catalog::CatalogManagerRef; use catalog::information_schema::NoopInformationExtension; use catalog::kvbackend::KvBackendCatalogManagerBuilder; use client::client_manager::NodeClients; @@ -29,9 +28,7 @@ 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, TableRouteCacheRef, TableSchemaCacheRef, -}; +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; @@ -61,10 +58,9 @@ use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef}; use object_store::ObjectStore; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::util::normalize_dir; -use partition::cache::PartitionInfoCacheRef; -use partition::manager::PartitionRuleManager; use query::QueryEngineFactory; use query::dummy_catalog::TableProviderFactoryRef; +use query::region_query::RegionQueryHandlerFactoryRef; use servers::server::ServerHandlers; use snafu::{OptionExt, ResultExt, ensure}; use store_api::logstore::LogStore; @@ -79,9 +75,8 @@ use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use crate::error::{ self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu, DataFusionSnafu, DuplicateRegionEngineConfigSnafu, GetMetadataSnafu, - InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, MissingNodeIdSnafu, - ObjectStoreWalNotStandaloneSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, - ShutdownServerSnafu, StartServerSnafu, + InvalidObjectStoreWalConfigSnafu, MissingNodeIdSnafu, ObjectStoreWalNotStandaloneSnafu, + OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu, }; use crate::event_listener::{ NoopRegionServerEventListener, RegionServerEventListenerRef, RegionServerEventReceiver, @@ -506,40 +501,16 @@ impl DatanodeBuilder { ) -> Result { let opts: &DatanodeOptions = &self.opts; - // The datanode executes the plans of the region queries it receives. Such a plan may - // contain a `MergeScan` node: the dist planner of the datanode plans it into a - // `MergeScanExec`, which resolves the metadata of the target table from the catalog and the - // routes of its regions from the partition rule manager, then queries their leaders through - // the datanode client. - // - // Notes: - // - The `DistExtensionPlanner` is registered as long as both the partition rule manager and - // the region query handler are given (see `DfQueryPlanner::new`), it doesn't depend on - // `with_dist_planner`. - // - `with_dist_planner` stays `false` on purpose: the datanode only executes the plans it - // receives, it must not rewrite its own plans into distributed ones. - let catalog_manager: CatalogManagerRef = KvBackendCatalogManagerBuilder::new( + // 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(); - let table_route_cache: TableRouteCacheRef = - cache_registry.get().context(MissingCacheSnafu)?; - let partition_info_cache: PartitionInfoCacheRef = - cache_registry.get().context(MissingCacheSnafu)?; - let partition_manager = Arc::new(PartitionRuleManager::new( - self.kv_backend.clone(), - table_route_cache, - partition_info_cache, - )); - - // Clients to the other datanodes. They are only used when a plan contains a `MergeScan` - // node, i.e. they stay idle for the plans that the datanode executes locally. - // - // A region query is a streaming request, so it must not be cut off by the request timeout: - // same as the frontend to datanode client. + // MergeScan region requests stream until completion, so they have no client timeout. let mut channel_config = ChannelConfig { timeout: None, ..Default::default() @@ -550,7 +521,11 @@ impl DatanodeBuilder { } let node_manager: NodeManagerRef = Arc::new(NodeClients::new(channel_config)); let region_query_handler = - DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager); + if let Some(factory) = self.plugins.get::() { + 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. @@ -1122,6 +1097,7 @@ mod tests { 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}; @@ -1131,6 +1107,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; @@ -1141,6 +1119,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) { @@ -1493,6 +1472,38 @@ mod tests { ); } + struct CountingRegionQueryHandlerFactory(Arc); + + 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"); diff --git a/src/datanode/src/region_query.rs b/src/datanode/src/region_query.rs index 0327395d1647..1642362ee1ae 100644 --- a/src/datanode/src/region_query.rs +++ b/src/datanode/src/region_query.rs @@ -12,13 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! A [`RegionQueryHandler`] implementation for the datanode. -//! -//! A datanode executes the plans it receives through region queries. Such a plan may contain a -//! `MergeScan` node (see [`query::dist_plan::MergeScanLogicalPlan`]) that is planned into a -//! `MergeScanExec` by the dist planner of the datanode, which then has to query the regions of -//! that `MergeScan` from the datanodes that host them. This handler provides the datanode with the -//! client side of those region queries. +//! The datanode's [`RegionQueryHandler`] implementation for querying remote regions in a +//! `MergeScan` plan. use std::sync::Arc; @@ -35,13 +30,8 @@ use session::ReadPreference; use snafu::ResultExt; use store_api::storage::RegionId; -/// Serves the region queries issued by the datanode itself. -/// -/// Known limitation: [`RegionQueryHandler::select_target`] resolves a region leader through the -/// version-checked table route cache, whose version is shared by the `CacheContainer`. An -/// unrelated `TableId` invalidation can therefore add route loads and cold-load latency; sustained -/// conflicts may exhaust the retry budget and fail the region query. -pub struct DatanodeRegionQueryHandler { +/// Serves region queries by resolving their leader and forwarding requests to its datanode. +pub(crate) struct DatanodeRegionQueryHandler { partition_manager: PartitionRuleManagerRef, node_manager: NodeManagerRef, } @@ -67,8 +57,6 @@ impl RegionQueryHandler for DatanodeRegionQueryHandler { ) -> QueryResult { // The leader of a region is the only peer that can serve a region query. // - // This lookup goes through the version checked table route cache: see the known limitation - // of `DatanodeRegionQueryHandler` about retries of unrelated invalidations. let peer = self .partition_manager .find_region_leader(region_id) @@ -93,9 +81,7 @@ impl RegionQueryHandler for DatanodeRegionQueryHandler { .context(RegionQuerySnafu) } - /// The remote dynamic filter of a nested merge scan is not supported yet (PoC): a datanode only - /// dispatches region queries for the `MergeScan` nodes of the plans it receives, it doesn't - /// produce remote dynamic filters of its own. + /// Datanodes forward received region queries and do not produce remote dynamic filters. async fn handle_remote_dyn_filter_update( &self, _target: &RegionQueryTarget, diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 7cebd73f6a68..7d5218f30b75 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -290,9 +290,10 @@ impl RegionServer { request: api::v1::region::QueryRequest, query_ctx: QueryContextRef, ) -> Result { - let permit = match &self.inner.parallelism { - Some(p) => Some(p.acquire().await?), - None => None, + let permit = if let Some(p) = &self.inner.parallelism { + Some(p.acquire().await?) + } else { + None }; let region_id = RegionId::from_u64(request.region_id); @@ -347,9 +348,10 @@ impl RegionServer { } async fn handle_read_inner(&self, request: QueryRequest) -> Result { - let permit = match &self.inner.parallelism { - Some(p) => Some(p.acquire().await?), - None => None, + let permit = if let Some(p) = &self.inner.parallelism { + Some(p.acquire().await?) + } else { + None }; let ctx = request.header.as_ref().map(|h| h.into()); diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index a01aed240afc..9e26252c47c5 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -16,8 +16,8 @@ use std::collections::{BTreeSet, HashMap}; use std::env; use std::net::TcpListener; use std::ops::RangeInclusive; +use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use std::time::Duration; @@ -88,28 +88,20 @@ use crate::test_util::{ create_tmp_dir_and_datanode_opts, }; -/// The gRPC requests another node sent to a datanode over its network address (see +/// Flight `DoGet` requests another node sent to a datanode over its network address (see /// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]). /// -/// The frontend and the datanodes inside the test process talk to a datanode through an -/// in-process client, so a request counted here comes from a remote node, e.g. the region query -/// that a datanode issues for the regions of another datanode. +/// The frontend and datanodes inside the test process use in-process clients, so counted requests +/// come from remote nodes querying regions owned by this datanode. #[derive(Debug, Default)] pub struct DatanodeRpcStats { - requests: AtomicUsize, - paths: Mutex>, + do_get_requests: AtomicUsize, } impl DatanodeRpcStats { - /// The number of gRPC requests another node sent to the datanode. + /// The number of Flight `DoGet` requests another node sent to the datanode. pub fn requests(&self) -> usize { - self.requests.load(Ordering::Relaxed) - } - - /// The gRPC paths of those requests, e.g. - /// `/arrow.flight.protocol.FlightService/DoGet` for a region query. - pub fn paths(&self) -> BTreeSet { - self.paths.lock().unwrap().clone() + self.do_get_requests.load(Ordering::Relaxed) } } @@ -121,11 +113,6 @@ pub struct GreptimeDbCluster { /// The stats of the network address of every datanode, empty unless the cluster serves the /// datanodes at their own addresses. pub datanode_rpc_stats: HashMap>, - /// The layered cache registry of every datanode, i.e. the registry that the datanode - /// invalidates locally after it received an invalidation instruction from metasrv. A test can - /// invalidate a cache ident (e.g. an unrelated `TableId`) exactly the way the datanode does - /// when it handles the broadcast of metasrv. - pub datanode_cache_registries: HashMap>, pub kv_backend: KvBackendRef, pub metasrv: Arc, pub frontend: Arc, @@ -392,7 +379,7 @@ impl GreptimeDbClusterBuilder { ) .await; - let (datanode_instances, datanode_cache_registries) = self + let datanode_instances = self .build_datanodes_with_options(&metasrv, &datanode_options) .await; @@ -431,7 +418,6 @@ impl GreptimeDbClusterBuilder { guards, datanode_instances, datanode_rpc_stats, - datanode_cache_registries, kv_backend: self.kv_backend.clone(), metasrv: metasrv.metasrv, frontend: Arc::new(frontend), @@ -514,22 +500,16 @@ impl GreptimeDbClusterBuilder { &self, metasrv: &MockInfo, options: &[DatanodeOptions], - ) -> ( - HashMap, - HashMap>, - ) { + ) -> HashMap { let mut instances = HashMap::with_capacity(options.len()); - let mut cache_registries = HashMap::with_capacity(options.len()); for opts in options { - let (datanode, cache_registry) = - self.create_datanode(opts.clone(), metasrv.clone()).await; + let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await; let datanode_id = opts.node_id.unwrap(); instances.insert(datanode_id, datanode); - cache_registries.insert(datanode_id, cache_registry); } - (instances, cache_registries) + instances } async fn wait_datanodes_alive( @@ -555,11 +535,7 @@ impl GreptimeDbClusterBuilder { panic!("Some Datanodes are not alive in 10 seconds!") } - async fn create_datanode( - &self, - opts: DatanodeOptions, - metasrv: MockInfo, - ) -> (Datanode, Arc) { + async fn create_datanode(&self, opts: DatanodeOptions, metasrv: MockInfo) -> Datanode { let mut meta_client = MetaClientBuilder::datanode_default_options(opts.node_id.unwrap()) .channel_manager(metasrv.channel_manager) .build(); @@ -568,8 +544,6 @@ impl GreptimeDbClusterBuilder { let meta_backend = new_read_only_meta_kv_backend(meta_client.clone()); - // Use the same layered registry as the datanode of a running instance, so the derived - // caches of the test are invalidated after the caches they are derived from. let layered_cache_registry = Arc::new(build_datanode_layered_cache_registry(meta_backend.clone())); @@ -581,10 +555,7 @@ impl GreptimeDbClusterBuilder { datanode.start_heartbeat().await.unwrap(); - // The datanode invalidates this registry locally when it receives an invalidation - // instruction, so a test can drive the same path: see - // [`GreptimeDbCluster::datanode_cache_registries`]. - (datanode, layered_cache_registry) + datanode } async fn build_frontend( @@ -840,7 +811,7 @@ async fn create_datanode_client( ) } -/// Counts the gRPC requests that the wrapped server receives. +/// Counts Flight `DoGet` requests received by the wrapped server. #[derive(Clone)] struct CountRequests { inner: S, @@ -860,13 +831,10 @@ where } fn call(&mut self, request: http::Request) -> Self::Future { - if let Some(rpc_stats) = &self.rpc_stats { - rpc_stats.requests.fetch_add(1, Ordering::Relaxed); - rpc_stats - .paths - .lock() - .unwrap() - .insert(request.uri().path().to_string()); + if request.uri().path() == "/arrow.flight.protocol.FlightService/DoGet" + && let Some(rpc_stats) = &self.rpc_stats + { + rpc_stats.do_get_requests.fetch_add(1, Ordering::Relaxed); } self.inner.call(request) diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index bf0986391008..9be69a645a0d 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -12,56 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Capability test of the nested `MergeScan` execution of a datanode. -//! -//! A datanode executes the plans it receives. Such a plan may carry a nested plan in a `MergeScan` -//! node, i.e. a `MergeScan` that is not a placeholder, and the datanode then plans that node into -//! a `MergeScanExec` which queries the regions of the *inner* table from the datanodes that own -//! them (see `DatanodeRegionQueryHandler`). That is what lets a distributed sub-plan of a join run -//! on the datanode that owns a region of the probe table while the build side of the join is read -//! from the datanodes that own the regions of the build table. -//! -//! The plans below are built **by hand** and handed to a datanode by calling -//! `RegionServer::handle_remote_read` in process, so these tests exercise the execution capability -//! of the datanode only: no rewriter, hint or switch of this repository takes part in producing the -//! plan, and the outer request itself never goes over the wire. The region queries that the -//! datanode then issues for the regions of the inner table do: the tests build the cluster with -//! `GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`, so every datanode is reachable at the -//! address it registered and the inner `MergeScan` reaches the datanodes that own the regions of -//! the inner table over a real Flight connection. -//! -//! The plan of one probe region is `Join(, MergeScan())`: -//! -//! * the probe side is decoded with the region aware catalog list of the request -//! (`NameAwareCatalogList`), so it reads the region of the probe table that the request names, -//! * the `MergeScan` payload is decoded with the catalog of the query engine (see -//! `query::query_engine::default_serializer`), so its table scan refers to the *build table* of -//! the cluster and not to the region of the request. The two tables of the test share no column -//! name besides their join key on purpose: a payload that was bound to the region of the request -//! would decode with a schema that does not even match the payload. -//! -//! The checks are: -//! -//! * every region of the probe table is asked to execute the nested plan and the union of the -//! results is the same multiset of rows as the result of the equivalent join SQL on the frontend -//! (the build table has duplicate join keys, so a probe row joins several build rows: a lost row -//! and a row that was counted twice both fail the comparison), -//! * the datanode really queried the *other* datanode for the regions of the build table that it -//! owns (see [`assert_cross_datanode_region_query`]), -//! * a nested plan whose inner relation does not exist, and a nested plan with an inner region that -//! is not served any more, both make the whole query fail: the datanode never returns the rows of -//! the inner regions that happen to succeed, -//! * a cold load of the route of the build table returns the rows of the join after an unrelated -//! table id is invalidated once (see -//! [`test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation`]). +//! Exercises hand-built nested `MergeScan` plans at datanode execution, not optimizer behavior. +//! The outer `handle_remote_read` request is in-process; inner region reads use real TCP Flight. +//! The probe scan resolves against the request's region catalog, while the nested payload resolves +//! against the engine catalog; distinct column names ensure these scopes cannot be confused. +//! Success compares the complete duplicate-sensitive result multiset with frontend SQL and known +//! rows. Missing inner tables and closed remote inner regions must each fail the whole query. use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use api::v1::region::{QueryRequest, RegionRequestHeader}; -use common_meta::cache::TableRouteCacheRef; -use common_meta::cache_invalidator::{CacheInvalidator, Context}; -use common_meta::instruction::CacheIdent; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::info; @@ -75,30 +36,20 @@ use query::query_engine::DefaultSerializer; use servers::query_handler::sql::SqlQueryHandler; use session::context::{QueryContext, QueryContextRef}; use store_api::region_request::{RegionCloseRequest, RegionRequest}; -use store_api::storage::{RegionId, TableId}; +use store_api::storage::RegionId; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; use crate::test_util::execute_sql; -/// The probe table of the nested join: partitioned, so that its regions are spread over both -/// datanodes. +/// Partitioned probe table. const PROBE_TABLE: &str = "nested_cap_probe"; -/// The build table of the nested join: partitioned as well, and its join keys are duplicated. +/// Partitioned build table with duplicate join keys. const BUILD_TABLE: &str = "nested_cap_build"; -/// The build table of the error path: no such table exists, so the inner relation of the nested -/// plan cannot be resolved. +/// Missing inner table used to verify plan-resolution failure. const MISSING_INNER_TABLE: &str = "nested_cap_missing_build"; -/// The table id that no table of the cluster uses. The cold-load test invalidates this unrelated -/// table once before it executes nested queries for [`BUILD_TABLE`]. -const UNRELATED_TABLE_ID: TableId = 999_999; - -/// `(a_id, probe_key, probe_v)` of the probe table: `a_id` is the primary key and the partition -/// column, `probe_key` is the join key and `probe_v` is the value of the row. -/// -/// The probe table and the build table share no column name besides their join key on purpose, see -/// the module documentation. +/// `(a_id, probe_key, probe_v)` fixture rows; `a_id` is the partition key. const PROBE_ROWS: &[(i32, i32, i32)] = &[ (1, 10, 100), (2, 20, 200), @@ -107,8 +58,7 @@ const PROBE_ROWS: &[(i32, i32, i32)] = &[ (121, 10, 400), ]; -/// `(b_id, build_key, build_v)` of the build table: the join key `10` appears twice and `40` -/// matches no row of the probe table. +/// `(b_id, build_key, build_v)` fixture rows; key `10` duplicates and `40` is unmatched. const BUILD_ROWS: &[(i32, i32, i32)] = &[ (1, 10, 1), (2, 10, 2), @@ -118,18 +68,6 @@ const BUILD_ROWS: &[(i32, i32, i32)] = &[ (121, 40, 5), ]; -/// The datanodes that have to serve a region query of the nested plan of a probe region. -#[derive(Debug, Default)] -struct CrossDatanodeExpectation { - /// Datanode id -> the number of (probe region, build region) pairs that the nested plans have - /// to read from that datanode over the network. - remote_datanodes: BTreeMap, -} - -/// The capability of a datanode to execute a hand-built nested plan: every region of the probe -/// table is executed with -/// `Join(, MergeScan())`, and -/// the datanodes fetch the regions of the build table from the datanode that owns them. #[tokio::test(flavor = "multi_thread")] async fn test_nested_merge_scan_capability_join_across_datanodes() { common_telemetry::init_default_ut_logging(); @@ -141,11 +79,6 @@ async fn test_nested_merge_scan_capability_join_across_datanodes() { let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; - info!("region leaders of {PROBE_TABLE}: {probe_leaders:?}"); - info!("region leaders of {BUILD_TABLE}: {build_leaders:?}"); - // Every table is partitioned into two regions. Both tables span both datanodes, so a nested - // plan of a probe region always has to read a region of the build table from the other - // datanode over the network. assert_eq!( 2, datanodes(&probe_leaders).len(), @@ -159,149 +92,37 @@ async fn test_nested_merge_scan_capability_join_across_datanodes() { leaders: {build_leaders:?}" ); - // The reference: the frontend executes the same join with its default distributed plan, i.e. - // both sides of the join are read with a `MergeScan` and the join runs on the frontend. let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; - info!("join result of the frontend:\n{reference}"); assert_eq!( expected_join_rows(), multiset(table_cells(&reference)), "unexpected result of the join on the frontend:\n{reference}" ); - // The nested plan is the same for every probe region: the region of the request decides which - // rows of the probe table the plan reads, the inner `MergeScan` reads all regions of the build - // table. let query_ctx = query_ctx(); let nested_plan = nested_plan(&frontend, &query_ctx).await; let plan = encode(&nested_plan); - // The requests of the remote nodes are counted per datanode: a nested plan of a probe region - // that a datanode `d` owns has to query every datanode that owns a region of the build table - // *other* than `d`. - let expectation = cross_datanode_expectation(&probe_leaders, &build_leaders); - let baseline_requests = expectation - .remote_datanodes - .keys() + let remote_datanodes = datanodes(&build_leaders); + let baseline_requests = remote_datanodes + .iter() .map(|datanode| (*datanode, rpc_requests(&cluster, *datanode))) .collect::>(); - let mut rows = Vec::new(); - for (probe_region, (probe_datanode, _)) in &probe_leaders { - let stream = region_server(&cluster, *probe_datanode) - .handle_remote_read( - QueryRequest { - header: Some(region_request_header(&query_ctx)), - region_id: *probe_region, - plan: plan.clone(), - }, - query_ctx.clone(), - ) - .await - .unwrap_or_else(|e| { - panic!( - "the datanode {probe_datanode} must execute the nested plan of the region \ - {probe_region}: {e}" - ) - }); - let batches = RecordBatches::try_collect(stream) - .await - .unwrap_or_else(|e| { - panic!( - "the datanode {probe_datanode} must execute the nested plan of the region \ - {probe_region} completely: {e}" - ) - }); - let actual = batches - .pretty_print() - .expect("the result of the datanode must be printable"); - info!("join result of the nested plan of the region {probe_region}:\n{actual}"); - rows.extend(table_cells(&actual)); - } + let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; - // The nested plan of a probe region reads that region of the probe table, so the union of the - // results of all probe regions is the result of the whole join. The join key `10` appears twice - // in the build table, so the probe rows of that key appear twice as well: comparing the rows as - // a multiset catches both a lost row and a row that was counted twice. + let actual_rows = multiset(rows); assert_eq!( expected_join_rows(), - multiset(rows.clone()), - "the nested plans of the datanodes returned an unexpected multiset of rows: {rows:?}" + actual_rows, + "the nested plans of the datanodes returned an unexpected multiset of rows" ); assert_eq!( multiset(table_cells(&reference)), - multiset(rows), - "the nested plans of the datanodes returned a different multiset of rows than the \ - frontend:\nfrontend:\n{reference}" + actual_rows, + "frontend result differs:\n{reference}" ); - - assert_cross_datanode_region_query(&cluster, &expectation, &baseline_requests); -} - -/// The route of the build table is cold before the nested queries run. An unrelated table id is -/// invalidated once before those queries, and the nested plan must still return the join rows. -#[tokio::test(flavor = "multi_thread")] -async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation() { - common_telemetry::init_default_ut_logging(); - - let cluster = - build_cluster("test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation") - .await; - let frontend = cluster.fe_instance().clone(); - - prepare_tables(&frontend).await; - - let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; - let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; - info!("region leaders of {PROBE_TABLE}: {probe_leaders:?}"); - info!("region leaders of {BUILD_TABLE}: {build_leaders:?}"); - // Both datanodes own a probe region, so both of them resolve the leader of the regions of the - // build table when they execute the nested plan of their probe region. - assert_eq!( - 2, - datanodes(&probe_leaders).len(), - "expected the regions of {PROBE_TABLE} to be spread over both datanodes, actual region \ - leaders: {probe_leaders:?}" - ); - assert_eq!( - 2, - datanodes(&build_leaders).len(), - "expected the regions of {BUILD_TABLE} to be spread over both datanodes, actual region \ - leaders: {build_leaders:?}" - ); - - let build_table_id = table_id(&build_leaders); - - let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; - info!("join result of the frontend:\n{reference}"); - assert_eq!( - expected_join_rows(), - multiset(table_cells(&reference)), - "unexpected result of the join on the frontend:\n{reference}" - ); - - let query_ctx = query_ctx(); - let plan = encode(&nested_plan(&frontend, &query_ctx).await); - - invalidate_table_ids(&cluster, &[build_table_id]).await; - assert_route_cold(&cluster, build_table_id); - invalidate_table_ids(&cluster, &[UNRELATED_TABLE_ID]).await; - - let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; - assert_eq!( - expected_join_rows(), - multiset(rows), - "the nested plans returned an unexpected multiset of rows after the route of {BUILD_TABLE} \ - was dropped and unrelated table id {UNRELATED_TABLE_ID} was invalidated" - ); - - for datanode in datanodes(&probe_leaders) { - assert!( - table_route_cache(&cluster, datanode).contains_key(&build_table_id), - "the datanode {datanode} must hold the route of {BUILD_TABLE} after it executed the \ - nested plan of its probe region" - ); - } + assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests); } /// A failing inner region makes the whole query fail: the datanode must not return the rows of the @@ -318,11 +139,7 @@ async fn test_nested_merge_scan_capability_inner_failure_fails_query() { let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; - info!("region leaders of {PROBE_TABLE}: {probe_leaders:?}"); - info!("region leaders of {BUILD_TABLE}: {build_leaders:?}"); - // The request goes to a region of the probe table that the datanode `1` owns, and the nested - // plan reads the regions of the build table, so one of them comes from the other datanode. let (probe_region, (probe_datanode, _)) = probe_leaders .iter() .find(|(_, (datanode, _))| *datanode == 1) @@ -332,7 +149,7 @@ async fn test_nested_merge_scan_capability_inner_failure_fails_query() { {probe_leaders:?}" ) }); - let (build_region, (build_datanode, build_addr)) = build_leaders + let (build_region, (build_datanode, _)) = build_leaders .iter() .find(|(_, (datanode, _))| *datanode != *probe_datanode) .unwrap_or_else(|| { @@ -342,9 +159,7 @@ async fn test_nested_merge_scan_capability_inner_failure_fails_query() { ) }); - // 1. The inner relation of the nested plan does not exist. The payload of the `MergeScan` is - // resolved with the catalog of the datanode, so the datanode fails to decode the plan instead - // of binding the inner table scan to the region of the request. + // The nested payload resolves against the engine catalog, not the request's region catalog. let query_ctx = query_ctx(); let missing_inner_plan = nested_plan_with_inner_table(&frontend, &query_ctx, MISSING_INNER_TABLE).await; @@ -364,21 +179,14 @@ async fn test_nested_merge_scan_capability_inner_failure_fails_query() { // The top level error of the datanode only names the failing step (`Failed to decode logical // plan`): the reason is the source of it, which names the table that could not be resolved. let chain = error_chain(&error); - info!("the datanode rejected the nested plan with the missing inner relation: {chain}"); assert!( chain.contains(MISSING_INNER_TABLE), "expected the datanode to fail while resolving the inner relation {MISSING_INNER_TABLE}, \ actual error: {chain}" ); - // 2. An inner region of the nested plan fails: close the region of the build table that the - // other datanode owns. The nested plan of the probe region reads every region of the build - // table, so it has to query that datanode, which then cannot serve the region any more. The - // query of the probe region must fail as a whole: a partial result would hide the failure of - // the inner region, because the region of the build table that the datanode owns keeps serving - // its rows. - let baseline = rpc_requests(&cluster, *build_datanode); - let closed = region_server(&cluster, *build_datanode) + // A closed remote inner region must fail the whole query, not return a partial result. + region_server(&cluster, *build_datanode) .handle_request( RegionId::from_u64(*build_region), RegionRequest::Close(RegionCloseRequest { @@ -387,13 +195,10 @@ async fn test_nested_merge_scan_capability_inner_failure_fails_query() { ) .await .expect("the datanode must close the region of the build table"); - info!( - "closed the region {build_region} of {BUILD_TABLE} on the datanode {build_datanode} \ - ({closed:?}), which is reachable at {build_addr}" - ); let nested_plan = nested_plan(&frontend, &query_ctx).await; let plan = encode(&nested_plan); + let baseline = rpc_requests(&cluster, *build_datanode); let stream = region_server(&cluster, *probe_datanode) .handle_remote_read( QueryRequest { @@ -410,27 +215,11 @@ async fn test_nested_merge_scan_capability_inner_failure_fails_query() { .expect("the query must fail when one of its inner regions fails"); info!("the nested plan failed as a whole: {error}"); - // The failure comes from the region query of the *other* datanode: it reached that datanode - // over the network, i.e. the query failed while reading a region of the inner table, not while - // planning the local datanode. - let stats = cluster - .datanode_rpc_stats(*build_datanode) - .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {build_datanode}")); - info!( - "the datanode {build_datanode} served {} remote gRPC requests, paths: {:?}", - stats.requests(), - stats.paths() - ); - assert!( - stats.requests() > baseline, - "expected the failed nested plan to query the datanode {build_datanode} for the region \ - {build_region} it owns, but the requests of the datanode did not grow" - ); + // The increased DoGet count proves execution reached the remote region owner. + let actual = rpc_requests(&cluster, *build_datanode); assert!( - stats.paths().iter().any(|path| path.contains("DoGet")), - "expected the datanode {build_datanode} to serve the region query of the nested plan, \ - actual remote gRPC paths: {:?}", - stats.paths() + actual > baseline, + "expected the failed nested plan to reach datanode {build_datanode} for region {build_region} over Flight DoGet (count {actual}, baseline {baseline})" ); } @@ -447,81 +236,26 @@ async fn build_cluster(test_name: &str) -> GreptimeDbCluster { .await } -/// The `DoGet` requests that the nested plans of the probe regions have to send to the datanodes -/// that own the regions of the build table. -fn cross_datanode_expectation( - probe_leaders: &BTreeMap, - build_leaders: &BTreeMap, -) -> CrossDatanodeExpectation { - let mut expectation = CrossDatanodeExpectation::default(); - for (probe_region, (probe_datanode, probe_addr)) in probe_leaders { - for (build_region, (build_datanode, build_addr)) in build_leaders { - if probe_datanode == build_datanode { - continue; - } - *expectation - .remote_datanodes - .entry(*build_datanode) - .or_default() += 1; - info!( - "the datanode {probe_datanode} ({probe_addr}) executing the region {probe_region} \ - of {PROBE_TABLE} queries the region {build_region} of {BUILD_TABLE} from the \ - datanode {build_datanode} ({build_addr})" - ); - } - } - expectation -} - -/// Asserts that the datanodes issued a real region query to every datanode that owns a region of -/// the build table. -/// -/// The cluster serves every datanode at its own address, while the frontend and the local -/// bookkeeping of the test process reach a datanode through an in-process client, so a request -/// recorded at the address of a datanode comes from a remote node: a request of the `DoGet` (a -/// region query) or of the region service proves the datanode-to-datanode region query. fn assert_cross_datanode_region_query( cluster: &GreptimeDbCluster, - expectation: &CrossDatanodeExpectation, + remote_datanodes: &BTreeSet, baseline_requests: &BTreeMap, ) { - assert!( - !expectation.remote_datanodes.is_empty(), - "expected the nested plans of the probe regions to read a region of the build table from \ - another datanode, but the regions of {BUILD_TABLE} are not spread over both datanodes" + assert_eq!( + 2, + remote_datanodes.len(), + "build table must span both datanodes" ); - - for (datanode, probe_regions) in &expectation.remote_datanodes { - let stats = cluster - .datanode_rpc_stats(*datanode) - .unwrap_or_else(|| panic!("expected the rpc stats of the datanode {datanode}")); - let baseline = baseline_requests.get(datanode).copied().unwrap_or_default(); - info!( - "the datanode {datanode} served {} remote gRPC requests (baseline {baseline}), \ - paths: {:?}", - stats.requests(), - stats.paths() - ); + for datanode in remote_datanodes { + let actual = rpc_requests(cluster, *datanode); + let baseline = baseline_requests[datanode]; assert!( - stats.requests() > baseline, - "expected the nested plans of {PROBE_TABLE} to query the datanode {datanode} \ - {probe_regions} times for the regions of {BUILD_TABLE} it owns, but the datanode \ - served no additional remote request ({} requests, baseline {baseline})", - stats.requests() - ); - assert!( - stats - .paths() - .iter() - .any(|path| path.contains("DoGet") || path.contains("Region")), - "expected the datanode {datanode} to serve the region query of the nested plan, actual \ - remote gRPC paths: {:?}", - stats.paths() + actual > baseline, + "expected a nested Flight DoGet at datanode {datanode}, count {actual}, baseline {baseline}" ); } } -/// The number of gRPC requests that the remote nodes sent to the datanode so far. fn rpc_requests(cluster: &GreptimeDbCluster, datanode_id: u64) -> usize { cluster .datanode_rpc_stats(datanode_id) @@ -529,7 +263,6 @@ fn rpc_requests(cluster: &GreptimeDbCluster, datanode_id: u64) -> usize { .requests() } -/// Returns the region server of the datanode `datanode_id`. fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer { cluster .datanode_instances @@ -538,8 +271,6 @@ fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer .region_server() } -/// Executes the nested plan of every probe region on the datanode that owns the region, consumes -/// every stream completely and returns the union of the rows of the results. async fn run_nested_queries( cluster: &GreptimeDbCluster, probe_leaders: &BTreeMap, @@ -575,71 +306,12 @@ async fn run_nested_queries( let actual = batches .pretty_print() .expect("the result of the datanode must be printable"); - info!("join result of the nested plan of the region {probe_region}:\n{actual}"); rows.extend(table_cells(&actual)); } rows } -/// The id of the single table that the regions of `leaders` belong to. -fn table_id(leaders: &BTreeMap) -> TableId { - let table_ids = leaders - .keys() - .map(|region_id| RegionId::from_u64(*region_id).table_id()) - .collect::>(); - assert_eq!( - 1, - table_ids.len(), - "expected the regions to belong to a single table, actual region leaders: {leaders:?}" - ); - - *table_ids.first().expect("there is at least one region") -} - -/// Returns the table route cache of the datanode `datanode_id`: the cache that -/// `DatanodeRegionQueryHandler::select_target` reads to resolve the leader of a region. -fn table_route_cache(cluster: &GreptimeDbCluster, datanode_id: u64) -> TableRouteCacheRef { - cluster - .datanode_cache_registries - .get(&datanode_id) - .unwrap_or_else(|| panic!("expected the cache registry of the datanode {datanode_id}")) - .get::() - .unwrap_or_else(|| panic!("expected the table route cache of the datanode {datanode_id}")) -} - -/// Invalidates `table_ids` on every datanode, the way a datanode handles the invalidation -/// instruction of metasrv (see [`GreptimeDbCluster::datanode_cache_registries`]). -async fn invalidate_table_ids(cluster: &GreptimeDbCluster, table_ids: &[TableId]) { - let idents = table_ids - .iter() - .copied() - .map(CacheIdent::TableId) - .collect::>(); - for (datanode_id, registry) in &cluster.datanode_cache_registries { - registry - .invalidate(&Context::default(), &idents) - .await - .unwrap_or_else(|e| { - panic!("the datanode {datanode_id} must invalidate {idents:?}: {e}") - }); - } -} - -/// Asserts that no datanode holds the route of `table_id`, i.e. the route is cold. -fn assert_route_cold(cluster: &GreptimeDbCluster, table_id: TableId) { - let cached = cluster - .datanode_cache_registries - .keys() - .filter(|datanode_id| table_route_cache(cluster, **datanode_id).contains_key(&table_id)) - .collect::>(); - assert!( - cached.is_empty(), - "the datanodes {cached:?} must drop the route of the table {table_id} after the \ - invalidation" - ); -} - /// Creates the probe and build tables of the tests and inserts the rows. async fn prepare_tables(frontend: &Arc) { execute_sql( @@ -713,7 +385,6 @@ async fn prepare_tables(frontend: &Arc) { .await; } -/// Encodes a plan the way a frontend sends it to a datanode. fn encode(plan: &LogicalPlan) -> Vec { DFLogicalSubstraitConvertor .encode(plan, DefaultSerializer) @@ -721,41 +392,11 @@ fn encode(plan: &LogicalPlan) -> Vec { .to_vec() } -/// Builds the nested plan of the tests by hand: -/// `Join(, MergeScan())`. -/// -/// The plan is the remote input of a `MergeScan` that a frontend sends to the datanode that owns a -/// region of the probe table: the probe side is decoded with the region aware catalog list of the -/// request, so it reads the region of the request, and the `MergeScan` is not a placeholder, so the -/// datanode plans it into a `MergeScanExec` over the regions of the build table. -/// -/// The table scans are planned by the frontend (a `SELECT` of the columns of the join), which gives -/// them the schema of the tables of the cluster. The join and the nested `MergeScan` around the -/// build side are assembled here: no rewrite of the frontend is involved. +/// Plans the hand-built nested join sent in the region request. async fn nested_plan(frontend: &Arc, query_ctx: &QueryContextRef) -> LogicalPlan { - let probe = plan_sql( - frontend, - &format!("SELECT a_id, probe_key, probe_v FROM {PROBE_TABLE}"), - query_ctx, - ) - .await; - let build = plan_sql( - frontend, - &format!("SELECT build_key, build_v FROM {BUILD_TABLE}"), - query_ctx, - ) - .await; - - nested_plan_of(&probe, build) + nested_plan_with_inner_table(frontend, query_ctx, BUILD_TABLE).await } -/// Same as [`nested_plan`], but the inner `MergeScan` scans `inner_table` instead of -/// [`BUILD_TABLE`]. -/// -/// The build side of the nested plan is the build table with the name of its scan replaced: the -/// scan keeps the schema and the source of the build table, so the payload of the `MergeScan` looks -/// exactly like the payload of [`nested_plan`] to the encoder, while the table it names does not -/// exist. async fn nested_plan_with_inner_table( frontend: &Arc, query_ctx: &QueryContextRef, @@ -778,7 +419,6 @@ async fn nested_plan_with_inner_table( nested_plan_of(&probe, build) } -/// Assembles `Join(probe, MergeScan(build))` and logs both sides. fn nested_plan_of(probe: &LogicalPlan, build: LogicalPlan) -> LogicalPlan { info!("probe side of the nested plan:\n{probe}"); info!("build side of the nested plan:\n{build}"); @@ -799,8 +439,6 @@ fn nested_plan_of(probe: &LogicalPlan, build: LogicalPlan) -> LogicalPlan { plan } -/// Returns `plan` with every table scan renamed to `table_name`, keeping the schema and the source -/// of the scan. fn rename_scans(plan: LogicalPlan, table_name: &str) -> LogicalPlan { use datafusion::common::TableReference; use datafusion::common::tree_node::{Transformed, TreeNode}; @@ -820,7 +458,6 @@ fn rename_scans(plan: LogicalPlan, table_name: &str) -> LogicalPlan { .data } -/// Plans `sql` on the frontend and returns the logical plan. async fn plan_sql(frontend: &Arc, sql: &str, query_ctx: &QueryContextRef) -> LogicalPlan { let stmt = QueryLanguageParser::parse_sql(sql, query_ctx) .unwrap_or_else(|e| panic!("failed to parse `{sql}`: {e}")); @@ -831,8 +468,7 @@ async fn plan_sql(frontend: &Arc, sql: &str, query_ctx: &QueryContextR .unwrap_or_else(|e| panic!("failed to plan `{sql}`: {e}")) } -/// The join of the tests, over the columns of the nested plan of [`nested_plan`] in order: -/// `(a_id, probe_key, probe_v, build_key, build_v)`. +/// Frontend reference join over the nested plan's output columns. fn join_sql() -> String { format!( "SELECT p.a_id, p.probe_key, p.probe_v, b.build_key, b.build_v @@ -840,7 +476,7 @@ fn join_sql() -> String { ) } -/// The rows of [`join_sql`] as a multiset: every probe row joins every build row of its join key. +/// Expected complete join result, retaining duplicate rows. fn expected_join_rows() -> Vec> { let mut rows = Vec::new(); for (a_id, probe_key, probe_v) in PROBE_ROWS { @@ -859,7 +495,6 @@ fn expected_join_rows() -> Vec> { multiset(rows) } -/// Returns the region leaders of `table`: region id -> (datanode id, datanode address). async fn region_leaders(frontend: &Arc, table: &str) -> BTreeMap { use datatypes::arrow::array::AsArray; use datatypes::arrow::datatypes::UInt64Type; @@ -897,12 +532,10 @@ async fn region_leaders(frontend: &Arc, table: &str) -> BTreeMap) -> BTreeSet { leaders.values().map(|(id, _)| *id).collect() } -/// Collects the record batches of a query output. async fn output_batches(output: Output) -> RecordBatches { match output.data { common_query::OutputData::Stream(stream) => RecordBatches::try_collect(stream) @@ -915,7 +548,7 @@ async fn output_batches(output: Output) -> RecordBatches { } } -/// Consumes `stream` and returns the error it failed with, or `None` if it returned completely. +/// Consumes the full stream to detect errors after any partial rows. async fn collect_error(stream: SendableRecordBatchStream) -> Option { use futures::TryStreamExt; @@ -932,8 +565,6 @@ async fn collect_error(stream: SendableRecordBatchStream) -> Option { } } -/// Concatenates the message of every error of the source chain of `error`, so that the failure of -/// an outer step (e.g. decoding a plan) does not hide its reason (e.g. an unknown table). fn error_chain(error: &(dyn std::error::Error + 'static)) -> String { let mut messages = vec![error.to_string()]; let mut source = error.source(); @@ -944,7 +575,6 @@ fn error_chain(error: &(dyn std::error::Error + 'static)) -> String { messages.join(": ") } -/// Runs `sql` on the frontend and returns the pretty printed result. async fn query_pretty(frontend: &Arc, sql: &str, query_ctx: QueryContextRef) -> String { run_sql(frontend, sql, query_ctx) .await @@ -960,17 +590,13 @@ async fn run_sql(frontend: &Arc, sql: &str, query_ctx: QueryContextRef .unwrap() } -/// The query context of a new session, with a single target partition so that the plans of the -/// tests are stable. The datanode builds the query context of its own execution from the header of -/// the request, so the request of the tests carries the same query context. +/// Uses one target partition to keep the plans stable. fn query_ctx() -> QueryContextRef { let mut query_ctx = QueryContext::with_db_name(None); query_ctx.set_extension(QUERY_PARALLELISM_HINT, "1"); Arc::new(query_ctx) } -/// The header of the region query of the tests, carrying the query context of the session of the -/// plan. fn region_request_header(query_ctx: &QueryContextRef) -> RegionRequestHeader { RegionRequestHeader { query_context: Some(query_ctx.as_ref().into()), @@ -978,7 +604,6 @@ fn region_request_header(query_ctx: &QueryContextRef) -> RegionRequestHeader { } } -/// Parses the cells of the data rows of a pretty printed table. fn table_cells(pretty: &str) -> Vec> { let mut lines = pretty.lines().filter(|line| line.starts_with('|')); // The first `|` line is the header of the table. @@ -994,7 +619,6 @@ fn table_cells(pretty: &str) -> Vec> { .collect() } -/// Sorts the rows of a result, so that two results can be compared as multisets. fn multiset(mut rows: Vec>) -> Vec> { rows.sort(); rows From b1880fce06e4b675388a0cb2ae7552dab45b6d81 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:42:23 +0800 Subject: [PATCH 25/37] fix(datanode): retain required cache error import Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datanode/src/datanode.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index e6f82602a4eb..a4e283350687 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -75,8 +75,9 @@ use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig}; use crate::error::{ self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu, DataFusionSnafu, DuplicateRegionEngineConfigSnafu, GetMetadataSnafu, - InvalidObjectStoreWalConfigSnafu, MissingNodeIdSnafu, ObjectStoreWalNotStandaloneSnafu, - OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, ShutdownServerSnafu, StartServerSnafu, + InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, MissingNodeIdSnafu, + ObjectStoreWalNotStandaloneSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu, + ShutdownServerSnafu, StartServerSnafu, }; use crate::event_listener::{ NoopRegionServerEventListener, RegionServerEventListenerRef, RegionServerEventReceiver, From 85c5d6d425965ff67e6205c51ec5cabf9987dd35 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:42:23 +0800 Subject: [PATCH 26/37] refactor(query): reuse registered state for nested scan decoding Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/query_engine/default_serializer.rs | 545 +++++++++++------- 1 file changed, 329 insertions(+), 216 deletions(-) diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 31b9e59be6f5..6e48adf34cd7 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -107,10 +107,9 @@ impl SerializerRegistry for DefaultSerializer { /// Async `MergeScan` payload decoding over DataFusion's default Substrait consumer. struct MergeScanSubstraitConsumer<'a> { inner: DefaultSubstraitConsumer<'a>, - /// State of the level being decoded; payload states are rebuilt from it. session_state: &'a SessionState, - /// Engine catalog for payload tables; absent for manually built session states. catalog_manager: Option, + payload: bool, } impl MergeScanSubstraitConsumer<'_> { @@ -121,43 +120,8 @@ impl MergeScanSubstraitConsumer<'_> { .unwrap_or_else(QueryContext::arc) } - /// Rebuilds payload state with request defaults and re-registers Greptime function bindings. - fn payload_state(&self, catalog_list: Arc) -> Result { - let query_ctx = self.query_ctx(); - let mut config = self.session_state.config().clone(); - { - let catalog_options = &mut config.options_mut().catalog; - catalog_options.default_catalog = query_ctx.current_catalog().to_string(); - catalog_options.default_schema = query_ctx.current_schema(); - } - - let mut state = SessionStateBuilder::new_from_existing(self.session_state.clone()) - .with_config(config) - .with_catalog_list(catalog_list) - .build(); - register_greptime_functions(&mut state, &query_ctx).map_err(DataFusionError::from)?; - - Ok(state) - } - - /// Decodes payload tables through the engine catalog, not the region-bound request catalog. - /// - /// Do not fall back after an engine-catalog failure: it can bind payload tables to the request - /// region with the wrong schema. Manually built states lack that catalog and use the request list. async fn decode_payload(&self, payload: Vec) -> Result { - let state = if let Some(catalog_manager) = &self.catalog_manager { - let engine_catalog = Arc::new( - catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( - catalog_manager.clone(), - self.query_ctx(), - ), - ); - self.payload_state(engine_catalog)? - } else { - self.payload_state(self.session_state.catalog_list().clone())? - }; - - decode_plan(Bytes::from(payload), state) + decode_plan(Bytes::from(payload), self.session_state, true) .await .map_err(|e| DataFusionError::External(Box::new(e))) } @@ -169,7 +133,50 @@ impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> { &self, table_ref: &TableReference, ) -> Result>> { - self.inner.resolve_table_ref(table_ref).await + if !self.payload { + return self.inner.resolve_table_ref(table_ref).await; + } + + let query_ctx = self.query_ctx(); + let resolved = table_ref + .clone() + .resolve(query_ctx.current_catalog(), &query_ctx.current_schema()); + // Manual states retain DataFusion's own catalog and information-schema behavior. + let Some(manager) = &self.catalog_manager else { + return self + .inner + .resolve_table_ref(&TableReference::from(resolved)) + .await; + }; + let catalog_list: Arc = Arc::new( + catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( + manager.clone(), + query_ctx, + ), + ); + let schema_name = resolved.schema.as_ref(); + let catalog_name = resolved.catalog.as_ref(); + let schema = if self.session_state.config().information_schema() + && schema_name == datafusion::catalog::information_schema::INFORMATION_SCHEMA + { + Arc::new( + datafusion::catalog::information_schema::InformationSchemaProvider::new( + catalog_list.clone(), + ) + .with_table_functions(self.session_state.table_functions().clone()), + ) as Arc + } else { + catalog_list + .catalog(catalog_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("failed to resolve catalog: {catalog_name}")) + })? + .schema(schema_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("failed to resolve schema: {schema_name}")) + })? + }; + schema.table(&resolved.table).await } fn get_extensions(&self) -> &Extensions { @@ -270,7 +277,11 @@ impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> { } /// Decode with async MergeScan handling, preserving the default consumer's extension checks. -async fn decode_plan(bytes: Bytes, state: SessionState) -> substrait::error::Result { +async fn decode_plan( + bytes: Bytes, + state: &SessionState, + payload: bool, +) -> substrait::error::Result { let plan = Plan::decode(bytes).context(DecodeRelSnafu)?; let extensions = Extensions::try_from(&plan.extensions).context(DecodeDfPlanSnafu)?; if !extensions.type_variations.is_empty() { @@ -278,19 +289,19 @@ async fn decode_plan(bytes: Bytes, state: SessionState) -> substrait::error::Res .context(DecodeDfPlanSnafu); } - // Payloads need the engine catalog rather than the request's region-bound catalog. let catalog_manager = state .config() .get_extension::() .map(|engine_state| engine_state.catalog_manager().clone()); let consumer = MergeScanSubstraitConsumer { - inner: DefaultSubstraitConsumer::new(&extensions, &state), - session_state: &state, + inner: DefaultSubstraitConsumer::new(&extensions, state), + session_state: state, catalog_manager, + payload, }; DFLogicalSubstraitConvertor - .decode_with_consumer(&plan, &state, &consumer) + .decode_with_consumer(&plan, state, &consumer) .await } @@ -312,10 +323,10 @@ impl DefaultPlanDecoder { } } -/// Re-registers Greptime functions after every state build. +/// Registers Greptime functions after building the request decoder state. /// /// DataFusion aliases can rebind Greptime `date_format` to built-in `to_char`; registration must -/// happen after `SessionStateBuilder::build` on the state that decodes the plan. +/// happen after `SessionStateBuilder::build` on the state that decodes the request plan. fn register_greptime_functions( session_state: &mut SessionState, query_ctx: &QueryContextRef, @@ -429,7 +440,7 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { // Re-register after the build to avoid Greptime UDF alias collisions. register_greptime_functions(&mut session_state, &self.query_ctx)?; - let logical_plan = decode_plan(message, session_state) + let logical_plan = decode_plan(message, &session_state, false) .await .map_err(BoxedError::new) .context(common_query::error::DecodePlanSnafu)?; @@ -986,57 +997,54 @@ mod tests { ); } - /// Rebuilding state must retain the Greptime `date_format` binding. + /// Request state registration keeps Greptime `date_format` ahead of DataFusion aliases. #[tokio::test] async fn test_serializer_decode_keeps_greptime_date_format() { - // Repeat to expose nondeterministic function-registration order. - for _ in 0..8 { - let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); - let factory = QueryEngineFactory::new( - catalog_list, - None, - None, - None, - None, - false, - QueryOptions::default(), - ); - let engine = factory.query_engine(); + let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); + let factory = QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); - let table_provider = Arc::new(mock_table_provider(1.into())); - let query_ctx = date_format_query_ctx(); - let plan = LogicalPlanBuilder::scan( - "devices", - Arc::new(LogicalTableSource::new(table_provider.schema().clone())), - None, + let table_provider = Arc::new(mock_table_provider(1.into())); + let query_ctx = date_format_query_ctx(); + let plan = LogicalPlanBuilder::scan( + "devices", + Arc::new(LogicalTableSource::new(table_provider.schema().clone())), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))]) + .unwrap() + .build() + .unwrap(); + + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .new_plan_decoder() + .unwrap(); + let decoded = plan_decoder + .decode( + bytes, + Arc::new(DummyCatalogList::with_table_provider(table_provider)), + false, ) - .unwrap() - .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))]) - .unwrap() - .build() + .await .unwrap(); - let bytes = DFLogicalSubstraitConvertor - .encode(&plan, DefaultSerializer) - .unwrap(); - let plan_decoder = engine - .engine_context(query_ctx.clone()) - .new_plan_decoder() - .unwrap(); - let decoded = plan_decoder - .decode( - bytes, - Arc::new(DummyCatalogList::with_table_provider(table_provider)), - false, - ) - .await - .unwrap(); - - assert_only_function_is_greptime_date_format(&decoded, &query_ctx); - } + assert_only_function_is_greptime_date_format(&decoded, &query_ctx); } - /// Payload-state rebuilding must retain the Greptime `date_format` binding. + /// Nested payload decoding retains the Greptime `date_format` binding. #[tokio::test] async fn test_serializer_decode_merge_scan_payload_keeps_greptime_date_format() { let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); @@ -1050,67 +1058,64 @@ mod tests { }) .unwrap(); - for _ in 0..8 { - let factory = QueryEngineFactory::new( - catalog_manager.clone(), - None, - None, - None, - None, - false, - QueryOptions::default(), - ); - let engine = factory.query_engine(); - let query_ctx = date_format_query_ctx(); + let factory = QueryEngineFactory::new( + catalog_manager.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + let query_ctx = date_format_query_ctx(); - let input = LogicalPlanBuilder::scan( - NUMBERS_TABLE_NAME, - Arc::new(LogicalTableSource::new( - NumbersTable::schema().arrow_schema().clone(), - )), - None, - ) - .unwrap() - .project(vec![greptime_date_format_expr( - &query_ctx, - Expr::Cast(Cast::new( - Box::new(col("number")), - ArrowDataType::Timestamp(TimeUnit::Millisecond, None), - )), - )]) - .unwrap() - .build() - .unwrap(); - let plan = - MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); + let input = LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr( + &query_ctx, + Expr::Cast(Cast::new( + Box::new(col("number")), + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + )), + )]) + .unwrap() + .build() + .unwrap(); + let plan = MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); - let bytes = DFLogicalSubstraitConvertor - .encode(&plan, DefaultSerializer) - .unwrap(); - let plan_decoder = engine - .engine_context(query_ctx.clone()) - .new_plan_decoder() - .unwrap(); - let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( - mock_table_provider(1.into()), - ))); + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .new_plan_decoder() + .unwrap(); + let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))); - let decoded = plan_decoder - .decode(bytes, catalog_list, false) - .await - .unwrap(); + let decoded = plan_decoder + .decode(bytes, catalog_list, false) + .await + .unwrap(); - let LogicalPlan::Extension(extension) = &decoded else { - panic!("Expect a MergeScan plan, got: {decoded}"); - }; - let merge_scan = extension - .node - .as_any() - .downcast_ref::() - .expect("Expect a MergeScan plan node"); + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); - assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx); - } + assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx); } /// Registers a `numbers` table in a memory catalog manager. @@ -1159,17 +1164,119 @@ mod tests { .unwrap() } - /// Request catalog that resolves any table to a provider the payload must not use. + /// Request catalog that resolves any table to a provider payload resolution must not use. fn request_catalog_list() -> Arc { Arc::new(DummyCatalogList::with_table_provider(Arc::new( mock_table_provider(1.into()), ))) } + /// Payload resolution preserves context and explicit catalog/schema components. + #[tokio::test] + async fn test_serializer_decode_payload_observes_catalog_and_query_context() { + let observed = Arc::new(std::sync::Mutex::new(Vec::new())); + let inner = catalog::memory::new_memory_catalog_manager().unwrap(); + inner.register_catalog_sync("request_cat").unwrap(); + inner.register_catalog_sync("explicit_catalog").unwrap(); + for (catalog, schema) in [ + ("request_cat", "request_schema"), + ("request_cat", "explicit_schema"), + ("explicit_catalog", "full_schema"), + ] { + inner + .register_schema_sync(catalog::RegisterSchemaRequest { + catalog: catalog.to_string(), + schema: schema.to_string(), + }) + .unwrap(); + inner + .register_table_sync(RegisterTableRequest { + catalog: catalog.to_string(), + schema: schema.to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + } + let manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager { + inner, + resolution: TableResolution::ObserveContext(observed.clone()), + }); + let mut resolver_ctx = QueryContext::arc(); + Arc::make_mut(&mut resolver_ctx).set_current_catalog("request_cat"); + Arc::make_mut(&mut resolver_ctx).set_current_schema("request_schema"); + let engine = QueryEngineFactory::new( + manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let mut state = engine.engine_context(resolver_ctx.clone()).state().clone(); + state.config_mut().set_extension(resolver_ctx); + let extensions = Extensions::default(); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, &state), + session_state: &state, + catalog_manager: state + .config() + .get_extension::() + .map(|s| s.catalog_manager().clone()), + payload: true, + }; + for reference in [ + TableReference::bare(NUMBERS_TABLE_NAME), + TableReference::partial("explicit_schema", NUMBERS_TABLE_NAME), + TableReference::full( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + NUMBERS_TABLE_NAME, + ), + ] { + consumer + .resolve_table_ref(&reference) + .await + .unwrap() + .unwrap(); + } + assert_eq!( + *observed.lock().unwrap(), + vec![ + ( + "request_cat".into(), + "request_schema".into(), + NUMBERS_TABLE_NAME.into(), + "request_cat".into(), + "request_schema".into() + ), + ( + "request_cat".into(), + "explicit_schema".into(), + NUMBERS_TABLE_NAME.into(), + "request_cat".into(), + "request_schema".into() + ), + ( + DEFAULT_CATALOG_NAME.into(), + DEFAULT_SCHEMA_NAME.into(), + NUMBERS_TABLE_NAME.into(), + "request_cat".into(), + "request_schema".into() + ), + ], + ); + } + /// How [`InterceptingCatalogManager`] resolves payload tables. enum TableResolution { /// Fails with the frontend's cross-catalog access error. AccessDenied, + /// Records catalog lookup and request context values. + ObserveContext(Arc>>), /// Resolves once another task on the caller runtime signals. AwaitSignal { notify: Arc, @@ -1248,6 +1355,17 @@ mod tests { query_ctx: Option<&QueryContext>, ) -> CatalogResult> { match &self.resolution { + TableResolution::ObserveContext(observed) => { + let query_ctx = query_ctx.unwrap(); + observed.lock().unwrap().push(( + catalog.to_string(), + schema.to_string(), + table_name.to_string(), + query_ctx.current_catalog().to_string(), + query_ctx.current_schema(), + )); + Ok(Some(NumbersTable::table(NUMBERS_TABLE_ID))) + } TableResolution::AccessDenied => QueryAccessDeniedSnafu { catalog: catalog.to_string(), schema: schema.to_string(), @@ -1328,68 +1446,6 @@ mod tests { .query_engine() } - /// Decoding must not need a Tokio runtime: the public path is only `async`. - #[test] - fn test_serializer_decode_nested_merge_scan_without_tokio_runtime() { - let engine = QueryEngineFactory::new( - numbers_catalog_manager(), - None, - None, - None, - None, - false, - QueryOptions::default(), - ) - .query_engine(); - let plan = merge_scan(merge_scan(numbers_scan(), false), true); - let bytes = encode_plan(&plan); - - let decoded = futures::executor::block_on(plan_decoder(engine.as_ref()).decode( - bytes, - request_catalog_list(), - false, - )) - .unwrap(); - - assert_eq!(decoded.to_string(), plan.to_string()); - assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); - } - - /// Payload decoding must not block a multi-thread runtime either. - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn test_serializer_decode_deeply_nested_merge_scan_without_blocking() { - let engine = QueryEngineFactory::new( - numbers_catalog_manager(), - None, - None, - None, - None, - false, - QueryOptions::default(), - ) - .query_engine(); - let plan = merge_scan( - merge_scan( - LogicalPlanBuilder::from(merge_scan(numbers_scan(), false)) - .filter(col("number").lt(lit(10u32))) - .unwrap() - .build() - .unwrap(), - false, - ), - false, - ); - let bytes = encode_plan(&plan); - - let decoded = plan_decoder(engine.as_ref()) - .decode(bytes, request_catalog_list(), false) - .await - .unwrap(); - - assert_eq!(decoded.to_string(), plan.to_string()); - assert_eq!(decoded.to_string().matches("MergeScan [").count(), 3); - } - /// A payload lookup parked on a caller-runtime task must keep the decode async. /// /// The engine catalog resolves only once a task spawned on this runtime signals it, so a @@ -1495,11 +1551,36 @@ mod tests { ); } - /// Manually built states need ordinary extension decoding without an engine catalog. + /// Manual states use their own catalog and query-context defaults for payloads. #[tokio::test] async fn test_serializer_decode_manual_state_keeps_extension_registry() { - let state = SessionStateBuilder::new().with_default_features().build(); - let decoder = DefaultPlanDecoder::new(state, &QueryContext::arc()).unwrap(); + let manager = catalog::memory::new_memory_catalog_manager().unwrap(); + manager.register_catalog_sync("manual_cat").unwrap(); + manager + .register_schema_sync(catalog::RegisterSchemaRequest { + catalog: "manual_cat".to_string(), + schema: "manual_schema".to_string(), + }) + .unwrap(); + manager + .register_table_sync(RegisterTableRequest { + catalog: "manual_cat".to_string(), + schema: "manual_schema".to_string(), + table_name: NUMBERS_TABLE_NAME.to_string(), + table_id: NUMBERS_TABLE_ID, + table: NumbersTable::table(NUMBERS_TABLE_ID), + }) + .unwrap(); + let query_ctx = Arc::new(QueryContext::with("manual_cat", "manual_schema")); + let mut config = datafusion::execution::context::SessionConfig::new() + .with_information_schema(true) + .with_default_catalog_and_schema("manual_cat", "manual_schema"); + config.set_extension(query_ctx.clone()); + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(config) + .build(); + let decoder = DefaultPlanDecoder::new(state, &query_ctx).unwrap(); let plan = LogicalPlan::Extension(Extension { node: Arc::new(SeriesNormalize::new( 0, @@ -1509,9 +1590,11 @@ mod tests { merge_scan(numbers_scan(), false), )), }); - let request_catalog = Arc::new(DummyCatalogList::with_table_provider(Arc::new( - MemTable::try_new(NumbersTable::schema().arrow_schema().clone(), vec![vec![]]).unwrap(), - ))); + let request_catalog = Arc::new( + catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx( + manager, query_ctx, + ), + ); let decoded = decoder .decode(encode_plan(&plan), request_catalog, false) .await @@ -1521,6 +1604,30 @@ mod tests { assert_eq!(decoded.schema().as_arrow(), plan.schema().as_arrow()); } + /// Manual information-schema resolution uses the SessionState catalog and table functions. + #[tokio::test] + async fn test_serializer_decode_manual_state_information_schema() { + let query_ctx = Arc::new(QueryContext::with("manual_cat", "manual_schema")); + let mut config = datafusion::execution::context::SessionConfig::new() + .with_information_schema(true) + .with_default_catalog_and_schema("manual_cat", "manual_schema"); + config.set_extension(query_ctx.clone()); + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(config) + .with_table_function_list(datafusion::functions_table::all_default_table_functions()) + .build(); + let functions = state.table_functions().clone(); + assert!(functions.contains_key("generate_series")); + let bytes = encode_plan(&merge_scan(numbers_scan(), false)); + let decoder = DefaultPlanDecoder::new(state, &query_ctx).unwrap(); + let catalog = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + MemTable::try_new(NumbersTable::schema().arrow_schema().clone(), vec![vec![]]).unwrap(), + ))); + let decoded = decoder.decode(bytes, catalog, false).await.unwrap(); + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 1); + } + /// `MergeScan` decodes under an ordinary single-input extension parent. #[tokio::test] async fn test_serializer_decode_merge_scan_under_series_normalize() { @@ -1684,11 +1791,17 @@ mod tests { QueryOptions::default(), ) .query_engine(); - let state = engine.engine_context(QueryContext::arc()).state().clone(); + let mut state = SessionStateBuilder::new_from_existing( + engine.engine_context(QueryContext::arc()).state().clone(), + ) + .with_serializer_registry(Arc::new(DefaultSerializer)) + .build(); + register_greptime_functions(&mut state, &QueryContext::arc()).unwrap(); let consumer = MergeScanSubstraitConsumer { inner: DefaultSubstraitConsumer::new(&extensions, &state), session_state: &state, catalog_manager: Some(catalog_manager), + payload: false, }; let decoded = consumer.decode_payload(bytes.to_vec()).await.unwrap(); From 015298074b9e50e810f318240c9f56cc0e6dbe48 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:55:21 +0800 Subject: [PATCH 27/37] test(query): verify payload information schema table functions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/query_engine/default_serializer.rs | 142 ++++++++++-------- tests-integration/src/cluster.rs | 2 +- 2 files changed, 79 insertions(+), 65 deletions(-) diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 6e48adf34cd7..9a6809730edd 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -1176,29 +1176,6 @@ mod tests { async fn test_serializer_decode_payload_observes_catalog_and_query_context() { let observed = Arc::new(std::sync::Mutex::new(Vec::new())); let inner = catalog::memory::new_memory_catalog_manager().unwrap(); - inner.register_catalog_sync("request_cat").unwrap(); - inner.register_catalog_sync("explicit_catalog").unwrap(); - for (catalog, schema) in [ - ("request_cat", "request_schema"), - ("request_cat", "explicit_schema"), - ("explicit_catalog", "full_schema"), - ] { - inner - .register_schema_sync(catalog::RegisterSchemaRequest { - catalog: catalog.to_string(), - schema: schema.to_string(), - }) - .unwrap(); - inner - .register_table_sync(RegisterTableRequest { - catalog: catalog.to_string(), - schema: schema.to_string(), - table_name: NUMBERS_TABLE_NAME.to_string(), - table_id: NUMBERS_TABLE_ID, - table: NumbersTable::table(NUMBERS_TABLE_ID), - }) - .unwrap(); - } let manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager { inner, resolution: TableResolution::ObserveContext(observed.clone()), @@ -1246,27 +1223,27 @@ mod tests { assert_eq!( *observed.lock().unwrap(), vec![ - ( - "request_cat".into(), - "request_schema".into(), - NUMBERS_TABLE_NAME.into(), - "request_cat".into(), - "request_schema".into() - ), - ( - "request_cat".into(), - "explicit_schema".into(), - NUMBERS_TABLE_NAME.into(), - "request_cat".into(), - "request_schema".into() - ), - ( - DEFAULT_CATALOG_NAME.into(), - DEFAULT_SCHEMA_NAME.into(), - NUMBERS_TABLE_NAME.into(), - "request_cat".into(), - "request_schema".into() - ), + [ + "request_cat".to_string(), + "request_schema".to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], + [ + "request_cat".to_string(), + "explicit_schema".to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], + [ + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + NUMBERS_TABLE_NAME.to_string(), + "request_cat".to_string(), + "request_schema".to_string(), + ], ], ); } @@ -1276,7 +1253,7 @@ mod tests { /// Fails with the frontend's cross-catalog access error. AccessDenied, /// Records catalog lookup and request context values. - ObserveContext(Arc>>), + ObserveContext(Arc>>), /// Resolves once another task on the caller runtime signals. AwaitSignal { notify: Arc, @@ -1357,13 +1334,13 @@ mod tests { match &self.resolution { TableResolution::ObserveContext(observed) => { let query_ctx = query_ctx.unwrap(); - observed.lock().unwrap().push(( + observed.lock().unwrap().push([ catalog.to_string(), schema.to_string(), table_name.to_string(), query_ctx.current_catalog().to_string(), query_ctx.current_schema(), - )); + ]); Ok(Some(NumbersTable::table(NUMBERS_TABLE_ID))) } TableResolution::AccessDenied => QueryAccessDeniedSnafu { @@ -1604,28 +1581,65 @@ mod tests { assert_eq!(decoded.schema().as_arrow(), plan.schema().as_arrow()); } - /// Manual information-schema resolution uses the SessionState catalog and table functions. + /// Payload information-schema resolution exposes registered table functions as routines. #[tokio::test] - async fn test_serializer_decode_manual_state_information_schema() { + async fn test_serializer_payload_information_schema_exposes_table_functions() { let query_ctx = Arc::new(QueryContext::with("manual_cat", "manual_schema")); - let mut config = datafusion::execution::context::SessionConfig::new() - .with_information_schema(true) - .with_default_catalog_and_schema("manual_cat", "manual_schema"); - config.set_extension(query_ctx.clone()); - let state = SessionStateBuilder::new() - .with_default_features() + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let state = engine.engine_context(query_ctx).state().clone(); + let config = state.config().clone().with_information_schema(true); + let state = SessionStateBuilder::new_from_existing(state) .with_config(config) .with_table_function_list(datafusion::functions_table::all_default_table_functions()) .build(); - let functions = state.table_functions().clone(); - assert!(functions.contains_key("generate_series")); - let bytes = encode_plan(&merge_scan(numbers_scan(), false)); - let decoder = DefaultPlanDecoder::new(state, &query_ctx).unwrap(); - let catalog = Arc::new(DummyCatalogList::with_table_provider(Arc::new( - MemTable::try_new(NumbersTable::schema().arrow_schema().clone(), vec![vec![]]).unwrap(), - ))); - let decoded = decoder.decode(bytes, catalog, false).await.unwrap(); - assert_eq!(decoded.to_string().matches("MergeScan [").count(), 1); + + let extensions = Extensions::default(); + let consumer = MergeScanSubstraitConsumer { + inner: DefaultSubstraitConsumer::new(&extensions, &state), + session_state: &state, + catalog_manager: state + .config() + .get_extension::() + .map(|s| s.catalog_manager().clone()), + payload: true, + }; + let provider = consumer + .resolve_table_ref(&TableReference::partial("information_schema", "routines")) + .await + .unwrap() + .unwrap(); + let plan = provider.scan(&state, None, &[], None).await.unwrap(); + let batches = datafusion::physical_plan::collect(plan, state.task_ctx()) + .await + .unwrap(); + let schema = provider.schema(); + let routine_name_index = schema.index_of("routine_name").unwrap(); + let function_type_index = schema.index_of("function_type").unwrap(); + assert!(batches.iter().any(|batch| { + let routine_names = batch + .column(routine_name_index) + .as_any() + .downcast_ref::() + .unwrap(); + let function_types = batch + .column(function_type_index) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).any(|row| { + routine_names.value(row) == "generate_series" + && function_types.value(row) == "TABLE" + }) + })); } /// `MergeScan` decodes under an ordinary single-input extension parent. diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 9e26252c47c5..095968743b26 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -41,7 +41,7 @@ use common_datasource::object_store::LocalFileAccess; use common_event_recorder::EventRecorderOptions; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::DatanodeId; -use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder}; +use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::chroot::ChrootKvBackend; use common_meta::kv_backend::etcd::EtcdStore; From b0826aad9be45cf33bff4d02ea67e667cf415322 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:08:55 +0800 Subject: [PATCH 28/37] test(query): retain nested scan runtime and cold-route guards Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/query_engine/default_serializer.rs | 250 +++++++++++------- tests-integration/src/cluster.rs | 27 +- .../nested_merge_scan_capability_test.rs | 115 +++++++- 3 files changed, 293 insertions(+), 99 deletions(-) diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 9a6809730edd..98f265c99e9f 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -1000,48 +1000,51 @@ mod tests { /// Request state registration keeps Greptime `date_format` ahead of DataFusion aliases. #[tokio::test] async fn test_serializer_decode_keeps_greptime_date_format() { - let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); - let factory = QueryEngineFactory::new( - catalog_list, - None, - None, - None, - None, - false, - QueryOptions::default(), - ); - let engine = factory.query_engine(); - - let table_provider = Arc::new(mock_table_provider(1.into())); - let query_ctx = date_format_query_ctx(); - let plan = LogicalPlanBuilder::scan( - "devices", - Arc::new(LogicalTableSource::new(table_provider.schema().clone())), - None, - ) - .unwrap() - .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))]) - .unwrap() - .build() - .unwrap(); - - let bytes = DFLogicalSubstraitConvertor - .encode(&plan, DefaultSerializer) - .unwrap(); - let plan_decoder = engine - .engine_context(query_ctx.clone()) - .new_plan_decoder() - .unwrap(); - let decoded = plan_decoder - .decode( - bytes, - Arc::new(DummyCatalogList::with_table_provider(table_provider)), + // Repeat to expose nondeterministic function-registration order. + for _ in 0..8 { + let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap(); + let factory = QueryEngineFactory::new( + catalog_list, + None, + None, + None, + None, false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + + let table_provider = Arc::new(mock_table_provider(1.into())); + let query_ctx = date_format_query_ctx(); + let plan = LogicalPlanBuilder::scan( + "devices", + Arc::new(LogicalTableSource::new(table_provider.schema().clone())), + None, ) - .await + .unwrap() + .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))]) + .unwrap() + .build() .unwrap(); - assert_only_function_is_greptime_date_format(&decoded, &query_ctx); + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .new_plan_decoder() + .unwrap(); + let decoded = plan_decoder + .decode( + bytes, + Arc::new(DummyCatalogList::with_table_provider(table_provider)), + false, + ) + .await + .unwrap(); + + assert_only_function_is_greptime_date_format(&decoded, &query_ctx); + } } /// Nested payload decoding retains the Greptime `date_format` binding. @@ -1058,64 +1061,67 @@ mod tests { }) .unwrap(); - let factory = QueryEngineFactory::new( - catalog_manager.clone(), - None, - None, - None, - None, - false, - QueryOptions::default(), - ); - let engine = factory.query_engine(); - let query_ctx = date_format_query_ctx(); - - let input = LogicalPlanBuilder::scan( - NUMBERS_TABLE_NAME, - Arc::new(LogicalTableSource::new( - NumbersTable::schema().arrow_schema().clone(), - )), - None, - ) - .unwrap() - .project(vec![greptime_date_format_expr( - &query_ctx, - Expr::Cast(Cast::new( - Box::new(col("number")), - ArrowDataType::Timestamp(TimeUnit::Millisecond, None), - )), - )]) - .unwrap() - .build() - .unwrap(); - let plan = MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); + for _ in 0..8 { + let factory = QueryEngineFactory::new( + catalog_manager.clone(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ); + let engine = factory.query_engine(); + let query_ctx = date_format_query_ctx(); - let bytes = DFLogicalSubstraitConvertor - .encode(&plan, DefaultSerializer) - .unwrap(); - let plan_decoder = engine - .engine_context(query_ctx.clone()) - .new_plan_decoder() + let input = LogicalPlanBuilder::scan( + NUMBERS_TABLE_NAME, + Arc::new(LogicalTableSource::new( + NumbersTable::schema().arrow_schema().clone(), + )), + None, + ) + .unwrap() + .project(vec![greptime_date_format_expr( + &query_ctx, + Expr::Cast(Cast::new( + Box::new(col("number")), + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + )), + )]) + .unwrap() + .build() .unwrap(); - let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( - mock_table_provider(1.into()), - ))); + let plan = + MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); - let decoded = plan_decoder - .decode(bytes, catalog_list, false) - .await - .unwrap(); + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let plan_decoder = engine + .engine_context(query_ctx.clone()) + .new_plan_decoder() + .unwrap(); + let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( + mock_table_provider(1.into()), + ))); - let LogicalPlan::Extension(extension) = &decoded else { - panic!("Expect a MergeScan plan, got: {decoded}"); - }; - let merge_scan = extension - .node - .as_any() - .downcast_ref::() - .expect("Expect a MergeScan plan node"); + let decoded = plan_decoder + .decode(bytes, catalog_list, false) + .await + .unwrap(); + + let LogicalPlan::Extension(extension) = &decoded else { + panic!("Expect a MergeScan plan, got: {decoded}"); + }; + let merge_scan = extension + .node + .as_any() + .downcast_ref::() + .expect("Expect a MergeScan plan node"); - assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx); + assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx); + } } /// Registers a `numbers` table in a memory catalog manager. @@ -1423,6 +1429,68 @@ mod tests { .query_engine() } + /// Decoding must not need a Tokio runtime: the public path is only `async`. + #[test] + fn test_serializer_decode_nested_merge_scan_without_tokio_runtime() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let plan = merge_scan(merge_scan(numbers_scan(), false), true); + let bytes = encode_plan(&plan); + + let decoded = futures::executor::block_on(plan_decoder(engine.as_ref()).decode( + bytes, + request_catalog_list(), + false, + )) + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + } + + /// Payload decoding must not block a multi-thread runtime either. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_serializer_decode_deeply_nested_merge_scan_without_blocking() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let plan = merge_scan( + merge_scan( + LogicalPlanBuilder::from(merge_scan(numbers_scan(), false)) + .filter(col("number").lt(lit(10u32))) + .unwrap() + .build() + .unwrap(), + false, + ), + false, + ); + let bytes = encode_plan(&plan); + + let decoded = plan_decoder(engine.as_ref()) + .decode(bytes, request_catalog_list(), false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 3); + } + /// A payload lookup parked on a caller-runtime task must keep the decode async. /// /// The engine catalog resolves only once a task spawned on this runtime signals it, so a diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 095968743b26..b7089c336710 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -41,7 +41,7 @@ use common_datasource::object_store::LocalFileAccess; use common_event_recorder::EventRecorderOptions; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::DatanodeId; -use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; +use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder}; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::chroot::ChrootKvBackend; use common_meta::kv_backend::etcd::EtcdStore; @@ -113,6 +113,8 @@ pub struct GreptimeDbCluster { /// The stats of the network address of every datanode, empty unless the cluster serves the /// datanodes at their own addresses. pub datanode_rpc_stats: HashMap>, + /// Cache registries used by datanodes to apply cache invalidations. + pub datanode_cache_registries: HashMap>, pub kv_backend: KvBackendRef, pub metasrv: Arc, pub frontend: Arc, @@ -379,7 +381,7 @@ impl GreptimeDbClusterBuilder { ) .await; - let datanode_instances = self + let (datanode_instances, datanode_cache_registries) = self .build_datanodes_with_options(&metasrv, &datanode_options) .await; @@ -418,6 +420,7 @@ impl GreptimeDbClusterBuilder { guards, datanode_instances, datanode_rpc_stats, + datanode_cache_registries, kv_backend: self.kv_backend.clone(), metasrv: metasrv.metasrv, frontend: Arc::new(frontend), @@ -500,16 +503,22 @@ impl GreptimeDbClusterBuilder { &self, metasrv: &MockInfo, options: &[DatanodeOptions], - ) -> HashMap { + ) -> ( + HashMap, + HashMap>, + ) { let mut instances = HashMap::with_capacity(options.len()); + let mut cache_registries = HashMap::with_capacity(options.len()); for opts in options { - let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await; + let (datanode, cache_registry) = + self.create_datanode(opts.clone(), metasrv.clone()).await; let datanode_id = opts.node_id.unwrap(); instances.insert(datanode_id, datanode); + cache_registries.insert(datanode_id, cache_registry); } - instances + (instances, cache_registries) } async fn wait_datanodes_alive( @@ -535,7 +544,11 @@ impl GreptimeDbClusterBuilder { panic!("Some Datanodes are not alive in 10 seconds!") } - async fn create_datanode(&self, opts: DatanodeOptions, metasrv: MockInfo) -> Datanode { + async fn create_datanode( + &self, + opts: DatanodeOptions, + metasrv: MockInfo, + ) -> (Datanode, Arc) { let mut meta_client = MetaClientBuilder::datanode_default_options(opts.node_id.unwrap()) .channel_manager(metasrv.channel_manager) .build(); @@ -555,7 +568,7 @@ impl GreptimeDbClusterBuilder { datanode.start_heartbeat().await.unwrap(); - datanode + (datanode, layered_cache_registry) } async fn build_frontend( diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 9be69a645a0d..110600b4460e 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -23,6 +23,9 @@ use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use api::v1::region::{QueryRequest, RegionRequestHeader}; +use common_meta::cache::TableRouteCacheRef; +use common_meta::cache_invalidator::{CacheInvalidator, Context}; +use common_meta::instruction::CacheIdent; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::info; @@ -36,7 +39,7 @@ use query::query_engine::DefaultSerializer; use servers::query_handler::sql::SqlQueryHandler; use session::context::{QueryContext, QueryContextRef}; use store_api::region_request::{RegionCloseRequest, RegionRequest}; -use store_api::storage::RegionId; +use store_api::storage::{RegionId, TableId}; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; @@ -48,6 +51,7 @@ const PROBE_TABLE: &str = "nested_cap_probe"; const BUILD_TABLE: &str = "nested_cap_build"; /// Missing inner table used to verify plan-resolution failure. const MISSING_INNER_TABLE: &str = "nested_cap_missing_build"; +const UNRELATED_TABLE_ID: TableId = 10_000_000; /// `(a_id, probe_key, probe_v)` fixture rows; `a_id` is the partition key. const PROBE_ROWS: &[(i32, i32, i32)] = &[ @@ -125,6 +129,57 @@ async fn test_nested_merge_scan_capability_join_across_datanodes() { assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests); } +/// Explicit route eviction and unrelated invalidation finish before execution; nested reads refill +/// the route and return complete rows. This does not exercise a concurrent invalidation race. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation() { + common_telemetry::init_default_ut_logging(); + + let cluster = + build_cluster("test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation") + .await; + let frontend = cluster.fe_instance().clone(); + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; + assert_eq!(2, datanodes(&probe_leaders).len()); + assert_eq!(2, datanodes(&build_leaders).len()); + let build_table_id = table_id(&build_leaders); + + let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; + let expected = expected_join_rows(); + assert_eq!(expected.clone(), multiset(table_cells(&reference))); + + let query_ctx = query_ctx(); + let plan = encode(&nested_plan(&frontend, &query_ctx).await); + invalidate_table_ids(&cluster, &[build_table_id]).await; + assert_route_cold(&cluster, build_table_id); + invalidate_table_ids(&cluster, &[UNRELATED_TABLE_ID]).await; + assert_route_cold(&cluster, build_table_id); + + let remote_datanodes = datanodes(&build_leaders); + let baseline_requests = remote_datanodes + .iter() + .map(|datanode| (*datanode, rpc_requests(&cluster, *datanode))) + .collect::>(); + let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; + let actual = multiset(rows); + assert_eq!( + expected, actual, + "unexpected nested join rows after cold load" + ); + assert_eq!(multiset(table_cells(&reference)), actual); + assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests); + + for datanode in datanodes(&probe_leaders) { + assert!( + table_route_cache(&cluster, datanode).contains_key(&build_table_id), + "datanode {datanode} must refill the route of build table {build_table_id}" + ); + } +} + /// A failing inner region makes the whole query fail: the datanode must not return the rows of the /// inner regions that succeed. #[tokio::test(flavor = "multi_thread")] @@ -271,6 +326,64 @@ fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer .region_server() } +/// The id of the single table that the regions of `leaders` belong to. +fn table_id(leaders: &BTreeMap) -> TableId { + let table_ids = leaders + .keys() + .map(|region_id| RegionId::from_u64(*region_id).table_id()) + .collect::>(); + assert_eq!( + 1, + table_ids.len(), + "expected the regions to belong to a single table, actual region leaders: {leaders:?}" + ); + + *table_ids.first().expect("there is at least one region") +} + +/// Returns the table route cache of the datanode `datanode_id`: the cache that +/// `DatanodeRegionQueryHandler::select_target` reads to resolve the leader of a region. +fn table_route_cache(cluster: &GreptimeDbCluster, datanode_id: u64) -> TableRouteCacheRef { + cluster + .datanode_cache_registries + .get(&datanode_id) + .unwrap_or_else(|| panic!("expected the cache registry of the datanode {datanode_id}")) + .get::() + .unwrap_or_else(|| panic!("expected the table route cache of the datanode {datanode_id}")) +} + +/// Invalidates `table_ids` on every datanode, the way a datanode handles the invalidation +/// instruction of metasrv (see [`GreptimeDbCluster::datanode_cache_registries`]). +async fn invalidate_table_ids(cluster: &GreptimeDbCluster, table_ids: &[TableId]) { + let idents = table_ids + .iter() + .copied() + .map(CacheIdent::TableId) + .collect::>(); + for (datanode_id, registry) in &cluster.datanode_cache_registries { + registry + .invalidate(&Context::default(), &idents) + .await + .unwrap_or_else(|e| { + panic!("the datanode {datanode_id} must invalidate {idents:?}: {e}") + }); + } +} + +/// Asserts that no datanode holds the route of `table_id`, i.e. the route is cold. +fn assert_route_cold(cluster: &GreptimeDbCluster, table_id: TableId) { + let cached = cluster + .datanode_cache_registries + .keys() + .filter(|datanode_id| table_route_cache(cluster, **datanode_id).contains_key(&table_id)) + .collect::>(); + assert!( + cached.is_empty(), + "the datanodes {cached:?} must drop the route of the table {table_id} after the \ + invalidation" + ); +} + async fn run_nested_queries( cluster: &GreptimeDbCluster, probe_leaders: &BTreeMap, From a7db8480991794142d6890c9e2d9a0ab53f96f0e Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 9 Oct 2026 12:35:13 +0800 Subject: [PATCH 29/37] refactor(query): address nested scan review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cache/src/lib.rs | 45 +- src/datanode/src/datanode.rs | 2 - .../src/query_engine/default_serializer.rs | 429 +++--------------- tests-integration/src/cluster.rs | 27 +- tests-integration/src/standalone.rs | 2 - .../nested_merge_scan_capability_test.rs | 115 +---- 6 files changed, 75 insertions(+), 545 deletions(-) diff --git a/src/cache/src/lib.rs b/src/cache/src/lib.rs index 58ad4c14e5fc..c7304ae60096 100644 --- a/src/cache/src/lib.rs +++ b/src/cache/src/lib.rs @@ -53,19 +53,10 @@ 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 the layered cache registry for datanode. +/// Builds the datanode cache registry with base metadata caches before derived caches. /// -/// Layers separate caches derived from kv-backed metadata. The first layer holds the schema, table -/// id-to-schema-name, table info, table name and table route caches; the second holds the table cache -/// (derived from table info and name) and partition info cache (derived from table routes). -/// -/// [LayeredCacheRegistry] invalidates a layer only after the previous layer finished invalidating -/// (see [LayeredCacheRegistryBuilder::add_cache_registry]), so the derived caches are invalidated -/// after the caches they are derived from. A flat [CacheRegistry] invalidates its caches -/// concurrently: a derived cache could be invalidated first and then refilled by a concurrent query -/// from a base cache that still holds the old value, and the stale metadata would stay in the -/// derived cache until the entry expires. This is the only datanode cache registry: production code -/// and test infrastructure build the same one. +/// 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( @@ -243,33 +234,3 @@ pub fn with_default_composite_cache_registry( Ok(builder.add_cache_registry(registry)) } - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use catalog::kvbackend::TableCacheRef; - use common_meta::cache::{ - SchemaCacheRef, TableInfoCacheRef, TableNameCacheRef, TableRouteCacheRef, - TableSchemaCacheRef, - }; - use common_meta::kv_backend::memory::MemoryKvBackend; - use partition::cache::PartitionInfoCacheRef; - - use super::*; - - #[test] - fn test_datanode_layered_cache_registry_holds_all_caches() { - let registry = build_datanode_layered_cache_registry(Arc::new(MemoryKvBackend::< - common_meta::error::Error, - >::new())); - - assert!(registry.get::().is_some()); - assert!(registry.get::().is_some()); - assert!(registry.get::().is_some()); - assert!(registry.get::().is_some()); - assert!(registry.get::().is_some()); - assert!(registry.get::().is_some()); - assert!(registry.get::().is_some()); - } -} diff --git a/src/datanode/src/datanode.rs b/src/datanode/src/datanode.rs index a4e283350687..2e0300a54a30 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -1149,8 +1149,6 @@ mod tests { mock_region_server.register_engine(mock_region.clone()); let kv_backend = Arc::new(MemoryKvBackend::new()); - // Use the same layered registry as the datanode of a running instance, so the derived - // caches of the test are invalidated after the caches they are derived from. let layered_cache_registry = Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 98f265c99e9f..b1d168a5a5d6 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -156,26 +156,15 @@ impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> { ); let schema_name = resolved.schema.as_ref(); let catalog_name = resolved.catalog.as_ref(); - let schema = if self.session_state.config().information_schema() - && schema_name == datafusion::catalog::information_schema::INFORMATION_SCHEMA - { - Arc::new( - datafusion::catalog::information_schema::InformationSchemaProvider::new( - catalog_list.clone(), - ) - .with_table_functions(self.session_state.table_functions().clone()), - ) as Arc - } else { - catalog_list - .catalog(catalog_name) - .ok_or_else(|| { - DataFusionError::Plan(format!("failed to resolve catalog: {catalog_name}")) - })? - .schema(schema_name) - .ok_or_else(|| { - DataFusionError::Plan(format!("failed to resolve schema: {schema_name}")) - })? - }; + let schema = catalog_list + .catalog(catalog_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("failed to resolve catalog: {catalog_name}")) + })? + .schema(schema_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("failed to resolve schema: {schema_name}")) + })?; schema.table(&resolved.table).await } @@ -458,7 +447,7 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { #[cfg(test)] mod tests { use std::any::Any; - use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicBool, Ordering}; use catalog::error::{QueryAccessDeniedSnafu, Result as CatalogResult}; use catalog::{CatalogManager, RegisterTableRequest}; @@ -724,55 +713,19 @@ mod tests { #[tokio::test] async fn test_serializer_decode_merge_scan() { - let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); - // Payload tables must be registered in the engine catalog. - catalog_manager - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: NUMBERS_TABLE_NAME.to_string(), - table_id: NUMBERS_TABLE_ID, - table: NumbersTable::table(NUMBERS_TABLE_ID), - }) - .unwrap(); - let factory = QueryEngineFactory::new( - catalog_manager, + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), None, None, None, None, false, QueryOptions::default(), - ); - let engine = factory.query_engine(); - - let input = LogicalPlanBuilder::scan( - NUMBERS_TABLE_NAME, - Arc::new(LogicalTableSource::new( - NumbersTable::schema().arrow_schema().clone(), - )), - None, ) - .unwrap() - .build() - .unwrap(); - let plan = - MergeScanLogicalPlan::new(input.clone(), true, Default::default()).into_logical_plan(); - - let bytes = DFLogicalSubstraitConvertor - .encode(&plan, DefaultSerializer) - .unwrap(); - let plan_decoder = engine - .engine_context(QueryContext::arc()) - .new_plan_decoder() - .unwrap(); - // Must not bind payload tables through the request catalog. - let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( - mock_table_provider(1.into()), - ))); - - let decoded = plan_decoder - .decode(bytes, catalog_list, false) + .query_engine(); + let plan = merge_scan(numbers_scan(), true); + let decoded = plan_decoder(engine.as_ref()) + .decode(encode_plan(&plan), request_catalog_list(), false) .await .unwrap(); @@ -785,7 +738,7 @@ mod tests { .downcast_ref::() .expect("Expect a MergeScan plan node"); assert!(merge_scan.is_placeholder()); - assert_eq!(merge_scan.input().to_string(), input.to_string()); + assert_eq!(merge_scan.input().to_string(), numbers_scan().to_string()); // `PbMergeScan` doesn't carry `partition_cols`, so decoded nodes have an empty mapping. assert!(merge_scan.partition_cols().is_empty()); } @@ -793,52 +746,22 @@ mod tests { /// Payload table binding uses the engine catalog, not the request catalog. #[tokio::test] async fn test_serializer_decode_merge_scan_with_engine_catalog() { - let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); - catalog_manager - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: NUMBERS_TABLE_NAME.to_string(), - table_id: NUMBERS_TABLE_ID, - table: NumbersTable::table(NUMBERS_TABLE_ID), - }) - .unwrap(); - let factory = QueryEngineFactory::new( - catalog_manager, + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), None, None, None, None, false, QueryOptions::default(), - ); - let engine = factory.query_engine(); - - let input = LogicalPlanBuilder::scan( - NUMBERS_TABLE_NAME, - Arc::new(LogicalTableSource::new( - NumbersTable::schema().arrow_schema().clone(), - )), - None, ) - .unwrap() - .build() - .unwrap(); - let plan = MergeScanLogicalPlan::new(input, false, Default::default()).into_logical_plan(); - let bytes = DFLogicalSubstraitConvertor - .encode(&plan, DefaultSerializer) - .unwrap(); - - // This request catalog deliberately has a different schema. - let table_provider = Arc::new(mock_table_provider(1.into())); - let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider)); - - let plan_decoder = engine - .engine_context(QueryContext::arc()) - .new_plan_decoder() - .unwrap(); - let decoded = plan_decoder - .decode(bytes, catalog_list, false) + .query_engine(); + let decoded = plan_decoder(engine.as_ref()) + .decode( + encode_plan(&merge_scan(numbers_scan(), false)), + request_catalog_list(), + false, + ) .await .unwrap(); @@ -861,62 +784,26 @@ mod tests { /// Nested `MergeScan` payloads decode recursively. #[tokio::test] async fn test_serializer_decode_nested_merge_scan() { - let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap(); - catalog_manager - .register_table_sync(RegisterTableRequest { - catalog: DEFAULT_CATALOG_NAME.to_string(), - schema: DEFAULT_SCHEMA_NAME.to_string(), - table_name: NUMBERS_TABLE_NAME.to_string(), - table_id: NUMBERS_TABLE_ID, - table: NumbersTable::table(NUMBERS_TABLE_ID), - }) - .unwrap(); - let factory = QueryEngineFactory::new( - catalog_manager, + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), None, None, None, None, false, QueryOptions::default(), - ); - let engine = factory.query_engine(); - - let numbers_scan = || { - LogicalPlanBuilder::scan( - NUMBERS_TABLE_NAME, - Arc::new(LogicalTableSource::new( - NumbersTable::schema().arrow_schema().clone(), - )), - None, - ) - .unwrap() - .build() - .unwrap() - }; - let inner = MergeScanLogicalPlan::new(numbers_scan(), false, Default::default()) - .into_logical_plan(); + ) + .query_engine(); + let inner = merge_scan(numbers_scan(), false); let input = LogicalPlanBuilder::from(inner) .filter(col("number").lt(lit(10u32))) .unwrap() .build() .unwrap(); - let plan = - MergeScanLogicalPlan::new(input.clone(), false, Default::default()).into_logical_plan(); - - let bytes = DFLogicalSubstraitConvertor - .encode(&plan, DefaultSerializer) - .unwrap(); - let plan_decoder = engine - .engine_context(QueryContext::arc()) - .new_plan_decoder() - .unwrap(); - let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new( - mock_table_provider(1.into()), - ))); + let plan = merge_scan(input, false); - let decoded = plan_decoder - .decode(bytes, catalog_list, false) + let decoded = plan_decoder(engine.as_ref()) + .decode(encode_plan(&plan), request_catalog_list(), false) .await .unwrap(); @@ -1265,11 +1152,6 @@ mod tests { notify: Arc, entered: Arc, }, - /// Never resolves, holding `live` until the caller drops the decode future. - Pending { - live: Arc, - entered: Arc, - }, } /// Engine catalog manager that takes over `table` resolution while delegating the rest. @@ -1361,12 +1243,6 @@ mod tests { .table(catalog, schema, table_name, query_ctx) .await } - TableResolution::Pending { live, entered } => { - entered.store(true, Ordering::SeqCst); - let _guard = LiveGuard::new(live); - std::future::pending::<()>().await; - unreachable!("pending resolution never completes") - } } } @@ -1393,23 +1269,6 @@ mod tests { } } - /// Counts live payload resolutions, so a dropped decode leaves no work behind. - struct LiveGuard(Arc); - - impl LiveGuard { - fn new(live: &Arc) -> Self { - let live = live.clone(); - live.fetch_add(1, Ordering::SeqCst); - Self(live) - } - } - - impl Drop for LiveGuard { - fn drop(&mut self) { - self.0.fetch_sub(1, Ordering::SeqCst); - } - } - fn engine_with_intercepting_catalog( resolution: TableResolution, ) -> Arc { @@ -1429,68 +1288,6 @@ mod tests { .query_engine() } - /// Decoding must not need a Tokio runtime: the public path is only `async`. - #[test] - fn test_serializer_decode_nested_merge_scan_without_tokio_runtime() { - let engine = QueryEngineFactory::new( - numbers_catalog_manager(), - None, - None, - None, - None, - false, - QueryOptions::default(), - ) - .query_engine(); - let plan = merge_scan(merge_scan(numbers_scan(), false), true); - let bytes = encode_plan(&plan); - - let decoded = futures::executor::block_on(plan_decoder(engine.as_ref()).decode( - bytes, - request_catalog_list(), - false, - )) - .unwrap(); - - assert_eq!(decoded.to_string(), plan.to_string()); - assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); - } - - /// Payload decoding must not block a multi-thread runtime either. - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn test_serializer_decode_deeply_nested_merge_scan_without_blocking() { - let engine = QueryEngineFactory::new( - numbers_catalog_manager(), - None, - None, - None, - None, - false, - QueryOptions::default(), - ) - .query_engine(); - let plan = merge_scan( - merge_scan( - LogicalPlanBuilder::from(merge_scan(numbers_scan(), false)) - .filter(col("number").lt(lit(10u32))) - .unwrap() - .build() - .unwrap(), - false, - ), - false, - ); - let bytes = encode_plan(&plan); - - let decoded = plan_decoder(engine.as_ref()) - .decode(bytes, request_catalog_list(), false) - .await - .unwrap(); - - assert_eq!(decoded.to_string(), plan.to_string()); - assert_eq!(decoded.to_string().matches("MergeScan [").count(), 3); - } - /// A payload lookup parked on a caller-runtime task must keep the decode async. /// /// The engine catalog resolves only once a task spawned on this runtime signals it, so a @@ -1522,78 +1319,41 @@ mod tests { assert_eq!(decoded.to_string(), plan.to_string()); } - /// Dropping a decode parked on the catalog cancels it; no detached job keeps running. - #[tokio::test(flavor = "current_thread")] - async fn test_serializer_decode_dropped_pending_payload_leaves_no_job() { - let live = Arc::new(AtomicUsize::new(0)); - let entered = Arc::new(AtomicBool::new(false)); - let engine = engine_with_intercepting_catalog(TableResolution::Pending { - live: live.clone(), - entered: entered.clone(), - }); - let bytes = encode_plan(&merge_scan(numbers_scan(), false)); - - let decoder = plan_decoder(engine.as_ref()); - { - let mut decode = Box::pin(decoder.decode(bytes, request_catalog_list(), false)); - assert!(futures::poll!(decode.as_mut()).is_pending()); - assert_eq!(live.load(Ordering::SeqCst), 1); - assert!(entered.load(Ordering::SeqCst)); - } - - assert_eq!( - live.load(Ordering::SeqCst), - 0, - "dropping the decode future must cancel its payload work" - ); - } - - /// A payload table missing from the engine catalog must not fall back to the request catalog. + /// Missing and denied payload tables do not fall back to the usable request catalog. #[tokio::test] - async fn test_serializer_decode_payload_missing_in_engine_catalog() { - // The engine catalog deliberately has no `numbers` table. - let catalog_manager: CatalogManagerRef = + async fn test_serializer_decode_payload_negative_engine_catalog_resolution() { + let missing_manager: CatalogManagerRef = catalog::memory::new_memory_catalog_manager().unwrap(); - let engine = QueryEngineFactory::new( - catalog_manager, - None, - None, - None, - None, - false, - QueryOptions::default(), - ) - .query_engine(); - let bytes = encode_plan(&merge_scan(numbers_scan(), false)); - - let err = plan_decoder(engine.as_ref()) - .decode(bytes, request_catalog_list(), false) - .await - .unwrap_err(); - - let err = format!("{err:?}"); - assert!( - err.contains("Table not found: greptime.public.numbers"), - "unexpected error: {err}" - ); - } - - /// A payload table denied by the engine catalog must not fall back to the request catalog. - #[tokio::test] - async fn test_serializer_decode_payload_access_denied_in_engine_catalog() { - let engine = engine_with_intercepting_catalog(TableResolution::AccessDenied); + let cases = [ + ( + QueryEngineFactory::new( + missing_manager, + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(), + "Table not found: greptime.public.numbers", + ), + ( + engine_with_intercepting_catalog(TableResolution::AccessDenied), + "Illegal access to catalog", + ), + ]; let bytes = encode_plan(&merge_scan(numbers_scan(), false)); - let err = plan_decoder(engine.as_ref()) - .decode(bytes, request_catalog_list(), false) - .await - .unwrap_err(); + for (engine, expected) in cases { + let err = plan_decoder(engine.as_ref()) + .decode(bytes.clone(), request_catalog_list(), false) + .await + .unwrap_err(); - let err = format!("{err:?}"); - assert!( - err.contains("Illegal access to catalog"), - "unexpected error: {err}" - ); + let err = format!("{err:?}"); + assert!(err.contains(expected), "expected {expected:?}, got: {err}"); + } } /// Manual states use their own catalog and query-context defaults for payloads. @@ -1649,67 +1409,6 @@ mod tests { assert_eq!(decoded.schema().as_arrow(), plan.schema().as_arrow()); } - /// Payload information-schema resolution exposes registered table functions as routines. - #[tokio::test] - async fn test_serializer_payload_information_schema_exposes_table_functions() { - let query_ctx = Arc::new(QueryContext::with("manual_cat", "manual_schema")); - let engine = QueryEngineFactory::new( - numbers_catalog_manager(), - None, - None, - None, - None, - false, - QueryOptions::default(), - ) - .query_engine(); - let state = engine.engine_context(query_ctx).state().clone(); - let config = state.config().clone().with_information_schema(true); - let state = SessionStateBuilder::new_from_existing(state) - .with_config(config) - .with_table_function_list(datafusion::functions_table::all_default_table_functions()) - .build(); - - let extensions = Extensions::default(); - let consumer = MergeScanSubstraitConsumer { - inner: DefaultSubstraitConsumer::new(&extensions, &state), - session_state: &state, - catalog_manager: state - .config() - .get_extension::() - .map(|s| s.catalog_manager().clone()), - payload: true, - }; - let provider = consumer - .resolve_table_ref(&TableReference::partial("information_schema", "routines")) - .await - .unwrap() - .unwrap(); - let plan = provider.scan(&state, None, &[], None).await.unwrap(); - let batches = datafusion::physical_plan::collect(plan, state.task_ctx()) - .await - .unwrap(); - let schema = provider.schema(); - let routine_name_index = schema.index_of("routine_name").unwrap(); - let function_type_index = schema.index_of("function_type").unwrap(); - assert!(batches.iter().any(|batch| { - let routine_names = batch - .column(routine_name_index) - .as_any() - .downcast_ref::() - .unwrap(); - let function_types = batch - .column(function_type_index) - .as_any() - .downcast_ref::() - .unwrap(); - (0..batch.num_rows()).any(|row| { - routine_names.value(row) == "generate_series" - && function_types.value(row) == "TABLE" - }) - })); - } - /// `MergeScan` decodes under an ordinary single-input extension parent. #[tokio::test] async fn test_serializer_decode_merge_scan_under_series_normalize() { diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index b7089c336710..095968743b26 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -41,7 +41,7 @@ use common_datasource::object_store::LocalFileAccess; use common_event_recorder::EventRecorderOptions; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_meta::DatanodeId; -use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder}; +use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder}; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::chroot::ChrootKvBackend; use common_meta::kv_backend::etcd::EtcdStore; @@ -113,8 +113,6 @@ pub struct GreptimeDbCluster { /// The stats of the network address of every datanode, empty unless the cluster serves the /// datanodes at their own addresses. pub datanode_rpc_stats: HashMap>, - /// Cache registries used by datanodes to apply cache invalidations. - pub datanode_cache_registries: HashMap>, pub kv_backend: KvBackendRef, pub metasrv: Arc, pub frontend: Arc, @@ -381,7 +379,7 @@ impl GreptimeDbClusterBuilder { ) .await; - let (datanode_instances, datanode_cache_registries) = self + let datanode_instances = self .build_datanodes_with_options(&metasrv, &datanode_options) .await; @@ -420,7 +418,6 @@ impl GreptimeDbClusterBuilder { guards, datanode_instances, datanode_rpc_stats, - datanode_cache_registries, kv_backend: self.kv_backend.clone(), metasrv: metasrv.metasrv, frontend: Arc::new(frontend), @@ -503,22 +500,16 @@ impl GreptimeDbClusterBuilder { &self, metasrv: &MockInfo, options: &[DatanodeOptions], - ) -> ( - HashMap, - HashMap>, - ) { + ) -> HashMap { let mut instances = HashMap::with_capacity(options.len()); - let mut cache_registries = HashMap::with_capacity(options.len()); for opts in options { - let (datanode, cache_registry) = - self.create_datanode(opts.clone(), metasrv.clone()).await; + let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await; let datanode_id = opts.node_id.unwrap(); instances.insert(datanode_id, datanode); - cache_registries.insert(datanode_id, cache_registry); } - (instances, cache_registries) + instances } async fn wait_datanodes_alive( @@ -544,11 +535,7 @@ impl GreptimeDbClusterBuilder { panic!("Some Datanodes are not alive in 10 seconds!") } - async fn create_datanode( - &self, - opts: DatanodeOptions, - metasrv: MockInfo, - ) -> (Datanode, Arc) { + async fn create_datanode(&self, opts: DatanodeOptions, metasrv: MockInfo) -> Datanode { let mut meta_client = MetaClientBuilder::datanode_default_options(opts.node_id.unwrap()) .channel_manager(metasrv.channel_manager) .build(); @@ -568,7 +555,7 @@ impl GreptimeDbClusterBuilder { datanode.start_heartbeat().await.unwrap(); - (datanode, layered_cache_registry) + datanode } async fn build_frontend( diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 3feaca375f9c..59a41038ea98 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -216,8 +216,6 @@ impl GreptimeDbStandaloneBuilder { ) -> GreptimeDbStandalone { let plugins = self.plugin.clone().unwrap_or_default(); - // Use the same layered registry as the datanode of a running instance, so the derived - // caches of the test are invalidated after the caches they are derived from. let layered_cache_registry = Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 110600b4460e..9be69a645a0d 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -23,9 +23,6 @@ use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use api::v1::region::{QueryRequest, RegionRequestHeader}; -use common_meta::cache::TableRouteCacheRef; -use common_meta::cache_invalidator::{CacheInvalidator, Context}; -use common_meta::instruction::CacheIdent; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::info; @@ -39,7 +36,7 @@ use query::query_engine::DefaultSerializer; use servers::query_handler::sql::SqlQueryHandler; use session::context::{QueryContext, QueryContextRef}; use store_api::region_request::{RegionCloseRequest, RegionRequest}; -use store_api::storage::{RegionId, TableId}; +use store_api::storage::RegionId; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; @@ -51,7 +48,6 @@ const PROBE_TABLE: &str = "nested_cap_probe"; const BUILD_TABLE: &str = "nested_cap_build"; /// Missing inner table used to verify plan-resolution failure. const MISSING_INNER_TABLE: &str = "nested_cap_missing_build"; -const UNRELATED_TABLE_ID: TableId = 10_000_000; /// `(a_id, probe_key, probe_v)` fixture rows; `a_id` is the partition key. const PROBE_ROWS: &[(i32, i32, i32)] = &[ @@ -129,57 +125,6 @@ async fn test_nested_merge_scan_capability_join_across_datanodes() { assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests); } -/// Explicit route eviction and unrelated invalidation finish before execution; nested reads refill -/// the route and return complete rows. This does not exercise a concurrent invalidation race. -#[tokio::test(flavor = "multi_thread")] -async fn test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation() { - common_telemetry::init_default_ut_logging(); - - let cluster = - build_cluster("test_nested_merge_scan_capability_cold_load_with_unrelated_invalidation") - .await; - let frontend = cluster.fe_instance().clone(); - prepare_tables(&frontend).await; - - let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; - let build_leaders = region_leaders(&frontend, BUILD_TABLE).await; - assert_eq!(2, datanodes(&probe_leaders).len()); - assert_eq!(2, datanodes(&build_leaders).len()); - let build_table_id = table_id(&build_leaders); - - let reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; - let expected = expected_join_rows(); - assert_eq!(expected.clone(), multiset(table_cells(&reference))); - - let query_ctx = query_ctx(); - let plan = encode(&nested_plan(&frontend, &query_ctx).await); - invalidate_table_ids(&cluster, &[build_table_id]).await; - assert_route_cold(&cluster, build_table_id); - invalidate_table_ids(&cluster, &[UNRELATED_TABLE_ID]).await; - assert_route_cold(&cluster, build_table_id); - - let remote_datanodes = datanodes(&build_leaders); - let baseline_requests = remote_datanodes - .iter() - .map(|datanode| (*datanode, rpc_requests(&cluster, *datanode))) - .collect::>(); - let rows = run_nested_queries(&cluster, &probe_leaders, &plan, &query_ctx).await; - let actual = multiset(rows); - assert_eq!( - expected, actual, - "unexpected nested join rows after cold load" - ); - assert_eq!(multiset(table_cells(&reference)), actual); - assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests); - - for datanode in datanodes(&probe_leaders) { - assert!( - table_route_cache(&cluster, datanode).contains_key(&build_table_id), - "datanode {datanode} must refill the route of build table {build_table_id}" - ); - } -} - /// A failing inner region makes the whole query fail: the datanode must not return the rows of the /// inner regions that succeed. #[tokio::test(flavor = "multi_thread")] @@ -326,64 +271,6 @@ fn region_server(cluster: &GreptimeDbCluster, datanode_id: u64) -> RegionServer .region_server() } -/// The id of the single table that the regions of `leaders` belong to. -fn table_id(leaders: &BTreeMap) -> TableId { - let table_ids = leaders - .keys() - .map(|region_id| RegionId::from_u64(*region_id).table_id()) - .collect::>(); - assert_eq!( - 1, - table_ids.len(), - "expected the regions to belong to a single table, actual region leaders: {leaders:?}" - ); - - *table_ids.first().expect("there is at least one region") -} - -/// Returns the table route cache of the datanode `datanode_id`: the cache that -/// `DatanodeRegionQueryHandler::select_target` reads to resolve the leader of a region. -fn table_route_cache(cluster: &GreptimeDbCluster, datanode_id: u64) -> TableRouteCacheRef { - cluster - .datanode_cache_registries - .get(&datanode_id) - .unwrap_or_else(|| panic!("expected the cache registry of the datanode {datanode_id}")) - .get::() - .unwrap_or_else(|| panic!("expected the table route cache of the datanode {datanode_id}")) -} - -/// Invalidates `table_ids` on every datanode, the way a datanode handles the invalidation -/// instruction of metasrv (see [`GreptimeDbCluster::datanode_cache_registries`]). -async fn invalidate_table_ids(cluster: &GreptimeDbCluster, table_ids: &[TableId]) { - let idents = table_ids - .iter() - .copied() - .map(CacheIdent::TableId) - .collect::>(); - for (datanode_id, registry) in &cluster.datanode_cache_registries { - registry - .invalidate(&Context::default(), &idents) - .await - .unwrap_or_else(|e| { - panic!("the datanode {datanode_id} must invalidate {idents:?}: {e}") - }); - } -} - -/// Asserts that no datanode holds the route of `table_id`, i.e. the route is cold. -fn assert_route_cold(cluster: &GreptimeDbCluster, table_id: TableId) { - let cached = cluster - .datanode_cache_registries - .keys() - .filter(|datanode_id| table_route_cache(cluster, **datanode_id).contains_key(&table_id)) - .collect::>(); - assert!( - cached.is_empty(), - "the datanodes {cached:?} must drop the route of the table {table_id} after the \ - invalidation" - ); -} - async fn run_nested_queries( cluster: &GreptimeDbCluster, probe_leaders: &BTreeMap, From b442339cf476d8116ba2d7e2a4ce38b9299c7c7a Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:34:14 +0800 Subject: [PATCH 30/37] fix(session): preserve read preference across region query hops Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + src/datanode/src/error.rs | 10 +- src/datanode/src/region_server.rs | 22 ++- src/flow/src/adapter/flownode_impl.rs | 11 +- src/servers/src/grpc/greptime_handler.rs | 9 +- src/servers/src/http/hints.rs | 9 +- src/session/Cargo.toml | 3 + src/session/src/context.rs | 176 ++++++++++++++++-- src/session/src/hints.rs | 5 +- .../nested_merge_scan_capability_test.rs | 124 ++++++++++++ tests/compatibility/README.md | 13 +- .../wire/read_preference/README.md | 45 +++++ .../read_preference/legacy_query_context.pb | 3 + 13 files changed, 406 insertions(+), 25 deletions(-) create mode 100644 tests/compatibility/wire/read_preference/README.md create mode 100644 tests/compatibility/wire/read_preference/legacy_query_context.pb diff --git a/Cargo.lock b/Cargo.lock index c16572c2fcc9..5986a2348143 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -13696,6 +13696,7 @@ dependencies = [ "datafusion-common", "derive_builder 0.20.2", "derive_more", + "prost 0.14.1", "snafu 0.8.6", "sql", "uuid", diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index 539712711b71..e9b5d5ec1dca 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -145,6 +145,13 @@ pub enum Error { source: Box, }, + #[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 { .. } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 7d5218f30b75..ef333cff04d8 100644 --- a/src/datanode/src/region_server.rs +++ b/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(); diff --git a/src/flow/src/adapter/flownode_impl.rs b/src/flow/src/adapter/flownode_impl.rs index 62e8597dbb67..cf4b6fd2d349 100644 --- a/src/flow/src/adapter/flownode_impl.rs +++ b/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), diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index ad0a693ac598..c29ecc6493a8 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/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 diff --git a/src/servers/src/http/hints.rs b/src/servers/src/http/hints.rs index 09e4879dd50f..7ef9401cf185 100644 --- a/src/servers/src/http/hints.rs +++ b/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") diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index 5485198c266e..93c46ff7ff72 100644 --- a/src/session/Cargo.toml +++ b/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 diff --git a/src/session/src/context.rs b/src/session/src/context.rs index c1448053f352..2f81dfabffbf 100644 --- a/src/session/src/context.rs +++ b/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; @@ -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 { 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 for QueryContext { - fn from(ctx: api::v1::QueryContext) -> Self { +impl TryFrom for QueryContext { + type Error = crate::session_config::Error; + + fn try_from(mut ctx: api::v1::QueryContext) -> Result { + 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 for QueryContext { .unwrap_or_default(), ))) .explain_options(ctx.explain) - .build() + .read_preference(read_preference) + .build()) } } @@ -234,6 +254,17 @@ impl From 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() diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index 971426f4cc03..3ed4fd8dc2e0 100644 --- a/src/session/src/hints.rs +++ b/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)); } } diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 9be69a645a0d..298bf0e41776 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -23,6 +23,11 @@ use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use api::v1::region::{QueryRequest, RegionRequestHeader}; +use client::Client; +use client::region::RegionRequester; +use common_error::ext::ErrorExt; +use common_error::status_code::StatusCode; +use common_grpc::channel_manager::ChannelManager; use common_query::Output; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::info; @@ -125,6 +130,125 @@ async fn test_nested_merge_scan_capability_join_across_datanodes() { assert_cross_datanode_region_query(&cluster, &remote_datanodes, &baseline_requests); } +/// Internal read-preference wire metadata is rejected before nested execution on both inbound paths. +#[tokio::test(flavor = "multi_thread")] +async fn test_nested_merge_scan_capability_rejects_internal_read_preference() { + use common_meta::node_manager::Datanode; + use common_query::request::QueryRequest as LogicalQueryRequest; + use futures::TryStreamExt; + use session::hints::READ_PREFERENCE_EXTENSION_KEY; + + common_telemetry::init_default_ut_logging(); + + let cluster = build_cluster("test_nested_merge_scan_internal_read_preference").await; + let frontend = cluster.fe_instance().clone(); + prepare_tables(&frontend).await; + + let probe_leaders = region_leaders(&frontend, PROBE_TABLE).await; + let query_ctx = query_ctx(); + let plan = nested_plan(&frontend, &query_ctx).await; + let (region_id, (datanode_id, address)) = probe_leaders.iter().next().unwrap(); + let target = region_server(&cluster, *datanode_id); + // Confirm the same plan completes on the valid default-Leader path before asserting rejection. + let valid_rows = run_nested_queries(&cluster, &probe_leaders, &encode(&plan), &query_ctx).await; + assert_eq!(expected_join_rows(), multiset(valid_rows)); + let mut baseline = cluster + .datanode_instances + .keys() + .map(|id| (*id, rpc_requests(&cluster, *id))) + .collect::>(); + + for value in ["", "invalid", "follower"] { + let mut header = region_request_header(&query_ctx); + header + .query_context + .as_mut() + .unwrap() + .extensions + .insert(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string()); + let header = Some(header); + let error = target + .handle_read(LogicalQueryRequest { + header: header.clone(), + region_id: RegionId::from_u64(*region_id), + plan: plan.clone(), + }) + .await + .err() + .expect("internal read preference must be rejected at logical ingress"); + assert_eq!( + StatusCode::InvalidArguments, + error.status_code(), + "value={value:?}" + ); + let logical_message = error.to_string().to_ascii_lowercase(); + assert!( + logical_message.contains("read_preference"), + "expected an invalid query-context read-preference error, got: {logical_message}" + ); + if !value.is_empty() { + assert!( + logical_message.contains(value), + "unexpected logical error: {logical_message}" + ); + } + for id in cluster.datanode_instances.keys() { + assert_eq!( + baseline[id], + rpc_requests(&cluster, *id), + "logical invalid context caused DoGet at datanode {id}" + ); + } + + let client = Client::with_query_and_control_managers( + ChannelManager::new(), + ChannelManager::new(), + [address.as_str()], + ); + let requester = RegionRequester::new(client, false, false); + let request = LogicalQueryRequest { + header, + region_id: RegionId::from_u64(*region_id), + plan: plan.clone(), + }; + let flight_message = match Datanode::handle_query(&requester, request).await { + Err(error) => { + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + error.to_string().to_ascii_lowercase() + } + Ok(stream) => { + let error = stream + .try_collect::>() + .await + .expect_err("Flight must reject invalid internal read preference"); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + error.to_string().to_ascii_lowercase() + } + }; + assert!( + flight_message.contains("read_preference"), + "expected an invalid-parameter read-preference error, got: {flight_message}" + ); + if !value.is_empty() { + assert!( + flight_message.contains(value), + "unexpected Flight error: {flight_message}" + ); + } + for id in cluster.datanode_instances.keys() { + let expected = baseline[id] + u8::from(*id == *datanode_id) as usize; + assert_eq!( + expected, + rpc_requests(&cluster, *id), + "invalid Flight context caused nested DoGet at datanode {id}" + ); + } + for id in cluster.datanode_instances.keys() { + baseline.insert(*id, rpc_requests(&cluster, *id)); + } + } +} + /// A failing inner region makes the whole query fail: the datanode must not return the rows of the /// inner regions that succeed. #[tokio::test(flavor = "multi_thread")] diff --git a/tests/compatibility/README.md b/tests/compatibility/README.md index 542f10debfd5..404b6332b947 100644 --- a/tests/compatibility/README.md +++ b/tests/compatibility/README.md @@ -2,7 +2,18 @@ Compatibility tests verify that one GreptimeDB version can restart on state written by another version. -Tests are run via `cargo sqlness compat` and reuse the sqlness-runner infrastructure. +SQL restart tests are run via `cargo sqlness compat` and reuse the sqlness-runner infrastructure. + +## Wire fixtures + +Frozen protobuf fixtures are exercised by the Rust tests that own their context +conversion, not by SQLness case discovery. The +[legacy query-context fixture](wire/read_preference/README.md) checks omitted +read-preference defaulting and unrelated metadata preservation: + +```shell +cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)' +``` ## Quick Start diff --git a/tests/compatibility/wire/read_preference/README.md b/tests/compatibility/wire/read_preference/README.md new file mode 100644 index 000000000000..9e7b9ba74c10 --- /dev/null +++ b/tests/compatibility/wire/read_preference/README.md @@ -0,0 +1,45 @@ +# Legacy query-context wire compatibility + +`legacy_query_context.pb` is a generated legacy-shaped protobuf fixture, not a +capture from a released binary. It uses the unchanged `greptime.v1.QueryContext` +schema at greptime-proto revision +`9533078d6fd4bdaddbca0750cb9bf7ad64b2e568` and omits the internal +`query.read_preference` extension. + +Logical contents: + +- Catalog `c1`, schema `s1`, timezone `UTC`. +- Unrelated extension `flow.return_region_seq = true`. +- Channel `Internal` (`255`); no snapshot sequences or explain options. + +Run the real protobuf decode and session conversion regression: + +```shell +cargo nextest run -p session -E 'test(test_legacy_query_context_wire_defaults_to_leader)' +``` + +The test checks that the omitted preference defaults to Leader, that the original +fields and unrelated extension survive conversion and re-encoding, and that +Leader does not introduce a preference carrier. Existing session tests cover +present invalid or unsupported values. This fixture does not demonstrate that +an old receiver honors a non-Leader preference sent by an updated node. + +## Reproduce the fixture + +With `PROTO_ROOT` pointing to the `proto` directory of the revision above: + +```shell +printf '%s\n' \ + 'current_catalog: "c1"' \ + 'current_schema: "s1"' \ + 'timezone: "UTC"' \ + 'extensions { key: "flow.return_region_seq" value: "true" }' \ + 'channel: 255' | + protoc --proto_path="$PROTO_ROOT" --encode=greptime.v1.QueryContext \ + greptime/v1/common.proto > legacy_query_context.pb +``` + +Generated with `libprotoc 25.1`. The 48-byte fixture has SHA-256 +`4127ebbee8cdb1375ef8ad52754be56ed206c12513e5f61c4354c04587cba638`. +The test never rewrites it and does not require protobuf map ordering to remain +stable. diff --git a/tests/compatibility/wire/read_preference/legacy_query_context.pb b/tests/compatibility/wire/read_preference/legacy_query_context.pb new file mode 100644 index 000000000000..1990ede94119 --- /dev/null +++ b/tests/compatibility/wire/read_preference/legacy_query_context.pb @@ -0,0 +1,3 @@ + +c1s1"UTC* +flow.return_region_seqtrue0ÿ \ No newline at end of file From 6fcfd44b5fbf7b4b40e181cfc981b6cdd212b6e0 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:53:31 +0800 Subject: [PATCH 31/37] test(session): verify strict context rejection at Flow ingress Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- tests-integration/src/cluster.rs | 30 ++++++++++++++++++- .../nested_merge_scan_capability_test.rs | 6 ++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 095968743b26..840f389052a1 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -1115,13 +1115,41 @@ mod tests { // Assert: the inactive peer was rejected and the production Flow service handled the RPC. assert_eq!( - (vec![active_peer.addr], vec![inactive_peer.addr]), + (vec![active_peer.addr.clone()], vec![inactive_peer.addr]), client.peer_addresses_by_state() ); assert_eq!( StatusCode::InvalidArguments, response.unwrap_err().status_code() ); + + for value in ["", "unknown", "follower"] { + let mut query_context = + api::v1::QueryContext::from(session::context::QueryContext::arc().as_ref()); + query_context.extensions.insert( + session::hints::READ_PREFERENCE_EXTENSION_KEY.to_string(), + value.to_string(), + ); + let error = clients + .flownode(&active_peer) + .await + .handle(FlowRequest { + header: Some(api::v1::flow::FlowRequestHeader { + tracing_context: Default::default(), + query_context: Some(query_context), + }), + body: None, + }) + .await + .unwrap_err(); + assert_eq!(StatusCode::InvalidArguments, error.status_code()); + let message = error.output_msg(); + assert!(message.contains("read_preference"), "{message}"); + if !value.is_empty() { + assert!(message.contains(value), "{message}"); + } + } + flownode.shutdown().await.unwrap(); } } diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 298bf0e41776..2c86589112d8 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -181,7 +181,7 @@ async fn test_nested_merge_scan_capability_rejects_internal_read_preference() { error.status_code(), "value={value:?}" ); - let logical_message = error.to_string().to_ascii_lowercase(); + let logical_message = error.output_msg().to_ascii_lowercase(); assert!( logical_message.contains("read_preference"), "expected an invalid query-context read-preference error, got: {logical_message}" @@ -214,7 +214,7 @@ async fn test_nested_merge_scan_capability_rejects_internal_read_preference() { let flight_message = match Datanode::handle_query(&requester, request).await { Err(error) => { assert_eq!(StatusCode::InvalidArguments, error.status_code()); - error.to_string().to_ascii_lowercase() + error.output_msg().to_ascii_lowercase() } Ok(stream) => { let error = stream @@ -222,7 +222,7 @@ async fn test_nested_merge_scan_capability_rejects_internal_read_preference() { .await .expect_err("Flight must reject invalid internal read preference"); assert_eq!(StatusCode::InvalidArguments, error.status_code()); - error.to_string().to_ascii_lowercase() + error.output_msg().to_ascii_lowercase() } }; assert!( From 2c016aa0ca4fca6ab862b84aabcb0402d0a04863 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:54:26 +0800 Subject: [PATCH 32/37] docs(cmd): correct regression runner parser spelling Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cmd/src/bin/query_regression_runner/measure.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cmd/src/bin/query_regression_runner/measure.rs b/src/cmd/src/bin/query_regression_runner/measure.rs index a217d7a64497..a3ef70487006 100644 --- a/src/cmd/src/bin/query_regression_runner/measure.rs +++ b/src/cmd/src/bin/query_regression_runner/measure.rs @@ -355,7 +355,7 @@ fn executed_plan_preflight(sample: &mut Value, operator: &str) -> Option } /// The `(stage, plan)` of one executed TEXT DistAnalyze row, or `Ok(None)` for the explicit -/// null/null `Total rows:` trailer row. A row with a missing, non-null, or unparseable stage or +/// null/null `Total rows:` trailer row. A row with a missing, non-null, or unparsable stage or /// node, or without non-empty plan text, is malformed. fn plan_row(row: &Value) -> std::result::Result, String> { // Only a row whose `stage` and `node` are explicitly null is the trailer; a row with missing From 53acb9829556d1334abc5c1947709c766d9dc6c4 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:27:47 +0800 Subject: [PATCH 33/37] fix(query): only estimate rows when experimental_dist_join is enabled The default query path must not receive row statistics that could flip DataFusion JoinSelection from Partitioned to CollectLeft. Only compute row estimates when the session explicitly opts in with SET experimental_dist_join = true. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/planner.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index 319e95ccc33d..00213d880b0b 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -227,7 +227,13 @@ impl ExtensionPlanner for DistExtensionPlanner { .config() .get_extension() .unwrap_or_else(QueryContext::arc); - let row_estimate = self.ordinary_scan_row_estimate(input_plan, ®ions).await; + // Only estimate rows for the nested broadcast join; the default query path must not + // receive statistics that could flip JoinSelection (e.g. Partitioned → CollectLeft). + let row_estimate = if query_ctx.configuration_parameter().experimental_dist_join() { + self.ordinary_scan_row_estimate(input_plan, ®ions).await + } else { + None + }; let merge_scan_plan = MergeScanExec::new( session_state, table_name, From 5f4aa4c17e0bb38c07c11ecdc5b44b20f67492c9 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 12 Oct 2026 20:24:12 +0800 Subject: [PATCH 34/37] test: regenerate sqlness results for dist join planner defaults dist_join.result: the default (non-opt-in) physical plan now stays a partitioned hash join (RepartitionExec on both sides) after row estimates are only computed when experimental_dist_join is enabled; drop the CollectLeft-shaped expectation recorded while estimates were injected unconditionally. explain.result: DistJoinPlanner is a registered analyzer rule, so EXPLAIN VERBOSE now prints its 'logical_plan after DistJoinPlanner' stage (same pattern as DistPlannerAnalyzer/JsonSchemaConcretizeRule). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../cases/distributed/explain/dist_join.result | 18 ++++++++---------- .../common/tql-explain-analyze/explain.result | 3 +++ 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/tests/cases/distributed/explain/dist_join.result b/tests/cases/distributed/explain/dist_join.result index 31ea97a4959e..e53966b0ee3f 100644 --- a/tests/cases/distributed/explain/dist_join.result +++ b/tests/cases/distributed/explain/dist_join.result @@ -130,15 +130,14 @@ ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; |_|_TableScan: dist_join_build, partial_filters=[dist_join_build.b_key IS NOT NULL]_| |_| ]]_| | physical_plan | SortPreservingMergeExec: [p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST]_| -|_|_ProjectionExec: expr=[p_id@3 as p_id, p_key@4 as p_key, p_val@5 as p_val, b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| -|_|_SortExec: expr=[p_val@5 ASC NULLS LAST, p_id@3 ASC NULLS LAST, b_val@2 ASC NULLS LAST, b_id@0 ASC NULLS LAST], preserve_partitioning=[true] | +|_|_SortExec: expr=[p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST], preserve_partitioning=[true] | |_|_REDACTED -|_|_CoalescePartitionsExec_| -|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| -|_|_MergeScanExec: REDACTED |_|_RepartitionExec: partitioning=REDACTED |_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val]_| |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED +|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| +|_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -240,15 +239,14 @@ ORDER BY p.p_val, p.p_id, b.b_val, b.b_id; |_|_TableScan: dist_join_build, partial_filters=[dist_join_build.b_key IS NOT NULL]_| |_| ]]_| | physical_plan | SortPreservingMergeExec: [p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST]_| -|_|_ProjectionExec: expr=[p_id@3 as p_id, p_key@4 as p_key, p_val@5 as p_val, b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| -|_|_SortExec: expr=[p_val@5 ASC NULLS LAST, p_id@3 ASC NULLS LAST, b_val@2 ASC NULLS LAST, b_id@0 ASC NULLS LAST], preserve_partitioning=[true] | +|_|_SortExec: expr=[p_val@2 ASC NULLS LAST, p_id@0 ASC NULLS LAST, b_val@5 ASC NULLS LAST, b_id@3 ASC NULLS LAST], preserve_partitioning=[true] | |_|_REDACTED -|_|_CoalescePartitionsExec_| -|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| -|_|_MergeScanExec: REDACTED |_|_RepartitionExec: partitioning=REDACTED |_|_ProjectionExec: expr=[p_id@0 as p_id, p_key@1 as p_key, p_val@2 as p_val]_| |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED +|_|_ProjectionExec: expr=[b_id@0 as b_id, b_key@1 as b_key, b_val@2 as b_val]_| +|_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/standalone/common/tql-explain-analyze/explain.result b/tests/cases/standalone/common/tql-explain-analyze/explain.result index 3e8b0ec21334..8a940e4bc522 100644 --- a/tests/cases/standalone/common/tql-explain-analyze/explain.result +++ b/tests/cases/standalone/common/tql-explain-analyze/explain.result @@ -100,6 +100,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test; |_|_Filter: test.j >= TimestampMillisecond(-299999, None) AND test.j <= TimestampMillisecond(10000, None)_| |_|_TableScan: test, partial_filters=[test.j >= TimestampMillisecond(-299999, None), test.j <= TimestampMillisecond(10000, None)] | |_| ]]_| +| logical_plan after DistJoinPlanner_| SAME TEXT AS ABOVE_| | logical_plan after JsonSchemaConcretizeRule_| SAME TEXT AS ABOVE_| | logical_plan after FixStateUdafOrderingAnalyzer_| SAME TEXT AS ABOVE_| | analyzed_logical_plan_| SAME TEXT AS ABOVE_| @@ -254,6 +255,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test AS series; |_|_Filter: test.j >= TimestampMillisecond(-299999, None) AND test.j <= TimestampMillisecond(10000, None)_| |_|_TableScan: test, partial_filters=[test.j >= TimestampMillisecond(-299999, None), test.j <= TimestampMillisecond(10000, None)] | |_| ]]_| +| logical_plan after DistJoinPlanner_| SAME TEXT AS ABOVE_| | logical_plan after JsonSchemaConcretizeRule_| SAME TEXT AS ABOVE_| | logical_plan after FixStateUdafOrderingAnalyzer_| SAME TEXT AS ABOVE_| | analyzed_logical_plan_| SAME TEXT AS ABOVE_| @@ -440,6 +442,7 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test_nano; |_|_Filter: test_nano.j > TimestampNanosecond(-300000000000, None) AND test_nano.j <= TimestampNanosecond(10000000000, None)_| |_|_TableScan: test_nano, partial_filters=[test_nano.j > TimestampNanosecond(-300000000000, None), test_nano.j <= TimestampNanosecond(10000000000, None)] | |_| ]]_| +| logical_plan after DistJoinPlanner_| SAME TEXT AS ABOVE_| | logical_plan after JsonSchemaConcretizeRule_| SAME TEXT AS ABOVE_| | logical_plan after FixStateUdafOrderingAnalyzer_| SAME TEXT AS ABOVE_| | analyzed_logical_plan_| SAME TEXT AS ABOVE_| From 47f3b2998b1d207dca5a5977e45b969d2df6178b Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:02:28 +0800 Subject: [PATCH 35/37] fix(query): close volatile-side and wire-flag review findings for dist join Reviewer findings on the opt-in nested broadcast join: P1: the rewritten plan re-executes the SMALL-side subplan on every BIG-side region, so volatile expressions (e.g. random()) on the SMALL side change result semantics within one query. Add a volatility guard: contains_volatile_expr inspects the stripped small-side plan (projections, filters, table-scan filters, recursively) and refuses the rewrite. Conservative fallback only; a shared-materialization implementation is a separate stack PR. P2: experimental_dist_join never crossed the QueryContext protobuf roundtrip, so the datanode's nested MergeScan planning always skipped row estimates and JoinSelection could flip to BIG-build when only the append-only BIG side had exact statistics. Follow the read_preference extension-key precedent: encode query.experimental_dist_join in the protobuf extensions when enabled, restore on decode. Tests: volatile-small-side rejection unit test, protobuf roundtrip unit tests, and a real-cluster regression test asserting every dispatched BIG region's HashJoinExec builds the 2-row SMALL side and probes the append-only BIG side. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/dist_join_planner.rs | 84 +++++++++- src/session/src/context.rs | 52 +++++- src/session/src/hints.rs | 4 +- .../nested_merge_scan_capability_test.rs | 148 +++++++++++++++++- 4 files changed, 278 insertions(+), 10 deletions(-) diff --git a/src/query/src/dist_plan/dist_join_planner.rs b/src/query/src/dist_plan/dist_join_planner.rs index 638e3869e75c..27e2334f5894 100644 --- a/src/query/src/dist_plan/dist_join_planner.rs +++ b/src/query/src/dist_plan/dist_join_planner.rs @@ -358,7 +358,10 @@ impl NestedBroadcastJoinRewriter<'_> { } let (large_plan, large_merge_scan) = strip_merge_scan(&join.left)?; - let (_, small_merge_scan) = strip_merge_scan(&join.right)?; + let (small_plan, small_merge_scan) = strip_merge_scan(&join.right)?; + if contains_volatile_expr(&small_plan) { + return None; + } if large_merge_scan.is_placeholder() || small_merge_scan.is_placeholder() { return None; } @@ -493,6 +496,23 @@ fn restore_join_schema(join: Join, schema: datafusion_common::DFSchemaRef) -> Op ))) } +/// Whether `plan` contains a volatile expression. +/// +/// The small side is re-executed on every large-side region, so volatile expressions there +/// would change result semantics within a single query. +fn contains_volatile_expr(plan: &LogicalPlan) -> bool { + if plan.expressions().iter().any(Expr::is_volatile) { + return true; + } + if let LogicalPlan::TableScan(scan) = plan + && scan.filters.iter().any(Expr::is_volatile) + { + return true; + } + + plan.inputs().into_iter().any(contains_volatile_expr) +} + /// Whether `plan` contains a visible `MergeScan` node. fn contains_merge_scan(plan: &LogicalPlan) -> bool { if let LogicalPlan::Extension(extension) = plan @@ -558,7 +578,10 @@ mod tests { use common_meta::rpc::router::Region; use datafusion::datasource::DefaultTableSource; use datafusion_common::{DFSchema, DFSchemaRef, JoinConstraint, NullEquality, TableReference}; - use datafusion_expr::{Expr, LogicalPlanBuilder, build_join_schema, col, lit}; + use datafusion_expr::{ + Expr, LogicalPlanBuilder, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, build_join_schema, col, lit, + }; use datatypes::data_type::ConcreteDataType; use datatypes::extension::json::JsonExtensionType; use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; @@ -721,6 +744,35 @@ mod tests { join } + #[derive(Debug, PartialEq, Eq, Hash)] + struct VolatileTestUdf; + + impl ScalarUDFImpl for VolatileTestUdf { + fn name(&self) -> &str { + "volatile_test_udf" + } + + fn signature(&self) -> &Signature { + static SIGNATURE: std::sync::LazyLock = std::sync::LazyLock::new(|| { + Signature::uniform(1, vec![DataType::UInt32], Volatility::Volatile) + }); + &SIGNATURE + } + + fn return_type(&self, _arg_types: &[DataType]) -> datafusion_common::Result { + Ok(DataType::Float64) + } + + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + Ok(datafusion_expr::ColumnarValue::Scalar( + datafusion_common::ScalarValue::Float64(Some(0.5)), + )) + } + } + /// Without the build table option the join stays above the two `MergeScan`s. #[test] fn nested_broadcast_join_rewrite_disabled_by_default() { @@ -734,6 +786,34 @@ mod tests { assert_eq!(2, merge_scans(&result).len()); } + #[test] + fn nested_broadcast_join_rewrite_rejects_volatile_small_side() { + let volatile = Expr::ScalarFunction(datafusion_expr::expr::ScalarFunction::new_udf( + Arc::new(ScalarUDF::new_from_impl(VolatileTestUdf)), + vec![lit(1u32)], + )); + let small = LogicalPlanBuilder::from(table_scan("t2", 2, "t2")) + .project(vec![col("t2.number"), volatile]) + .unwrap() + .build() + .unwrap(); + let plan = DistPlannerAnalyzer {} + .analyze( + inner_join( + table_scan("t1", 1, "t1"), + small, + col("t1.number"), + col("t2.number"), + ), + &ConfigOptions::default(), + ) + .unwrap(); + + let result = rewrite(plan.clone()); + assert_eq!(plan.to_string(), result.to_string()); + assert!(matches!(result, LogicalPlan::Join(_))); + } + /// When statistics favor the small right side as build, its `MergeScan` becomes the /// inner scan, while the large left side probes locally on regions selected by its routes. /// diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 515ad44cd1de..839137978e6a 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -35,8 +35,9 @@ use derive_builder::Builder; use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect}; pub use crate::hints::{ - LIVE_ANALYZE_METRICS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY, - REMOTE_QUERY_ID_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + EXPERIMENTAL_DIST_JOIN_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; @@ -204,6 +205,11 @@ impl TryFrom for QueryContext { type Error = crate::session_config::Error; fn try_from(mut ctx: api::v1::QueryContext) -> Result { + let experimental_dist_join = ctx + .extensions + .remove(EXPERIMENTAL_DIST_JOIN_EXTENSION_KEY) + .as_deref() + == Some("true"); let read_preference = ctx.extensions.remove(READ_PREFERENCE_EXTENSION_KEY); let read_preference = match read_preference { Some(value) => ReadPreference::from_str(&value).map_err(|_| { @@ -217,7 +223,7 @@ impl TryFrom for QueryContext { None => ReadPreference::default(), }; let sequences = ctx.snapshot_seqs.as_ref(); - Ok(QueryContextBuilder::default() + let query_context = QueryContextBuilder::default() .current_catalog(ctx.current_catalog) .current_schema(ctx.current_schema) .timezone(parse_timezone(Some(&ctx.timezone))) @@ -235,7 +241,13 @@ impl TryFrom for QueryContext { ))) .explain_options(ctx.explain) .read_preference(read_preference) - .build()) + .build(); + if experimental_dist_join { + query_context + .configuration_parameter() + .set_experimental_dist_join(true); + } + Ok(query_context) } } @@ -249,6 +261,7 @@ impl From for api::v1::QueryContext { snapshot_seqs, sst_min_sequences, mutable_query_context_data, + configuration_parameter, .. }: QueryContext, ) -> Self { @@ -256,6 +269,13 @@ impl From for api::v1::QueryContext { let mutable_inner = mutable_inner.read().unwrap(); let mut extensions = extensions; extensions.remove(READ_PREFERENCE_EXTENSION_KEY); + extensions.remove(EXPERIMENTAL_DIST_JOIN_EXTENSION_KEY); + if configuration_parameter.experimental_dist_join() { + extensions.insert( + EXPERIMENTAL_DIST_JOIN_EXTENSION_KEY.to_string(), + "true".to_string(), + ); + } if mutable_inner.read_preference != ReadPreference::Leader { extensions.insert( READ_PREFERENCE_EXTENSION_KEY.to_string(), @@ -895,6 +915,30 @@ mod test { ); } + #[test] + fn test_experimental_dist_join_query_context_roundtrip() { + let enabled = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + enabled + .configuration_parameter() + .set_experimental_dist_join(true); + let wire: api::v1::QueryContext = enabled.into(); + assert_eq!( + Some(&"true".to_string()), + wire.extensions.get(EXPERIMENTAL_DIST_JOIN_EXTENSION_KEY) + ); + let restored = QueryContext::try_from(wire).unwrap(); + assert!(restored.configuration_parameter().experimental_dist_join()); + + let wire: api::v1::QueryContext = QueryContext::with(DEFAULT_CATALOG_NAME, "public").into(); + assert!( + !wire + .extensions + .contains_key(EXPERIMENTAL_DIST_JOIN_EXTENSION_KEY) + ); + let restored = QueryContext::try_from(wire).unwrap(); + assert!(!restored.configuration_parameter().experimental_dist_join()); + } + #[test] fn test_write_admission_is_local_and_database_scoped() { let ctx = QueryContext::with_channel("greptime", "public", Channel::Otlp); diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index 3ed4fd8dc2e0..b57ad1fbab58 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -28,12 +28,14 @@ pub const INSERT_SKIP_WAL_HINT: &str = "insert_skip_wal"; pub const READ_PREFERENCE_HINT: &str = "read_preference"; pub const READ_PREFERENCE_EXTENSION_KEY: &str = "query.read_preference"; -pub const RESERVED_EXTENSION_KEYS: [&str; 5] = [ +pub const EXPERIMENTAL_DIST_JOIN_EXTENSION_KEY: &str = "query.experimental_dist_join"; +pub const RESERVED_EXTENSION_KEYS: [&str; 6] = [ 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, + EXPERIMENTAL_DIST_JOIN_EXTENSION_KEY, ]; /// Deprecated, use `HINTS_KEY` instead. diff --git a/tests-integration/src/tests/nested_merge_scan_capability_test.rs b/tests-integration/src/tests/nested_merge_scan_capability_test.rs index 9b38169728ac..1da9a4802ab4 100644 --- a/tests-integration/src/tests/nested_merge_scan_capability_test.rs +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -394,6 +394,139 @@ async fn test_nested_broadcast_join_rewrite_executes_on_probe_regions() { ); } +/// The append-only BIG scan reports exact row counts to DataFusion; the SMALL side must remain +/// the hash-build input after the experimental setting crosses the nested-plan wire boundary. +#[tokio::test(flavor = "multi_thread")] +async fn test_experimental_dist_join_roundtrip_keeps_small_hash_build() { + common_telemetry::init_default_ut_logging(); + + let cluster = + build_cluster("test_experimental_dist_join_roundtrip_keeps_small_hash_build").await; + let frontend = cluster.fe_instance().clone(); + let big = "nested_cap_stats_probe"; + let small = "nested_cap_stats_build"; + + execute_sql( + &frontend, + &format!( + "CREATE TABLE {big} ( + big_id INT, + join_key INT, + payload STRING, + big_ts TIMESTAMP, + TIME INDEX (big_ts) + ) PARTITION ON COLUMNS (big_id) ( + big_id < 100, + big_id >= 100 + ) ENGINE=mito" + ), + ) + .await; + execute_sql( + &frontend, + &format!( + "CREATE TABLE {small} ( + small_id INT, + join_key INT, + small_ts TIMESTAMP, + TIME INDEX (small_ts), + PRIMARY KEY (small_id) + ) PARTITION ON COLUMNS (small_id) ( + small_id < 100, + small_id >= 100 + ) ENGINE=mito" + ), + ) + .await; + execute_sql( + &frontend, + &format!( + "INSERT INTO {big}(big_id, join_key, payload, big_ts) VALUES + (1, 10, repeat('large-big-payload-', 80), 1000), + (2, 20, repeat('large-big-payload-', 80), 2000), + (3, 30, repeat('large-big-payload-', 80), 3000), + (4, 40, repeat('large-big-payload-', 80), 4000), + (120, 10, repeat('large-big-payload-', 80), 5000), + (121, 20, repeat('large-big-payload-', 80), 6000), + (122, 30, repeat('large-big-payload-', 80), 7000), + (123, 40, repeat('large-big-payload-', 80), 8000), + (5, 10, repeat('large-big-payload-', 80), 9000), + (6, 20, repeat('large-big-payload-', 80), 10000), + (124, 10, repeat('large-big-payload-', 80), 11000), + (125, 20, repeat('large-big-payload-', 80), 12000), + (126, 10, repeat('large-big-payload-', 80), 13000), + (127, 20, repeat('large-big-payload-', 80), 14000), + (128, 10, repeat('large-big-payload-', 80), 15000), + (129, 20, repeat('large-big-payload-', 80), 16000)" + ), + ) + .await; + execute_sql( + &frontend, + &format!( + "INSERT INTO {small}(small_id, join_key, small_ts) VALUES + (1, 10, 1000), (2, 20, 2000)" + ), + ) + .await; + + let (big_bytes, small_bytes) = wait_for_named_region_statistics(&frontend, big, small, 2).await; + assert!( + 2 * small_bytes < big_bytes, + "the planner cost heuristic must favor SMALL: BIG={big_bytes} bytes, SMALL={small_bytes} bytes" + ); + let big_leaders = region_leaders(&frontend, big).await; + assert_eq!(2, big_leaders.len()); + assert_eq!(2, datanodes(&big_leaders).len()); + + let query_ctx = query_ctx(); + run_sql( + &frontend, + "SET experimental_dist_join = true", + query_ctx.clone(), + ) + .await; + let sql = format!( + "SELECT b.join_key, s.join_key FROM {big} b JOIN {small} s ON b.join_key = s.join_key" + ); + let actual = run_sql(&frontend, &sql, query_ctx).await; + let physical = actual + .meta + .plan + .clone() + .expect("query output must include its physical plan"); + let merge_scan = find_merge_scan_exec(&physical) + .unwrap_or_else(|| panic!("expected nested MergeScan, got {physical:?}")); + let rows = output_batches(actual).await; + let total_rows: usize = rows.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(12, total_rows); + + let stages = merge_scan.sub_stage_metrics(); + let hash_joins = stages + .iter() + .flat_map(|stage| &stage.plan_metrics) + .filter(|metric| metric.plan_name == "HashJoinExec") + .collect::>(); + assert_eq!( + merge_scan.regions().len(), + hash_joins.len(), + "each dispatched BIG region must report its DN HashJoinExec: {stages:?}" + ); + assert!( + hash_joins.iter().all(|metric| { + metric + .metrics + .iter() + .any(|(name, value)| *name == "build_input_rows" && *value == 2) + && metric + .metrics + .iter() + .any(|(name, value)| *name == "input_rows" && *value > 2) + }), + "each DN join must build the two SMALL rows and probe more BIG rows: {hash_joins:?}" + ); +} + /// Complete expected output for the unfiltered broadcast SQL, in its projected column order. fn expected_full_broadcast_rows() -> Vec> { let mut rows = Vec::new(); @@ -511,16 +644,25 @@ async fn region_statistics(frontend: &Arc, table: &str) -> (u64, u64, async fn wait_for_region_statistics( frontend: &Arc, expected_build_rows: u64, +) -> (u64, u64) { + wait_for_named_region_statistics(frontend, PROBE_TABLE, BUILD_TABLE, expected_build_rows).await +} + +async fn wait_for_named_region_statistics( + frontend: &Arc, + probe_table: &str, + build_table: &str, + expected_build_rows: u64, ) -> (u64, u64) { let deadline = tokio::time::Instant::now() + Duration::from_secs(120); loop { let mut stats = Vec::new(); - for table in [PROBE_TABLE, BUILD_TABLE] { + for table in [probe_table, build_table] { stats.push(region_statistics(frontend, table).await); } let probe = stats[0]; let build = stats[1]; - info!("region statistics of {PROBE_TABLE}: {probe:?}, of {BUILD_TABLE}: {build:?}"); + info!("region statistics of {probe_table}: {probe:?}, of {build_table}: {build:?}"); if probe.0 == 2 && probe.2 > 0 && probe.3 > 0 @@ -532,7 +674,7 @@ async fn wait_for_region_statistics( } assert!( tokio::time::Instant::now() < deadline, - "the region statistics of {PROBE_TABLE} {probe:?} and {BUILD_TABLE} {build:?} must be \ + "the region statistics of {probe_table} {probe:?} and {build_table} {build:?} must be \ reported" ); tokio::time::sleep(Duration::from_secs(1)).await; From f08a8b79aa682e3c8b52916510583f2c78e4c9cd Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:01:23 +0800 Subject: [PATCH 36/37] fix(query): extend volatile guard to JOIN key expressions The P1 guard rejected volatile expressions inside the small-side sub-plan but missed volatile expressions that appear only in join.on keys. ExtractEquijoinPredicate can extract volatile expressions into join keys (e.g. `ON b.k = CASE WHEN random() < 0.5 THEN s.k ELSE s.k + 1 END`), and the physical-plan-time Projection that computes such keys would be replicated per BIG region after the broadcast rewrite. Extend the guard to also check small-side join key expressions for volatility, and add a unit test covering this shape. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/dist_join_planner.rs | 28 +++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/src/query/src/dist_plan/dist_join_planner.rs b/src/query/src/dist_plan/dist_join_planner.rs index 27e2334f5894..e23f35855154 100644 --- a/src/query/src/dist_plan/dist_join_planner.rs +++ b/src/query/src/dist_plan/dist_join_planner.rs @@ -359,7 +359,12 @@ impl NestedBroadcastJoinRewriter<'_> { let (large_plan, large_merge_scan) = strip_merge_scan(&join.left)?; let (small_plan, small_merge_scan) = strip_merge_scan(&join.right)?; - if contains_volatile_expr(&small_plan) { + // The small sub-plan is re-executed per large region, so reject volatile + // expressions there. Join keys are not part of either input plan, so + // check the small-side keys explicitly as well. + if contains_volatile_expr(&small_plan) + || join.on.iter().any(|(_, small_key)| small_key.is_volatile()) + { return None; } if large_merge_scan.is_placeholder() || small_merge_scan.is_placeholder() { @@ -814,6 +819,27 @@ mod tests { assert!(matches!(result, LogicalPlan::Join(_))); } + #[test] + fn nested_broadcast_join_rewrite_rejects_volatile_join_key() { + let volatile = Expr::ScalarFunction(datafusion_expr::expr::ScalarFunction::new_udf( + Arc::new(ScalarUDF::new_from_impl(VolatileTestUdf)), + vec![col("t2.number")], + )); + // The small-side plan has no volatile expressions — the volatile + // expression only appears in the JOIN key. + let small = table_scan("t2", 2, "t2"); + let plan = DistPlannerAnalyzer {} + .analyze( + inner_join(table_scan("t1", 1, "t1"), small, col("t1.number"), volatile), + &ConfigOptions::default(), + ) + .unwrap(); + + let result = rewrite(plan.clone()); + assert_eq!(plan.to_string(), result.to_string()); + assert!(matches!(result, LogicalPlan::Join(_))); + } + /// When statistics favor the small right side as build, its `MergeScan` becomes the /// inner scan, while the large left side probes locally on regions selected by its routes. /// From 4b317f7e168ed69b3192710e823bc20e74f70318 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:14:17 +0800 Subject: [PATCH 37/37] test(query): strengthen volatile JOIN-key rejection test premise Declare the volatile test UDF as non-nullable via return_field_from_args so FilterNullJoinKeys does not inject a volatile IS-NOT-NULL filter into the small sub-plan. Add assertions that verify the test premise: the small plan has no volatile expressions and the volatile expression exists only in the JOIN key, ensuring the test exercises the join-key guard specifically. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/query/src/dist_plan/dist_join_planner.rs | 31 ++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/src/query/src/dist_plan/dist_join_planner.rs b/src/query/src/dist_plan/dist_join_planner.rs index e23f35855154..a78d014025bb 100644 --- a/src/query/src/dist_plan/dist_join_planner.rs +++ b/src/query/src/dist_plan/dist_join_planner.rs @@ -768,6 +768,20 @@ mod tests { Ok(DataType::Float64) } + fn return_field_from_args( + &self, + args: datafusion_expr::ReturnFieldArgs<'_>, + ) -> datafusion_common::Result { + let _ = args; + Ok(std::sync::Arc::new( + datafusion_common::arrow::datatypes::Field::new( + self.name(), + DataType::Float64, + false, // non-nullable + ), + )) + } + fn invoke_with_args( &self, _args: ScalarFunctionArgs, @@ -835,6 +849,23 @@ mod tests { ) .unwrap(); + // Verify the test premise: after the analyzer, the small sub-plan must + // NOT contain volatile expressions (they must only be in the JOIN key). + // This ensures the test exercises the join-key guard, not the small-plan guard. + if let LogicalPlan::Join(join) = &plan { + let (small_plan, _) = strip_merge_scan(&join.right).unwrap(); + assert!( + !contains_volatile_expr(&small_plan), + "small plan must not contain volatile expressions" + ); + assert!( + join.on.iter().any(|(_, key)| key.is_volatile()), + "join key must contain volatile expression" + ); + } else { + panic!("expected Join plan after analyzer"); + } + let result = rewrite(plan.clone()); assert_eq!(plan.to_string(), result.to_string()); assert!(matches!(result, LogicalPlan::Join(_)));