From 00c2d7a1a32b32cd3badb253f81cf79dd6445e6f Mon Sep 17 00:00:00 2001 From: Morax <100508620+fzlzjerry@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:04:51 +0000 Subject: [PATCH] fix(query): respect query timezone in timestamp casts (#8859) * fix(query): respect timezone in insert values Signed-off-by: Morax * fix(query): normalize timestamp casts with query timezone Signed-off-by: Morax * fix(query): scope timestamp conversion to insert assignments Signed-off-by: Morax * style(query): simplify Arc usage Signed-off-by: Morax --------- Signed-off-by: Morax (cherry picked from commit f479ddb38f3bdaf29fcb2e174bdfb8f5a13c8e19) (cherry picked from commit a6c99a5334af982888f83b17e592316161b21daf) --- src/query/src/optimizer/type_conversion.rs | 284 +++++++++++++++--- src/query/src/planner.rs | 70 ++++- .../insert/insert_default_timezone.result | 39 +++ .../common/insert/insert_default_timezone.sql | 18 ++ 4 files changed, 370 insertions(+), 41 deletions(-) diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 1702aacce83..366a7d8a31d 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -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 { 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::>(); + let expr = plan + .expressions_consider_join() + .into_iter() + .map(|e| e.rewrite(&mut converter).map(|x| x.data)) + .collect::>>()?; + + 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::>(); 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::>(); 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 { if let Expr::Column(_) = expr && let Ok(v) = expr.get_type(&self.schema) @@ -162,7 +175,7 @@ impl TypeConverter { ) -> Result { 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 { + 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::>>()?; + + plan.with_new_exprs(expr, vec![input]) +} + +fn rewrite_insert_values(plan: LogicalPlan, query_ctx: QueryContextRef) -> Result { + 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::>>() + }) + .collect::>>()?; + + Ok(LogicalPlan::Values(values)) +} + +fn projected_literals(plan: &LogicalPlan) -> Option<(DFSchemaRef, Vec>)> { + 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 { + 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>)>, +} + +impl InsertAssignmentConverter { + fn assignment_literal(&self, expr: &Expr) -> Option { + 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> { + 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 { +fn cast_string_to_timestamp( + string: &str, + target_type: &DataType, + timezone: Option<&Timezone>, +) -> Result { + let parsed = parse_string_to_timestamp(string, timezone)?; + cast_timestamp(parsed, target_type) +} + +fn parse_string_to_timestamp(string: &str, timezone: Option<&Timezone>) -> Result { 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 { + 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()), diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index 44b3a4604dc..c333eb5a69d 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -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( diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index 9b12e4a15d0..da8b1261e7a 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -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 + diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index 2473cdcee49..fba8c134fa0 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -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;