fix(query): respect query timezone in timestamp casts (#8859)

* fix(query): respect timezone in insert values Signed-off-by: Morax <james20081204@gmail.com> * fix(query): normalize timestamp casts with query timezone Signed-off-by: Morax <james20081204@gmail.com> * fix(query): scope timestamp conversion to insert assignments Signed-off-by: Morax <james20081204@gmail.com> * style(query): simplify Arc usage Signed-off-by: Morax <james20081204@gmail.com> --------- Signed-off-by: Morax <james20081204@gmail.com> (cherry picked from commit f479ddb38f3bdaf29fcb2e174bdfb8f5a13c8e19) (cherry picked from commit a6c99a5334af982888f83b17e592316161b21daf)
2026-10-13 07:22:32 +00:00 · 2026-08-20 14:45:10 +08:00
parent a024250fdb
commit 00c2d7a1a3
4 changed files with 370 additions and 41 deletions
@@ -12,14 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

+use std::sync::Arc;
+
 use common_time::Timezone;
 use common_time::timestamp::{TimeUnit, Timestamp};
 use datafusion::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
 use datafusion_common::{DFSchemaRef, DataFusionError, Result, ScalarValue};
-use datafusion_expr::expr::InList;
+use datafusion_expr::expr::{Cast, InList};
 use datafusion_expr::{
-    Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan,
+    Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan, WriteOp,
 };
 use datatypes::arrow::compute;
 use datatypes::arrow::datatypes::DataType;
@@ -32,7 +34,7 @@ use crate::plan::ExtractExpr;
 /// TypeConversionRule converts some literal values in logical plan to other types according
 /// to data type of corresponding columns.
 /// Specifically:
-/// - string literal of timestamp is converted to `Expr::Literal(ScalarValue::TimestampMillis)`
+/// - string literal of timestamp is converted to the target timestamp type
 /// - string literal of boolean is converted to `Expr::Literal(ScalarValue::Boolean)`
 pub struct TypeConversionRule;

@@ -45,10 +47,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
    ) -> Result<LogicalPlan> {
        plan.transform_up_with_subqueries(|plan| match plan {
            LogicalPlan::Filter(filter) => {
-                let mut converter = TypeConverter {
-                    schema: filter.input.schema().clone(),
-                    query_ctx: ctx.query_ctx(),
-                };
+                let mut converter =
+                    TypeConverter::new(filter.input.schema().clone(), ctx.query_ctx());
                let rewritten = filter.predicate.clone().rewrite(&mut converter)?.data;
                Ok(Transformed::yes(LogicalPlan::Filter(Filter::try_new(
                    rewritten,
@@ -63,10 +63,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
                filters,
                fetch,
            }) => {
-                let mut converter = TypeConverter {
-                    schema: projected_schema.clone(),
-                    query_ctx: ctx.query_ctx(),
-                };
+                let mut converter = TypeConverter::new(projected_schema.clone(), ctx.query_ctx());
                let rewrite_filters = filters
                    .into_iter()
                    .map(|e| e.rewrite(&mut converter).map(|x| x.data))
@@ -80,8 +77,18 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
                    fetch,
                })))
            }
-            LogicalPlan::Projection { .. }
-            | LogicalPlan::Window { .. }
+            LogicalPlan::Projection { .. } => {
+                let mut converter = TypeConverter::new(plan.schema().clone(), ctx.query_ctx());
+                let inputs = plan.inputs().into_iter().cloned().collect::<Vec<_>>();
+                let expr = plan
+                    .expressions_consider_join()
+                    .into_iter()
+                    .map(|e| e.rewrite(&mut converter).map(|x| x.data))
+                    .collect::<Result<Vec<_>>>()?;
+
+                plan.with_new_exprs(expr, inputs).map(Transformed::yes)
+            }
+            LogicalPlan::Window { .. }
            | LogicalPlan::Aggregate { .. }
            | LogicalPlan::Repartition { .. }
            | LogicalPlan::Extension { .. }
@@ -89,10 +96,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
            | LogicalPlan::Union { .. }
            | LogicalPlan::Values { .. }
            | LogicalPlan::Analyze { .. } => {
-                let mut converter = TypeConverter {
-                    schema: plan.schema().clone(),
-                    query_ctx: ctx.query_ctx(),
-                };
+                let mut converter = TypeConverter::new(plan.schema().clone(), ctx.query_ctx());
                let inputs = plan.inputs().into_iter().cloned().collect::<Vec<_>>();
                let expr = plan
                    .expressions_consider_join()
@@ -107,10 +111,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
                let Ok(schema) = join.left.schema().join(join.right.schema()) else {
                    return Ok(Transformed::no(LogicalPlan::Join(join)));
                };
-                let mut converter = TypeConverter {
-                    schema: std::sync::Arc::new(schema),
-                    query_ctx: ctx.query_ctx(),
-                };
+                let mut converter = TypeConverter::new(Arc::new(schema), ctx.query_ctx());
                let plan = LogicalPlan::Join(join);
                let inputs = plan.inputs().into_iter().cloned().collect::<Vec<_>>();
                let expr = plan
@@ -122,6 +123,14 @@ impl ExtensionAnalyzerRule for TypeConversionRule {
                plan.with_new_exprs(expr, inputs).map(Transformed::yes)
            }

+            LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => {
+                dml.input = Arc::new(rewrite_insert_assignments(
+                    dml.input.as_ref().clone(),
+                    ctx.query_ctx(),
+                )?);
+                Ok(Transformed::yes(LogicalPlan::Dml(dml)))
+            }
+
            LogicalPlan::Distinct { .. }
            | LogicalPlan::Limit { .. }
            | LogicalPlan::Subquery { .. }
@@ -146,6 +155,10 @@ struct TypeConverter {
 }

 impl TypeConverter {
+    fn new(schema: DFSchemaRef, query_ctx: QueryContextRef) -> Self {
+        Self { query_ctx, schema }
+    }
+
    fn column_type(&self, expr: &Expr) -> Option<DataType> {
        if let Expr::Column(_) = expr
            && let Ok(v) = expr.get_type(&self.schema)
@@ -162,7 +175,7 @@ impl TypeConverter {
    ) -> Result<ScalarValue> {
        match (target_type, value) {
            (DataType::Timestamp(_, _), ScalarValue::Utf8(Some(v))) => {
-                string_to_timestamp_ms(v, Some(&self.query_ctx.timezone()))
+                parse_string_to_timestamp(v, Some(&self.query_ctx.timezone()))
            }
            (DataType::Boolean, ScalarValue::Utf8(Some(v))) => match v.to_lowercase().as_str() {
                "true" => Ok(ScalarValue::Boolean(Some(true))),
@@ -294,6 +307,120 @@ impl TreeNodeRewriter for TypeConverter {
    }
 }

+/// Rewrites timestamp assignment casts at the INSERT boundary. The source query
+/// is intentionally left untouched so explicit casts keep DataFusion semantics.
+fn rewrite_insert_assignments(
+    plan: LogicalPlan,
+    query_ctx: QueryContextRef,
+) -> Result<LogicalPlan> {
+    let LogicalPlan::Projection(projection) = plan else {
+        return Ok(plan);
+    };
+
+    // VALUES type coercion is performed inside the Values node. Only rewrite
+    // the direct source of the assignment projection, not nested query plans.
+    let input = rewrite_insert_values(projection.input.as_ref().clone(), query_ctx.clone())?;
+    let input_literals = projected_literals(&input);
+    let mut converter = InsertAssignmentConverter {
+        query_ctx,
+        input_literals,
+    };
+    let plan = LogicalPlan::Projection(projection);
+    let expr = plan
+        .expressions_consider_join()
+        .into_iter()
+        .map(|expr| expr.rewrite(&mut converter).map(|x| x.data))
+        .collect::<Result<Vec<_>>>()?;
+
+    plan.with_new_exprs(expr, vec![input])
+}
+
+fn rewrite_insert_values(plan: LogicalPlan, query_ctx: QueryContextRef) -> Result<LogicalPlan> {
+    let LogicalPlan::Values(mut values) = plan else {
+        return Ok(plan);
+    };
+    let mut converter = InsertAssignmentConverter {
+        query_ctx,
+        input_literals: None,
+    };
+    values.values = values
+        .values
+        .into_iter()
+        .map(|row| {
+            row.into_iter()
+                .map(|expr| expr.rewrite(&mut converter).map(|x| x.data))
+                .collect::<Result<Vec<_>>>()
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    Ok(LogicalPlan::Values(values))
+}
+
+fn projected_literals(plan: &LogicalPlan) -> Option<(DFSchemaRef, Vec<Option<ScalarValue>>)> {
+    let LogicalPlan::Projection(projection) = plan else {
+        return None;
+    };
+    let literals = projection.expr.iter().map(projected_literal).collect();
+    Some((projection.schema.clone(), literals))
+}
+
+fn projected_literal(expr: &Expr) -> Option<ScalarValue> {
+    match expr {
+        Expr::Literal(value, _) => Some(value.clone()),
+        Expr::Alias(alias) => projected_literal(&alias.expr),
+        _ => None,
+    }
+}
+
+struct InsertAssignmentConverter {
+    query_ctx: QueryContextRef,
+    input_literals: Option<(DFSchemaRef, Vec<Option<ScalarValue>>)>,
+}
+
+impl InsertAssignmentConverter {
+    fn assignment_literal(&self, expr: &Expr) -> Option<ScalarValue> {
+        let Expr::Column(column) = expr else {
+            return expr.as_literal().cloned();
+        };
+        let (schema, literals) = self.input_literals.as_ref()?;
+        let index = schema.maybe_index_of_column(column)?;
+        literals.get(index)?.clone()
+    }
+}
+
+impl TreeNodeRewriter for InsertAssignmentConverter {
+    type Node = Expr;
+
+    fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
+        let Expr::Cast(Cast {
+            expr: inner,
+            data_type: target_type,
+        }) = expr
+        else {
+            return Ok(Transformed::no(expr));
+        };
+        let original = || {
+            Expr::Cast(Cast {
+                expr: inner.clone(),
+                data_type: target_type.clone(),
+            })
+        };
+        if !matches!(&target_type, DataType::Timestamp(_, _)) {
+            return Ok(Transformed::no(original()));
+        }
+        let Some(ScalarValue::Utf8(Some(value))) = self.assignment_literal(inner.as_ref()) else {
+            return Ok(Transformed::no(original()));
+        };
+
+        match cast_string_to_timestamp(&value, &target_type, Some(&self.query_ctx.timezone())) {
+            Ok(value) if !value.is_null() => Ok(Transformed::yes(Expr::Literal(value, None))),
+            // Preserve DataFusion's existing parser for unsupported timestamp
+            // syntax and out-of-range values.
+            Ok(_) | Err(_) => Ok(Transformed::no(original())),
+        }
+    }
+}
+
 fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr {
    let timestamp = match unit {
        TimeUnit::Second => val * 1_000,
@@ -308,18 +435,34 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr {
    )
 }

-fn string_to_timestamp_ms(string: &str, timezone: Option<&Timezone>) -> Result<ScalarValue> {
+fn cast_string_to_timestamp(
+    string: &str,
+    target_type: &DataType,
+    timezone: Option<&Timezone>,
+) -> Result<ScalarValue> {
+    let parsed = parse_string_to_timestamp(string, timezone)?;
+    cast_timestamp(parsed, target_type)
+}
+
+fn parse_string_to_timestamp(string: &str, timezone: Option<&Timezone>) -> Result<ScalarValue> {
    let ts = Timestamp::from_str(string, timezone)
        .map_err(|e| DataFusionError::External(Box::new(e)))?;

    let value = Some(ts.value());
-    let scalar = match ts.unit() {
+    Ok(match ts.unit() {
        TimeUnit::Second => ScalarValue::TimestampSecond(value, None),
        TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(value, None),
        TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(value, None),
        TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(value, None),
-    };
-    Ok(scalar)
+    })
+}
+
+fn cast_timestamp(parsed: ScalarValue, target_type: &DataType) -> Result<ScalarValue> {
+    let parsed = parsed.to_array()?;
+    let casted = compute::cast(&parsed, target_type)
+        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
+
+    ScalarValue::try_from_array(&casted, 0)
 }

 #[cfg(test)]
@@ -337,19 +480,23 @@ mod tests {
    use super::*;

    #[test]
-    fn test_string_to_timestamp_ms() {
+    fn test_cast_string_to_timestamp() {
+        use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit;
+
+        let target_type = DataType::Timestamp(ArrowTimeUnit::Second, None);
        assert_eq!(
-            string_to_timestamp_ms("2022-02-02 19:00:00+08:00", None).unwrap(),
+            cast_string_to_timestamp("2022-02-02 19:00:00+08:00", &target_type, None).unwrap(),
            ScalarValue::TimestampSecond(Some(1643799600), None)
        );
        assert_eq!(
-            string_to_timestamp_ms("2009-02-13 23:31:30Z", None).unwrap(),
+            cast_string_to_timestamp("2009-02-13 23:31:30Z", &target_type, None).unwrap(),
            ScalarValue::TimestampSecond(Some(1234567890), None)
        );

        assert_eq!(
-            string_to_timestamp_ms(
+            cast_string_to_timestamp(
                "2009-02-13 23:31:30",
+                &target_type,
                Some(&Timezone::from_tz_string("Asia/Shanghai").unwrap())
            )
            .unwrap(),
@@ -357,13 +504,24 @@ mod tests {
        );

        assert_eq!(
-            string_to_timestamp_ms(
+            cast_string_to_timestamp(
                "2009-02-13 23:31:30",
+                &target_type,
                Some(&Timezone::from_tz_string("-8:00").unwrap())
            )
            .unwrap(),
            ScalarValue::TimestampSecond(Some(1234567890 + 8 * 3600), None)
        );
+
+        assert_eq!(
+            cast_string_to_timestamp(
+                "2009-02-13 23:31:30.123456789Z",
+                &DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
+                None,
+            )
+            .unwrap(),
+            ScalarValue::TimestampNanosecond(Some(1_234_567_890_123_456_789), None)
+        );
    }

    #[test]
@@ -421,14 +579,11 @@ mod tests {
            )
            .unwrap(),
        );
-        let mut converter = TypeConverter {
-            schema,
-            query_ctx: QueryContext::arc(),
-        };
+        let mut converter = TypeConverter::new(schema, QueryContext::arc());

        assert_eq!(
            Expr::Column(Column::from_name("ts")).gt(ScalarValue::TimestampSecond(
-                Some(1599514949),
+                Some(1_599_514_949),
                None
            )
            .lit()),
@@ -439,6 +594,58 @@ mod tests {
        );
    }

+    #[test]
+    fn test_insert_assignment_uses_query_timezone_and_target_precision() {
+        use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit;
+
+        let query_ctx = QueryContext::arc();
+        query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap());
+        let mut converter = InsertAssignmentConverter {
+            query_ctx,
+            input_literals: None,
+        };
+
+        let expr = Expr::Cast(Cast {
+            expr: Box::new("2009-02-13 23:31:30.123456789".lit()),
+            data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
+        });
+        assert_eq!(
+            converter.f_up(expr).unwrap().data,
+            ScalarValue::TimestampNanosecond(Some(1_234_539_090_123_456_789), None).lit()
+        );
+    }
+
+    #[test]
+    fn test_insert_assignment_falls_back_for_unsupported_literal() {
+        use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit;
+
+        let mut converter = InsertAssignmentConverter {
+            query_ctx: QueryContext::arc(),
+            input_literals: None,
+        };
+        for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] {
+            let expr = Expr::Cast(Cast {
+                expr: Box::new(literal.lit()),
+                data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None),
+            });
+
+            assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr);
+        }
+    }
+
+    #[test]
+    fn test_type_converter_leaves_explicit_timestamp_cast_unchanged() {
+        use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit;
+
+        let mut converter = TypeConverter::new(Arc::new(DFSchema::empty()), QueryContext::arc());
+        let expr = Expr::Cast(Cast {
+            expr: Box::new("2009-02-13 23:31:30".lit()),
+            data_type: DataType::Timestamp(ArrowTimeUnit::Millisecond, None),
+        });
+
+        assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr);
+    }
+
    #[test]
    fn test_convert_bool() {
        let col_name = "is_valid";
@@ -452,10 +659,7 @@ mod tests {
            )
            .unwrap(),
        );
-        let mut converter = TypeConverter {
-            schema,
-            query_ctx: QueryContext::arc(),
-        };
+        let mut converter = TypeConverter::new(schema, QueryContext::arc());

        assert_eq!(
            Expr::Column(Column::from_name(col_name)).eq(true.lit()),
@@ -650,9 +650,10 @@ mod tests {
    use catalog::RegisterTableRequest;
    use catalog::memory::MemoryCatalogManager;
    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+    use common_time::Timezone;
    use datatypes::prelude::ConcreteDataType;
    use datatypes::schema::{ColumnSchema, Schema};
-    use session::context::QueryContext;
+    use session::context::{QueryContext, QueryContextBuilder};
    use store_api::metric_engine_consts::{
        DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY,
        METRIC_ENGINE_NAME,
@@ -683,6 +684,36 @@ mod tests {
        crate::tests::new_query_engine_with_table(table)
    }

+    async fn create_timestamp_test_engine() -> QueryEngineRef {
+        let columns = vec![
+            ColumnSchema::new(
+                "ts",
+                ConcreteDataType::timestamp_millisecond_datatype(),
+                false,
+            )
+            .with_time_index(true),
+            ColumnSchema::new(
+                "st",
+                ConcreteDataType::timestamp_millisecond_datatype(),
+                false,
+            ),
+        ];
+        let schema = Arc::new(Schema::new(columns));
+        let table_meta = TableMetaBuilder::empty()
+            .schema(schema)
+            .primary_key_indices(vec![])
+            .value_indices(vec![0, 1])
+            .next_column_id(1024)
+            .build()
+            .unwrap();
+        let table_info = TableInfoBuilder::new("timestamps", table_meta)
+            .build()
+            .unwrap();
+        let table = EmptyTable::from_table_info(&table_info);
+
+        crate::tests::new_query_engine_with_table(table)
+    }
+
    fn create_promql_test_engine() -> QueryEngineRef {
        let catalog_manager = MemoryCatalogManager::with_default_setup();
        let physical_table_name = "phy";
@@ -815,6 +846,43 @@ mod tests {
        engine.planner().plan(&stmt, query_ctx).await.unwrap()
    }

+    #[tokio::test]
+    async fn test_insert_timestamp_literals_use_query_timezone() {
+        let query_ctx = Arc::new(
+            QueryContextBuilder::default()
+                .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
+                .build(),
+        );
+        let engine = create_timestamp_test_engine().await;
+
+        for (sql, expected_timestamp) in [
+            (
+                "INSERT INTO timestamps (ts, st) \
+                 VALUES ('2026-08-02 12:00:00.001', now())",
+                1_785_643_200_001_i64,
+            ),
+            (
+                "INSERT INTO timestamps (ts, st) \
+                 SELECT '2026-08-03 12:00:00.001', now()",
+                1_785_729_600_001_i64,
+            ),
+        ] {
+            let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
+            let plan = engine
+                .planner()
+                .plan(&stmt, query_ctx.clone())
+                .await
+                .unwrap()
+                .display_indent()
+                .to_string();
+
+            assert!(
+                plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")),
+                "{plan}"
+            );
+        }
+    }
+
    #[tokio::test]
    async fn test_extract_placeholder_cast_types_multiple() {
        let plan = parse_sql_to_plan(
@@ -52,6 +52,41 @@ SELECT * FROM test1 ORDER BY j;
 | 4 | 2025-02-01T00:01:01 |
 +---+---------------------+

+CREATE TABLE test3 (ts TIMESTAMP TIME INDEX, st TIMESTAMP, ts_ns TIMESTAMP(9));
+
+Affected Rows: 0
+
+INSERT INTO test3 (ts) VALUES ('2026-08-01 12:00:00.001');
+
+Affected Rows: 1
+
+INSERT INTO test3 (ts, st) VALUES ('2026-08-02 12:00:00.001', now());
+
+Affected Rows: 1
+
+INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now();
+
+Affected Rows: 1
+
+INSERT INTO test3 (ts, st, ts_ns) VALUES (
+    '2026-08-04 12:00:00.001',
+    now(),
+    '2026-08-04 12:00:00.123456789'
+);
+
+Affected Rows: 1
+
+SELECT ts, ts_ns FROM test3 ORDER BY ts;
+
+-------------------------+-------------------------------+
+| ts                      | ts_ns                         |
+-------------------------+-------------------------------+
+| 2026-08-01T04:00:00.001 |                               |
+| 2026-08-02T04:00:00.001 |                               |
+| 2026-08-03T04:00:00.001 |                               |
+| 2026-08-04T04:00:00.001 | 2026-08-04T04:00:00.123456789 |
+-------------------------+-------------------------------+
+
 SET time_zone = 'UTC';

 Affected Rows: 0
@@ -64,3 +99,7 @@ DROP TABLE test2;

 Affected Rows: 0

+DROP TABLE test3;
+
+Affected Rows: 0
+
@@ -16,8 +16,26 @@ SELECT * FROM test2 ORDER BY j;

 SELECT * FROM test1 ORDER BY j;

+CREATE TABLE test3 (ts TIMESTAMP TIME INDEX, st TIMESTAMP, ts_ns TIMESTAMP(9));
+
+INSERT INTO test3 (ts) VALUES ('2026-08-01 12:00:00.001');
+
+INSERT INTO test3 (ts, st) VALUES ('2026-08-02 12:00:00.001', now());
+
+INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now();
+
+INSERT INTO test3 (ts, st, ts_ns) VALUES (
+    '2026-08-04 12:00:00.001',
+    now(),
+    '2026-08-04 12:00:00.123456789'
+);
+
+SELECT ts, ts_ns FROM test3 ORDER BY ts;
+
 SET time_zone = 'UTC';

 DROP TABLE test1;

 DROP TABLE test2;
+
+DROP TABLE test3;