fix: align EXPLAIN output with actual partition pruning statistics (#17491)

* fix: apply pruning pipeline in EXPLAIN to avoid OOM

* fix: fix some stats missing

* fixup

* fixup ut

* fix possible explain hang

* fix possible explain hang

---------

Co-authored-by: Winter Zhang <[email protected]>

Author: dqhl76

src/query/catalog/src/table_context.rs

@@ -71,6 +71,7 @@ use crate::lock::LockTableOption;
 use crate::merge_into_join::MergeIntoJoin;
 use crate::plan::DataSourcePlan;
 use crate::plan::PartInfoPtr;
+use crate::plan::PartStatistics;
 use crate::plan::Partitions;
 use crate::query_kind::QueryKind;
 use crate::runtime_filter_info::RuntimeFilterInfo;
@@ -408,6 +409,13 @@ pub trait TableContext: Send + Sync {
     fn get_consume_streams(&self, _query: bool) -> Result<Vec<Arc<dyn Table>>> {
         unimplemented!()
     }
+
+    fn get_pruned_partitions_stats(&self) -> Option<PartStatistics> {
+        unimplemented!()
+    }
+    fn set_pruned_partitions_stats(&self, _partitions: PartStatistics) {
+        unimplemented!()
+    }
 }
 
 pub type AbortChecker = Arc<dyn CheckAbort + Send + Sync>;

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

@@ -21,6 +21,8 @@ use databend_common_ast::ast::FormatTreeNode;
 use databend_common_base::runtime::profile::get_statistics_desc;
 use databend_common_base::runtime::profile::ProfileDesc;
 use databend_common_base::runtime::profile::ProfileStatisticsName;
+use databend_common_catalog::plan::DataSourcePlan;
+use databend_common_catalog::plan::PartStatistics;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -31,6 +33,7 @@ use databend_common_expression::FromData;
 use databend_common_pipeline_core::always_callback;
 use databend_common_pipeline_core::processors::PlanProfile;
 use databend_common_pipeline_core::ExecutionInfo;
+use databend_common_pipeline_core::Pipeline;
 use databend_common_sql::binder::ExplainConfig;
 use databend_common_sql::executor::format_partial_tree;
 use databend_common_sql::executor::MutationBuildInfo;
@@ -39,6 +42,8 @@ use databend_common_sql::BindContext;
 use databend_common_sql::ColumnSet;
 use databend_common_sql::FormatOptions;
 use databend_common_sql::MetadataRef;
+use databend_common_storages_fuse::FuseLazyPartInfo;
+use databend_common_storages_fuse::FuseTable;
 use databend_common_storages_result_cache::gen_result_cache_key;
 use databend_common_storages_result_cache::ResultCacheReader;
 use databend_common_users::UserApiProvider;
@@ -54,6 +59,7 @@ use crate::interpreters::Interpreter;
 use crate::pipelines::executor::ExecutorSettings;
 use crate::pipelines::executor::PipelineCompleteExecutor;
 use crate::pipelines::executor::PipelinePullingExecutor;
+use crate::pipelines::executor::QueryPipelineExecutor;
 use crate::pipelines::PipelineBuildResult;
 use crate::schedulers::build_query_pipeline;
 use crate::schedulers::Fragmenter;
@@ -152,7 +158,8 @@ impl Interpreter for ExplainInterpreter {
                         schema.clone(),
                         metadata.clone(),
                     )?;
-                    let plan = interpreter.build_physical_plan(&mutation, true).await?;
+                    let mut plan = interpreter.build_physical_plan(&mutation, true).await?;
+                    self.inject_pruned_partitions_stats(&mut plan, metadata)?;
                     self.explain_physical_plan(&plan, metadata, &None).await?
                 }
                 _ => self.explain_plan(&self.plan)?,
@@ -541,7 +548,8 @@ impl ExplainInterpreter {
         // we should not use `dry_run` mode to build the physical plan.
         // It's because we need to get the same partitions as the original selecting plan.
         let mut builder = PhysicalPlanBuilder::new(metadata.clone(), ctx, formatted_ast.is_none());
-        let plan = builder.build(s_expr, bind_context.column_set()).await?;
+        let mut plan = builder.build(s_expr, bind_context.column_set()).await?;
+        self.inject_pruned_partitions_stats(&mut plan, metadata)?;
         self.explain_physical_plan(&plan, metadata, formatted_ast)
             .await
     }
