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