feat: add partial truncate (#6602)

* feat: add partial truncate Signed-off-by: discord9 <discord9@163.com> * fix: per review Signed-off-by: discord9 <discord9@163.com> * feat: add proto partial truncate kind Signed-off-by: discord9 <discord9@163.com> * chore: clippy Signed-off-by: discord9 <discord9@163.com> * chore: update branched proto Signed-off-by: discord9 <discord9@163.com> * feat: grpc support truncate WIP sql support Signed-off-by: discord9 <discord9@163.com> * wip: parse truncate range Signed-off-by: discord9 <discord9@163.com> * feat: truncate by range Signed-off-by: discord9 <discord9@163.com> * fix: truncate range display Signed-off-by: discord9 <discord9@163.com> * chore: resolve todo Signed-off-by: discord9 <discord9@163.com> * refactor: per review Signed-off-by: discord9 <discord9@163.com> * test: more invalid parse Signed-off-by: discord9 <discord9@163.com> * chore: per review Signed-off-by: discord9 <discord9@163.com> * refactor: per review Signed-off-by: discord9 <discord9@163.com> * chore: unused Signed-off-by: discord9 <discord9@163.com> * chore: per review Signed-off-by: discord9 <discord9@163.com> * chore: update branch Signed-off-by: discord9 <discord9@163.com> --------- Signed-off-by: discord9 <discord9@163.com>
2026-05-28 19:00:39 +00:00 · 2025-08-04 18:50:27 +08:00
parent 414101fafa
commit 1afa0afc67
25 changed files with 938 additions and 102 deletions
--- a/src/operator/src/statement.rs
+++ b/src/operator/src/statement.rs
@@ -47,6 +47,7 @@ use common_telemetry::tracing;
 use common_time::range::TimestampRange;
 use common_time::Timestamp;
 use datafusion_expr::LogicalPlan;
+use datatypes::prelude::ConcreteDataType;
 use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef};
 use query::parser::QueryStatement;
 use query::QueryEngineRef;
@@ -73,8 +74,8 @@ use self::set::{
 };
 use crate::error::{
    self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidSqlSnafu, NotSupportedSnafu,
-    PlanStatementSnafu, Result, SchemaNotFoundSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu,
-    UpgradeCatalogManagerRefSnafu,
+    PlanStatementSnafu, Result, SchemaNotFoundSnafu, SqlCommonSnafu, TableMetadataManagerSnafu,
+    TableNotFoundSnafu, UnexpectedSnafu, UpgradeCatalogManagerRefSnafu,
 };
 use crate::insert::InserterRef;
 use crate::statement::copy_database::{COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY};