@@ -572,4 +580,54 @@ impl ExplainInterpreter {
         let formatted_plan = StringType::from_data(line_split_result);
         Ok(vec![DataBlock::new_from_columns(vec![formatted_plan])])
     }
+
+    fn inject_pruned_partitions_stats(
+        &self,
+        plan: &mut PhysicalPlan,
+        metadata: &MetadataRef,
+    ) -> Result<()> {
+        let mut sources = vec![];
+        plan.get_all_data_source(&mut sources);
+        for (id, source) in sources {
+            if let Some(stat) = self.prune_lazy_parts(metadata, &source)? {
+                plan.set_pruning_stats(id, stat);
+            }
+        }
+        Ok(())
+    }
+
+    fn prune_lazy_parts(
+        &self,
+        metadata: &MetadataRef,
+        source: &DataSourcePlan,
+    ) -> Result<Option<PartStatistics>> {
+        let partitions = source.parts.partitions.first();
+        if partitions.is_none_or(|part| part.as_any().downcast_ref::<FuseLazyPartInfo>().is_none())
+        {
+            return Ok(None);
+        }
+        let meta = metadata.read();
+        let max_threads = self.ctx.get_settings().get_max_threads()?;
+        let table_entry = meta.table(source.scan_id);
+        if let Some(fuse_table) = table_entry.table().as_any().downcast_ref::<FuseTable>() {
+            let mut dummy_pipeline = Pipeline::create();
+            let prune_pipeline = fuse_table.do_build_prune_pipeline(
+                self.ctx.clone(),
+                source,
+                &mut dummy_pipeline,
+                true,
+            )?;
+            if let Some(mut pipeline) = prune_pipeline {
+                pipeline.set_max_threads(max_threads as usize);
+                let settings = ExecutorSettings::try_create(self.ctx.clone())?;
+                let executor = QueryPipelineExecutor::create(pipeline, settings)?;
+                executor.execute()?;
+            }
+            let stats = self.ctx.get_pruned_partitions_stats();
+            if stats.is_some() {
+                return Ok(stats);
+            }
+        }
+        Ok(None)
+    }
 }

src/query/service/src/sessions/query_ctx.rs

@@ -47,6 +47,7 @@ use databend_common_catalog::plan::DataSourceInfo;
 use databend_common_catalog::plan::DataSourcePlan;
 use databend_common_catalog::plan::ParquetReadOptions;
 use databend_common_catalog::plan::PartInfoPtr;
+use databend_common_catalog::plan::PartStatistics;
 use databend_common_catalog::plan::Partitions;
 use databend_common_catalog::plan::StageTableInfo;
 use databend_common_catalog::query_kind::QueryKind;
@@ -1826,6 +1827,14 @@ impl TableContext for QueryContext {
     async fn get_warehouse_cluster(&self) -> Result<Arc<Cluster>> {
         self.shared.get_warehouse_clusters().await
     }
+
+    fn get_pruned_partitions_stats(&self) -> Option<PartStatistics> {
+        self.shared.get_pruned_partitions_stats()
+    }
+
+    fn set_pruned_partitions_stats(&self, partitions: PartStatistics) {
+        self.shared.set_pruned_partitions_stats(partitions);
+    }
 }
 
 impl TrySpawn for QueryContext {

src/query/service/src/sessions/query_ctx_shared.rs

@@ -33,6 +33,7 @@ use databend_common_base::runtime::Runtime;
 use databend_common_catalog::catalog::Catalog;
 use databend_common_catalog::catalog::CatalogManager;
 use databend_common_catalog::merge_into_join::MergeIntoJoin;
+use databend_common_catalog::plan::PartStatistics;
 use databend_common_catalog::query_kind::QueryKind;
 use databend_common_catalog::runtime_filter_info::RuntimeFilterInfo;
 use databend_common_catalog::runtime_filter_info::RuntimeFilterReady;
@@ -167,6 +168,8 @@ pub struct QueryContextShared {
 
     // Used by hilbert clustering when do recluster.
     pub(in crate::sessions) selected_segment_locs: Arc<RwLock<HashSet<Location>>>,
+
+    pub(in crate::sessions) pruned_partitions_stats: Arc<RwLock<Option<PartStatistics>>>,
 }
 
 impl QueryContextShared {
@@ -232,6 +235,7 @@ impl QueryContextShared {
             mem_stat: Arc::new(RwLock::new(None)),
             node_memory_usage: Arc::new(RwLock::new(HashMap::new())),
             selected_segment_locs: Default::default(),
+            pruned_partitions_stats: Arc::new(RwLock::new(None)),
         }))
     }
 
