diff --git a/src/api/grpc/src/handler/grpc/flight/mod.rs b/src/api/grpc/src/handler/grpc/flight/mod.rs index c839b11435..43be4746d1 100644 --- a/src/api/grpc/src/handler/grpc/flight/mod.rs +++ b/src/api/grpc/src/handler/grpc/flight/mod.rs @@ -15,7 +15,6 @@ use std::{io::Cursor, sync::Arc}; -use arrow::ipc::{CompressionType, writer::IpcWriteOptions}; use arrow_flight::{ Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket, @@ -28,11 +27,14 @@ use datafusion::{ common::{DataFusionError, Result}, physical_plan::execute_stream_partitioned, }; -use flight::common::{MetricsInfo, PreCustomMessage}; +use flight::common::{MetricsInfo, NodeMessages}; use futures::{StreamExt, stream::BoxStream}; use futures_util::pin_mut; use prost::Message; -use search::inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}; +use search::{ + datafusion::distributed_plan::profile, + inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, +}; use search_service::{grpc::flight as grpcFlight, work_group::DeferredLock}; use tonic::{Request, Response, Status, Streaming}; use tracing::Instrument; @@ -45,11 +47,8 @@ use { use crate::handler::grpc::{ MetadataMap, flight::{ - stream::{FlightEncoderStreamBuilder, MAX_FLIGHT_DATA_SIZE}, - visitor::{ - get_cluster_metrics, get_partial_err, get_peak_memory, get_peak_memory_from_ctx, - get_scan_stats, - }, + stream::FlightEncoderStreamBuilder, + visitor::{get_peak_memory_from_ctx, node_role}, }, }; @@ -164,31 +163,15 @@ impl FlightService for FlightServiceImpl { } let start = std::time::Instant::now(); - let write_options: IpcWriteOptions = IpcWriteOptions::default() - .try_with_compression(Some(CompressionType::ZSTD)) - .map_err(|e| { - log::error!( - "[trace_id {trace_id}] flight->search: do_get create IPC write options error: {e:?}", - ); - Status::internal(e.to_string()) - })?; - - // used for EXPLAIN ANALYZE to collect metrics after stream is done - let metrics = req.search_info.is_analyze.then_some(MetricsInfo { - plan: plan.clone(), - is_super_cluster, - func: Box::new(super_cluster_enabled), - }); - - // Get the peak memory usage from the memory pool - // Note: We get peak memory after stream execution, so we pass it via a shared reference - let peak_memory = get_peak_memory_from_ctx(&ctx); - - // used for super cluster follower leader to get information from follower node - let scan_stats_ref = get_scan_stats(&plan); - let metrics_ref = get_cluster_metrics(&plan); - let peak_memory_ref = get_peak_memory(&plan); - let partial_err_ref = get_partial_err(&plan); + let node_messages = NodeMessages { + scan_stats, + peak_memory: get_peak_memory_from_ctx(&ctx), + profile: req.search_info.is_analyze.then(|| MetricsInfo { + plan: plan.clone(), + build_profile: profile::build_profile, + }), + role: node_role(&plan), + }; // One stream per output partition so they encode in parallel let streams = execute_stream_partitioned(plan, ctx.task_ctx()).map_err(|e| { @@ -198,23 +181,12 @@ impl FlightService for FlightServiceImpl { Status::internal(e.to_string()) })?; - let mut stream = - FlightEncoderStreamBuilder::new(write_options, MAX_FLIGHT_DATA_SIZE, session_guard) - .with_trace_id(trace_id.to_string()) - .with_is_super(is_super_cluster) - .with_defer_lock(lock) - .with_start(start) - .with_custom_message(PreCustomMessage::ScanStats(scan_stats)) - .with_custom_message(PreCustomMessage::ScanStatsRef(scan_stats_ref)) - .with_custom_message(PreCustomMessage::Metrics(metrics)) - .with_custom_message(PreCustomMessage::MetricsRef(metrics_ref)) - .with_custom_message(PreCustomMessage::PeakMemoryRef(Some(peak_memory))) - .with_custom_message(PreCustomMessage::PeakMemoryRef(peak_memory_ref)) - .with_custom_message(PreCustomMessage::PartialErrRefEarly( - partial_err_ref.clone(), - )) - .with_custom_message(PreCustomMessage::PartialErrRef(partial_err_ref)) - .build(streams, span); + let mut stream = FlightEncoderStreamBuilder::new(session_guard, node_messages) + .with_trace_id(trace_id.to_string()) + .with_is_super(is_super_cluster) + .with_defer_lock(lock) + .with_start(start) + .build(streams, span); let stream = async_stream::stream! { let timeout = tokio::time::sleep(tokio::time::Duration::from_secs(timeout)); @@ -367,16 +339,6 @@ fn clear_session_data(trace_id: &str) { log::info!("Cleared session for trace_id: {trace_id}"); } -#[cfg(feature = "enterprise")] -fn super_cluster_enabled() -> bool { - get_o2_config().super_cluster.enabled -} - -#[cfg(not(feature = "enterprise"))] -fn super_cluster_enabled() -> bool { - false -} - #[cfg(test)] mod tests { use super::*; @@ -385,9 +347,4 @@ mod tests { fn test_flight_service_impl_default() { let _impl = FlightServiceImpl; } - - #[test] - fn test_super_cluster_disabled_in_oss() { - assert!(!super_cluster_enabled()); - } } diff --git a/src/api/grpc/src/handler/grpc/flight/partition_encoder.rs b/src/api/grpc/src/handler/grpc/flight/partition_encoder.rs index 3f6d8725a0..756357fa59 100644 --- a/src/api/grpc/src/handler/grpc/flight/partition_encoder.rs +++ b/src/api/grpc/src/handler/grpc/flight/partition_encoder.rs @@ -23,6 +23,8 @@ use datafusion::execution::SendableRecordBatchStream; use flight::encoder::{ChunkCoalescer, FlightDataEncoder}; use futures::{Stream, StreamExt}; +use super::stream::MAX_FLIGHT_DATA_SIZE; + /// Coalesces and ZSTD-encodes one datafusion output partition into FlightData. do_get runs one per /// partition on its own tokio task, so partitions execute and encode in parallel. /// @@ -39,15 +41,11 @@ pub(super) struct PartitionEncoderStream { } impl PartitionEncoderStream { - pub(super) fn new( - inner: SendableRecordBatchStream, - options: IpcWriteOptions, - max_flight_data_size: usize, - ) -> Self { + pub(super) fn new(inner: SendableRecordBatchStream, options: IpcWriteOptions) -> Self { Self { inner, coalescer: ChunkCoalescer::new(), - encoder: FlightDataEncoder::new(options, max_flight_data_size), + encoder: FlightDataEncoder::new(options, MAX_FLIGHT_DATA_SIZE), out: VecDeque::new(), inner_done: false, } diff --git a/src/api/grpc/src/handler/grpc/flight/stream.rs b/src/api/grpc/src/handler/grpc/flight/stream.rs index 519c19ccfe..75ea5368f0 100644 --- a/src/api/grpc/src/handler/grpc/flight/stream.rs +++ b/src/api/grpc/src/handler/grpc/flight/stream.rs @@ -15,11 +15,11 @@ use std::{collections::VecDeque, pin::Pin, task::Poll}; -use arrow::ipc::writer::IpcWriteOptions; +use arrow::ipc::{CompressionType, writer::IpcWriteOptions}; use arrow_flight::{FlightData, error::FlightError}; use config::metrics; use datafusion::execution::SendableRecordBatchStream; -use flight::{common::PreCustomMessage, encoder::FlightDataEncoder}; +use flight::{common::NodeMessages, encoder::FlightDataEncoder}; use futures::{Stream, StreamExt}; use search_service::work_group::DeferredLock; use tokio::{sync::mpsc, task::JoinSet}; @@ -30,12 +30,10 @@ use crate::handler::grpc::flight::{SessionGuard, partition_encoder::PartitionEnc /// Cap on one encoded FlightData message. Mirrors the `ZO_GRPC_MAX_MESSAGE_SIZE` default, /// which is what the receiving channel accepts. -pub(crate) const MAX_FLIGHT_DATA_SIZE: usize = 32 * 1024 * 1024; +pub(super) const MAX_FLIGHT_DATA_SIZE: usize = 32 * 1024 * 1024; pub struct FlightEncoderStreamBuilder { - options: IpcWriteOptions, - max_flight_data_size: usize, - custom_messages: Vec, + node_messages: NodeMessages, // query context trace_id: String, is_super: bool, @@ -45,15 +43,9 @@ pub struct FlightEncoderStreamBuilder { } impl FlightEncoderStreamBuilder { - pub(crate) fn new( - options: IpcWriteOptions, - max_flight_data_size: usize, - session_guard: SessionGuard, - ) -> Self { + pub(crate) fn new(session_guard: SessionGuard, node_messages: NodeMessages) -> Self { Self { - options, - max_flight_data_size, - custom_messages: vec![], + node_messages, trace_id: String::new(), is_super: false, defer_lock: None, @@ -62,11 +54,6 @@ impl FlightEncoderStreamBuilder { } } - pub fn with_custom_message(mut self, message: PreCustomMessage) -> Self { - self.custom_messages.push(message); - self - } - pub fn with_trace_id(mut self, trace_id: String) -> Self { self.trace_id = trace_id; self @@ -98,12 +85,12 @@ impl FlightEncoderStreamBuilder { ) -> FlightEncoderStream { let child_span = info_span!("grpc:search:flight:execute_physical_plan"); let _ = child_span.set_parent(span.context()); + let options = write_options(); let (source, tasks) = if partitions.len() == 1 { // one partition: no parallelism to gain, poll inline (no task, no channel) let inner = partitions.into_iter().next().unwrap(); - let stream = - PartitionEncoderStream::new(inner, self.options.clone(), self.max_flight_data_size); + let stream = PartitionEncoderStream::new(inner, options.clone()); (FlightSource::Inline(Box::new(stream)), JoinSet::new()) } else { // one task per partition, merged over the channel @@ -112,11 +99,7 @@ impl FlightEncoderStreamBuilder { let (tx, rx) = mpsc::channel::, FlightError>>(cap); for inner in partitions { let tx = tx.clone(); - let mut stream = PartitionEncoderStream::new( - inner, - self.options.clone(), - self.max_flight_data_size, - ); + let mut stream = PartitionEncoderStream::new(inner, options.clone()); tasks.spawn(async move { while let Some(item) = stream.next().await { let is_err = item.is_err(); @@ -134,9 +117,9 @@ impl FlightEncoderStreamBuilder { FlightEncoderStream { source, - encoder: FlightDataEncoder::new(self.options, self.max_flight_data_size), + encoder: FlightDataEncoder::new(options, MAX_FLIGHT_DATA_SIZE), queue: VecDeque::new(), - custom_messages: self.custom_messages, + node_messages: self.node_messages, first_data: true, done: false, ctx: StreamContext { @@ -167,7 +150,7 @@ pub struct FlightEncoderStream { /// Only used to encode custom messages (scan stats, peak memory, ...). encoder: FlightDataEncoder, queue: VecDeque, - custom_messages: Vec, + node_messages: NodeMessages, first_data: bool, done: bool, /// query context + resources consumed by the stream's `Drop` teardown. @@ -232,31 +215,20 @@ impl EventLog { } impl FlightEncoderStream { - fn encode_custom_messages(&mut self) -> Result<(), FlightError> { - let custom_messages = std::mem::take(&mut self.custom_messages); - for message in custom_messages.into_iter() { - if let Some(message) = message.get_custom_message() { - let flight_data = self.encoder.encode_custom(&message)?; - self.queue.push_back(flight_data); - } + /// Queues the node's messages due now; a stream with no batch sends its first ones at the end. + fn encode_node_messages(&mut self, at_end: bool) -> Result<(), FlightError> { + let mut messages = Vec::new(); + if self.first_data { + messages.extend(self.node_messages.first()); } - Ok(()) - } - - fn encode_early_emit_messages(&mut self) -> Result<(), FlightError> { - let custom_messages = std::mem::take(&mut self.custom_messages); - let mut remainder_messages = Vec::new(); - for message in custom_messages.into_iter() { - if message.is_early_emit() { - if let Some(message) = message.get_custom_message() { - let flight_data = self.encoder.encode_custom(&message)?; - self.queue.push_back(flight_data); - } - } else { - remainder_messages.push(message); - } + if at_end { + messages.extend(self.node_messages.last()); + } + self.first_data = false; + for message in messages { + let flight_data = self.encoder.encode_custom(&message)?; + self.queue.push_back(flight_data); } - self.custom_messages = remainder_messages; Ok(()) } @@ -318,14 +290,10 @@ impl Stream for FlightEncoderStream { // 4. pull the next encoded FlightData from the partition encoder(s) match this.poll_source(cx) { Poll::Ready(Some(Ok(group))) => { - // before the first data, emit early-emit custom messages - if this.first_data { - this.first_data = false; - if !this.custom_messages.is_empty() - && let Err(e) = this.encode_early_emit_messages() - { - return this.fail(e); - } + if this.first_data + && let Err(e) = this.encode_node_messages(false) + { + return this.fail(e); } // enqueue the whole group so the chunk's dictionary and records stay contiguous this.queue.extend(group); @@ -345,7 +313,7 @@ impl Stream for FlightEncoderStream { if !this.ctx.tasks.is_empty() { return Poll::Pending; } - if let Err(e) = this.encode_custom_messages() { + if let Err(e) = this.encode_node_messages(true) { return this.fail(e); } this.done = true; @@ -393,27 +361,30 @@ impl Drop for FlightEncoderStream { } } +fn write_options() -> IpcWriteOptions { + IpcWriteOptions::default() + .try_with_compression(Some(CompressionType::ZSTD)) + // Cannot fail: compression only needs IPC metadata V5, which is the default. + .expect("default IPC options accept compression") +} + #[cfg(test)] mod tests { use std::sync::Arc; - use arrow::{ - array::{Int32Array, RecordBatch}, - ipc::writer::IpcWriteOptions, - }; + use arrow::array::{Int32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; use config::meta::search::ScanStats; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; - use flight::common::PreCustomMessage; + use flight::common::{CustomMessage, NodeRole}; use futures::TryStreamExt; use super::*; fn make_builder() -> FlightEncoderStreamBuilder { FlightEncoderStreamBuilder::new( - IpcWriteOptions::default(), - 4096, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) } @@ -424,7 +395,6 @@ mod tests { assert!(!b.is_super); assert!(b.defer_lock.is_none()); assert!(b.session_guard.is_some()); - assert!(b.custom_messages.is_empty()); } #[test] @@ -458,11 +428,53 @@ mod tests { let _ = b.start; } - #[test] - fn test_builder_with_custom_message() { - let msg = PreCustomMessage::ScanStats(ScanStats::new()); - let b = make_builder().with_custom_message(msg); - assert_eq!(b.custom_messages.len(), 1); + fn follower_messages() -> NodeMessages { + NodeMessages { + scan_stats: ScanStats::new(), + peak_memory: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + profile: None, + role: NodeRole::Follower, + } + } + + async fn message_kinds(partition: SendableRecordBatchStream) -> Vec<&'static str> { + let fds: Vec = make_builder() + .build(vec![partition], tracing::Span::none()) + .try_collect() + .await + .unwrap(); + fds.iter() + .map( + |fd| match config::utils::json::from_slice(&fd.app_metadata) { + Ok(CustomMessage::ScanStats(_)) => "ScanStats", + Ok(CustomMessage::PeakMemory(_)) => "PeakMemory", + Ok(_) => "other", + Err(_) => "data", + }, + ) + .collect() + } + + #[tokio::test] + async fn test_node_messages_surround_the_batches() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))]) + .unwrap(); + let with_rows: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema.clone(), + futures::stream::iter(vec![Ok(batch)]), + )); + let kinds = message_kinds(with_rows).await; + assert_eq!(kinds.first(), Some(&"ScanStats")); + assert_eq!(kinds.last(), Some(&"PeakMemory")); + assert!(kinds[1..kinds.len() - 1].iter().all(|kind| *kind == "data")); + + let empty: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::empty(), + )); + assert_eq!(message_kinds(empty).await, ["ScanStats", "PeakMemory"]); } #[test] @@ -503,13 +515,9 @@ mod tests { make_partition(total / 2..total), ]; - let options = IpcWriteOptions::default() - .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) - .unwrap(); let stream = FlightEncoderStreamBuilder::new( - options, - MAX_FLIGHT_DATA_SIZE, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) .build(partitions, tracing::Span::none()); @@ -518,6 +526,7 @@ mod tests { let dicts = std::collections::HashMap::new(); let mut decoded: Vec = fds .iter() + .filter(|fd| fd.app_metadata.is_empty()) .flat_map(|fd| { let b = arrow_flight::utils::flight_data_to_arrow_batch(fd, schema.clone(), &dicts) .unwrap(); @@ -562,9 +571,8 @@ mod tests { )); let stream = FlightEncoderStreamBuilder::new( - IpcWriteOptions::default(), - MAX_FLIGHT_DATA_SIZE, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) .build(vec![ok_partition, panic_partition], tracing::Span::none()); @@ -676,13 +684,9 @@ mod tests { let id_base_1 = n_per_partition; let partitions = vec![make_partition(0, 0), make_partition(1, id_base_1)]; - let options = IpcWriteOptions::default() - .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) - .unwrap(); let stream = FlightEncoderStreamBuilder::new( - options, - MAX_FLIGHT_DATA_SIZE, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) .build(partitions, tracing::Span::none()); let fds: Vec = stream.try_collect().await.unwrap(); @@ -749,13 +753,9 @@ mod tests { futures::stream::iter(batches.into_iter().map(Ok)), )); - let options = IpcWriteOptions::default() - .try_with_compression(Some(arrow::ipc::CompressionType::ZSTD)) - .unwrap(); let stream = FlightEncoderStreamBuilder::new( - options, - MAX_FLIGHT_DATA_SIZE, SessionGuard::new("test-trace".to_string()), + follower_messages(), ) .build(vec![inner], tracing::Span::none()); @@ -767,6 +767,7 @@ mod tests { let dicts = std::collections::HashMap::new(); let decoded: Vec = fds .iter() + .filter(|fd| fd.app_metadata.is_empty()) .flat_map(|fd| { let b = arrow_flight::utils::flight_data_to_arrow_batch(fd, schema.clone(), &dicts) .unwrap(); diff --git a/src/api/grpc/src/handler/grpc/flight/visitor.rs b/src/api/grpc/src/handler/grpc/flight/visitor.rs index 59567c423f..def7730613 100644 --- a/src/api/grpc/src/handler/grpc/flight/visitor.rs +++ b/src/api/grpc/src/handler/grpc/flight/visitor.rs @@ -13,163 +13,40 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::sync::{Arc, atomic::AtomicUsize}; +use std::sync::Arc; -use config::meta::search::ScanStats; use datafusion::{ - common::{ - Result, - tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}, - }, + common::tree_node::{TreeNode, TreeNodeRecursion}, physical_plan::ExecutionPlan, }; -use parking_lot::Mutex; +use flight::common::NodeRole; use search::datafusion::distributed_plan::remote_scan_exec::RemoteScanExec; -pub use search::datafusion::plan_metrics::{get_cluster_metrics, get_peak_memory_from_ctx}; +pub use search::datafusion::plan_metrics::get_peak_memory_from_ctx; -pub fn get_scan_stats(plan: &Arc) -> Option>> { - let mut visitor = RemoteScanVisitor::new(); - let _ = plan.visit(&mut visitor); - visitor.scan_stats -} - -struct RemoteScanVisitor { - scan_stats: Option>>, -} - -impl RemoteScanVisitor { - pub fn new() -> Self { - Self { scan_stats: None } - } -} - -impl<'n> TreeNodeVisitor<'n> for RemoteScanVisitor { - type Node = Arc; - - fn f_up(&mut self, node: &'n Self::Node) -> Result { - let name = node.name(); - if name == "RemoteScanExec" { - let remote_scan_exec = node.downcast_ref::().unwrap(); - self.scan_stats = Some(remote_scan_exec.scan_stats()); - Ok(TreeNodeRecursion::Stop) - } else { - Ok(TreeNodeRecursion::Continue) - } - } -} - -pub fn get_partial_err(plan: &Arc) -> Option>> { - let mut visitor = PartialErrVisitor::new(); - let _ = plan.visit(&mut visitor); - visitor.partial_err -} - -struct PartialErrVisitor { - partial_err: Option>>, -} - -impl PartialErrVisitor { - pub fn new() -> Self { - Self { partial_err: None } - } -} - -impl<'n> TreeNodeVisitor<'n> for PartialErrVisitor { - type Node = Arc; - - fn f_up(&mut self, node: &'n Self::Node) -> Result { - if node.name() == "RemoteScanExec" { - let remote_scan_exec = node.downcast_ref::().unwrap(); - self.partial_err = Some(remote_scan_exec.partial_err()); - Ok(TreeNodeRecursion::Stop) - } else { - Ok(TreeNodeRecursion::Continue) - } - } -} - -pub fn get_peak_memory(plan: &Arc) -> Option> { - let mut visitor = PeakMemoryVisitor::new(); - let _ = plan.visit(&mut visitor); - visitor.peak_memory -} - -struct PeakMemoryVisitor { - peak_memory: Option>, -} - -impl PeakMemoryVisitor { - pub fn new() -> Self { - Self { peak_memory: None } - } -} - -impl<'n> TreeNodeVisitor<'n> for PeakMemoryVisitor { - type Node = Arc; - - fn f_up(&mut self, node: &'n Self::Node) -> Result { - let name = node.name(); - if name == "RemoteScanExec" { - let remote_scan_exec = node.downcast_ref::().unwrap(); - self.peak_memory = Some(remote_scan_exec.peak_memory()); - } - Ok(TreeNodeRecursion::Continue) - } +/// Only a super cluster's cluster leader runs a plan with a remote scan in it. +pub fn node_role(plan: &Arc) -> NodeRole { + let mut role = NodeRole::Follower; + let _ = plan.apply(|node| { + let Some(remote_scan) = node.downcast_ref::() else { + return Ok(TreeNodeRecursion::Continue); + }; + role = NodeRole::ClusterLeader(remote_scan.state.clone()); + Ok(TreeNodeRecursion::Stop) + }); + role } #[cfg(test)] mod tests { - use std::sync::Arc; - use arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_plan::empty::EmptyExec; use super::*; - fn empty_plan() -> Arc { + #[test] + fn test_plan_without_remote_scan_is_follower() { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); - Arc::new(EmptyExec::new(schema)) - } - - #[test] - fn test_remote_scan_visitor_new_has_no_stats() { - let v = RemoteScanVisitor::new(); - assert!(v.scan_stats.is_none()); - } - - #[test] - fn test_partial_err_visitor_new_has_no_err() { - let v = PartialErrVisitor::new(); - assert!(v.partial_err.is_none()); - } - - #[test] - fn test_peak_memory_visitor_new_has_none() { - let v = PeakMemoryVisitor::new(); - assert!(v.peak_memory.is_none()); - } - - #[test] - fn test_get_scan_stats_no_remote_exec_returns_none() { - let plan = empty_plan(); - assert!(get_scan_stats(&plan).is_none()); - } - - #[test] - fn test_get_cluster_metrics_no_remote_exec_returns_empty() { - let plan = empty_plan(); - assert!(get_cluster_metrics(&plan).is_empty()); - } - - #[test] - fn test_get_partial_err_no_remote_exec_returns_none() { - let plan = empty_plan(); - assert!(get_partial_err(&plan).is_none()); - } - - #[test] - fn test_get_peak_memory_no_remote_exec_returns_none() { - let plan = empty_plan(); - assert!(get_peak_memory(&plan).is_none()); + let plan: Arc = Arc::new(EmptyExec::new(schema)); + assert!(matches!(node_role(&plan), NodeRole::Follower)); } } diff --git a/src/config/src/datafusion/request.rs b/src/config/src/datafusion/request.rs index 64f5014722..4ebfd4dbf3 100644 --- a/src/config/src/datafusion/request.rs +++ b/src/config/src/datafusion/request.rs @@ -39,6 +39,8 @@ pub struct Request { /// request because the query is planned again in spawned tasks, where /// that task-local is gone. Never derived from client input. pub chat_history_reader: bool, + /// Adds the structured `profile` column to EXPLAIN ANALYZE output. + pub analyze_profile: bool, } impl Default for Request { @@ -59,6 +61,7 @@ impl Default for Request { overwrite_cache: false, histogram_interval: 0, chat_history_reader: false, + analyze_profile: false, } } } @@ -92,6 +95,7 @@ impl Request { overwrite_cache, histogram_interval, chat_history_reader: false, + analyze_profile: false, } } @@ -112,6 +116,10 @@ impl Request { self.use_cache = use_cache; } + pub fn set_analyze_profile(&mut self, analyze_profile: bool) { + self.analyze_profile = analyze_profile; + } + /// Mark the request as the AI Chat API's history reader when — and only /// when — it is issued from inside that reader's scope. pub fn mark_chat_history_reader(&mut self) { @@ -138,6 +146,7 @@ impl From for Request { histogram_interval: req.search_info.histogram_interval, // Never carried across nodes. chat_history_reader: false, + analyze_profile: false, } } } diff --git a/src/config/src/meta/search.rs b/src/config/src/meta/search.rs index 43ae73516f..f8cdbe2c4d 100644 --- a/src/config/src/meta/search.rs +++ b/src/config/src/meta/search.rs @@ -198,6 +198,9 @@ pub struct Query { /// that don't carry their own 3rd timezone argument. None / "UTC" / "" => UTC. #[serde(default, skip_serializing_if = "Option::is_none")] pub timezone: Option, + /// Adds the structured `profile` column to EXPLAIN ANALYZE output. + #[serde(default)] + pub analyze_profile: bool, } fn default_size() -> i64 { @@ -223,6 +226,7 @@ impl Default for Query { streaming_id: None, histogram_interval: 0, timezone: None, + analyze_profile: false, } } } diff --git a/src/flight/src/common.rs b/src/flight/src/common.rs index 245802aed7..abaed37d97 100644 --- a/src/flight/src/common.rs +++ b/src/flight/src/common.rs @@ -14,7 +14,7 @@ // along with this program. If not, see . use std::{ - fmt, + collections::BTreeMap, sync::{ Arc, atomic::{AtomicUsize, Ordering}, @@ -31,14 +31,15 @@ use config::{ use datafusion::{ self, physical_plan::{ - DisplayFormatType, ExecutionPlan, - display::DisplayableExecutionPlan, + ExecutionPlan, metrics::{self, ExecutionPlanMetricsSet, MetricBuilder, MetricValue}, }, }; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; +use crate::profile::{ProfileNode, RemoteNodeProfile}; + #[derive(Debug, Clone)] pub struct RemoteScanMetrics { /// Time in nanos to execute child operator and fetch batches @@ -115,154 +116,117 @@ pub enum CustomMessage { #[derive(Default, Debug, Serialize, Deserialize, Clone)] pub struct Metrics { + /// Unused; kept, like `metrics`, so that older leaders can still parse the message. pub stage: i32, pub node_address: String, pub node_name: String, + /// Left empty: the leader renders the text plan from `profile`. pub metrics: String, + /// The plan this node ran; older nodes do not send it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub profile: Option, } -/// the CustomMessage is not ready, so we need to use this struct to store -/// the scan stats, and send it to the client when the scan stats is ready -pub enum PreCustomMessage { - ScanStats(ScanStats), - // use for super cluster follower leader - ScanStatsRef(Option>>), - // physical plan, is_super_cluster - Metrics(Option), - // use for super cluster follower leader - MetricsRef(Vec>>>), - // use for storing memory pool reference to extract peak later - PeakMemoryRef(Option>), - // sent before first batch (early signal for fast-failing nodes) - PartialErrRefEarly(Option>>), - // sent after all batches (final accumulated value) - PartialErrRef(Option>>), +/// What a node sends besides its record batches; the shared references are read when sent. +pub struct NodeMessages { + /// This node's own scan stats, known before execution. + pub scan_stats: ScanStats, + /// The peak of this node's memory pool. + pub peak_memory: Arc, + /// The plan to send back as a profile; set only for EXPLAIN ANALYZE. + pub profile: Option, + pub role: NodeRole, +} + +impl NodeMessages { + /// Sent before the first batch, so the caller sees scan stats and fast failures early. + pub fn first(&self) -> Vec { + let mut messages = vec![CustomMessage::ScanStats(self.reported_scan_stats())]; + messages.extend(self.partial_err()); + messages + } + + /// Sent after the last batch. + pub fn last(&self) -> Vec { + let mut messages = vec![CustomMessage::PeakMemory(self.reported_peak_memory())]; + messages.extend(self.partial_err()); + messages.extend( + self.profile + .as_ref() + .map(|info| CustomMessage::Metrics(collect_metrics(info))), + ); + messages + } + + /// Merged here rather than sent twice: the caller's add gives the same result either way. + fn reported_scan_stats(&self) -> ScanStats { + let mut stats = self.scan_stats; + if let NodeRole::ClusterLeader(remote) = &self.role { + stats.add(&remote.scan_stats.lock()); + } + stats + } + + fn reported_peak_memory(&self) -> usize { + let own = self.peak_memory.load(Ordering::Relaxed); + match &self.role { + NodeRole::Follower => own, + NodeRole::ClusterLeader(remote) => own.max(remote.peak_memory.load(Ordering::Relaxed)), + } + } + + /// Only a cluster leader has one: a follower's own failure ends its stream instead. + fn partial_err(&self) -> Option { + let NodeRole::ClusterLeader(remote) = &self.role else { + return None; + }; + let err = remote.partial_err.lock().clone(); + (!err.is_empty()).then_some(CustomMessage::PartialErr(err)) + } +} + +pub enum NodeRole { + /// Scans the data itself. + Follower, + /// A super cluster's cluster leader, which passes the query on to its own followers. + ClusterLeader(RemoteScanState), +} + +/// Per-partition profiles of one remote scan, keyed by partition. +pub type NodeProfiles = Arc>>; + +/// What a remote scan collects from the nodes it calls, shared by every copy of the scan. +#[derive(Debug, Clone, Default)] +pub struct RemoteScanState { + pub scan_stats: Arc>, + pub peak_memory: Arc, + pub partial_err: Arc>, + /// Filled only for EXPLAIN ANALYZE. + pub node_profiles: NodeProfiles, } pub struct MetricsInfo { pub plan: Arc, - pub is_super_cluster: bool, - pub func: Box bool + Send>, + /// Passed in because it lives in the search crate, which depends on this one. + pub build_profile: fn(&dyn ExecutionPlan) -> ProfileNode, } -impl PreCustomMessage { - pub fn is_early_emit(&self) -> bool { - matches!( - self, - PreCustomMessage::ScanStats(_) - | PreCustomMessage::ScanStatsRef(_) - | PreCustomMessage::PartialErrRefEarly(_) - ) - } - - pub fn get_custom_message(&self) -> Option { - match self { - PreCustomMessage::ScanStats(stats) => Some(CustomMessage::ScanStats(*stats)), - PreCustomMessage::ScanStatsRef(stats_ref) => stats_ref - .as_ref() - .map(|stats| CustomMessage::ScanStats(*stats.lock())), - PreCustomMessage::Metrics(metrics_ref) => metrics_ref - .as_ref() - .map(|metrics_info| CustomMessage::Metrics(collect_metrics(metrics_info))), - PreCustomMessage::MetricsRef(metrics) => { - let metrics: Vec = metrics.iter().flat_map(|m| m.lock().clone()).collect(); - (!metrics.is_empty()).then_some(CustomMessage::Metrics(metrics)) - } - PreCustomMessage::PeakMemoryRef(peak_memory_ref) => peak_memory_ref - .as_ref() - .map(|peak| CustomMessage::PeakMemory(peak.load(Ordering::Relaxed))), - PreCustomMessage::PartialErrRefEarly(ref_opt) - | PreCustomMessage::PartialErrRef(ref_opt) => ref_opt.as_ref().and_then(|r| { - let err = r.lock().clone(); - if err.is_empty() { - None - } else { - Some(CustomMessage::PartialErr(err)) - } - }), - } - } -} - -// collect metrics from physical plan fn collect_metrics(metrics_info: &MetricsInfo) -> Vec { - let plan = &metrics_info.plan; - let is_super_cluster = metrics_info.is_super_cluster; - let func = &metrics_info.func; - let plan_with_metrics = if is_super_cluster { - format_leader_plan(plan.as_ref()) - } else { - DisplayableExecutionPlan::with_metrics(plan.as_ref()) - .set_show_statistics(false) - .indent(false) - .to_string() - }; - let stage = if func() { - if is_super_cluster { 1 } else { 2 } - } else { - 1 - }; vec![Metrics { - stage, node_address: LOCAL_NODE.get_grpc_addr().to_string(), node_name: LOCAL_NODE.get_name().to_string(), - metrics: plan_with_metrics, + profile: Some((metrics_info.build_profile)(metrics_info.plan.as_ref())), + ..Default::default() }] } -fn format_leader_plan(plan: &dyn ExecutionPlan) -> String { - let mut output = String::new(); - write_plan(&mut output, plan, 0); - output -} - -struct PlanDisplay<'a>(&'a dyn ExecutionPlan); - -impl fmt::Display for PlanDisplay<'_> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt_as(DisplayFormatType::Default, f) - } -} - -fn write_plan(output: &mut String, plan: &dyn ExecutionPlan, indent: usize) { - use fmt::Write; - - write!(output, "{:indent$}", "", indent = indent * 2).expect("write to String"); - write!(output, "{}", PlanDisplay(plan)).expect("write to String"); - - if let Some(metrics) = plan.metrics() { - let metrics = metrics - .aggregate_by_name() - .sorted_for_display() - .timestamps_removed(); - write!(output, ", metrics=[{metrics}]").expect("write to String"); - } else { - write!(output, ", metrics=[]").expect("write to String"); - } - writeln!(output).expect("write to String"); - - if plan.name() == "RemoteScanExec" { - return; - } - - for child in plan.children() { - write_plan(output, child.as_ref(), indent + 1); - } -} - #[cfg(test)] mod tests { use std::sync::Arc; use arrow_schema::{DataType, Field, Schema}; use config::meta::search::ScanStats; - use datafusion::{ - common::Result, - execution::{SendableRecordBatchStream, TaskContext}, - physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, empty::EmptyExec, - }, - }; + use datafusion::physical_plan::{ExecutionPlan, empty::EmptyExec}; use parking_lot::Mutex; use super::*; @@ -285,70 +249,18 @@ mod tests { } } + fn test_profile(plan: &dyn ExecutionPlan) -> ProfileNode { + ProfileNode { + name: plan.name().to_string(), + ..Default::default() + } + } + fn create_test_execution_plan() -> Arc { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); Arc::new(EmptyExec::new(schema)) } - #[derive(Debug)] - struct NamedExec { - name: &'static str, - inner: Arc, - } - - impl DisplayAs for NamedExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "{}", self.name) - } - } - - impl ExecutionPlan for NamedExec { - fn apply_expressions( - &self, - _f: &mut dyn FnMut( - &Arc, - ) -> datafusion::common::Result< - datafusion::common::tree_node::TreeNodeRecursion, - >, - ) -> datafusion::common::Result { - Ok(datafusion::common::tree_node::TreeNodeRecursion::Continue) - } - - fn name(&self) -> &'static str { - self.name - } - - fn properties(&self) -> &Arc { - self.inner.properties() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.inner] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - unimplemented!() - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - unimplemented!() - } - } - - fn create_named_test_execution_plan( - name: &'static str, - inner: Arc, - ) -> Arc { - Arc::new(NamedExec { name, inner }) - } - #[test] fn test_metrics_default() { let metrics = Metrics::default(); @@ -358,225 +270,97 @@ mod tests { assert_eq!(metrics.metrics, ""); } - #[test] - fn test_pre_custom_message_is_early_emit() { - let scan_stats = create_test_scan_stats(); - - // Test ScanStats variant - let pre_msg = PreCustomMessage::ScanStats(scan_stats); - assert!(pre_msg.is_early_emit()); - - // Test ScanStatsRef variant with Some - let stats_ref = Arc::new(Mutex::new(scan_stats)); - let pre_msg = PreCustomMessage::ScanStatsRef(Some(stats_ref)); - assert!(pre_msg.is_early_emit()); - - // Test ScanStatsRef variant with None - let pre_msg = PreCustomMessage::ScanStatsRef(None); - assert!(pre_msg.is_early_emit()); - - // Test Metrics variant (should be false) - let pre_msg = PreCustomMessage::Metrics(None); - assert!(!pre_msg.is_early_emit()); - - // Test MetricsRef variant (should be false) - let pre_msg = PreCustomMessage::MetricsRef(vec![]); - assert!(!pre_msg.is_early_emit()); + fn describe(messages: Vec) -> Vec { + messages + .into_iter() + .map(|message| match message { + CustomMessage::ScanStats(stats) => format!("ScanStats({})", stats.files), + CustomMessage::PeakMemory(peak) => format!("PeakMemory({peak})"), + CustomMessage::PartialErr(err) => format!("PartialErr({err})"), + CustomMessage::Metrics(metrics) => { + format!("Metrics({})", metrics[0].profile.as_ref().unwrap().name) + } + }) + .collect() } #[test] - fn test_pre_custom_message_get_custom_message_scan_stats() { - let scan_stats = create_test_scan_stats(); - - let pre_msg = PreCustomMessage::ScanStats(scan_stats); - let custom_msg = pre_msg.get_custom_message().unwrap(); - - match custom_msg { - CustomMessage::ScanStats(stats) => { - assert_eq!(stats.files, 10); - assert_eq!(stats.records, 1000); - } - _ => panic!("Expected ScanStats variant"), - } - } - - #[test] - fn test_pre_custom_message_get_custom_message_scan_stats_ref() { - let scan_stats = create_test_scan_stats(); - let stats_ref = Arc::new(Mutex::new(scan_stats)); - - let pre_msg = PreCustomMessage::ScanStatsRef(Some(stats_ref)); - let custom_msg = pre_msg.get_custom_message().unwrap(); - - match custom_msg { - CustomMessage::ScanStats(stats) => { - assert_eq!(stats.files, 10); - assert_eq!(stats.records, 1000); - } - _ => panic!("Expected ScanStats variant"), - } - } - - #[test] - fn test_pre_custom_message_get_custom_message_scan_stats_ref_none() { - let pre_msg = PreCustomMessage::ScanStatsRef(None); - let custom_msg = pre_msg.get_custom_message(); - assert!(custom_msg.is_none()); - } - - #[test] - fn test_pre_custom_message_get_custom_message_metrics() { - let plan = create_test_execution_plan(); - let metrics_info = MetricsInfo { - plan, - is_super_cluster: false, - func: Box::new(|| true), + fn test_follower_sends_only_its_own_values() { + let messages = NodeMessages { + scan_stats: create_test_scan_stats(), + peak_memory: Arc::new(AtomicUsize::new(5)), + profile: None, + role: NodeRole::Follower, }; - - let pre_msg = PreCustomMessage::Metrics(Some(metrics_info)); - let custom_msg = pre_msg.get_custom_message().unwrap(); - - match custom_msg { - CustomMessage::Metrics(metrics) => { - assert_eq!(metrics.len(), 1); - assert_eq!(metrics[0].stage, 2); // func() returns true, is_super_cluster is false - assert!(!metrics[0].node_address.is_empty()); - assert!(!metrics[0].node_name.is_empty()); - assert!(!metrics[0].metrics.is_empty()); - } - _ => panic!("Expected Metrics variant"), - } + assert_eq!(describe(messages.first()), ["ScanStats(10)"]); + assert_eq!(describe(messages.last()), ["PeakMemory(5)"]); } #[test] - fn test_pre_custom_message_get_custom_message_metrics_none() { - let pre_msg = PreCustomMessage::Metrics(None); - let custom_msg = pre_msg.get_custom_message(); - assert!(custom_msg.is_none()); - } - - #[test] - fn test_pre_custom_message_get_custom_message_metrics_ref() { - let metrics1 = vec![Metrics { - stage: 1, - node_address: "node1".to_string(), - node_name: "Node 1".to_string(), - metrics: "metrics1".to_string(), - }]; - let metrics2 = vec![Metrics { - stage: 2, - node_address: "node2".to_string(), - node_name: "Node 2".to_string(), - metrics: "metrics2".to_string(), - }]; - - let metrics_refs = vec![ - Arc::new(Mutex::new(metrics1)), - Arc::new(Mutex::new(metrics2)), - ]; - - let pre_msg = PreCustomMessage::MetricsRef(metrics_refs); - let custom_msg = pre_msg.get_custom_message().unwrap(); - - match custom_msg { - CustomMessage::Metrics(metrics) => { - assert_eq!(metrics.len(), 2); - assert_eq!(metrics[0].node_address, "node1"); - assert_eq!(metrics[1].node_address, "node2"); - } - _ => panic!("Expected Metrics variant"), - } - } - - #[test] - fn test_pre_custom_message_get_custom_message_metrics_ref_empty() { - let metrics_refs = vec![Arc::new(Mutex::new(vec![]))]; - - let pre_msg = PreCustomMessage::MetricsRef(metrics_refs); - let custom_msg = pre_msg.get_custom_message(); - assert!(custom_msg.is_none()); - } - - #[test] - fn test_collect_metrics_super_cluster_true() { - let plan = create_test_execution_plan(); - let metrics_info = MetricsInfo { - plan, - is_super_cluster: true, - func: Box::new(|| true), + fn test_cluster_leader_merges_what_its_remote_scan_collected() { + let remote = RemoteScanState { + scan_stats: Arc::new(Mutex::new(ScanStats { + files: 3, + ..Default::default() + })), + peak_memory: Arc::new(AtomicUsize::new(7)), + ..Default::default() }; + let partial_err = remote.partial_err.clone(); + let messages = NodeMessages { + scan_stats: create_test_scan_stats(), + peak_memory: Arc::new(AtomicUsize::new(5)), + profile: Some(MetricsInfo { + plan: create_test_execution_plan(), + build_profile: test_profile, + }), + role: NodeRole::ClusterLeader(remote), + }; + assert_eq!(describe(messages.first()), ["ScanStats(13)"]); - let metrics = collect_metrics(&metrics_info); + *partial_err.lock() = "node down".to_string(); + assert_eq!( + describe(messages.first()), + ["ScanStats(13)", "PartialErr(node down)"] + ); + assert_eq!( + describe(messages.last()), + [ + "PeakMemory(7)", + "PartialErr(node down)", + "Metrics(EmptyExec)" + ] + ); + } + + #[test] + fn test_profile_carries_this_node() { + let metrics = collect_metrics(&MetricsInfo { + plan: create_test_execution_plan(), + build_profile: test_profile, + }); assert_eq!(metrics.len(), 1); - assert_eq!(metrics[0].stage, 1); // func() returns true, is_super_cluster is true assert!(!metrics[0].node_address.is_empty()); - assert!(!metrics[0].metrics.is_empty()); + assert!(!metrics[0].node_name.is_empty()); } #[test] - fn test_collect_metrics_super_cluster_false() { - let plan = create_test_execution_plan(); - let metrics_info = MetricsInfo { - plan, - is_super_cluster: false, - func: Box::new(|| true), - }; - - let metrics = collect_metrics(&metrics_info); - assert_eq!(metrics.len(), 1); - assert_eq!(metrics[0].stage, 2); // func() returns true, is_super_cluster is false + fn test_metrics_without_profile_deserializes() { + let json = r#"{"stage":1,"node_address":"a","node_name":"b","metrics":"m"}"#; + let metrics: Metrics = serde_json::from_str(json).unwrap(); + assert!(metrics.profile.is_none()); } #[test] - fn test_collect_metrics_func_returns_false() { - let plan = create_test_execution_plan(); - let metrics_info = MetricsInfo { - plan, - is_super_cluster: true, - func: Box::new(|| false), + fn test_metrics_keeps_fields_older_leaders_require() { + let metrics = Metrics { + profile: Some(ProfileNode::default()), + ..Default::default() }; - - let metrics = collect_metrics(&metrics_info); - assert_eq!(metrics.len(), 1); - assert_eq!(metrics[0].stage, 1); // func() returns false, stage defaults to 1 - } - - #[test] - fn test_super_cluster_stops_at_remote_scan() { - let follower_plan = - create_named_test_execution_plan("FollowerOnlyExec", create_test_execution_plan()); - let remote_scan_plan = create_named_test_execution_plan("RemoteScanExec", follower_plan); - let region_leader_plan = - create_named_test_execution_plan("SortPreservingMergeExec", remote_scan_plan); - let metrics_info = MetricsInfo { - plan: region_leader_plan, - is_super_cluster: true, - func: Box::new(|| true), - }; - - let metrics = collect_metrics(&metrics_info); - assert_eq!(metrics.len(), 1); - assert!(metrics[0].metrics.contains("SortPreservingMergeExec")); - assert!(metrics[0].metrics.contains("RemoteScanExec")); - assert!(!metrics[0].metrics.contains("FollowerOnlyExec")); - } - - #[test] - fn test_non_super_cluster_keeps_remote_scan_child_metrics() { - let follower_plan = - create_named_test_execution_plan("FollowerOnlyExec", create_test_execution_plan()); - let remote_scan_plan = create_named_test_execution_plan("RemoteScanExec", follower_plan); - let region_leader_plan = - create_named_test_execution_plan("SortPreservingMergeExec", remote_scan_plan); - let metrics_info = MetricsInfo { - plan: region_leader_plan, - is_super_cluster: false, - func: Box::new(|| true), - }; - - let metrics = collect_metrics(&metrics_info); - assert_eq!(metrics.len(), 1); - assert!(metrics[0].metrics.contains("FollowerOnlyExec, metrics=[]")); + let json = serde_json::to_value(&metrics).unwrap(); + for key in ["stage", "node_address", "node_name", "metrics"] { + assert!(json.get(key).is_some(), "missing {key}"); + } } #[test] @@ -603,6 +387,7 @@ mod tests { node_address: "test_node".to_string(), node_name: "Test Node".to_string(), metrics: "test_metrics".to_string(), + profile: None, }]; let custom_msg = CustomMessage::Metrics(metrics); diff --git a/src/flight/src/lib.rs b/src/flight/src/lib.rs index b0842e58f1..30fa19159d 100644 --- a/src/flight/src/lib.rs +++ b/src/flight/src/lib.rs @@ -16,3 +16,4 @@ pub mod common; pub mod decoder; pub mod encoder; +pub mod profile; diff --git a/src/flight/src/profile.rs b/src/flight/src/profile.rs new file mode 100644 index 0000000000..e6e6a8c9fb --- /dev/null +++ b/src/flight/src/profile.rs @@ -0,0 +1,251 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! Structured EXPLAIN ANALYZE profile: the executed plan with typed, per-node metrics. +//! +//! The profile is a tree that follows the leader's physical plan. Ordinary operators +//! recurse through `ProfileInputs::Children`. A `RemoteScanExec` stops the recursion +//! with `ProfileInputs::Remote` instead, because the subtree under it on the leader is +//! only the plan template sent to the remote nodes and never ran there. That list +//! holds one `RemoteNodeProfile` per node it called directly (one per partition), +//! including nodes skipped because they had no files. +//! +//! A `RemoteNodeProfile` has two parts: +//! - what the caller observed of that node: start and end time, rows, batches and bytes received +//! over Flight, decode time, and the node's scan stats; +//! - `plan`: the plan the node actually ran, with its own metrics, sent back by the node itself. +//! +//! A caller lists only the nodes it called directly; deeper nodes nest inside +//! `plan`. In a super cluster the cluster leader's plan has its own +//! `RemoteScanExec`, whose `remote` lists that cluster's queriers and ingesters: +//! +//! ```text +//! QueryProfile super leader +//! └─ plan … RemoteScanExec +//! └─ remote: [sc-a, sc-b, sc-c] cluster leaders +//! └─ plan … RemoteScanExec +//! └─ remote: [querier-0 … ingester-1] nodes of that cluster +//! └─ plan … +//! ``` +//! +//! A plan with several remote scans, such as a join of two streams, has one +//! `remote` list per scan, so the same node appears once under each of them. +//! +//! Each plan appears exactly once: a node sends back only its own plan, and its +//! caller stores it under the remote scan that made the call. +//! +//! The leader also renders the text rows of EXPLAIN ANALYZE from this tree, one per +//! node that sent back its plan, which is why `ProfileMetric::display` keeps each +//! value exactly as DataFusion prints it. + +use std::fmt::Write; + +use config::meta::search::ScanStats; +use serde::{Deserialize, Serialize}; + +/// The profile of one EXPLAIN ANALYZE run, as returned to the client. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct QueryProfile { + pub node_name: String, + pub node_address: String, + pub start_us: i64, + pub end_us: i64, + pub output_rows: usize, + pub plan: ProfileNode, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ProfileNode { + pub name: String, + pub detail: String, + pub partitions: usize, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub metrics: Vec, + #[serde(flatten)] + pub inputs: ProfileInputs, +} + +impl ProfileNode { + /// Renders the EXPLAIN ANALYZE text of the plan this node heads, stopping at remote scans. + pub fn to_text(&self) -> String { + let mut out = String::new(); + self.write_text(&mut out, 0); + out + } + + fn write_text(&self, out: &mut String, indent: usize) { + let metrics = self + .metrics + .iter() + .map(|m| format!("{}={}", m.name, m.display)) + .collect::>() + .join(", "); + writeln!( + out, + "{:indent$}{}, metrics=[{metrics}]", + "", + self.detail, + indent = indent * 2 + ) + .expect("write to String"); + if let ProfileInputs::Children(children) = &self.inputs { + for child in children { + child.write_text(out, indent + 1); + } + } + } +} + +/// Where a node's input came from; serialized as a `children` or a `remote` key. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ProfileInputs { + /// Operators that ran on this node; empty for a leaf. + Children(Vec), + /// Only on a remote scan: one entry per node it called directly. + Remote(Vec), +} + +impl Default for ProfileInputs { + fn default() -> Self { + Self::Children(Vec::new()) + } +} + +/// A metric summed across partitions. +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +pub struct ProfileMetric { + pub name: String, + pub kind: MetricKind, + pub value: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub total: Option, + /// The value as DataFusion prints it, such as `1.20 MB` or `100 total → 20 matched`. + #[serde(default)] + pub display: String, +} + +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum MetricKind { + Count, + /// Nanoseconds. + Time, + Bytes, + Gauge, + Ratio, + Pruning, + #[default] + #[serde(other)] + Other, +} + +/// What the caller saw of one remote node, plus the plan that node ran. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct RemoteNodeProfile { + pub partition: usize, + pub node_name: String, + pub node_address: String, + pub role: String, + pub cluster: String, + pub start_us: i64, + pub end_us: i64, + pub output_rows: u64, + pub output_batches: u64, + pub output_bytes: u64, + pub decode_ns: u64, + pub skipped: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scan_stats: Option, + /// The plan this node ran; its own remote scans nest the nodes below it the same way. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub plan: Option, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_profile_json_round_trip_and_unknown_kind() { + let node = ProfileNode { + name: "FilterExec".to_string(), + metrics: vec![ProfileMetric { + name: "output_rows".to_string(), + kind: MetricKind::Count, + value: 1, + total: None, + display: "1".to_string(), + }], + ..Default::default() + }; + let json = serde_json::to_string(&node).unwrap(); + assert!(json.contains(r#""children":[]"#) && !json.contains("remote")); + let back: ProfileNode = serde_json::from_str(&json).unwrap(); + assert_eq!(back.metrics, node.metrics); + assert!(matches!(back.inputs, ProfileInputs::Children(c) if c.is_empty())); + let m: ProfileMetric = + serde_json::from_str(r#"{"name":"x","kind":"future_kind","value":1}"#).unwrap(); + assert_eq!(m.kind, MetricKind::Other); + } + + #[test] + fn test_remote_inputs_serialize_under_remote_key() { + let node = ProfileNode { + name: "RemoteScanExec".to_string(), + inputs: ProfileInputs::Remote(vec![RemoteNodeProfile { + node_name: "q-0".to_string(), + ..Default::default() + }]), + ..Default::default() + }; + let json = serde_json::to_string(&node).unwrap(); + assert!(json.contains(r#""remote":[{"#) && !json.contains("children")); + let back: ProfileNode = serde_json::from_str(&json).unwrap(); + let ProfileInputs::Remote(remote) = back.inputs else { + panic!("expected remote inputs"); + }; + assert_eq!(remote[0].node_name, "q-0"); + } + + #[test] + fn test_to_text_indents_children_and_stops_at_remote() { + let metric = |name: &str, display: &str| ProfileMetric { + name: name.to_string(), + display: display.to_string(), + ..Default::default() + }; + let node = ProfileNode { + detail: "SortExec: expr=[a@0 ASC]".to_string(), + metrics: vec![metric("output_rows", "3"), metric("spilled_bytes", "0.0 B")], + inputs: ProfileInputs::Children(vec![ProfileNode { + detail: "RemoteScanExec: partitions=2".to_string(), + inputs: ProfileInputs::Remote(vec![RemoteNodeProfile { + plan: Some(ProfileNode { + detail: "FilterExec".to_string(), + ..Default::default() + }), + ..Default::default() + }]), + ..Default::default() + }]), + ..Default::default() + }; + assert_eq!( + node.to_text(), + "SortExec: expr=[a@0 ASC], metrics=[output_rows=3, spilled_bytes=0.0 B]\n RemoteScanExec: partitions=2, metrics=[]\n" + ); + } +} diff --git a/src/search/src/datafusion/distributed_plan/common.rs b/src/search/src/datafusion/distributed_plan/common.rs index f3f4470207..d021b36f2b 100644 --- a/src/search/src/datafusion/distributed_plan/common.rs +++ b/src/search/src/datafusion/distributed_plan/common.rs @@ -22,9 +22,10 @@ use datafusion::{ common::Result, execution::SendableRecordBatchStream, physical_plan::stream::RecordBatchStreamAdapter, }; -use flight::common::Metrics; use parking_lot::Mutex; +use super::profile::NodeProfileRecorder; + #[derive(Debug)] pub struct QueryContext { pub trace_id: String, @@ -33,13 +34,13 @@ pub struct QueryContext { pub is_querier: bool, pub scan_stats: Arc>, pub partial_err: Arc>, - pub cluster_metrics: Arc>>, pub peak_memory: Arc, pub start: std::time::Instant, pub num_rows: usize, pub req_id: u64, pub req_last_time: std::time::Instant, pub print_key_event: bool, + pub node_profile: NodeProfileRecorder, } impl QueryContext { @@ -50,7 +51,6 @@ impl QueryContext { is_querier: false, scan_stats: Arc::new(Mutex::new(ScanStats::new())), partial_err: Arc::new(Mutex::new(String::new())), - cluster_metrics: Arc::new(Mutex::new(Vec::new())), peak_memory: Arc::new(AtomicUsize::new(0)), trace_id: String::new(), start: std::time::Instant::now(), @@ -58,6 +58,7 @@ impl QueryContext { req_id: 0, req_last_time: std::time::Instant::now(), print_key_event: config::get_config().common.print_key_event, + node_profile: NodeProfileRecorder::default(), } } @@ -86,11 +87,6 @@ impl QueryContext { self } - pub fn with_cluster_metrics(mut self, metrics: Arc>>) -> Self { - self.cluster_metrics = metrics; - self - } - pub fn with_start_time(mut self, start: std::time::Instant) -> Self { self.start = start; self @@ -100,6 +96,11 @@ impl QueryContext { self.peak_memory = peak_memory; self } + + pub fn with_node_profile(mut self, node_profile: NodeProfileRecorder) -> Self { + self.node_profile = node_profile; + self + } } /// EmptyStream is a stream that returns an empty record batch diff --git a/src/search/src/datafusion/distributed_plan/decoder_stream.rs b/src/search/src/datafusion/distributed_plan/decoder_stream.rs index 1aa382da1f..000759e828 100644 --- a/src/search/src/datafusion/distributed_plan/decoder_stream.rs +++ b/src/search/src/datafusion/distributed_plan/decoder_stream.rs @@ -67,7 +67,7 @@ impl FlightDecoderStream { self.query_context.scan_stats.lock().add(&stats); } CustomMessage::Metrics(metrics) => { - self.query_context.cluster_metrics.lock().extend(metrics); + self.query_context.node_profile.record_plan(metrics); } CustomMessage::PeakMemory(peak_memory) => { self.query_context @@ -191,6 +191,10 @@ impl Stream for FlightDecoderStream { impl Drop for FlightDecoderStream { fn drop(&mut self) { + let scan_stats = self.scan_stats; + self.query_context + .node_profile + .finish(&self.metrics, scan_stats); let QueryContext { trace_id, node, @@ -206,7 +210,6 @@ impl Drop for FlightDecoderStream { } else { "follower".to_string() }; - let scan_stats = self.scan_stats; let batches = self.metrics.batches_received.value(); let bytes = self.metrics.bytes_received.value(); let avg_ms = self.metrics.avg_batch_time_ms(); diff --git a/src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs b/src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs index 5903d86def..1f7c28904e 100644 --- a/src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs +++ b/src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs @@ -16,12 +16,12 @@ use std::sync::Arc; use arrow::{ - array::{Int32Builder, StringBuilder}, + array::{ArrayBuilder, ArrayRef, Int32Builder, StringArray, StringBuilder}, datatypes::SchemaRef, record_batch::RecordBatch, }; use arrow_schema::{DataType, Field, Schema}; -use config::{cluster::LOCAL_NODE, meta::cluster::NodeInfo}; +use config::{cluster::LOCAL_NODE, meta::cluster::NodeInfo, utils::time::now_micros}; use datafusion::{ common::{DataFusionError, Result, internal_err, tree_node::TreeNodeRecursion}, execution::TaskContext, @@ -32,24 +32,22 @@ use datafusion::{ stream::RecordBatchStreamAdapter, }, }; -use flight::common::Metrics; +use flight::profile::{ProfileInputs, ProfileNode, QueryProfile, RemoteNodeProfile}; use futures::StreamExt; -use crate::datafusion::{ - distributed_plan::display::DisplayableExecutionPlan, plan_metrics::get_cluster_metrics, +use crate::datafusion::distributed_plan::{ + display::DisplayableExecutionPlan, profile::build_profile, }; -/// query EXPLAIN ANALYZE in distributed mode -/// the output recordbatch's schema is -/// phase: current we have three phase -/// node: the ip or partition id -/// plan: plan with metrics -#[derive(Debug)] +/// EXPLAIN ANALYZE across nodes: one text row per node that ran a plan, phase = its depth. +#[derive(Debug, Clone)] pub struct DistributeAnalyzeExec { /// Control how much extra to print verbose: bool, /// If statistics should be displayed show_statistics: bool, + /// Adds the `profile` column, the whole profile as JSON on the first row. + profile: bool, /// The input plan (the plan being analyzed) pub input: Arc, /// The output schema for RecordBatches of this exec node @@ -58,25 +56,90 @@ pub struct DistributeAnalyzeExec { } impl DistributeAnalyzeExec { - /// Create a new AnalyzeExec - pub fn new(verbose: bool, show_statistics: bool, input: Arc) -> Self { - let fields = vec![ + pub fn new( + verbose: bool, + show_statistics: bool, + profile: bool, + input: Arc, + ) -> Self { + let mut fields = vec![ Field::new("phase", DataType::Int32, false), Field::new("node_address", DataType::Utf8, false), Field::new("node_name", DataType::Utf8, false), Field::new("plan", DataType::Utf8, false), ]; + if profile { + fields.push(Field::new("profile", DataType::Utf8, true)); + } let schema = Arc::new(Schema::new(fields)); let cache = Self::compute_properties(&input, Arc::clone(&schema)); DistributeAnalyzeExec { verbose, show_statistics, + profile, input, schema, cache, } } + fn output_batch( + &self, + profile: &QueryProfile, + duration: std::time::Duration, + ) -> Result { + let mut phase_builder = Int32Builder::new(); + let mut address_builder = StringBuilder::with_capacity(1, 1024); + let mut name_builder = StringBuilder::with_capacity(1, 1024); + let mut plan_builder = StringBuilder::with_capacity(1, 1024); + + phase_builder.append_value(0); + address_builder.append_value(&profile.node_address); + name_builder.append_value(&profile.node_name); + let annotated_plan = DisplayableExecutionPlan::new(self.input.as_ref()) + .indent(self.verbose) + .to_string(); + plan_builder.append_value(annotated_plan); + + let mut remote_rows = Vec::new(); + collect_remote_rows(&profile.plan, 1, &mut remote_rows); + for (stage, node, plan) in remote_rows { + phase_builder.append_value(stage); + address_builder.append_value(&node.node_address); + name_builder.append_value(&node.node_name); + plan_builder.append_value(plan.to_text()); + } + + if self.verbose { + phase_builder.append_value(0); + address_builder.append_value("Output Rows"); + name_builder.append_value(""); + plan_builder.append_value(profile.output_rows.to_string()); + + phase_builder.append_value(0); + address_builder.append_value("Duration"); + name_builder.append_value(""); + plan_builder.append_value(format!("{duration:?}")); + } + + let rows = phase_builder.len(); + let mut columns: Vec = vec![ + Arc::new(phase_builder.finish()), + Arc::new(address_builder.finish()), + Arc::new(name_builder.finish()), + Arc::new(plan_builder.finish()), + ]; + if self.profile { + let json = + serde_json::to_string(profile).map_err(|e| DataFusionError::External(e.into()))?; + columns.push(Arc::new(StringArray::from_iter( + std::iter::once(Some(json)).chain(std::iter::repeat_n(None, rows - 1)), + ))); + } + + RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(DataFusionError::from) + } + /// This function creates the cache object that stores the plan properties such as schema, /// equivalence properties, ordering, partitioning, etc. fn compute_properties( @@ -130,6 +193,7 @@ impl ExecutionPlan for DistributeAnalyzeExec { Ok(Arc::new(Self::new( self.verbose, self.show_statistics, + self.profile, children.pop().unwrap(), ))) } @@ -147,10 +211,8 @@ impl ExecutionPlan for DistributeAnalyzeExec { // Create future that computes the final output let start = std::time::Instant::now(); - let captured_input = Arc::clone(&self.input); - let captured_schema = Arc::clone(&self.schema); - let verbose = self.verbose; - let show_statistics = self.show_statistics; + let start_us = now_micros(); + let exec = self.clone(); let mut stream = execute_stream(self.input.clone(), context.clone())?; let output = async move { @@ -158,16 +220,19 @@ impl ExecutionPlan for DistributeAnalyzeExec { while let Some(batch) = stream.next().await.transpose()? { total_rows += batch.num_rows(); } + // remote nodes publish their profiles when their streams are dropped + drop(stream); let duration = std::time::Instant::now() - start; - create_output_batch( - verbose, - show_statistics, - total_rows, - duration, - captured_input, - captured_schema, - ) + let profile = QueryProfile { + node_name: LOCAL_NODE.get_name(), + node_address: LOCAL_NODE.get_grpc_addr(), + start_us, + end_us: now_micros(), + output_rows: total_rows, + plan: build_profile(exec.input.as_ref()), + }; + exec.output_batch(&profile, duration) }; Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -177,75 +242,35 @@ impl ExecutionPlan for DistributeAnalyzeExec { } } -fn create_output_batch( - verbose: bool, - _show_statistics: bool, - total_rows: usize, - duration: std::time::Duration, - input: Arc, - schema: SchemaRef, -) -> Result { - let mut phase_builder = Int32Builder::new(); - let mut address_builder = StringBuilder::with_capacity(1, 1024); - let mut name_builder = StringBuilder::with_capacity(1, 1024); - let mut plan_builder = StringBuilder::with_capacity(1, 1024); - - phase_builder.append_value(0); // Phase 0 for the main plan - address_builder.append_value(LOCAL_NODE.get_grpc_addr()); - name_builder.append_value(LOCAL_NODE.get_name()); - - let annotated_plan = DisplayableExecutionPlan::new(input.as_ref()) - .indent(verbose) - .to_string(); - plan_builder.append_value(annotated_plan); - - // add metrics from remote scan - let metrics = collect_metrics(input); - - for m in metrics { - phase_builder.append_value(m.stage); - address_builder.append_value(m.node_address); - name_builder.append_value(m.node_name); - plan_builder.append_value(m.metrics); +/// Lists, depth first, each remote node that sent back its plan, with its stage below the leader. +fn collect_remote_rows<'a>( + node: &'a ProfileNode, + stage: i32, + rows: &mut Vec<(i32, &'a RemoteNodeProfile, &'a ProfileNode)>, +) { + match &node.inputs { + ProfileInputs::Children(children) => { + for child in children { + collect_remote_rows(child, stage, rows); + } + } + ProfileInputs::Remote(remotes) => { + for (remote, plan) in remotes + .iter() + .filter_map(|r| r.plan.as_ref().map(|plan| (r, plan))) + { + rows.push((stage, remote, plan)); + collect_remote_rows(plan, stage + 1, rows); + } + } } - - if verbose { - phase_builder.append_value(0); // Phase 0 for the main plan - address_builder.append_value("Output Rows"); - name_builder.append_value(""); - plan_builder.append_value(total_rows.to_string()); - - phase_builder.append_value(0); // Phase 0 for the main plan - address_builder.append_value("Duration"); - name_builder.append_value(""); - plan_builder.append_value(format!("{duration:?}")); - } - - RecordBatch::try_new( - schema, - vec![ - Arc::new(phase_builder.finish()), - Arc::new(address_builder.finish()), - Arc::new(name_builder.finish()), - Arc::new(plan_builder.finish()), - ], - ) - .map_err(DataFusionError::from) -} - -fn collect_metrics(plan: Arc) -> Vec { - let mut metrics = Vec::new(); - for m in get_cluster_metrics(&plan) { - let m = m.lock().clone(); - metrics.extend(m); - } - metrics } #[cfg(test)] mod tests { use std::sync::Arc; + use arrow::array::Array; use arrow_schema::{DataType, Field, Schema}; use datafusion::physical_plan::empty::EmptyExec; @@ -259,7 +284,7 @@ mod tests { #[test] fn test_distribute_analyze_exec_new_non_verbose() { let input = empty_input(); - let exec = DistributeAnalyzeExec::new(false, false, input); + let exec = DistributeAnalyzeExec::new(false, false, false, input); assert_eq!(exec.name(), "DistributeAnalyzeExec"); assert_eq!(exec.children().len(), 1); } @@ -267,16 +292,91 @@ mod tests { #[test] fn test_distribute_analyze_exec_new_verbose() { let input = empty_input(); - let exec = DistributeAnalyzeExec::new(true, true, input); + let exec = DistributeAnalyzeExec::new(true, true, false, input); assert_eq!(exec.name(), "DistributeAnalyzeExec"); } #[test] fn test_distribute_analyze_exec_schema_has_four_fields() { let input = empty_input(); - let exec = DistributeAnalyzeExec::new(false, false, input); + let exec = DistributeAnalyzeExec::new(false, false, false, input); assert_eq!(exec.schema().fields().len(), 4); assert_eq!(exec.schema().field(0).name(), "phase"); assert_eq!(exec.schema().field(1).name(), "node_address"); } + + #[test] + fn test_distribute_analyze_exec_schema_adds_profile_when_requested() { + let exec = DistributeAnalyzeExec::new(false, false, true, empty_input()); + assert_eq!(exec.schema().fields().len(), 5); + assert_eq!(exec.schema().field(4).name(), "profile"); + assert!(exec.schema().field(4).is_nullable()); + } + + #[tokio::test] + async fn test_distribute_analyze_exec_without_profile_emits_four_columns() { + let exec = Arc::new(DistributeAnalyzeExec::new( + true, + false, + false, + empty_input(), + )); + let batches = datafusion::physical_plan::collect(exec, Arc::new(TaskContext::default())) + .await + .unwrap(); + assert_eq!(batches[0].num_columns(), 4); + } + + #[tokio::test] + async fn test_distribute_analyze_exec_emits_profile_on_first_row() { + let exec = Arc::new(DistributeAnalyzeExec::new(true, false, true, empty_input())); + let batches = datafusion::physical_plan::collect(exec, Arc::new(TaskContext::default())) + .await + .unwrap(); + let profile = batches[0] + .column(4) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(profile.is_valid(0)); + assert!((1..profile.len()).all(|i| profile.is_null(i))); + let parsed: QueryProfile = serde_json::from_str(profile.value(0)).unwrap(); + assert_eq!(parsed.plan.name, "EmptyExec"); + assert!(parsed.end_us >= parsed.start_us); + } + + #[test] + fn test_collect_remote_rows_nests_stages_and_skips_nodes_without_plan() { + let node = |name: &str, plan: Option| RemoteNodeProfile { + node_name: name.to_string(), + plan, + ..Default::default() + }; + let scan = |remotes| ProfileNode { + name: "RemoteScanExec".to_string(), + inputs: ProfileInputs::Remote(remotes), + ..Default::default() + }; + let cluster_leader = node( + "b-querier-0", + Some(scan(vec![node( + "b-querier-1", + Some(ProfileNode::default()), + )])), + ); + let root = ProfileNode { + inputs: ProfileInputs::Children(vec![scan(vec![ + cluster_leader, + node("skipped", None), + ])]), + ..Default::default() + }; + let mut rows = Vec::new(); + collect_remote_rows(&root, 1, &mut rows); + let rows: Vec<_> = rows + .iter() + .map(|(stage, node, _)| (*stage, node.node_name.as_str())) + .collect(); + assert_eq!(rows, vec![(1, "b-querier-0"), (2, "b-querier-1")]); + } } diff --git a/src/search/src/datafusion/distributed_plan/mod.rs b/src/search/src/datafusion/distributed_plan/mod.rs index 8b455ee154..3a6323a177 100644 --- a/src/search/src/datafusion/distributed_plan/mod.rs +++ b/src/search/src/datafusion/distributed_plan/mod.rs @@ -39,6 +39,7 @@ pub mod enrichment_exec; pub mod metadata_count_exec; pub mod node; mod once_async; +pub mod profile; pub mod remote_scan_exec; pub mod rewrite; pub mod streaming_aggs_exec; diff --git a/src/search/src/datafusion/distributed_plan/profile.rs b/src/search/src/datafusion/distributed_plan/profile.rs new file mode 100644 index 0000000000..2849c21b2b --- /dev/null +++ b/src/search/src/datafusion/distributed_plan/profile.rs @@ -0,0 +1,380 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::fmt; + +use config::{ + meta::{cluster::NodeInfo, search::ScanStats}, + utils::time::now_micros, +}; +use datafusion::physical_plan::{ + DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + metrics::{MetricValue, MetricsSet}, +}; +use flight::{ + common::{Metrics, NodeProfiles, RemoteScanMetrics}, + profile::{MetricKind, ProfileInputs, ProfileMetric, ProfileNode, RemoteNodeProfile}, +}; + +use super::remote_scan_exec::RemoteScanExec; + +/// One remote node as the leader sees it, published to its scan on drop; the default records +/// nothing. +#[derive(Debug, Default)] +pub struct NodeProfileRecorder { + profiles: Option, + profile: RemoteNodeProfile, +} + +impl NodeProfileRecorder { + pub fn start(profiles: NodeProfiles, partition: usize, node: &dyn NodeInfo) -> Self { + let role = match (node.is_querier(), node.is_ingester()) { + (true, true) => "querier,ingester", + (true, false) => "querier", + (false, true) => "ingester", + (false, false) => "", + }; + Self { + profiles: Some(profiles), + profile: RemoteNodeProfile { + partition, + node_name: node.get_name(), + node_address: node.get_grpc_addr(), + role: role.to_string(), + cluster: node.get_cluster(), + start_us: now_micros(), + ..Default::default() + }, + } + } + + /// Keeps the plan the node sent back about itself. + pub fn record_plan(&mut self, metrics: Vec) { + if let Some(plan) = metrics.into_iter().find_map(|m| m.profile) { + self.profile.plan = Some(plan); + } + } + + pub fn mark_skipped(&mut self) { + self.profile.skipped = true; + } + + pub fn finish(&mut self, metrics: &RemoteScanMetrics, scan_stats: ScanStats) { + let p = &mut self.profile; + p.output_rows = metrics.output_rows.value() as u64; + p.output_batches = metrics.batches_received.value() as u64; + p.output_bytes = metrics.bytes_received.value() as u64; + p.decode_ns = metrics.decode_time.value() as u64; + p.scan_stats = Some(scan_stats); + } +} + +impl Drop for NodeProfileRecorder { + fn drop(&mut self) { + if let Some(profiles) = self.profiles.take() { + self.profile.end_us = now_micros(); + let profile = std::mem::take(&mut self.profile); + profiles.lock().insert(profile.partition, profile); + } + } +} + +/// A plan node's own line, exactly as EXPLAIN prints it. +struct PlanDetail<'a>(&'a dyn ExecutionPlan); + +impl fmt::Display for PlanDetail<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Default, f) + } +} + +/// Profile of an executed plan; a remote scan lists the nodes it called, as its children never ran. +pub fn build_profile(plan: &dyn ExecutionPlan) -> ProfileNode { + let inputs = match plan.downcast_ref::() { + Some(remote) => ProfileInputs::Remote( + remote + .state + .node_profiles + .lock() + .values() + .cloned() + .collect(), + ), + None => ProfileInputs::Children( + plan.children() + .into_iter() + .map(|child| build_profile(child.as_ref())) + .collect(), + ), + }; + ProfileNode { + name: plan.name().to_string(), + detail: PlanDetail(plan).to_string(), + partitions: plan.output_partitioning().partition_count(), + metrics: plan + .metrics() + .map(|m| convert_metrics(&m)) + .unwrap_or_default(), + inputs, + } +} + +/// Aggregates and orders the metrics the same way the text plan does, so `display` matches it. +fn convert_metrics(metrics: &MetricsSet) -> Vec { + metrics + .aggregate_by_name() + .sorted_for_display() + .timestamps_removed() + .iter() + .map(|metric| { + let value = metric.value(); + let (kind, number, total) = metric_value(value); + ProfileMetric { + name: value.name().to_string(), + kind, + value: number, + total, + display: value.to_string(), + } + }) + .collect() +} + +fn metric_value(value: &MetricValue) -> (MetricKind, u64, Option) { + let as_u64 = |v: usize| v as u64; + match value { + MetricValue::OutputRows(c) + | MetricValue::SpillCount(c) + | MetricValue::OutputBatches(c) + | MetricValue::SpilledRows(c) + | MetricValue::Count { count: c, .. } => (MetricKind::Count, as_u64(c.value()), None), + MetricValue::SpilledBytes(c) | MetricValue::OutputBytes(c) => { + (MetricKind::Bytes, as_u64(c.value()), None) + } + MetricValue::CurrentMemoryUsage(g) | MetricValue::PeakMemoryUsage { gauge: g, .. } => { + (MetricKind::Bytes, as_u64(g.value()), None) + } + MetricValue::ElapsedCompute(t) | MetricValue::Time { time: t, .. } => { + (MetricKind::Time, as_u64(t.value()), None) + } + MetricValue::Gauge { gauge: g, .. } => (MetricKind::Gauge, as_u64(g.value()), None), + MetricValue::PruningMetrics { + pruning_metrics, .. + } => { + let pruned = as_u64(pruning_metrics.pruned()); + let total = pruned + as_u64(pruning_metrics.matched()); + (MetricKind::Pruning, pruned, Some(total)) + } + MetricValue::Ratio { ratio_metrics, .. } => ( + MetricKind::Ratio, + as_u64(ratio_metrics.part()), + Some(as_u64(ratio_metrics.total())), + ), + // Timestamps never get here: convert_metrics removes them first. + MetricValue::Custom { .. } + | MetricValue::StartTimestamp(_) + | MetricValue::EndTimestamp(_) => (MetricKind::Other, as_u64(value.as_usize()), None), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Int32Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion::{ + physical_plan::{ + collect, + display::DisplayableExecutionPlan, + empty::EmptyExec, + metrics::{ExecutionPlanMetricsSet, MetricBuilder}, + union::UnionExec, + }, + prelude::SessionContext, + }; + + use super::*; + use crate::datafusion::distributed_plan::node::RemoteScanNode; + + #[derive(Debug)] + struct TestNode; + + impl NodeInfo for TestNode { + fn is_querier(&self) -> bool { + true + } + fn get_grpc_addr(&self) -> String { + "10.0.0.1:5081".to_string() + } + fn get_auth_token(&self) -> String { + String::new() + } + fn get_name(&self) -> String { + "q-0".to_string() + } + fn is_local(&self) -> bool { + false + } + } + + fn empty() -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + Arc::new(EmptyExec::new(schema)) + } + + fn batch(rows: i32) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from_iter_values(0..rows))], + ) + .unwrap() + } + + #[test] + fn test_recorder_publishes_one_node_on_drop() { + let profiles = NodeProfiles::default(); + let mut recorder = NodeProfileRecorder::start(profiles.clone(), 2, &TestNode); + let metrics = RemoteScanMetrics::new(2, &ExecutionPlanMetricsSet::new()); + metrics.record_output(&batch(3)); + metrics.record_output(&batch(4)); + recorder.record_plan(vec![Metrics { + profile: Some(ProfileNode { + name: "FilterExec".to_string(), + ..Default::default() + }), + ..Default::default() + }]); + recorder.finish( + &metrics, + ScanStats { + files: 5, + ..Default::default() + }, + ); + assert!(profiles.lock().is_empty()); + drop(recorder); + + let p = profiles.lock().get(&2).cloned().unwrap(); + assert_eq!(p.node_name, "q-0"); + assert_eq!(p.role, "querier"); + assert_eq!(p.output_rows, 7); + assert_eq!(p.output_batches, 2); + assert!(p.end_us >= p.start_us); + assert_eq!(p.scan_stats.unwrap().files, 5); + assert_eq!(p.plan.unwrap().name, "FilterExec"); + } + + #[test] + fn test_recorder_publishes_skipped_and_unfinished_nodes() { + let profiles = NodeProfiles::default(); + NodeProfileRecorder::start(profiles.clone(), 0, &TestNode).mark_skipped(); + drop(NodeProfileRecorder::start(profiles.clone(), 1, &TestNode)); + NodeProfileRecorder::default().mark_skipped(); + let profiles = profiles.lock(); + assert_eq!(profiles.len(), 2); + assert!(profiles[&0].skipped); + assert!(!profiles[&1].skipped && profiles[&1].scan_stats.is_none()); + } + + #[test] + fn test_convert_metrics_sums_partitions() { + let set = ExecutionPlanMetricsSet::new(); + MetricBuilder::new(&set).output_rows(0).add(10); + MetricBuilder::new(&set).output_rows(1).add(30); + MetricBuilder::new(&set) + .elapsed_compute(0) + .add_duration(std::time::Duration::from_nanos(5)); + let metrics = convert_metrics(&set.clone_inner()); + let rows = metrics.iter().find(|m| m.name == "output_rows").unwrap(); + assert_eq!(rows.kind, MetricKind::Count); + assert_eq!(rows.value, 40); + let time = metrics + .iter() + .find(|m| m.name == "elapsed_compute") + .unwrap(); + assert_eq!(time.kind, MetricKind::Time); + assert_eq!(time.value, 5); + assert_eq!(time.display, "5ns"); + } + + #[tokio::test] + async fn test_profile_text_matches_datafusion() { + let ctx = SessionContext::new(); + let sql = "SELECT a % 7 AS k, count(*) FROM (SELECT unnest(range(1, 50)) AS a) \ + WHERE a > 10 GROUP BY k ORDER BY k"; + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + collect(plan.clone(), ctx.task_ctx()).await.unwrap(); + let expected = DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .indent(false) + .to_string(); + let profile = build_profile(plan.as_ref()); + assert_eq!(profile.to_text(), expected); + assert!(expected.contains("selectivity="), "{expected}"); + } + + #[test] + fn test_memory_metrics_are_bytes() { + let set = ExecutionPlanMetricsSet::new(); + MetricBuilder::new(&set).mem_used(0).set(2048); + let metrics = convert_metrics(&set.clone_inner()); + assert_eq!(metrics[0].kind, MetricKind::Bytes); + assert_eq!(metrics[0].value, 2048); + assert_eq!(metrics[0].display, "2.0 KB"); + } + + #[test] + fn test_build_profile_walks_children() { + let plan = UnionExec::try_new(vec![empty(), empty()]).unwrap(); + let profile = build_profile(plan.as_ref()); + assert_eq!(profile.name, "UnionExec"); + let ProfileInputs::Children(children) = profile.inputs else { + panic!("expected children"); + }; + assert_eq!(children.len(), 2); + assert_eq!(children[0].name, "EmptyExec"); + } + + #[test] + fn test_build_profile_stops_at_remote() { + let node = RemoteScanNode { + nodes: vec![Arc::new(TestNode)], + opentelemetry_context: opentelemetry::Context::current(), + query_identifier: Default::default(), + search_infos: Default::default(), + index_info: Default::default(), + super_cluster_info: Default::default(), + }; + let remote = RemoteScanExec::new(empty(), node).unwrap(); + drop(NodeProfileRecorder::start( + remote.state.node_profiles.clone(), + 0, + &TestNode, + )); + let profile = build_profile(&remote); + assert_eq!(profile.name, "RemoteScanExec"); + let ProfileInputs::Remote(remote) = profile.inputs else { + panic!("expected remote inputs"); + }; + assert_eq!(remote[0].node_name, "q-0"); + } +} diff --git a/src/search/src/datafusion/distributed_plan/remote_scan_exec.rs b/src/search/src/datafusion/distributed_plan/remote_scan_exec.rs index 1ca5a00db9..6a014af2a9 100644 --- a/src/search/src/datafusion/distributed_plan/remote_scan_exec.rs +++ b/src/search/src/datafusion/distributed_plan/remote_scan_exec.rs @@ -13,14 +13,11 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::sync::{Arc, atomic::AtomicUsize}; +use std::sync::Arc; use arrow_schema::SchemaRef; use config::{ - meta::{ - inverted_index::IndexOptimizeMode, - search::{ScanStats, SearchEventType}, - }, + meta::{inverted_index::IndexOptimizeMode, search::SearchEventType}, utils::rand::generate_random_string, }; use datafusion::{ @@ -35,14 +32,13 @@ use datafusion::{ }, }; use datafusion_proto::bytes::physical_plan_to_bytes_with_extension_codec; -use flight::common::{Metrics, RemoteScanMetrics}; +use flight::common::{RemoteScanMetrics, RemoteScanState}; use futures::{StreamExt, TryStreamExt}; use futures_util::pin_mut; use opentelemetry::context::FutureExt as _; -use parking_lot::Mutex; use rand::prelude::SliceRandom; -use super::node::RemoteScanNode; +use super::{node::RemoteScanNode, profile::NodeProfileRecorder}; use crate::datafusion::{ distributed_plan::{ codec::get_physical_extension_codec, @@ -60,10 +56,7 @@ pub struct RemoteScanExec { remote_scan_node: RemoteScanNode, partitions: usize, cache: Arc, - pub scan_stats: Arc>, - pub peak_memory: Arc, - pub partial_err: Arc>, - pub cluster_metrics: Arc>>, + pub state: RemoteScanState, pub enrich_mode_node_idx: usize, pub metrics: ExecutionPlanMetricsSet, } @@ -101,53 +94,19 @@ impl RemoteScanExec { remote_scan_node, partitions: output_partitions, cache, - scan_stats: Arc::new(Mutex::new(ScanStats::default())), - partial_err: Arc::new(Mutex::new(String::new())), - cluster_metrics: Arc::new(Mutex::new(Vec::new())), - peak_memory: Arc::new(AtomicUsize::new(0)), + state: RemoteScanState::default(), enrich_mode_node_idx, metrics: ExecutionPlanMetricsSet::new(), }) } - pub fn scan_stats(&self) -> Arc> { - self.scan_stats.clone() - } - - pub fn partial_err(&self) -> Arc> { - self.partial_err.clone() - } - - pub fn cluster_metrics(&self) -> Arc>> { - self.cluster_metrics.clone() - } - - pub fn peak_memory(&self) -> Arc { - self.peak_memory.clone() - } - /// The input plan bytes exactly as every follower receives them. pub fn encoded_plan(&self) -> &[u8] { &self.remote_scan_node.search_infos.plan } - pub fn with_scan_stats(mut self, scan_stats: Arc>) -> Self { - self.scan_stats = scan_stats; - self - } - - pub fn with_partial_err(mut self, partial_err: Arc>) -> Self { - self.partial_err = partial_err; - self - } - - pub fn with_cluster_metrics(mut self, cluster_metrics: Arc>>) -> Self { - self.cluster_metrics = cluster_metrics; - self - } - - pub fn with_peak_memory(mut self, peak_memory: Arc) -> Self { - self.peak_memory = peak_memory; + pub fn with_state(mut self, state: RemoteScanState) -> Self { + self.state = state; self } @@ -254,10 +213,7 @@ impl ExecutionPlan for RemoteScanExec { return Ok(self); } let remote_scan = Self::new(children[0].clone(), self.remote_scan_node.clone())? - .with_scan_stats(self.scan_stats.clone()) - .with_partial_err(self.partial_err.clone()) - .with_cluster_metrics(self.cluster_metrics.clone()) - .with_peak_memory(self.peak_memory.clone()) + .with_state(self.state.clone()) .with_metrics(self.metrics.clone()); Ok(Arc::new(remote_scan)) } @@ -273,10 +229,7 @@ impl ExecutionPlan for RemoteScanExec { partition, partition == self.enrich_mode_node_idx, self.input.schema().clone(), - self.scan_stats(), - self.partial_err(), - self.cluster_metrics(), - self.peak_memory(), + self.state.clone(), baseline_metrics, ); let stream = futures::stream::once(fut).try_flatten(); @@ -297,12 +250,15 @@ async fn get_remote_batch( partition: usize, enrich_mode: bool, schema: SchemaRef, - scan_stats: Arc>, - partial_err: Arc>, - cluster_metrics: Arc>>, - peak_memory: Arc, + state: RemoteScanState, metrics: RemoteScanMetrics, ) -> Result { + let RemoteScanState { + scan_stats, + peak_memory, + partial_err, + node_profiles, + } = state; let start = std::time::Instant::now(); let cfg = config::get_config(); let trace_id = remote_scan_node.query_identifier.trace_id.clone(); @@ -335,6 +291,12 @@ async fn get_remote_batch( } } + let mut recorder = if remote_scan_node.search_infos.is_analyze { + NodeProfileRecorder::start(node_profiles, partition, node.as_ref()) + } else { + NodeProfileRecorder::default() + }; + let empty_stream = EmptyStream::new( &trace_id, schema.clone(), @@ -351,6 +313,7 @@ async fn get_remote_batch( && !enrich_mode && remote_scan_node.is_file_list_empty(partition) { + recorder.mark_skipped(); return Ok(get_empty_stream(empty_stream)); } @@ -412,9 +375,9 @@ async fn get_remote_batch( .with_is_querier(is_querier) .with_scan_stats(scan_stats) .with_partial_err(partial_err.clone()) - .with_cluster_metrics(cluster_metrics) .with_peak_memory(peak_memory) - .with_start_time(start); + .with_start_time(start) + .with_node_profile(recorder); let mut stream = FlightDecoderStream::new(stream, schema.clone(), metrics, query_context); let stream = async_stream::stream! { @@ -529,7 +492,7 @@ mod tests { #[test] fn test_remote_scan_exec_partial_err_empty() { let exec = make_remote_exec(); - assert!(exec.partial_err().lock().is_empty()); + assert!(exec.state.partial_err.lock().is_empty()); } #[test] @@ -541,11 +504,11 @@ mod tests { } #[test] - fn test_remote_scan_exec_with_partial_err() { - let exec = make_remote_exec(); - let new_err = Arc::new(parking_lot::Mutex::new("pre-set".to_string())); - let exec = exec.with_partial_err(new_err.clone()); - assert_eq!(*exec.partial_err().lock(), "pre-set"); + fn test_remote_scan_exec_with_state() { + let state = RemoteScanState::default(); + *state.partial_err.lock() = "pre-set".to_string(); + let exec = make_remote_exec().with_state(state); + assert_eq!(*exec.state.partial_err.lock(), "pre-set"); } #[test] diff --git a/src/search/src/datafusion/optimizer/mod.rs b/src/search/src/datafusion/optimizer/mod.rs index d1190adba5..10dd8109e9 100644 --- a/src/search/src/datafusion/optimizer/mod.rs +++ b/src/search/src/datafusion/optimizer/mod.rs @@ -297,6 +297,7 @@ pub fn generate_physical_optimizer_rules( pub async fn create_physical_plan( ctx: &SessionContext, sql: &str, + analyze_profile: bool, ) -> Result> { let plan = ctx.state().create_logical_plan(sql).await?; // Defense in depth: reject DDL/DML/COPY plans so search SQL cannot write files on the node. @@ -306,7 +307,7 @@ pub async fn create_physical_plan( .with_allow_statements(false) .verify_plan(&plan)?; let physical_plan = ctx.state().create_physical_plan(&plan).await?; - optimize_distribute_analyze(physical_plan) + optimize_distribute_analyze(physical_plan, analyze_profile) } #[cfg(test)] @@ -542,12 +543,24 @@ mod tests { "INSERT INTO t2 VALUES (1)", ] { assert!( - create_physical_plan(&ctx, sql).await.is_err(), + create_physical_plan(&ctx, sql, false).await.is_err(), "should reject: {sql}" ); } } + #[tokio::test] + async fn test_create_physical_plan_adds_profile_only_when_asked() { + let sql = "EXPLAIN ANALYZE SELECT a FROM t"; + let ctx = SessionContext::new(); + register_tables(&ctx); + let plan = create_physical_plan(&ctx, sql, false).await.unwrap(); + assert_eq!(plan.schema().fields().len(), 4); + + let plan = create_physical_plan(&ctx, sql, true).await.unwrap(); + assert_eq!(plan.schema().field(4).name(), "profile"); + } + #[tokio::test] async fn test_struct_field_access_matches_datafusion() { let mut sql = test_sql(); diff --git a/src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs b/src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs index 4f316888e9..4db2a1b3e7 100644 --- a/src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs +++ b/src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs @@ -28,11 +28,15 @@ use crate::datafusion::distributed_plan::{ }; // replace the AnalyzeExec to DistributeAnalyzeExec -pub fn optimize_distribute_analyze(plan: Arc) -> Result> { +pub fn optimize_distribute_analyze( + plan: Arc, + profile: bool, +) -> Result> { if let Some(analyze) = plan.downcast_ref::() { let distribute_analyze = Arc::new(DistributeAnalyzeExec::new( analyze.verbose(), analyze.show_statistics(), + profile, analyze.input().clone(), )) as Arc; let mut rewriter = AnalyzeRewrite::new(); @@ -78,9 +82,8 @@ mod tests { ])) } - #[test] - fn test_optimize_distribute_analyze() { - let plan = Arc::new( + fn analyze_plan() -> Arc { + Arc::new( AnalyzeExec::builder( false, false, @@ -89,9 +92,20 @@ mod tests { ) .with_metric_types(vec![]) .build(), - ); - let optimized_plan = optimize_distribute_analyze(plan).unwrap(); + ) + } + + #[test] + fn test_optimize_distribute_analyze() { + let optimized_plan = optimize_distribute_analyze(analyze_plan(), false).unwrap(); assert_eq!(optimized_plan.name(), "DistributeAnalyzeExec"); + assert_eq!(optimized_plan.schema().fields().len(), 4); + } + + #[test] + fn test_optimize_distribute_analyze_with_profile() { + let optimized_plan = optimize_distribute_analyze(analyze_plan(), true).unwrap(); + assert_eq!(optimized_plan.schema().fields().len(), 5); } #[test] diff --git a/src/search/src/datafusion/plan_metrics.rs b/src/search/src/datafusion/plan_metrics.rs index a81c74daad..f6600a4c4e 100644 --- a/src/search/src/datafusion/plan_metrics.rs +++ b/src/search/src/datafusion/plan_metrics.rs @@ -15,41 +15,9 @@ use std::sync::{Arc, atomic::AtomicUsize}; -use datafusion::{ - common::{ - Result, - tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}, - }, - physical_plan::ExecutionPlan, - prelude::SessionContext, -}; -use flight::common::Metrics; -use parking_lot::Mutex; +use datafusion::prelude::SessionContext; -use super::{distributed_plan::remote_scan_exec::RemoteScanExec, peak_memory_pool::PeakMemoryPool}; - -pub fn get_cluster_metrics(plan: &Arc) -> Vec>>> { - let mut visitor = MetricsVisitor::default(); - let _ = plan.visit(&mut visitor); - visitor.metrics -} - -#[derive(Default)] -struct MetricsVisitor { - metrics: Vec>>>, -} - -impl<'n> TreeNodeVisitor<'n> for MetricsVisitor { - type Node = Arc; - - fn f_up(&mut self, node: &'n Self::Node) -> Result { - if node.name() == "RemoteScanExec" { - let remote_scan_exec = node.downcast_ref::().unwrap(); - self.metrics.push(remote_scan_exec.cluster_metrics()); - } - Ok(TreeNodeRecursion::Continue) - } -} +use super::peak_memory_pool::PeakMemoryPool; /// Get the peak memory from the SessionContext's memory pool. /// This should be called after the query execution to get the peak memory usage. diff --git a/src/search/src/utils.rs b/src/search/src/utils.rs index 367957170e..c5556658e3 100644 --- a/src/search/src/utils.rs +++ b/src/search/src/utils.rs @@ -87,17 +87,17 @@ impl ExecutionPlanVisitor for ScanStatsVisitor { let mayby_remote_scan_exec = plan.downcast_ref::(); if let Some(remote_scan_exec) = mayby_remote_scan_exec { { - let guard = remote_scan_exec.scan_stats.lock(); + let guard = remote_scan_exec.state.scan_stats.lock(); let stats = *guard; self.scan_stats.add(&stats); } { - let guard = remote_scan_exec.partial_err.lock(); + let guard = remote_scan_exec.state.partial_err.lock(); let err = (*guard).clone(); self.partial_err.push_str(&err); } { - let peak_memory = remote_scan_exec.peak_memory().load(Ordering::Relaxed); + let peak_memory = remote_scan_exec.state.peak_memory.load(Ordering::Relaxed); self.peak_memory = peak_memory.max(self.peak_memory); } } diff --git a/src/search_service/src/cluster/flight.rs b/src/search_service/src/cluster/flight.rs index d285b6d0c4..ba7c447513 100644 --- a/src/search_service/src/cluster/flight.rs +++ b/src/search_service/src/cluster/flight.rs @@ -397,10 +397,12 @@ pub async fn run_datafusion( register_table(&ctx, &sql).await?; // create physical plan - let physical_plan = create_physical_plan(&ctx, &sql.sql).await.map_err(|e| { - log::error!("[trace_id {trace_id}] flight->search: create physical plan error: {e}"); - e - })?; + let physical_plan = create_physical_plan(&ctx, &sql.sql, req.analyze_profile) + .await + .map_err(|e| { + log::error!("[trace_id {trace_id}] flight->search: create physical plan error: {e}"); + e + })?; if cfg.common.print_key_sql { log::info!("[trace_id {trace_id}] leader physical plan"); diff --git a/src/search_service/src/lib.rs b/src/search_service/src/lib.rs index d6fd7926a8..e1bb5da309 100644 --- a/src/search_service/src/lib.rs +++ b/src/search_service/src/lib.rs @@ -163,6 +163,7 @@ pub async fn search( request.set_local_mode(Some(v)); } request.set_use_cache(in_req.use_cache); + request.set_analyze_profile(in_req.query.analyze_profile); // Before the query is planned again in spawned tasks (cluster::http). request.mark_chat_history_reader(); let meta = Sql::new_from_req(&request, &query).await?; diff --git a/src/search_service/src/super_cluster/leader.rs b/src/search_service/src/super_cluster/leader.rs index f8d126ffe3..8e57a0985c 100644 --- a/src/search_service/src/super_cluster/leader.rs +++ b/src/search_service/src/super_cluster/leader.rs @@ -243,7 +243,7 @@ async fn run_datafusion( register_table(&ctx, &sql).await?; // create physical plan - let physical_plan = create_physical_plan(&ctx, &sql.sql).await?; + let physical_plan = create_physical_plan(&ctx, &sql.sql, req.analyze_profile).await?; if cfg.common.print_key_sql { log::info!("[trace_id {trace_id}] super cluster leader: physical plan");