starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-05 00:09:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: opt-in structured profile for distributed EXPLAIN ANALYZE (#16293)

Browse Source
This commit is contained in:
Huaijin
2026-11-04 12:19:36 +00:00
committed by GitHub
parent 08f0d75c68
commit 81ec3d8bec
22 changed files with 1239 additions and 910 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/api/grpc/src/handler/grpc/flight/mod.rs
+22 -65
View File
@@ -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());
}
}
src/api/grpc/src/handler/grpc/flight/partition_encoder.rs
+4 -6
View File
@@ -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,
}
src/api/grpc/src/handler/grpc/flight/stream.rs
+92 -91
View File
@@ -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<PreCustomMessage>,
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::<Result<Vec<FlightData>, 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<FlightData>,
custom_messages: Vec<PreCustomMessage>,
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<FlightData> = 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<i32> = 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<FlightData> = 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<i32> = 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();
src/api/grpc/src/handler/grpc/flight/visitor.rs
+19 -142
View File
@@ -13,163 +13,40 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
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<dyn ExecutionPlan>) -> Option<Arc<Mutex<ScanStats>>> {
let mut visitor = RemoteScanVisitor::new();
let _ = plan.visit(&mut visitor);
visitor.scan_stats
}
struct RemoteScanVisitor {
scan_stats: Option<Arc<Mutex<ScanStats>>>,
}
impl RemoteScanVisitor {
pub fn new() -> Self {
Self { scan_stats: None }
}
}
impl<'n> TreeNodeVisitor<'n> for RemoteScanVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
let name = node.name();
if name == "RemoteScanExec" {
let remote_scan_exec = node.downcast_ref::<RemoteScanExec>().unwrap();
self.scan_stats = Some(remote_scan_exec.scan_stats());
Ok(TreeNodeRecursion::Stop)
} else {
Ok(TreeNodeRecursion::Continue)
}
}
}
pub fn get_partial_err(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<Mutex<String>>> {
let mut visitor = PartialErrVisitor::new();
let _ = plan.visit(&mut visitor);
visitor.partial_err
}
struct PartialErrVisitor {
partial_err: Option<Arc<Mutex<String>>>,
}
impl PartialErrVisitor {
pub fn new() -> Self {
Self { partial_err: None }
}
}
impl<'n> TreeNodeVisitor<'n> for PartialErrVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if node.name() == "RemoteScanExec" {
let remote_scan_exec = node.downcast_ref::<RemoteScanExec>().unwrap();
self.partial_err = Some(remote_scan_exec.partial_err());
Ok(TreeNodeRecursion::Stop)
} else {
Ok(TreeNodeRecursion::Continue)
}
}
}
pub fn get_peak_memory(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<AtomicUsize>> {
let mut visitor = PeakMemoryVisitor::new();
let _ = plan.visit(&mut visitor);
visitor.peak_memory
}
struct PeakMemoryVisitor {
peak_memory: Option<Arc<AtomicUsize>>,
}
impl PeakMemoryVisitor {
pub fn new() -> Self {
Self { peak_memory: None }
}
}
impl<'n> TreeNodeVisitor<'n> for PeakMemoryVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
let name = node.name();
if name == "RemoteScanExec" {
let remote_scan_exec = node.downcast_ref::<RemoteScanExec>().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<dyn ExecutionPlan>) -> NodeRole {
let mut role = NodeRole::Follower;
let _ = plan.apply(|node| {
let Some(remote_scan) = node.downcast_ref::<RemoteScanExec>() 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<dyn datafusion::physical_plan::ExecutionPlan> {
#[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<dyn ExecutionPlan> = Arc::new(EmptyExec::new(schema));
assert!(matches!(node_role(&plan), NodeRole::Follower));
}
}
src/config/src/datafusion/request.rs
+9
View File
@@ -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<FlightSearchRequest> for Request {
histogram_interval: req.search_info.histogram_interval,
// Never carried across nodes.
chat_history_reader: false,
analyze_profile: false,
}
}
}
src/config/src/meta/search.rs
+4
View File
@@ -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<String>,
/// 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,
}
}
}
src/flight/src/common.rs
+171 -386
View File
@@ -14,7 +14,7 @@
// along with this program. If not, see <http://www.gnu.org/licenses/>.
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<ProfileNode>,
}
/// 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<Arc<Mutex<ScanStats>>>),
// physical plan, is_super_cluster
Metrics(Option<MetricsInfo>),
// use for super cluster follower leader
MetricsRef(Vec<Arc<Mutex<Vec<Metrics>>>>),
// use for storing memory pool reference to extract peak later
PeakMemoryRef(Option<Arc<AtomicUsize>>),
// sent before first batch (early signal for fast-failing nodes)
PartialErrRefEarly(Option<Arc<Mutex<String>>>),
// sent after all batches (final accumulated value)
PartialErrRef(Option<Arc<Mutex<String>>>),
/// 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<AtomicUsize>,
/// The plan to send back as a profile; set only for EXPLAIN ANALYZE.
pub profile: Option<MetricsInfo>,
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<CustomMessage> {
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<CustomMessage> {
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<CustomMessage> {
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<Mutex<BTreeMap<usize, RemoteNodeProfile>>>;
/// 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<Mutex<ScanStats>>,
pub peak_memory: Arc<AtomicUsize>,
pub partial_err: Arc<Mutex<String>>,
/// Filled only for EXPLAIN ANALYZE.
pub node_profiles: NodeProfiles,
}
pub struct MetricsInfo {
pub plan: Arc<dyn ExecutionPlan>,
pub is_super_cluster: bool,
pub func: Box<dyn Fn() -> 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<CustomMessage> {
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> = 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<Metrics> {
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<dyn ExecutionPlan> {
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<dyn ExecutionPlan>,
}
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<dyn datafusion::physical_expr::PhysicalExpr>,
) -> datafusion::common::Result<
datafusion::common::tree_node::TreeNodeRecursion,
>,
) -> datafusion::common::Result<datafusion::common::tree_node::TreeNodeRecursion> {
Ok(datafusion::common::tree_node::TreeNodeRecursion::Continue)
}
fn name(&self) -> &'static str {
self.name
}
fn properties(&self) -> &Arc<PlanProperties> {
self.inner.properties()
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![&self.inner]
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
) -> Result<Arc<dyn ExecutionPlan>> {
unimplemented!()
}
fn execute(
&self,
_partition: usize,
_context: Arc<TaskContext>,
) -> Result<SendableRecordBatchStream> {
unimplemented!()
}
}
fn create_named_test_execution_plan(
name: &'static str,
inner: Arc<dyn ExecutionPlan>,
) -> Arc<dyn ExecutionPlan> {
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<CustomMessage>) -> Vec<String> {
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);
src/flight/src/lib.rs
+1
View File
@@ -16,3 +16,4 @@
pub mod common;
pub mod decoder;
pub mod encoder;
pub mod profile;
src/flight/src/profile.rs
+251
View File
@@ -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 <http://www.gnu.org/licenses/>.
//! 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<ProfileMetric>,
#[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::<Vec<_>>()
.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<ProfileNode>),
/// Only on a remote scan: one entry per node it called directly.
Remote(Vec<RemoteNodeProfile>),
}
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<u64>,
/// 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<ScanStats>,
/// 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<ProfileNode>,
}
#[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"
);
}
}
src/search/src/datafusion/distributed_plan/common.rs
+9 -8
View File
@@ -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<Mutex<ScanStats>>,
pub partial_err: Arc<Mutex<String>>,
pub cluster_metrics: Arc<Mutex<Vec<Metrics>>>,
pub peak_memory: Arc<AtomicUsize>,
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<Mutex<Vec<Metrics>>>) -> 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
src/search/src/datafusion/distributed_plan/decoder_stream.rs
+5 -2
View File
@@ -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();
src/search/src/datafusion/distributed_plan/distribute_analyze_exec.rs
+191 -91
View File
@@ -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<dyn ExecutionPlan>,
/// 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<dyn ExecutionPlan>) -> Self {
let fields = vec![
pub fn new(
verbose: bool,
show_statistics: bool,
profile: bool,
input: Arc<dyn ExecutionPlan>,
) -> 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<RecordBatch> {
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<ArrayRef> = 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<dyn ExecutionPlan>,
schema: SchemaRef,
) -> Result<RecordBatch> {
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<dyn ExecutionPlan>) -> Vec<Metrics> {
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::<StringArray>()
.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<ProfileNode>| 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")]);
}
}
src/search/src/datafusion/distributed_plan/mod.rs
+1
View File
@@ -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;
src/search/src/datafusion/distributed_plan/profile.rs
+380
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<NodeProfiles>,
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<Metrics>) {
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::<RemoteScanExec>() {
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<ProfileMetric> {
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<u64>) {
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<dyn ExecutionPlan> {
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");
}
}
src/search/src/datafusion/distributed_plan/remote_scan_exec.rs
+32 -69
View File
@@ -13,14 +13,11 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
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<PlanProperties>,
pub scan_stats: Arc<Mutex<ScanStats>>,
pub peak_memory: Arc<AtomicUsize>,
pub partial_err: Arc<Mutex<String>>,
pub cluster_metrics: Arc<Mutex<Vec<Metrics>>>,
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<Mutex<ScanStats>> {
self.scan_stats.clone()
}
pub fn partial_err(&self) -> Arc<Mutex<String>> {
self.partial_err.clone()
}
pub fn cluster_metrics(&self) -> Arc<Mutex<Vec<Metrics>>> {
self.cluster_metrics.clone()
}
pub fn peak_memory(&self) -> Arc<AtomicUsize> {
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<Mutex<ScanStats>>) -> Self {
self.scan_stats = scan_stats;
self
}
pub fn with_partial_err(mut self, partial_err: Arc<Mutex<String>>) -> Self {
self.partial_err = partial_err;
self
}
pub fn with_cluster_metrics(mut self, cluster_metrics: Arc<Mutex<Vec<Metrics>>>) -> Self {
self.cluster_metrics = cluster_metrics;
self
}
pub fn with_peak_memory(mut self, peak_memory: Arc<AtomicUsize>) -> 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<Mutex<ScanStats>>,
partial_err: Arc<Mutex<String>>,
cluster_metrics: Arc<Mutex<Vec<Metrics>>>,
peak_memory: Arc<AtomicUsize>,
state: RemoteScanState,
metrics: RemoteScanMetrics,
) -> Result<SendableRecordBatchStream> {
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]
src/search/src/datafusion/optimizer/mod.rs
+15 -2
View File
@@ -297,6 +297,7 @@ pub fn generate_physical_optimizer_rules(
pub async fn create_physical_plan(
ctx: &SessionContext,
sql: &str,
analyze_profile: bool,
) -> Result<Arc<dyn ExecutionPlan>> {
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();
src/search/src/datafusion/optimizer/physical_optimizer/distribute_analyze.rs
+20 -6
View File
@@ -28,11 +28,15 @@ use crate::datafusion::distributed_plan::{
};
// replace the AnalyzeExec to DistributeAnalyzeExec
pub fn optimize_distribute_analyze(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
pub fn optimize_distribute_analyze(
plan: Arc<dyn ExecutionPlan>,
profile: bool,
) -> Result<Arc<dyn ExecutionPlan>> {
if let Some(analyze) = plan.downcast_ref::<AnalyzeExec>() {
let distribute_analyze = Arc::new(DistributeAnalyzeExec::new(
analyze.verbose(),
analyze.show_statistics(),
profile,
analyze.input().clone(),
)) as Arc<dyn ExecutionPlan>;
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<dyn ExecutionPlan> {
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]
src/search/src/datafusion/plan_metrics.rs
+2 -34
View File
@@ -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<dyn ExecutionPlan>) -> Vec<Arc<Mutex<Vec<Metrics>>>> {
let mut visitor = MetricsVisitor::default();
let _ = plan.visit(&mut visitor);
visitor.metrics
}
#[derive(Default)]
struct MetricsVisitor {
metrics: Vec<Arc<Mutex<Vec<Metrics>>>>,
}
impl<'n> TreeNodeVisitor<'n> for MetricsVisitor {
type Node = Arc<dyn ExecutionPlan>;
fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
if node.name() == "RemoteScanExec" {
let remote_scan_exec = node.downcast_ref::<RemoteScanExec>().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.
src/search/src/utils.rs
+3 -3
View File
@@ -87,17 +87,17 @@ impl ExecutionPlanVisitor for ScanStatsVisitor {
let mayby_remote_scan_exec = plan.downcast_ref::<RemoteScanExec>();
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);
}
}
src/search_service/src/cluster/flight.rs
+6 -4
View File
@@ -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");
src/search_service/src/lib.rs
+1
View File
@@ -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?;
src/search_service/src/super_cluster/leader.rs
+1 -1
View File
@@ -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");