@@ -765,6 +769,15 @@ impl QueryContextShared {
     pub fn get_table_meta_timestamps(&self) -> Arc<Mutex<HashMap<u64, TableMetaTimestamps>>> {
         self.table_meta_timestamps.clone()
     }
+
+    pub fn get_pruned_partitions_stats(&self) -> Option<PartStatistics> {
+        self.pruned_partitions_stats.read().clone()
+    }
+
+    pub fn set_pruned_partitions_stats(&self, stats: PartStatistics) {
+        let mut guard = self.pruned_partitions_stats.write();
+        *guard = Some(stats);
+    }
 }
 
 impl Drop for QueryContextShared {

src/query/service/tests/it/storages/fuse/operations/internal_column.rs

@@ -12,24 +12,31 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::collections::HashSet;
 
 use databend_common_base::base::tokio;
 use databend_common_catalog::plan::InternalColumn;
 use databend_common_catalog::plan::InternalColumnMeta;
-use databend_common_catalog::plan::Partitions;
+use databend_common_catalog::plan::PartInfoPtr;
 use databend_common_exception::Result;
+use databend_common_expression::block_debug::pretty_format_blocks;
 use databend_common_expression::DataBlock;
+use databend_common_expression::FieldIndex;
 use databend_common_expression::BLOCK_NAME_COL_NAME;
 use databend_common_expression::ROW_ID_COL_NAME;
 use databend_common_expression::SEGMENT_NAME_COL_NAME;
 use databend_common_expression::SNAPSHOT_NAME_COL_NAME;
+use databend_common_pipeline_core::Pipeline;
 use databend_common_sql::binder::INTERNAL_COLUMN_FACTORY;
+use databend_common_sql::executor::table_read_plan::ToReadDataSourcePlan;
 use databend_common_sql::Planner;
 use databend_common_storages_fuse::io::MetaReaders;
 use databend_common_storages_fuse::FuseBlockPartInfo;
 use databend_common_storages_fuse::FuseTable;
 use databend_query::interpreters::InterpreterFactory;
+use databend_query::pipelines::executor::ExecutorSettings;
+use databend_query::pipelines::executor::QueryPipelineExecutor;
 use databend_query::test_kits::*;
 use databend_storages_common_cache::LoadParams;
 use databend_storages_common_table_meta::meta::SegmentInfo;
@@ -39,11 +46,11 @@ use databend_storages_common_table_meta::table::OPT_KEY_SNAPSHOT_LOCATION;
 use futures::TryStreamExt;
 
 fn expected_data_block(
-    parts: &Partitions,
+    parts: &Vec<PartInfoPtr>,
     internal_columns: &Vec<InternalColumn>,
 ) -> Result<Vec<DataBlock>> {
-    let mut data_blocks = Vec::with_capacity(parts.partitions.len());
-    for part in &parts.partitions {
+    let mut data_blocks = Vec::with_capacity(parts.len());
+    for part in parts {
         let fuse_part = FuseBlockPartInfo::from_part(part)?;
         let num_rows = fuse_part.nums_rows;
         let block_meta = fuse_part.block_meta_index.as_ref().unwrap();
@@ -65,28 +72,22 @@ fn expected_data_block(
         }
         data_blocks.push(DataBlock::new(columns, num_rows));
     }
-    data_blocks.reverse();
 
     Ok(data_blocks)
 }
 
 fn check_data_block(expected: Vec<DataBlock>, blocks: Vec<DataBlock>) -> Result<()> {
-    let expected_data_block = DataBlock::concat(&expected)?.consume_convert_to_full();
-    let data_block = DataBlock::concat(&blocks)?.consume_convert_to_full();
-
-    for (expected_column, column) in expected_data_block
-        .columns()
-        .iter()
-        .zip(data_block.columns())
-    {
-        assert_eq!(expected_column.data_type, column.data_type);
-        assert_eq!(expected_column.value, column.value);
-    }
+    let expected_blocks = pretty_format_blocks(&expected)?;
+    let expected_str: Vec<&str> = expected_blocks.split('\n').collect();
+    databend_common_expression::block_debug::assert_blocks_sorted_eq(
+        expected_str,
+        blocks.as_slice(),
+    );
 
     Ok(())
 }
 
-async fn check_partitions(parts: &Partitions, fixture: &TestFixture) -> Result<()> {
+async fn check_partitions(parts: &Vec<PartInfoPtr>, fixture: &TestFixture) -> Result<()> {
     let mut segment_name = HashSet::new();
     let mut block_name = HashSet::new();
 
@@ -129,7 +130,7 @@ async fn check_partitions(parts: &Partitions, fixture: &TestFixture) -> Result<(
         }
     }
 
-    for part in &parts.partitions {
+    for part in parts {
         let fuse_part = FuseBlockPartInfo::from_part(part)?;
         let block_meta = fuse_part.block_meta_index.as_ref().unwrap();
         assert_eq!(
@@ -166,6 +167,11 @@ async fn test_internal_column() -> Result<()> {
             .get_internal_column(BLOCK_NAME_COL_NAME)
             .unwrap(),
     ];
+    let internal_columns_map = internal_columns
+        .iter()
+        .enumerate()
+        .map(|(i, col)| (i, col.clone()))
+        .collect::<BTreeMap<FieldIndex, InternalColumn>>();
 
     // insert 5 times
     let n = 5;
@@ -188,7 +194,34 @@ async fn test_internal_column() -> Result<()> {
     let blocks = res.try_collect::<Vec<DataBlock>>().await?;
 
     let table = fixture.latest_default_table().await?;
-    let (_, parts) = table.read_partitions(ctx.clone(), None, true).await?;
+    let data_source_plan = table
+        .read_plan(
+            ctx.clone(),
+            None,
+            Some(internal_columns_map.clone()),
+            false,
+            false,
+        )
+        .await?;
+
+    let mut dummy_pipeline = Pipeline::create();
+    let parts = if let Some(mut prune_pipeline) =
+        table.build_prune_pipeline(ctx.clone(), &data_source_plan, &mut dummy_pipeline)?
+    {
+        let fuse_table = FuseTable::try_from_table(table.as_ref())?;
+        let rx = fuse_table.pruned_result_receiver.lock().clone().unwrap();
+        prune_pipeline.set_max_threads(1);
+        let settings = ExecutorSettings::try_create(ctx.clone())?;
+        let executor = QueryPipelineExecutor::create(prune_pipeline, settings)?;
+        executor.execute()?;
+        let mut parts = Vec::new();
+        while let Ok(Ok(segment)) = rx.recv().await {
+            parts.push(segment);
+        }
+        parts
+    } else {
+        data_source_plan.parts.partitions.clone()
+    };
     let expected = expected_data_block(&parts, &internal_columns)?;
     check_partitions(&parts, &fixture).await?;
     check_data_block(expected, blocks)?;
@@ -213,7 +246,35 @@ async fn test_internal_column() -> Result<()> {
     let blocks = res.try_collect::<Vec<DataBlock>>().await?;
 
     let table = fixture.latest_default_table().await?;
-    let (_, parts) = table.read_partitions(ctx.clone(), None, true).await?;
+    let data_source_plan = table
+        .read_plan(
+            ctx.clone(),
+            None,
+            Some(internal_columns_map.clone()),
+            false,
+            false,
+        )
+        .await?;
+
+    let mut dummy_pipeline = Pipeline::create();
+    let parts = if let Some(mut prune_pipeline) =
+        table.build_prune_pipeline(ctx.clone(), &data_source_plan, &mut dummy_pipeline)?
+    {
+        let fuse_table = FuseTable::try_from_table(table.as_ref())?;
+        let rx = fuse_table.pruned_result_receiver.lock().clone().unwrap();
+        prune_pipeline.set_max_threads(1);
+        let settings = ExecutorSettings::try_create(ctx.clone())?;
+        let executor = QueryPipelineExecutor::create(prune_pipeline, settings)?;
+        executor.execute()?;
+        let mut parts = Vec::new();
+        while let Ok(Ok(segment)) = rx.recv().await {
+            parts.push(segment);
+        }
+        parts
+    } else {
+        data_source_plan.parts.partitions.clone()
+    };
+
     let expected = expected_data_block(&parts, &internal_columns)?;
     check_partitions(&parts, &fixture).await?;
     check_data_block(expected, blocks)?;

src/query/service/tests/it/storages/fuse/pruning_pipeline.rs

@@ -95,6 +95,8 @@ async fn apply_snapshot_pruning(
         segment_rx,
         res_tx,
         cache_key,
+        segment_locs.len(),
+        false,
     )?;
     prune_pipeline.set_max_threads(1);
     prune_pipeline.set_on_init(move || {