fix: collected profiles lost in cluster mode (#18680)

* fix: collected profiles lost in cluster mode

* revert

* fix: collected profiles lost in cluster mode

* fix: collected profiles lost in cluster mode

* fix: collected pruned partitions lost in cluster mode

* test: add test for explain analyze in both standalone and cluster modes

* fixup

* fixup

* fixup

* fixup

* fixup

* fix test

* chore: bump sqllogictest to 0.38.4

Author: dqhl76

Cargo.lock

@@ -585,7 +585,7 @@ tracing-subscriber = { version = "0.3.17", features = ["env-filter", "json", "va
 
 # Databend Integration Test
 quickcheck = "1.0"
-sqllogictest = "0.28.0"
+sqllogictest = "0.28.4"
 
 [workspace.lints.rust]
 async_fn_in_trait = "allow"

Cargo.toml

@@ -466,6 +466,9 @@ pub trait TableContext: Send + Sync {
     fn get_running_query_execution_stats(&self) -> Vec<(String, ExecutorStatsSnapshot)> {
         unimplemented!()
     }
+    fn merge_pruned_partitions_stats(&self, _other: &HashMap<u32, PartStatistics>) {
+        unimplemented!()
+    }
 }
 
 pub type AbortChecker = Arc<dyn CheckAbort + Send + Sync>;

src/query/catalog/src/table_context.rs

@@ -539,15 +539,11 @@ impl ExplainInterpreter {
 
                 let executor = PipelineCompleteExecutor::from_pipelines(pipelines, settings)?;
                 executor.execute()?;
-                self.ctx
-                    .add_query_profiles(&executor.get_inner().fetch_profiling(false));
             }
             false => {
                 let mut executor = PipelinePullingExecutor::from_pipelines(build_res, settings)?;
                 executor.start();
                 while (executor.pull_data()?).is_some() {}
-                self.ctx
-                    .add_query_profiles(&executor.get_inner().fetch_profiling(false));
             }
         }
         Ok(self

src/query/service/src/interpreters/interpreter_explain.rs

@@ -24,6 +24,7 @@ use std::time::Duration;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::FlightData;
 use async_channel::Receiver;
+use databend_common_base::base::tokio::sync::oneshot;
 use databend_common_base::base::GlobalInstance;
 use databend_common_base::runtime::ExecutorStatsSnapshot;
 use databend_common_base::runtime::GlobalIORuntime;
@@ -35,6 +36,7 @@ use databend_common_config::GlobalConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_grpc::ConnectionFactory;
+use databend_common_pipeline_core::always_callback;
 use databend_common_pipeline_core::basic_callback;
 use databend_common_pipeline_core::ExecutionInfo;
 use fastrace::prelude::*;
@@ -941,6 +943,15 @@ impl QueryCoordinator {
             }
         }
 
+        let (finished_profiling_tx, finished_profiling_rx) = oneshot::channel();
+        if let Some(p) = pipelines.first_mut() {
+            p.set_on_finished(always_callback(move |info: &ExecutionInfo| {
+                let profiling = info.profiling.clone();
+                let _ = finished_profiling_tx.send(profiling);
+                Ok(())
+            }));
+        };
+
         let settings = ExecutorSettings::try_create(info.query_ctx.clone())?;
         let executor = PipelineCompleteExecutor::from_pipelines(pipelines, settings)?;
 
@@ -966,6 +977,7 @@ impl QueryCoordinator {
             request_server_exchange,
             executor.get_inner(),
             perf_guard,
+            finished_profiling_rx,
         );
 
         let span = if let Some(parent) = SpanContext::current_local_parent() {

src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs

@@ -130,6 +130,7 @@ impl BlockMetaTransform<ExchangeDeserializeMeta> for TransformExchangeDeserializ
             DataPacket::DataCacheMetrics(_) => unreachable!(),
             DataPacket::FragmentData(v) => Ok(vec![self.recv_data(meta.packet, v)?]),
             DataPacket::QueryPerf(_) => unreachable!(),
+            DataPacket::PartStatistics(_) => unreachable!(),
         }
     }
 }

src/query/service/src/servers/flight/v1/exchange/serde/exchange_deserializer.rs

@@ -187,6 +187,10 @@ impl StatisticsReceiver {
                 ctx.set_nodes_perf(source_target.to_string(), perf);
                 Ok(false)
             }
+            Ok(Some(DataPacket::PartStatistics(stat))) => {
+                ctx.merge_pruned_partitions_stats(&stat);
+                Ok(false)
+            }
         }
     }
 

src/query/service/src/servers/flight/v1/exchange/statistics_receiver.rs

@@ -12,10 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
 
 use async_channel::Sender;
+use databend_common_base::base::tokio::sync::oneshot;
 use databend_common_base::base::tokio::time::sleep;
 use databend_common_base::runtime::MemStat;
 use databend_common_base::runtime::QueryPerf;
@@ -25,6 +27,7 @@ use databend_common_base::JoinHandle;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_pipeline_core::PlanProfile;
 use databend_common_storage::MutationStatus;
 use futures_util::future::Either;
 use log::warn;
@@ -49,6 +52,7 @@ impl StatisticsSender {
         exchange: FlightExchange,
         executor: Arc<PipelineExecutor>,
         perf_guard: Option<QueryPerfGuard>,
+        profile_rx: oneshot::Receiver<HashMap<u32, PlanProfile>>,
     ) -> Self {
         let spawner = ctx.clone();
         let tx = exchange.convert_to_sender();
@@ -106,8 +110,8 @@ impl StatisticsSender {
                     }
                 }
 