@@ -306,7 +307,11 @@ impl StatementExecutor {
                        .map_err(BoxedError::new)
                        .context(ExternalSnafu)?;
                let table_name = TableName::new(catalog, schema, table);
-                self.truncate_table(table_name, query_ctx).await
+                let time_ranges = self
+                    .convert_truncate_time_ranges(&table_name, stmt.time_ranges(), &query_ctx)
+                    .await?;
+                self.truncate_table(table_name, time_ranges, query_ctx)
+                    .await
            }
            Statement::CreateDatabase(stmt) => {
                self.create_database(
@@ -530,6 +535,84 @@ impl StatementExecutor {
    pub fn cache_invalidator(&self) -> &CacheInvalidatorRef {
        &self.cache_invalidator
    }
+
+    /// Convert truncate time ranges for the given table from sql values to timestamps
+    ///
+    pub async fn convert_truncate_time_ranges(
+        &self,
+        table_name: &TableName,
+        sql_values_time_range: &[(sqlparser::ast::Value, sqlparser::ast::Value)],
+        query_ctx: &QueryContextRef,
+    ) -> Result<Vec<(Timestamp, Timestamp)>> {
+        if sql_values_time_range.is_empty() {
+            return Ok(vec![]);
+        }
+        let table = self.get_table(&table_name.table_ref()).await?;
+        let info = table.table_info();
+        let time_index_dt = info
+            .meta
+            .schema
+            .timestamp_column()
+            .context(UnexpectedSnafu {
+                violated: "Table must have a timestamp column",
+            })?;
+
+        let time_unit = time_index_dt
+            .data_type
+            .as_timestamp()
+            .with_context(|| UnexpectedSnafu {
+                violated: format!(
+                    "Table {}'s time index column must be a timestamp type, found: {:?}",
+                    table_name, time_index_dt
+                ),
+            })?
+            .unit();
+
+        let mut time_ranges = Vec::with_capacity(sql_values_time_range.len());
+        for (start, end) in sql_values_time_range {
+            let start = common_sql::convert::sql_value_to_value(
+                "range_start",
+                &ConcreteDataType::timestamp_datatype(time_unit),
+                start,
+                Some(&query_ctx.timezone()),
+                None,
+                false,
+            )
+            .context(SqlCommonSnafu)
+            .and_then(|v| {
+                if let datatypes::value::Value::Timestamp(t) = v {
+                    Ok(t)
+                } else {
+                    error::InvalidSqlSnafu {
+                        err_msg: format!("Expected a timestamp value, found {v:?}"),
+                    }
+                    .fail()
+                }
+            })?;
+
+            let end = common_sql::convert::sql_value_to_value(
+                "range_end",
+                &ConcreteDataType::timestamp_datatype(time_unit),
+                end,
+                Some(&query_ctx.timezone()),
+                None,
+                false,
+            )
+            .context(SqlCommonSnafu)
+            .and_then(|v| {
+                if let datatypes::value::Value::Timestamp(t) = v {
+                    Ok(t)
+                } else {
+                    error::InvalidSqlSnafu {
+                        err_msg: format!("Expected a timestamp value, found {v:?}"),
+                    }
+                    .fail()
+                }
+            })?;
+            time_ranges.push((start, end));
+        }
+        Ok(time_ranges)
+    }
 }

 fn to_copy_query_request(stmt: CopyQueryToArgument) -> Result<CopyQueryToRequest> {
--- a/src/operator/src/statement/ddl.rs
+++ b/src/operator/src/statement/ddl.rs
@@ -46,7 +46,7 @@ use common_meta::rpc::ddl::{
 use common_query::Output;
 use common_sql::convert::sql_value_to_value;
 use common_telemetry::{debug, info, tracing, warn};
-use common_time::Timezone;
+use common_time::{Timestamp, Timezone};
 use datafusion_common::tree_node::TreeNodeVisitor;
 use datafusion_expr::LogicalPlan;
 use datatypes::prelude::ConcreteDataType;
@@ -1151,6 +1151,7 @@ impl StatementExecutor {
    pub async fn truncate_table(
        &self,
        table_name: TableName,
+        time_ranges: Vec<(Timestamp, Timestamp)>,
        query_context: QueryContextRef,
    ) -> Result<Output> {
        ensure!(
@@ -1174,7 +1175,7 @@ impl StatementExecutor {
                table_name: table_name.to_string(),
            })?;
        let table_id = table.table_info().table_id();
-        self.truncate_table_procedure(&table_name, table_id, query_context)
+        self.truncate_table_procedure(&table_name, table_id, time_ranges, query_context)
            .await?;

        Ok(Output::new_with_affected_rows(0))
@@ -1490,6 +1491,7 @@ impl StatementExecutor {
        &self,
        table_name: &TableName,
        table_id: TableId,
+        time_ranges: Vec<(Timestamp, Timestamp)>,
        query_context: QueryContextRef,
    ) -> Result<SubmitDdlTaskResponse> {
        let request = SubmitDdlTaskRequest {
@@ -1499,6 +1501,7 @@ impl StatementExecutor {
                table_name.schema_name.to_string(),
                table_name.table_name.to_string(),
                table_id,
+                time_ranges,
            ),
        };