From 28697c9e550aa8dba3b39bbc873c68a96d96b047 Mon Sep 17 00:00:00 2001 From: discord9 Date: Sat, 10 Oct 2026 06:49:38 +0000 Subject: [PATCH] feat(datanode): execute nested MergeScan plans on the datanode (#9206) * feat(datanode): execute nested MergeScan plans on the datanode Give a datanode the ability to decode and execute a plan whose MergeScan carries a nested plan (e.g. an inner MergeScan over a build table), so a distributed sub-plan pushed to a datanode can itself fan out to other datanodes. This is the execution/codec capability only; it is not wired to any user-facing switch and produces no plan by itself. Codec (src/query/src/query_engine/default_serializer.rs): - MergeScanAwareSerializer keeps the request's SessionState and the engine catalog manager, and decodes a MergeScan payload's tables by (catalog, schema, table) name through the engine catalog rather than the region-aware request catalog list (which binds every table to the request's region and breaks on mismatched columns). Datanode (src/datanode): - New DatanodeRegionQueryHandler (region_query.rs) serves region queries for the MergeScan nodes of a received plan, resolving region leaders via PartitionRuleManager and dispatching to peer datanodes. - Wire KvBackendCatalogManager + PartitionRuleManager + NodeClients into the datanode's catalog/cache registry (with_dist_planner stays false: a datanode never rewrites its own plan). Bug fixes uncovered while wiring this up: - cache registry: register the table/table_info/table_name/table_route/ partition_info caches the datanode's catalog manager and partition rule manager require (previously missing -> panic on lookup). - default_serializer: decode_sub_plan bridge no longer drives a future on a current-thread runtime from within another executor (spawn a dedicated thread + runtime when not on a multi-thread runtime). Tested by query_engine::default_serializer unit tests: nested MergeScan round-trip and engine-catalog resolution with mismatched column names. No user-facing behavior change. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): keep GreptimeDB function bindings after rebuilding decode state SessionStateBuilder::new_from_existing collects the function registry (HashMap) into a Vec and build() re-registers each entry; register_udf registers aliases before the function's own name and the Vec order is unstable. Rebuilding the session state to install MergeScanAwareSerializer after the GreptimeDB functions were registered could therefore rebind a name/alias (e.g. date_format, an alias of DataFusion's to_char) back to the built-in function. This path runs for every decoded plan, not only nested MergeScan ones. Extract the GreptimeDB function registration into register_greptime_functions and re-apply it after the final build of both the top-level decode state and the MergeScan payload state. Add regression tests asserting date_format stays bound to the GreptimeDB implementation on both paths. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): address review on nested MergeScan capability Three follow-ups from review of the nested MergeScan execution capability: 1. Layer the datanode caches (P2b). table_cache derives from table_info_cache + table_name_cache, and partition_info_cache derives from table_route_cache, but they were all registered in a single flat CacheRegistry whose invalidate() runs concurrently with no ordering. A derived cache could be invalidated and refilled from a base cache that had not been invalidated yet, leaving stale metadata. Build the base caches and the derived caches as separate layers of a LayeredCacheRegistry (reusing the existing layering mechanism) so base caches finish invalidating before derived ones. Both the production path (cmd) and the test scaffolding now build the same layered registry; the flat build_datanode_cache_registry is removed. 2. Don't let a nested MergeScan's internal region query compete for the concurrency permit (P2a). An outer query holds its permit until its stream is fully consumed, and a nested MergeScan makes the datanode issue region queries to peer datanodes as execution stages of that same query. If such an inner stage also took a permit, the outer query waiting on it would deadlock against its own held permit (or time out when max_concurrent_queries > 0). Mark inner DN->DN region queries as internal stages: QueryRequest gains an flag carried over the wire in the request header's query_context extensions (reserved key, not injectable via hints), and the region server skips acquiring a concurrency permit for internal-stage requests on both acquire sites. Outer queries still acquire normally. 3. Add an end-to-end capability test (tests-integration/.../nested_merge_scan_capability_test.rs): hand-build Join(, MergeScan(build regions on both datanodes)) with mismatched column names and duplicate build keys, ship it to a datanode over a real TCP Flight endpoint, consume the stream fully and compare the row multiset against the frontend's reference join result, assert a real cross-datanode DoGet via an RPC counter, and assert an inner region failure fails the whole query. Also migrate the with_real_datanode_grpc_addr + DatanodeRpcStats test scaffolding. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): adapt date_format regression test to the updated DataFusion fork ScalarUDFImpl's trait bounds changed so its impl can no longer be downcast via as_any(); assert the function equals the GreptimeDB implementation instead, which to_char can never satisfy. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(datanode): cover the internal stage under the concurrency limit Add an end-to-end regression test for the nested MergeScan capability under max_concurrent_queries = 1: both datanodes' only permit is held by the outer queries' streams while their inner MergeScans fetch the build regions from the same datanodes over real Flight connections. The result consumption is bounded by a timeout, so a broken link anywhere in the chain (marking the inner request internal, carrying the marker in the request header, honoring it at the receiving acquire sites) fails the test instead of deadlocking. Also add a minimal with_datanode_options_override hook to the cluster builder, and fix the test module doc: the outer request is handed to RegionServer::handle_remote_read in process; only the inner region queries go over the wire. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(datanode): cover cold route loads under unrelated invalidations The table route cache uses InitStrategy::VersionChecked, whose version counter belongs to the whole CacheContainer rather than to a single table: any TableId invalidation bumps the shared version, so a cold load of one table's route retries when an unrelated table is invalidated concurrently. Document this as a known limitation on the region query handler, and pin the behavior with an integration test: run the nested plans while a background task invalidates an unrelated table id on every datanode, assert the results still match the expected multiset (retries only add work, they never return a stale route), and report the observed load amplification via the cache miss counter (2 loads without invalidations, 8 over three runs with 24 unrelated invalidations in this run). Also expose the datanode cache registries on the test cluster so a test can drive the local invalidation path that the heartbeat handler drives after a metasrv broadcast. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(query): trim nested plan decoder commentary Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor: remove redundant nested scan support plumbing Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): decode nested MergeScan payloads asynchronously Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): avoid unnecessary extension name allocation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): retain extension decoding for manual session states Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(datanode): keep nested scan query admission unchanged Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): honor query handler factories and trim nested scan fixtures Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datanode): retain required cache error import Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(query): reuse registered state for nested scan decoding Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): verify payload information schema table functions Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): retain nested scan runtime and cold-route guards Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(query): address nested scan review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(session): preserve read preference across region query hops Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(session): verify strict context rejection at Flow ingress Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(query): remove redundant serializer error case Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 4 + src/cache/src/lib.rs | 65 +- src/cmd/src/datanode/builder.rs | 11 +- src/common/substrait/src/df_substrait.rs | 19 +- src/common/substrait/src/lib.rs | 1 + src/datanode/Cargo.toml | 5 +- src/datanode/src/datanode.rs | 96 +- src/datanode/src/error.rs | 10 +- src/datanode/src/lib.rs | 1 + src/datanode/src/region_query.rs | 109 ++ src/datanode/src/region_server.rs | 22 +- src/flow/src/adapter/flownode_impl.rs | 11 +- .../src/query_engine/default_serializer.rs | 1371 +++++++++++++++-- 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 +- tests-integration/src/cluster.rs | 237 ++- tests-integration/src/standalone.rs | 9 +- tests-integration/src/tests.rs | 1 + .../nested_merge_scan_capability_test.rs | 749 +++++++++ tests/compatibility/README.md | 13 +- .../wire/read_preference/README.md | 45 + .../read_preference/legacy_query_context.pb | 3 + 25 files changed, 2790 insertions(+), 194 deletions(-) create mode 100644 src/datanode/src/region_query.rs create mode 100644 tests-integration/src/tests/nested_merge_scan_capability_test.rs 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 9b447169fbc..055bf712394 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4473,6 +4473,8 @@ dependencies = [ "async-trait", "bytes", "cache", + "catalog", + "client", "common-base", "common-catalog", "common-config", @@ -4480,6 +4482,7 @@ dependencies = [ "common-error", "common-function", "common-greptimedb-telemetry", + "common-grpc", "common-macro", "common-meta", "common-options", @@ -13663,6 +13666,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/cache/src/lib.rs b/src/cache/src/lib.rs index 1a234aa802e..c7304ae6009 100644 --- a/src/cache/src/lib.rs +++ b/src/cache/src/lib.rs @@ -20,8 +20,8 @@ use std::time::Duration; use catalog::kvbackend::new_table_cache; use common_meta::cache::{ - CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistryBuilder, new_schema_cache, - new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache, + CacheRegistry, CacheRegistryBuilder, LayeredCacheRegistry, LayeredCacheRegistryBuilder, + new_schema_cache, new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache, new_table_route_cache, new_table_schema_cache, new_view_info_cache, }; use common_meta::kv_backend::KvBackendRef; @@ -53,11 +53,11 @@ pub const TABLE_FLOWNODE_SET_CACHE_NAME: &str = "table_flownode_set_cache"; pub const TABLE_ROUTE_CACHE_NAME: &str = "table_route_cache"; pub const PARTITION_INFO_CACHE_NAME: &str = "partition_info_cache"; -/// Builds cache registry for datanode, including: -/// - Schema cache. -/// - Table id to schema name cache. -pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry { - // Builds table id schema name cache that never expires. +/// Builds the datanode cache registry with base metadata caches before derived caches. +/// +/// Invalidating the base layer first prevents a concurrent query from refilling a derived cache +/// from stale metadata after that derived cache has been invalidated. +pub fn build_datanode_layered_cache_registry(kv_backend: KvBackendRef) -> LayeredCacheRegistry { let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY).build(); let table_id_schema_cache = Arc::new(new_table_schema_cache( TABLE_SCHEMA_NAME_CACHE_NAME.to_string(), @@ -65,7 +65,6 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry kv_backend.clone(), )); - // Builds schema cache let cache = default_cache(); let schema_cache = Arc::new(new_schema_cache( SCHEMA_CACHE_NAME.to_string(), @@ -73,9 +72,57 @@ pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry kv_backend.clone(), )); - CacheRegistryBuilder::default() + let cache = default_cache(); + let table_info_cache = Arc::new(new_table_info_cache( + TABLE_INFO_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_name_cache = Arc::new(new_table_name_cache( + TABLE_NAME_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_route_cache = Arc::new(new_table_route_cache( + TABLE_ROUTE_CACHE_NAME.to_string(), + cache, + kv_backend.clone(), + )); + + let cache = default_cache(); + let table_cache = Arc::new(new_table_cache( + TABLE_CACHE_NAME.to_string(), + cache, + table_info_cache.clone(), + table_name_cache.clone(), + )); + + let cache = default_cache(); + let partition_info_cache = Arc::new(new_partition_info_cache( + PARTITION_INFO_CACHE_NAME.to_string(), + cache, + table_route_cache.clone(), + )); + + let base_registry = CacheRegistryBuilder::default() .add_cache(table_id_schema_cache) .add_cache(schema_cache) + .add_cache(table_info_cache) + .add_cache(table_name_cache) + .add_cache(table_route_cache) + .build(); + let derived_registry = CacheRegistryBuilder::default() + .add_cache(table_cache) + .add_cache(partition_info_cache) + .build(); + + LayeredCacheRegistryBuilder::default() + .add_cache_registry(base_registry) + .add_cache_registry(derived_registry) .build() } diff --git a/src/cmd/src/datanode/builder.rs b/src/cmd/src/datanode/builder.rs index 810dfeedc7b..6df7ac728e8 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/substrait/src/df_substrait.rs b/src/common/substrait/src/df_substrait.rs index 08eeedd8316..9570076ee7a 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 3eeb3c570e1..a6932451a62 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/datanode/Cargo.toml b/src/datanode/Cargo.toml index 7c126764723..4d4e470e665 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 fdac0597572..cfbcf948b96 100644 --- a/src/datanode/src/datanode.rs +++ b/src/datanode/src/datanode.rs @@ -19,17 +19,22 @@ use std::path::Path; use std::sync::Arc; use std::time::{Duration, Instant}; +use catalog::information_schema::NoopInformationExtension; +use catalog::kvbackend::KvBackendCatalogManagerBuilder; +use client::client_manager::NodeClients; use common_base::Plugins; use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE}; use common_datasource::object_store::LocalFileAccess; use common_error::ext::BoxedError; use common_greptimedb_telemetry::GreptimeDBTelemetryTask; +use common_grpc::channel_manager::ChannelConfig; use common_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::datanode::TopicStatsReporter; use common_meta::key::runtime_switch::RuntimeSwitchManager; use common_meta::key::{SchemaMetadataManager, SchemaMetadataManagerRef}; use common_meta::kv_backend::KvBackendRef; +use common_meta::node_manager::NodeManagerRef; pub use common_procedure::options::ProcedureConfig; use common_query::prelude::set_default_prefix; use common_stat::ResourceStatImpl; @@ -54,7 +59,8 @@ use object_store::ObjectStore; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::util::normalize_dir; use query::QueryEngineFactory; -use query::dummy_catalog::{DummyCatalogManager, TableProviderFactoryRef}; +use query::dummy_catalog::TableProviderFactoryRef; +use query::region_query::RegionQueryHandlerFactoryRef; use servers::server::ServerHandlers; use snafu::{OptionExt, ResultExt, ensure}; use store_api::logstore::LogStore; @@ -80,6 +86,7 @@ use crate::event_listener::{ use crate::greptimedb_telemetry::get_greptimedb_telemetry_task; use crate::heartbeat::HeartbeatTask; use crate::partition_expr_fetcher::MetaPartitionExprFetcher; +use crate::region_query::DatanodeRegionQueryHandler; use crate::region_server::{DummyTableProviderFactory, RegionServer}; use crate::store::{self, new_object_store_without_cache}; use crate::utils::{RegionOpenRequests, build_region_open_requests}; @@ -357,6 +364,7 @@ impl DatanodeBuilder { schema_metadata_manager, region_event_listener, file_ref_manager, + cache_registry.clone(), ) .await?; @@ -497,14 +505,41 @@ impl DatanodeBuilder { schema_metadata_manager: SchemaMetadataManagerRef, event_listener: RegionServerEventListenerRef, file_ref_manager: FileReferenceManagerRef, + cache_registry: Arc, ) -> Result { let opts: &DatanodeOptions = &self.opts; + // Keep `with_dist_planner` false: datanodes execute received MergeScan plans. + let catalog_manager = KvBackendCatalogManagerBuilder::new( + Arc::new(NoopInformationExtension), + self.kv_backend.clone(), + cache_registry.clone(), + ) + .build(); + let partition_manager = catalog_manager.partition_manager(); + + // MergeScan region requests stream until completion, so they have no client timeout. + let mut channel_config = ChannelConfig { + timeout: None, + ..Default::default() + }; + if opts.grpc.flight_compression.transport_compression() { + channel_config.accept_compression = true; + channel_config.send_compression = true; + } + let node_manager: NodeManagerRef = Arc::new(NodeClients::new(channel_config)); + let region_query_handler = + if let Some(factory) = self.plugins.get::() { + factory.build(partition_manager.clone(), node_manager) + } else { + DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager) + }; + let query_engine_factory = QueryEngineFactory::try_new_with_plugins( // query engine in datanode only executes plan with resolved table source. - DummyCatalogManager::arc(), - None, - None, + catalog_manager, + Some(partition_manager), + Some(region_query_handler), None, None, None, @@ -1060,17 +1095,17 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; - use cache::build_datanode_cache_registry; + use cache::build_datanode_layered_cache_registry; use common_base::Plugins; use common_base::readable_size::ReadableSize; use common_config::Configurable; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; - use common_meta::cache::LayeredCacheRegistryBuilder; use common_meta::key::RegionRoleSet; use common_meta::key::datanode_table::DatanodeTableManager; use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::node_manager::NodeManagerRef; use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir}; use common_wal::config::DatanodeWalConfig; use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION}; @@ -1080,6 +1115,8 @@ mod tests { use mito2::engine::MITO_ENGINE_NAME; use object_store::ObjectStore; use object_store::manager::ObjectStoreManager; + use partition::manager::PartitionRuleManagerRef; + use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef}; use store_api::logstore::LogStore; use store_api::region_request::RegionRequest; use store_api::storage::RegionId; @@ -1090,6 +1127,7 @@ mod tests { wal_object_store, }; use crate::error::{self, Error}; + use crate::region_query::DatanodeRegionQueryHandler; use crate::tests::{MockRegionEngine, mock_region_server}; async fn setup_table_datanode(kv: &KvBackendRef) { @@ -1118,11 +1156,8 @@ mod tests { mock_region_server.register_engine(mock_region.clone()); let kv_backend = Arc::new(MemoryKvBackend::new()); - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut builder = DatanodeBuilder::new( DatanodeOptions { @@ -1164,11 +1199,8 @@ mod tests { wal: DatanodeWalConfig, kv_backend: KvBackendRef, ) -> DatanodeBuilder { - let layered_cache_registry = Arc::new( - LayeredCacheRegistryBuilder::default() - .add_cache_registry(build_datanode_cache_registry(kv_backend.clone())) - .build(), - ); + let layered_cache_registry = + Arc::new(build_datanode_layered_cache_registry(kv_backend.clone())); let mut opts = DatanodeOptions { node_id: Some(0), wal, @@ -1446,6 +1478,38 @@ mod tests { ); } + struct CountingRegionQueryHandlerFactory(Arc); + + 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/error.rs b/src/datanode/src/error.rs index 539712711b7..e9b5d5ec1dc 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/lib.rs b/src/datanode/src/lib.rs index c14ec886fac..5a8948e2c84 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 00000000000..1642362ee1a --- /dev/null +++ b/src/datanode/src/region_query.rs @@ -0,0 +1,109 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The datanode's [`RegionQueryHandler`] implementation for querying remote regions in a +//! `MergeScan` plan. + +use std::sync::Arc; + +use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate}; +use async_trait::async_trait; +use common_error::ext::BoxedError; +use common_meta::node_manager::NodeManagerRef; +use common_query::request::QueryRequest; +use common_recordbatch::SendableRecordBatchStream; +use partition::manager::PartitionRuleManagerRef; +use query::error::{RegionQuerySnafu, Result as QueryResult, UnimplementedSnafu}; +use query::region_query::{RegionQueryHandler, RegionQueryTarget}; +use session::ReadPreference; +use snafu::ResultExt; +use store_api::storage::RegionId; + +/// Serves region queries by resolving their leader and forwarding requests to its datanode. +pub(crate) struct DatanodeRegionQueryHandler { + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, +} + +impl DatanodeRegionQueryHandler { + pub fn arc( + partition_manager: PartitionRuleManagerRef, + node_manager: NodeManagerRef, + ) -> Arc { + 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) + } + + /// Datanodes forward received region queries and do not produce remote dynamic filters. + async fn handle_remote_dyn_filter_update( + &self, + _target: &RegionQueryTarget, + _query_id: String, + _update: RemoteDynFilterUpdate, + ) -> QueryResult<()> { + UnimplementedSnafu { + operation: "remote dynamic filter update from a datanode", + } + .fail() + } + + /// See [`RegionQueryHandler::handle_remote_dyn_filter_update`]. + async fn handle_remote_dyn_filter_unregister( + &self, + _target: &RegionQueryTarget, + _query_id: String, + _unregister: RemoteDynFilterUnregister, + ) -> QueryResult<()> { + UnimplementedSnafu { + operation: "remote dynamic filter unregister from a datanode", + } + .fail() + } +} diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 7d5218f30b7..ef333cff04d 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 62e8597dbb6..cf4b6fd2d34 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/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index f4d6b29fa2c..c924a57a68e 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -14,19 +14,21 @@ 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; 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, @@ -44,12 +46,18 @@ 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::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; /// Extended [`substrait::extension_serializer::ExtensionSerializer`] but supports [`MergeScanLogicalPlan`] serialization. #[derive(Debug)] @@ -86,8 +94,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. + // This registry has no session state; the query engine decodes `MergeScan` payloads. Err(DataFusionError::Substrait(format!( "Unsupported plan node: {name}" ))) @@ -97,6 +104,196 @@ impl SerializerRegistry for DefaultSerializer { } } +/// Async `MergeScan` payload decoding over DataFusion's default Substrait consumer. +struct MergeScanSubstraitConsumer<'a> { + inner: DefaultSubstraitConsumer<'a>, + session_state: &'a SessionState, + catalog_manager: Option, + payload: bool, +} + +impl MergeScanSubstraitConsumer<'_> { + fn query_ctx(&self) -> QueryContextRef { + self.session_state + .config() + .get_extension() + .unwrap_or_else(QueryContext::arc) + } + + async fn decode_payload(&self, payload: Vec) -> Result { + decode_plan(Bytes::from(payload), self.session_state, true) + .await + .map_err(|e| DataFusionError::External(Box::new(e))) + } +} + +#[async_trait::async_trait] +impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> { + async fn resolve_table_ref( + &self, + table_ref: &TableReference, + ) -> Result>> { + 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 = 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 { + 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, + 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(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).await?; + + // `PbMergeScan` lacks `partition_cols`; decoded plans lose this optimization and may repartition. + Ok( + MergeScanLogicalPlan::new(input, merge_scan.is_placeholder, Default::default()) + .into_logical_plan(), + ) + } + + // 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 })) + } + + // 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); + } + 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, + 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() { + return not_impl_err!("Type variation extensions are not supported") + .context(DecodeDfPlanSnafu); + } + + 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, + payload, + }; + + DFLogicalSubstraitConvertor + .decode_with_consumer(&plan, state, &consumer) + .await +} + /// The datafusion `[LogicalPlan]` decoder. pub struct DefaultPlanDecoder { session_state: SessionState, @@ -115,6 +312,108 @@ impl DefaultPlanDecoder { } } +/// 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 request plan. +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( @@ -123,108 +422,14 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { catalog_list: Arc, optimize: bool, ) -> common_query::error::Result { - // The session_state already has the `DefaultSerialzier` as `SerializerRegistry`. let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone()) + .with_serializer_registry(Arc::new(DefaultSerializer)) .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() })?; - } + // Re-register after the build to avoid Greptime UDF alias collisions. + register_greptime_functions(&mut session_state, &self.query_ctx)?; - 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)); - } - - let logical_plan = DFLogicalSubstraitConvertor - .decode(message, session_state) + let logical_plan = decode_plan(message, &session_state, false) .await .map_err(BoxedError::new) .context(common_query::error::DecodePlanSnafu)?; @@ -241,18 +446,35 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { #[cfg(test)] mod tests { + use std::any::Any; + use std::sync::atomic::{AtomicBool, Ordering}; + + use catalog::error::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::ScalarFunction; - use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, 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::*; use crate::QueryEngineFactory; @@ -488,4 +710,935 @@ 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 engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .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(); + + 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(), numbers_scan().to_string()); + // `PbMergeScan` doesn't carry `partition_cols`, so decoded nodes have an empty mapping. + assert!(merge_scan.partition_cols().is_empty()); + } + + /// Payload table binding uses the engine catalog, not the request catalog. + #[tokio::test] + async fn test_serializer_decode_merge_scan_with_engine_catalog() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .query_engine(); + let decoded = plan_decoder(engine.as_ref()) + .decode( + encode_plan(&merge_scan(numbers_scan(), false)), + request_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()); + }; + 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"); + } + + /// Nested `MergeScan` payloads decode recursively. + #[tokio::test] + async fn test_serializer_decode_nested_merge_scan() { + let engine = QueryEngineFactory::new( + numbers_catalog_manager(), + None, + None, + None, + None, + false, + QueryOptions::default(), + ) + .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 = merge_scan(input, false); + + let decoded = plan_decoder(engine.as_ref()) + .decode(encode_plan(&plan), request_catalog_list(), false) + .await + .unwrap(); + + assert_eq!(decoded.to_string(), plan.to_string()); + assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2); + } + + 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() + } + + /// 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) + } + + 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")], + }) + } + + /// Includes `MergeScan` payload expressions, which are not node 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 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, + ) { + 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_eq!( + udf.as_ref(), + greptime_date_format_udf(query_ctx).as_ref(), + "`date_format` is not bound to the implementation of GreptimeDB: {plan}" + ); + } + + /// 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 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); + } + } + + /// 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(); + 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(); + + 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 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); + } + } + + /// 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 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(); + 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".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(), + ], + ], + ); + } + + /// How [`InterceptingCatalogManager`] resolves payload tables. + enum TableResolution { + /// Records catalog lookup and request context values. + ObserveContext(Arc>>), + /// Resolves once another task on the caller runtime signals. + AwaitSignal { + notify: 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::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::AwaitSignal { notify, entered } => { + entered.store(true, Ordering::SeqCst); + notify.notified().await; + self.inner + .table(catalog, schema, table_name, query_ctx) + .await + } + } + } + + 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) + } + } + + 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() + } + + /// 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()); + } + + /// A missing engine-catalog table does not fall back to the usable request catalog. + #[tokio::test] + async fn test_serializer_decode_missing_engine_catalog_table() { + let missing_manager: CatalogManagerRef = + catalog::memory::new_memory_catalog_manager().unwrap(); + let engine = QueryEngineFactory::new( + missing_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}" + ); + } + + /// 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 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, + "number", + false, + Vec::new(), + merge_scan(numbers_scan(), false), + )), + }); + 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 + .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() { + let outer = merge_scan(numbers_scan(), false); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 0, + "number", + 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 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(); + + 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}" + ); + } } diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index ad0a693ac59..c29ecc6493a 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 09e4879dd50..7ef9401cf18 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 5485198c266..93c46ff7ff7 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 c1448053f35..2f81dfabffb 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 971426f4cc0..3ed4fd8dc2e 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/cluster.rs b/tests-integration/src/cluster.rs index 76f198c9707..bd357033a7e 100644 --- a/tests-integration/src/cluster.rs +++ b/tests-integration/src/cluster.rs @@ -17,13 +17,15 @@ use std::env; use std::net::TcpListener; use std::ops::RangeInclusive; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +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,31 @@ use crate::test_util::{ create_tmp_dir_and_datanode_opts, }; +/// Flight `DoGet` requests another node sent to a datanode over its network address (see +/// [`GreptimeDbClusterBuilder::with_real_datanode_grpc_addr`]). +/// +/// 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 { + do_get_requests: AtomicUsize, +} + +impl DatanodeRpcStats { + /// The number of Flight `DoGet` requests another node sent to the datanode. + pub fn requests(&self) -> usize { + self.do_get_requests.load(Ordering::Relaxed) + } +} + 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 +123,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(); @@ -177,6 +208,7 @@ pub struct GreptimeDbClusterBuilder { plugins: Plugins, local_file_access: LocalFileAccess, event_recorder_options: EventRecorderOptions, + real_datanode_grpc_addr: bool, } impl GreptimeDbClusterBuilder { @@ -214,6 +246,7 @@ impl GreptimeDbClusterBuilder { plugins: Plugins::default(), local_file_access: LocalFileAccess::default(), event_recorder_options: EventRecorderOptions::default(), + real_datanode_grpc_addr: false, } } @@ -299,6 +332,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, @@ -346,7 +394,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; @@ -365,6 +428,7 @@ impl GreptimeDbClusterBuilder { datanode_options, guards, datanode_instances, + datanode_rpc_stats, kv_backend: self.kv_backend.clone(), metasrv: metasrv.metasrv, frontend: Arc::new(frontend), @@ -421,12 +485,28 @@ 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, @@ -436,7 +516,8 @@ impl GreptimeDbClusterBuilder { for opts in options { let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await; - instances.insert(opts.node_id.unwrap(), datanode); + let datanode_id = opts.node_id.unwrap(); + instances.insert(datanode_id, datanode); } instances @@ -474,15 +555,12 @@ 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(), - ); + 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 - .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(); @@ -630,18 +708,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") @@ -652,8 +746,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) @@ -669,12 +764,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( @@ -699,6 +822,52 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) { ) } +/// Counts Flight `DoGet` requests received by the wrapped server. +#[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 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) + } +} + +#[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)] @@ -993,13 +1162,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/standalone.rs b/tests-integration/src/standalone.rs index 98eb336d9cc..241d4daf5b5 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; @@ -209,11 +209,8 @@ 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(), - ); + 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 d70af021651..4a5b2f0c6c5 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 00000000000..2c86589112d --- /dev/null +++ b/tests-integration/src/tests/nested_merge_scan_capability_test.rs @@ -0,0 +1,749 @@ +// 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. + +//! 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 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; +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; + +/// Partitioned probe table. +const PROBE_TABLE: &str = "nested_cap_probe"; +/// Partitioned build table with duplicate join keys. +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"; + +/// `(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), + // The second partition of the probe table. + (120, 30, 300), + (121, 10, 400), +]; + +/// `(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), + (3, 20, 3), + // The second partition of the build table. + (120, 30, 4), + (121, 40, 5), +]; + +#[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; + 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 reference = query_pretty(&frontend, &join_sql(), query_ctx()).await; + 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 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_rows = multiset(rows); + assert_eq!( + expected_join_rows(), + actual_rows, + "the nested plans of the datanodes returned an unexpected multiset of rows" + ); + assert_eq!( + multiset(table_cells(&reference)), + actual_rows, + "frontend result differs:\n{reference}" + ); + 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.output_msg().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.output_msg().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.output_msg().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")] +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; + + 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_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:?}" + ) + }); + + // 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; + 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); + assert!( + chain.contains(MISSING_INNER_TABLE), + "expected the datanode to fail while resolving the inner relation {MISSING_INNER_TABLE}, \ + actual error: {chain}" + ); + + // 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 { + flush_on_close: false, + }), + ) + .await + .expect("the datanode must close the region of the build table"); + + 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 { + 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 increased DoGet count proves execution reached the remote region owner. + let actual = rpc_requests(&cluster, *build_datanode); + assert!( + actual > baseline, + "expected the failed nested plan to reach datanode {build_datanode} for region {build_region} over Flight DoGet (count {actual}, baseline {baseline})" + ); +} + +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 +} + +fn assert_cross_datanode_region_query( + cluster: &GreptimeDbCluster, + remote_datanodes: &BTreeSet, + baseline_requests: &BTreeMap, +) { + assert_eq!( + 2, + remote_datanodes.len(), + "build table must span both datanodes" + ); + for datanode in remote_datanodes { + let actual = rpc_requests(cluster, *datanode); + let baseline = baseline_requests[datanode]; + assert!( + actual > baseline, + "expected a nested Flight DoGet at datanode {datanode}, count {actual}, baseline {baseline}" + ); + } +} + +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() +} + +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() +} + +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"); + rows.extend(table_cells(&actual)); + } + + rows +} + +/// 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; +} + +fn encode(plan: &LogicalPlan) -> Vec { + DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .expect("the plan must be encodable") + .to_vec() +} + +/// Plans the hand-built nested join sent in the region request. +async fn nested_plan(frontend: &Arc, query_ctx: &QueryContextRef) -> LogicalPlan { + nested_plan_with_inner_table(frontend, query_ctx, BUILD_TABLE).await +} + +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) +} + +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 +} + +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 +} + +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}")) +} + +/// 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 + FROM {PROBE_TABLE} p JOIN {BUILD_TABLE} b ON p.probe_key = b.build_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 { + 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) +} + +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 +} + +fn datanodes(leaders: &BTreeMap) -> BTreeSet { + leaders.values().map(|(id, _)| *id).collect() +} + +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 the full stream to detect errors after any partial rows. +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)), + } +} + +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(": ") +} + +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() +} + +/// 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) +} + +fn region_request_header(query_ctx: &QueryContextRef) -> RegionRequestHeader { + RegionRequestHeader { + query_context: Some(query_ctx.as_ref().into()), + ..Default::default() + } +} + +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() +} + +fn multiset(mut rows: Vec>) -> Vec> { + rows.sort(); + rows +} diff --git a/tests/compatibility/README.md b/tests/compatibility/README.md index 542f10debfd..404b6332b94 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 00000000000..9e7b9ba74c1 --- /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 00000000000..1990ede9411 --- /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