-                if let Err(error) = Self::send_profile(&executor, &tx, true).await {
-                    warn!("Profiles send has error, cause: {:?}.", error);
+                if let Err(error) = Self::send_final_profile(profile_rx, &tx).await {
+                    warn!("Final profiles send has error, cause: {:?}.", error);
                 }
 
                 if let Err(error) = Self::send_copy_status(&ctx, &tx).await {
@@ -125,6 +129,10 @@ impl StatisticsSender {
                 if let Err(error) = Self::send_perf(&perf_guard, &tx).await {
                     warn!("Perf send has error, cause: {:?}.", error);
                 }
+
+                if let Err(error) = Self::send_part_statistics(&ctx, &tx).await {
+                    warn!("PartStatistics send has error, cause: {:?}.", error);
+                }
             }
         });
 
@@ -205,6 +213,7 @@ impl StatisticsSender {
         Ok(())
     }
 
+    #[async_backtrace::framed]
     async fn send_profile(
         executor: &PipelineExecutor,
         tx: &FlightSender,
@@ -220,6 +229,36 @@ impl StatisticsSender {
         Ok(())
     }
 
+    #[async_backtrace::framed]
+    async fn send_part_statistics(ctx: &Arc<QueryContext>, tx: &FlightSender) -> Result<()> {
+        let part_stats = ctx.get_pruned_partitions_stats();
+
+        if !part_stats.is_empty() {
+            let data_packet = DataPacket::PartStatistics(part_stats);
+            tx.send(data_packet).await?;
+        }
+
+        Ok(())
+    }
+
+    #[async_backtrace::framed]
+    async fn send_final_profile(
+        mut rx: oneshot::Receiver<HashMap<u32, PlanProfile>>,
+        tx: &FlightSender,
+    ) -> Result<()> {
+        // The plans_profile comes from the executor's on_finish callback.
+        // We use try_recv() instead of blocking recv() because the execution order
+        // guarantees that on_finish is called before the statistics sender shuts down.
+        if let Ok(plans_profile) = rx.try_recv() {
+            if !plans_profile.is_empty() {
+                let data_packet = DataPacket::QueryProfiles(plans_profile);
+                tx.send(data_packet).await?;
+            }
+        }
+
+        Ok(())
+    }
+
     #[async_backtrace::framed]
     async fn send_scan_cache_metrics(
         ctx: &Arc<QueryContext>,
@@ -272,8 +311,4 @@ impl StatisticsSender {
         }
         progress_info
     }
-
-    // fn fetch_profiling(ctx: &Arc<QueryContext>) -> Result<Vec<PlanProfile>> {
-    //     // ctx.get_exchange_manager()
-    // }
 }

src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs

@@ -23,6 +23,7 @@ use byteorder::BigEndian;
 use byteorder::ReadBytesExt;
 use byteorder::WriteBytesExt;
 use bytes::Bytes;
+use databend_common_catalog::plan::PartStatistics;
 use databend_common_catalog::statistics::data_cache_statistics::DataCacheMetricValues;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -64,6 +65,7 @@ pub enum DataPacket {
     CopyStatus(CopyStatus),
     MutationStatus(MutationStatus),
     DataCacheMetrics(DataCacheMetricValues),
+    PartStatistics(HashMap<u32, PartStatistics>),
     QueryPerf(String),
 }
 
@@ -83,6 +85,7 @@ impl DataPacket {
             DataPacket::QueryProfiles(_) => 0,
             DataPacket::DataCacheMetrics(_) => 0,
             DataPacket::QueryPerf(_) => 0,
+            DataPacket::PartStatistics(_) => 0,
         }
     }
 }
@@ -149,6 +152,12 @@ impl TryFrom<DataPacket> for FlightData {
                 data_header: Default::default(),
                 flight_descriptor: None,
             },
+            DataPacket::PartStatistics(stat) => FlightData {
+                app_metadata: vec![0x10].into(),
+                data_body: serde_json::to_vec(&stat)?.into(),
+                data_header: Default::default(),
+                flight_descriptor: None,
+            },
         })
     }
 }
@@ -215,6 +224,11 @@ impl TryFrom<FlightData> for DataPacket {
                     .map_err(|_| ErrorCode::BadBytes("Invalid UTF-8 in query performance data."))?;
                 Ok(DataPacket::QueryPerf(query_perf))
             }
+            0x10 => {
+                let stat =
+                    serde_json::from_slice::<HashMap<u32, PartStatistics>>(&flight_data.data_body)?;
+                Ok(DataPacket::PartStatistics(stat))
+            }
             _ => Err(ErrorCode::BadBytes("Unknown flight data packet type.")),
         }
     }

src/query/service/src/servers/flight/v1/packets/packet_data.rs

@@ -2036,6 +2036,10 @@ impl TableContext for QueryContext {
         self.shared.set_pruned_partitions_stats(plan_id, stats);
     }
 
+    fn merge_pruned_partitions_stats(&self, other: &HashMap<u32, PartStatistics>) {
+        self.shared.merge_pruned_partitions_stats(other);
+    }
+
     fn get_next_broadcast_id(&self) -> u32 {
         self.shared
             .next_broadcast_id