fix: vacuum drop table with limit does not work (#18460)

* fix: vacuum drop table with limit does not work

* add result set

* fix test

* fix test

Author: SkyFan2002

src/meta/api/src/schema_api_impl.rs

@@ -3703,31 +3703,40 @@ async fn get_history_tables_for_gc(
     }
 
     let mut filter_tb_infos = vec![];
+    const BATCH_SIZE: usize = 1000;
 
-    let limited_args = &args[..std::cmp::min(limit, args.len())];
-    let table_id_idents = limited_args.iter().map(|(table_id, _)| table_id.clone());
+    // Process in batches to avoid performance issues
+    for chunk in args.chunks(BATCH_SIZE) {
+        // Get table metadata for current batch
+        let table_id_idents = chunk.iter().map(|(table_id, _)| table_id.clone());
+        let seq_metas = kv_api.get_pb_values_vec(table_id_idents).await?;
 
-    let seq_metas = kv_api.get_pb_values_vec(table_id_idents).await?;
+        // Filter by drop_time_range for current batch
+        for (seq_meta, (table_id, table_name)) in seq_metas.into_iter().zip(chunk.iter()) {
+            let Some(seq_meta) = seq_meta else {
+                error!(
+                    "batch_filter_table_info cannot find {:?} table_meta",
+                    table_id
+                );
+                continue;
+            };
 
-    for (seq_meta, (table_id, table_name)) in seq_metas.into_iter().zip(limited_args.iter()) {
-        let Some(seq_meta) = seq_meta else {
-            error!(
-                "batch_filter_table_info cannot find {:?} table_meta",
-                table_id
-            );
-            continue;
-        };
+            if !drop_time_range.contains(&seq_meta.data.drop_on) {
+                info!("table {:?} is not in drop_time_range", seq_meta.data);
+                continue;
+            }
 
-        if !drop_time_range.contains(&seq_meta.data.drop_on) {
-            info!("table {:?} is not in drop_time_range", seq_meta.data);
-            continue;
-        }
+            filter_tb_infos.push(TableNIV::new(
+                DBIdTableName::new(db_id, table_name.clone()),
+                table_id.clone(),
+                seq_meta,
+            ));
 
-        filter_tb_infos.push(TableNIV::new(
-            DBIdTableName::new(db_id, table_name.clone()),
-            table_id.clone(),
-            seq_meta,
-        ));
+            // Check if we have reached the limit
+            if filter_tb_infos.len() >= limit {
+                return Ok(filter_tb_infos);
+            }
+        }
     }
 
     Ok(filter_tb_infos)

src/query/service/src/interpreters/interpreter_vacuum_drop_tables.rs

@@ -182,6 +182,8 @@ impl Interpreter for VacuumDropTablesInterpreter {
             tables.len()
         );
 
+        let tables_count = tables.len();
+
         let handler = get_vacuum_handler();
         let threads_nums = self.ctx.get_settings().get_max_threads()? as usize;
         let (files_opt, failed_tables) = handler
@@ -230,7 +232,10 @@ impl Interpreter for VacuumDropTablesInterpreter {
         }
 
         match files_opt {
-            None => return Ok(PipelineBuildResult::create()),
+            None => PipelineBuildResult::from_blocks(vec![DataBlock::new_from_columns(vec![
+                UInt64Type::from_data(vec![tables_count as u64 - failed_tables.len() as u64]),
+                UInt64Type::from_data(vec![failed_tables.len() as u64]),
+            ])]),
             Some(purge_files) => {
                 let mut len = min(purge_files.len(), DRY_RUN_LIMIT);
                 if let Some(limit) = self.plan.option.limit {

src/query/sql/src/planner/plans/ddl/table.rs

@@ -168,7 +168,16 @@ impl VacuumDropTablePlan {
                 ]))
             }
         } else {
-            Arc::new(DataSchema::empty())
+            Arc::new(DataSchema::new(vec![
+                DataField::new(
+                    "success_tables_count",
+                    DataType::Number(NumberDataType::UInt64),
+                ),
+                DataField::new(
+                    "failed_tables_count",
+                    DataType::Number(NumberDataType::UInt64),
+                ),
+            ]))
         }
     }
 }

tests/suites/5_ee/01_vacuum/01_003_vacuum_drop_table_continue.result

@@ -12,7 +12,6 @@
 >>>> insert into test_vacuum_drop_table_continue.d values (1)
 1
 >>>> drop database test_vacuum_drop_table_continue
->>>> set data_retention_time_in_days=0; vacuum drop table
 >>>> undrop database test_vacuum_drop_table_continue
 >>>> use test_vacuum_drop_table_continue;show tables
 a

tests/suites/5_ee/01_vacuum/01_003_vacuum_drop_table_continue.sh

@@ -29,7 +29,7 @@ chmod 444 /tmp/test_vacuum_drop_table_continue/
 stmt "drop database test_vacuum_drop_table_continue" 
 
 # can't vacuum files of table a, but can go on vacuum other tables
-stmt "set data_retention_time_in_days=0; vacuum drop table" 
+stmt "set data_retention_time_in_days=0; vacuum drop table" > /dev/null 2>&1
 
 chmod 755 /tmp/test_vacuum_drop_table_continue/
 find /tmp/test_vacuum_drop_table_continue/ -type d -exec chmod 755 {} +

tests/suites/5_ee/04_attach_read_only/02_0006_attach_table_ro_issue_16121.sh

@@ -13,7 +13,7 @@ echo "attach table issue_16121.attach_read_only 's3://testbucket/admin/$storage_
 echo "drop table issue_16121.base;" | $BENDSQL_CLIENT_CONNECT
 
 # purge base table data
-echo "set data_retention_time_in_days=0;vacuum drop table from issue_16121;" | $BENDSQL_CLIENT_CONNECT
+echo "set data_retention_time_in_days=0;vacuum drop table from issue_16121;" | $BENDSQL_CLIENT_CONNECT > /dev/null
 
 echo "expects no error(nothing outputs)"
 echo "drop table issue_16121.attach_read_only" | $BENDSQL_CLIENT_CONNECT

tests/suites/5_ee/04_attach_read_only/04_0001_check_mutations.sh

@@ -25,7 +25,7 @@ echo "VACUUM TABLE test_attach_only.test_json_read_only;" | $BENDSQL_CLIENT_CONN
 echo "vacuum drop table from db should not include the read_only attach table"
 # drop & vacuum
 echo "drop table test_attach_only.test_json_read_only" | $BENDSQL_CLIENT_CONNECT
-echo "vacuum drop table from test_attach_only" | $BENDSQL_CLIENT_CONNECT
+echo "vacuum drop table from test_attach_only" | $BENDSQL_CLIENT_CONNECT > /dev/null
 # attach it back
 echo "attach table test_attach_only.test_json_read_only 's3://testbucket/admin/data/$storage_prefix' connection=(access_key_id ='minioadmin' secret_access_key ='minioadmin' endpoint_url='${STORAGE_S3_ENDPOINT_URL}')" | $BENDSQL_CLIENT_CONNECT
 echo "expect table data still there"