refactor: streamline system.temp_files table scan (#16659)

* refactor: streamline system.temp_files result set

Changed the implementation to return the result set as a stream instead
of a single block, providing better responsiveness and memory usage.

* chore: remove old implementation

* fix: incorrect finish processing

should not stop processing if only step_limit is reached

* refactor: refine using stream chunk

* refactor: extract ListerStreamSourceBuilder

* chore: rename methods

---------

Co-authored-by: Bohu <[email protected]>

Author: dantengsky

src/query/storages/system/src/lib.rs

@@ -66,6 +66,7 @@ mod user_functions_table;
 mod users_table;
 mod util;
 mod virtual_columns_table;
+
 pub use background_jobs_table::BackgroundJobTable;
 pub use background_tasks_table::BackgroundTaskTable;
 pub use backtrace_table::BacktraceTable;

src/query/storages/system/src/temp_files_table.rs

@@ -12,8 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::any::Any;
+use std::future::Future;
 use std::sync::Arc;
 
+use databend_common_catalog::plan::DataSourcePlan;
+use databend_common_catalog::plan::PartStatistics;
+use databend_common_catalog::plan::Partitions;
+use databend_common_catalog::plan::PartitionsShuffleKind;
 use databend_common_catalog::plan::PushDownInfo;
 use databend_common_catalog::table::Table;
 use databend_common_catalog::table_context::TableContext;
@@ -27,99 +33,85 @@ use databend_common_expression::BlockEntry;
 use databend_common_expression::DataBlock;
 use databend_common_expression::FromData;
 use databend_common_expression::Scalar;
+use databend_common_expression::SendableDataBlockStream;
 use databend_common_expression::TableDataType;
 use databend_common_expression::TableField;
 use databend_common_expression::TableSchemaRefExt;
 use databend_common_expression::Value;
 use databend_common_meta_app::schema::TableIdent;
 use databend_common_meta_app::schema::TableInfo;
 use databend_common_meta_app::schema::TableMeta;
+use databend_common_pipeline_core::processors::OutputPort;
+use databend_common_pipeline_core::processors::ProcessorPtr;
 use databend_common_pipeline_core::query_spill_prefix;
+use databend_common_pipeline_core::Pipeline;
+use databend_common_pipeline_sources::EmptySource;
+use databend_common_pipeline_sources::StreamSource;
 use databend_common_storage::DataOperator;
+use futures::stream;
+use futures::stream::Chunks;
+use futures::stream::Take;
 use futures::StreamExt;
-use futures::TryStreamExt;
+use opendal::operator_futures::FutureLister;
+use opendal::Entry;
+use opendal::Lister;
 use opendal::Metakey;
 
-use crate::table::AsyncOneBlockSystemTable;
-use crate::table::AsyncSystemTable;
+use crate::table::SystemTablePart;
 
 pub struct TempFilesTable {
     table_info: TableInfo,
 }
 
 #[async_trait::async_trait]
-impl AsyncSystemTable for TempFilesTable {
-    const NAME: &'static str = "system.temp_files";
+impl Table for TempFilesTable {
+    fn is_local(&self) -> bool {
+        // Follow the practice of `SyncOneBlockSystemTable::is_local`
+        false
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
 
     fn get_table_info(&self) -> &TableInfo {
         &self.table_info
     }
 
     #[async_backtrace::framed]
-    async fn get_full_data(
+    async fn read_partitions(
         &self,
-        ctx: Arc<dyn TableContext>,
-        push_downs: Option<PushDownInfo>,
-    ) -> Result<DataBlock> {
-        let tenant = ctx.get_tenant();
-        let operator = DataOperator::instance().operator();
-
-        let mut temp_files_name: Vec<String> = vec![];
-        let mut temp_files_content_length = vec![];
-        let mut temp_files_last_modified = vec![];
+        _ctx: Arc<dyn TableContext>,
+        _push_downs: Option<PushDownInfo>,
+        _dry_run: bool,
+    ) -> Result<(PartStatistics, Partitions)> {
+        Ok((
+            PartStatistics::default(),
+            Partitions::create(PartitionsShuffleKind::Seq, vec![Arc::new(Box::new(
+                SystemTablePart,
+            ))]),
+        ))
+    }
 
-        let location_prefix = format!("{}/", query_spill_prefix(tenant.tenant_name(), ""));
-        if let Ok(lister) = operator
-            .lister_with(&location_prefix)
-            .recursive(true)
-            .metakey(Metakey::LastModified | Metakey::ContentLength)
-            .await
-        {
-            let limit = push_downs.and_then(|x| x.limit).unwrap_or(usize::MAX);
-            let mut lister = lister.take(limit);
-
-            while let Some(entry) = lister.try_next().await? {
-                let metadata = entry.metadata();
-
-                if metadata.is_file() {
-                    temp_files_name.push(
-                        entry
-                            .path()
-                            .trim_start_matches(&location_prefix)
-                            .to_string(),
-                    );
-
-                    temp_files_last_modified
-                        .push(metadata.last_modified().map(|x| x.timestamp_micros()));
-                    temp_files_content_length.push(metadata.content_length());
-                }
-            }
+    fn read_data(
+        &self,
+        ctx: Arc<dyn TableContext>,
+        plan: &DataSourcePlan,
+        pipeline: &mut Pipeline,
+        _put_cache: bool,
+    ) -> Result<()> {
+        // avoid duplicate read in cluster mode.
+        if plan.parts.partitions.is_empty() {
+            pipeline.add_source(EmptySource::create, 1)?;
+            return Ok(());
         }
 
-        let num_rows = temp_files_name.len();
-        let data_block = DataBlock::new(
-            vec![
-                BlockEntry::new(
-                    DataType::String,
-                    Value::Scalar(Scalar::String("Spill".to_string())),
-                ),
-                BlockEntry::new(
-                    DataType::String,
-                    Value::Column(StringType::from_data(temp_files_name)),
-                ),
-                BlockEntry::new(
-                    DataType::Number(NumberDataType::UInt64),
-                    Value::Column(NumberType::from_data(temp_files_content_length)),
-                ),
-                BlockEntry::new(
-                    DataType::Timestamp.wrap_nullable(),
-                    Value::Column(TimestampType::from_opt_data(temp_files_last_modified)),
-                ),
-            ],
-            num_rows,
-        );
+        pipeline.add_source(
+            |output| TempFilesTable::create_source(ctx.clone(), output, plan.push_downs.clone()),
+            1,
+        )?;
 
-        Ok(data_block.consume_convert_to_full())
+        Ok(())
     }
 }
 
@@ -151,6 +143,184 @@ impl TempFilesTable {
             ..Default::default()
         };
 
-        AsyncOneBlockSystemTable::create(Self { table_info })
+        Arc::new(Self { table_info })
+    }
+
+    pub fn create_source(
+        ctx: Arc<dyn TableContext>,
+        output: Arc<OutputPort>,
+        push_downs: Option<PushDownInfo>,
+    ) -> Result<ProcessorPtr> {
+        let tenant = ctx.get_tenant();
+        let location_prefix = format!("{}/", query_spill_prefix(tenant.tenant_name(), ""));
+        let limit = push_downs.as_ref().and_then(|x| x.limit);
+
+        let operator = DataOperator::instance().operator();
+        let lister = operator
+            .lister_with(&location_prefix)
+            .recursive(true)
+            .metakey(Metakey::LastModified | Metakey::ContentLength);
+
+        let stream = {
+            let prefix = location_prefix.clone();
+            let mut counter = 0;
+            let ctx = ctx.clone();
+            let builder = ListerStreamSourceBuilder::with_lister_fut(lister);
+            builder
+                .limit_opt(limit)
+                .chunk_size(MAX_BATCH_SIZE)
+                .build(move |entries| {
+                    counter += entries.len();
+                    let block = Self::block_from_entries(&prefix, entries)?;
+                    ctx.set_status_info(format!("{} entries processed", counter).as_str());
+                    Ok(block)
+                })?
+        };
+
+        StreamSource::create(ctx, Some(stream), output)
+    }
+
+    fn build_block(
+        names: Vec<String>,
+        file_lens: Vec<u64>,
+        file_last_modifieds: Vec<Option<i64>>,
+    ) -> DataBlock {
+        let row_number = names.len();
+        DataBlock::new(
+            vec![
+                BlockEntry::new(
+                    DataType::String,
+                    Value::Scalar(Scalar::String("Spill".to_string())),
+                ),
+                BlockEntry::new(
+                    DataType::String,
+                    Value::Column(StringType::from_data(names)),
+                ),
+                BlockEntry::new(
+                    DataType::Number(NumberDataType::UInt64),
+                    Value::Column(NumberType::from_data(file_lens)),
+                ),
+                BlockEntry::new(
+                    DataType::Timestamp.wrap_nullable(),
+                    Value::Column(TimestampType::from_opt_data(file_last_modifieds)),
+                ),
+            ],
+            row_number,
+        )
+    }
+
+    fn block_from_entries(location_prefix: &str, entries: Vec<Entry>) -> Result<DataBlock> {
+        let num_items = entries.len();
+        let mut temp_files_name: Vec<String> = Vec::with_capacity(num_items);
+        let mut temp_files_content_length = Vec::with_capacity(num_items);
+        let mut temp_files_last_modified = Vec::with_capacity(num_items);
+        for entry in entries {
+            let metadata = entry.metadata();
+            if metadata.is_file() {
+                temp_files_name.push(entry.path().trim_start_matches(location_prefix).to_string());
+
+                temp_files_last_modified
+                    .push(metadata.last_modified().map(|x| x.timestamp_micros()));
+                temp_files_content_length.push(metadata.content_length());
+            }
+        }
+
+        let data_block = TempFilesTable::build_block(
+            temp_files_name,
+            temp_files_content_length,
+            temp_files_last_modified,
+        );
+        Ok(data_block)
+    }
+}
+
+const MAX_BATCH_SIZE: usize = 1000;
+
+pub struct ListerStreamSourceBuilder<T>
+where T: Future<Output = opendal::Result<Lister>> + Send + 'static
+{
+    lister_fut: FutureLister<T>,
+    limit: Option<usize>,
+    chunk_size: usize,
+}
+
+impl<T> ListerStreamSourceBuilder<T>
+where T: Future<Output = opendal::Result<Lister>> + Send + 'static
+{
+    pub fn with_lister_fut(lister_fut: FutureLister<T>) -> Self {
+        Self {
+            lister_fut,
+            limit: None,
+            chunk_size: MAX_BATCH_SIZE,
+        }
+    }
+
+    pub fn limit(mut self, limit: usize) -> Self {
+        self.limit = Some(limit);
+        self
     }
+
+    pub fn limit_opt(mut self, limit: Option<usize>) -> Self {
+        self.limit = limit;
+        self
+    }
+
+    pub fn chunk_size(mut self, chunk_size: usize) -> Self {
+        self.chunk_size = chunk_size;
+        self
+    }
+
+    pub fn build(
+        self,
+        block_builder: (impl FnMut(Vec<Entry>) -> Result<DataBlock> + Sync + Send + 'static),
+    ) -> Result<SendableDataBlockStream> {
+        stream_source_from_entry_lister_with_chunk_size(
+            self.lister_fut,
+            self.limit,
+            self.chunk_size,
+            block_builder,
+        )
+    }
+}
+
+fn stream_source_from_entry_lister_with_chunk_size<T>(
+    lister_fut: FutureLister<T>,
+    limit: Option<usize>,
+    chunk_size: usize,
+    block_builder: (impl FnMut(Vec<Entry>) -> Result<DataBlock> + Sync + Send + 'static),
+) -> Result<SendableDataBlockStream>
+where
+    T: Future<Output = opendal::Result<Lister>> + Send + 'static,
+{
+    enum ListerState<U: Future<Output = opendal::Result<Lister>> + Send + 'static> {
+        Uninitialized(FutureLister<U>),
+        Initialized(Chunks<Take<Lister>>),
+    }
+
+    let state = ListerState::<T>::Uninitialized(lister_fut);
+
+    let stream = stream::try_unfold(
+        (state, block_builder),
+        move |(mut state, mut builder)| async move {
+            let mut lister = {
+                match state {
+                    ListerState::Uninitialized(fut) => {
+                        let lister = fut.await?;
+                        lister.take(limit.unwrap_or(usize::MAX)).chunks(chunk_size)
+                    }
+                    ListerState::Initialized(l) => l,
+                }
+            };
+            if let Some(entries) = lister.next().await {
+                let entries = entries.into_iter().collect::<opendal::Result<Vec<_>>>()?;
+                let data_block = builder(entries)?;
+                state = ListerState::Initialized(lister);
+                Ok(Some((data_block, (state, builder))))
+            } else {
+                Ok(None)
+            }
+        },
+    );
+
+    Ok(stream.boxed())
 }