diff --git a/src/common/time/src/timezone.rs b/src/common/time/src/timezone.rs index 41cc1f7842..dbd8dffcd4 100644 --- a/src/common/time/src/timezone.rs +++ b/src/common/time/src/timezone.rs @@ -108,6 +108,15 @@ impl Timezone { } } + /// A named zone merely sitting at +00:00 today is not UTC: it may have been + /// elsewhere at the timestamp being converted. + pub fn is_utc(&self) -> bool { + match self { + Self::Offset(offset) => offset.local_minus_utc() == 0, + Self::Named(tz) => matches!(tz, Tz::UTC), + } + } + /// Returns the number of seconds to add to convert from UTC to the local time. pub fn local_minus_utc(&self) -> i64 { match self { diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index 8827d48ed6..a6b36e6fda 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -17,6 +17,7 @@ pub mod constant_term; pub mod count_nest_aggr; pub mod count_wildcard; pub mod global_limit; +pub(crate) mod insert_assignment; pub(crate) mod json_type_concretize; pub mod parallelize_scan; pub mod pass_distribution; diff --git a/src/query/src/optimizer/insert_assignment.rs b/src/query/src/optimizer/insert_assignment.rs new file mode 100644 index 0000000000..ad77736ce0 --- /dev/null +++ b/src/query/src/optimizer/insert_assignment.rs @@ -0,0 +1,307 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use common_time::Timezone; +use datafusion::config::ConfigOptions; +use datafusion_common::{DFSchemaRef, Result, ScalarValue}; +use datafusion_expr::expr::{Alias, Cast}; +use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Values}; +use datafusion_optimizer::analyzer::AnalyzerRule; +use datafusion_optimizer::analyzer::type_coercion::TypeCoercion; +use datatypes::arrow::datatypes::{DataType, TimeUnit}; +use session::context::QueryContextRef; + +use crate::optimizer::type_conversion::cast_string_to_timestamp; + +/// Interprets strings assigned to timestamp columns at an `INSERT` boundary +/// using the session timezone. `plan` is the assignment projection under a +/// `WriteOp::Insert`. +/// +/// DataFusion plans `INSERT` as a projection casting each source column to its +/// target column type, and that cast reads a naive string as UTC. Arrow does +/// apply a timezone when the cast target carries one, so the assignment is +/// routed through `Timestamp(unit, Some(tz))` and back. Stripping the timezone +/// afterwards is value-preserving — arrow only shifts values in the opposite +/// direction. +/// +/// The source query is left untouched. Reinterpreting a value where it is +/// *produced* would change what the source query means: pushing the conversion +/// below a `UNION`'s `DISTINCT`, for instance, moves the dedup key from the raw +/// strings to parsed instants and silently drops rows. +/// +/// # Why `TypeCoercion` runs here +/// +/// The rewrite reads source types, and those are only settled once a `UNION`'s +/// branch types have been reconciled: before coercion a union carries its loose +/// schema (the first branch's types), so a mixed +/// `SELECT 'string' UNION ALL SELECT CAST(.. AS TIMESTAMP)` still looks like a +/// string. Retargeting that cast would leave `Timestamp(None) -> +/// Timestamp(Some(tz))` behind once coercion retypes the union — the one +/// direction in which arrow shifts the value instead of relabelling it. +/// +/// Coercing here rather than deferring to the analyzer is forced by where an +/// INSERT is still identifiable: `exec_dml_statement` strips the `Dml` node and +/// executes its input, so by the time the analyzer runs, an assignment +/// projection is indistinguishable from any other projection. +/// +/// Explicit casts stay out of this: the SQL layer turns a user's +/// `CAST(x AS TIMESTAMP)` into an `arrow_cast` call, which only becomes an +/// `Expr::Cast` in the optimizer's `SimplifyExpressions`. Assignment casts are +/// therefore the only `Expr::Cast` reaching a timestamp column here. +/// +/// # Reach +/// +/// The emitted cast only carries its timezone where the expression is evaluated +/// on this node. Substrait drops the timezone name when a plan is pushed down — +/// it encodes any zoned timestamp as `PrecisionTimestampTz` and decodes it back +/// as UTC — so a source reading from a table falls back to UTC, the behaviour it +/// had before this rule existed. Sources that never leave this node (literals, +/// `VALUES`, and `UNION`s of them) keep the session timezone, and those are what +/// an INSERT's timestamp assignment is in practice. +pub(crate) fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: &QueryContextRef, + config: &ConfigOptions, +) -> Result { + let Some(timezone) = session_timezone(query_ctx) else { + return Ok(plan); + }; + + let plan = TypeCoercion::new().analyze(plan, config)?; + rewrite_assignment(plan, &timezone) +} + +/// Session timezone, in both forms the rewrite needs. +struct SessionTimezone { + /// Parses literals, matching the plain `INSERT ... VALUES` path. + parsed: Timezone, + /// Names the intermediate arrow cast target. + name: Arc, +} + +fn session_timezone(query_ctx: &QueryContextRef) -> Option { + let parsed = query_ctx.timezone(); + + // A UTC session already gets UTC semantics from the plain assignment cast. + if parsed.is_utc() { + return None; + } + + Some(SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + }) +} + +fn rewrite_assignment(plan: LogicalPlan, timezone: &SessionTimezone) -> Result { + let LogicalPlan::Projection(assignment) = plan else { + return Ok(plan); + }; + + let mut exprs = assignment.expr.clone(); + let mut changed = false; + for expr in &mut exprs { + changed |= retarget_assignment_cast( + expr, + assignment.input.schema(), + Some(assignment.input.as_ref()), + timezone, + )?; + } + + // The planner types `VALUES` against the target table, so the assignment + // cast lands inside the `Values` rows instead of on the projection above. + let mut input = assignment.input.clone(); + if let LogicalPlan::Values(values) = assignment.input.as_ref() + && let Some(rewritten) = rewrite_values(values, timezone)? + { + input = Arc::new(LogicalPlan::Values(rewritten)); + changed = true; + } + + if !changed { + return Ok(LogicalPlan::Projection(assignment)); + } + Projection::try_new(exprs, input).map(LogicalPlan::Projection) +} + +fn rewrite_values(values: &Values, timezone: &SessionTimezone) -> Result> { + let mut rewritten = values.clone(); + let mut changed = false; + for row in &mut rewritten.values { + for expr in row.iter_mut() { + changed |= retarget_assignment_cast(expr, &values.schema, None, timezone)?; + } + } + + Ok(changed.then_some(rewritten)) +} + +/// Reinterprets one assignment cast, returning whether it was rewritten. +/// +/// `source_plan` is the projection's input, used to resolve a literal behind a +/// column reference; `Values` rows carry their expression inline and pass `None`. +fn retarget_assignment_cast( + expr: &mut Expr, + schema: &DFSchemaRef, + source_plan: Option<&LogicalPlan>, + timezone: &SessionTimezone, +) -> Result { + let expr = unalias_mut(expr); + let Expr::Cast(Cast { + expr: source, + data_type: DataType::Timestamp(unit, None), + }) = expr + else { + return Ok(false); + }; + let unit = *unit; + + if !matches!( + source.get_type(schema)?, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + return Ok(false); + } + + // Fold literals with the same parser the plain `INSERT ... VALUES` path + // uses, so a given string means the same thing however it reaches a column. + // The parsers disagree on ambiguous local times: this one resolves them, + // arrow rejects them. + let folded = source_literal(source.as_ref(), source_plan) + .and_then(|literal| convert_literal(&literal, unit, &timezone.parsed)); + if let Some(folded) = folded { + *expr = folded; + return Ok(true); + } + + let source = source.as_ref().clone(); + *expr = Expr::Cast(Cast::new( + Box::new(Expr::Cast(Cast::new( + Box::new(source), + DataType::Timestamp(unit, Some(timezone.name.clone())), + ))), + DataType::Timestamp(unit, None), + )); + Ok(true) +} + +fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option { + match source { + Expr::Literal(value, _) => Some(value.clone()), + Expr::Column(column) => { + let plan = source_plan?; + let index = plan.schema().maybe_index_of_column(column)?; + lineage_literal(plan, index).cloned() + } + _ => None, + } +} + +/// Resolves a literal when every row carries the same value at `output_idx`. +/// +/// Read-only: the literal is folded into the assignment above, so nodes that +/// drop, reorder or deduplicate rows can be traversed — none of them changes +/// the value a surviving row carries, and folding above them leaves their keys +/// on the original strings. +fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { + if output_idx >= plan.schema().fields().len() { + return None; + } + + match plan { + LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => Some(value), + Expr::Column(column) => { + let input_idx = projection.input.schema().maybe_index_of_column(column)?; + lineage_literal(projection.input.as_ref(), input_idx) + } + _ => None, + }, + LogicalPlan::Filter(_) + | LogicalPlan::Sort(_) + | LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return None; + }; + lineage_literal(input, output_idx) + } + _ => None, + } +} + +fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option { + let ScalarValue::Utf8(Some(value)) = value else { + return None; + }; + cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone)) + .ok() + .filter(|value| !value.is_null()) + .map(|value| Expr::Literal(value, None)) +} + +fn unalias(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias(expr), + expr => expr, + } +} + +fn unalias_mut(expr: &mut Expr) -> &mut Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias_mut(expr), + expr => expr, + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::DFSchema; + use datafusion_expr::expr::Placeholder; + + use super::*; + + fn shanghai() -> SessionTimezone { + let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + } + } + + /// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped + /// placeholder, which must survive for parameter substitution. + #[test] + fn test_untyped_placeholder_assignment_is_left_alone() { + let schema = Arc::new(DFSchema::empty()); + let mut expr = Expr::Cast(Cast::new( + Box::new(Expr::Placeholder(Placeholder::new_with_field( + "$1".to_string(), + None, + ))), + DataType::Timestamp(TimeUnit::Millisecond, None), + )); + let original = expr.clone(); + + let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap(); + + assert!(!changed); + assert_eq!(expr, original); + } +} diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 3941fedfe8..40bb43e791 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod insert_assignment; - use std::sync::Arc; use common_time::Timezone; @@ -31,7 +29,7 @@ use session::context::QueryContextRef; use crate::QueryEngineContext; use crate::optimizer::ExtensionAnalyzerRule; -use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments; +use crate::optimizer::insert_assignment::rewrite_insert_assignments; use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according @@ -46,7 +44,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { &self, plan: LogicalPlan, ctx: &QueryEngineContext, - _config: &ConfigOptions, + config: &ConfigOptions, ) -> Result { plan.transform_up_with_subqueries(|plan| match plan { LogicalPlan::Filter(filter) => { @@ -129,7 +127,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule { 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(), + &ctx.query_ctx(), + config, )?); Ok(Transformed::yes(LogicalPlan::Dml(dml))) } @@ -324,7 +323,7 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { ) } -fn cast_string_to_timestamp( +pub(crate) fn cast_string_to_timestamp( string: &str, target_type: &DataType, timezone: Option<&Timezone>, diff --git a/src/query/src/optimizer/type_conversion/insert_assignment.rs b/src/query/src/optimizer/type_conversion/insert_assignment.rs deleted file mode 100644 index 11ab7e7829..0000000000 --- a/src/query/src/optimizer/type_conversion/insert_assignment.rs +++ /dev/null @@ -1,336 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::sync::Arc; - -use datafusion_common::{Column, Result, ScalarValue}; -use datafusion_expr::expr::{Alias, Cast}; -use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values}; -use datatypes::arrow::datatypes::DataType; -use session::context::QueryContextRef; - -use crate::optimizer::type_conversion::cast_string_to_timestamp; -use crate::plan::ExtractExpr; - -/// Rewrites string literals that feed timestamp columns at an INSERT boundary. -/// Constants are folded at the assignment to avoid changing source types; -/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts -/// stay on DataFusion's existing path. -pub(super) fn rewrite_insert_assignments( - plan: LogicalPlan, - query_ctx: QueryContextRef, -) -> Result { - let LogicalPlan::Projection(assignment) = plan else { - return Ok(plan); - }; - - let converter = InsertAssignmentConverter { query_ctx }; - let mut exprs = assignment.expr.clone(); - let mut input = assignment.input.as_ref().clone(); - let mut changed = false; - for (output_idx, expr) in assignment.expr.iter().enumerate() { - let target_type = assignment.schema.field(output_idx).data_type(); - if !matches!(target_type, DataType::Timestamp(_, _)) { - continue; - } - - let Some(column) = assignment_input_column(expr) else { - continue; - }; - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - continue; - }; - - let literal = lineage_literal(&input, input_idx).cloned(); - if let Some(literal) = literal - && let Some(folded) = converter.convert_literal(&literal, target_type) - { - let (qualifier, field) = assignment.schema.qualified_field(output_idx); - exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name()); - changed = true; - continue; - } - - // A hand-built DML plan (e.g. from flow) may share one source column - // between targets; rewriting it in place would retype every reader. - if assignment - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - continue; - } - if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? { - input = rewritten; - changed = true; - } - } - - if !changed { - return Ok(LogicalPlan::Projection(assignment)); - } - Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection) -} - -/// Resolves a literal when every row carries the same value at `output_idx`. -fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { - if output_idx >= plan.schema().fields().len() { - return None; - } - - match plan { - LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => Some(value), - Expr::Column(column) => { - let input_idx = projection.input.schema().maybe_index_of_column(column)?; - lineage_literal(projection.input.as_ref(), input_idx) - } - _ => None, - }, - // These nodes drop, reorder or deduplicate rows without touching the - // value a surviving row carries, and their schema stays positional. - LogicalPlan::Filter(_) - | LogicalPlan::Sort(_) - | LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return None; - }; - lineage_literal(input, output_idx) - } - _ => None, - } -} - -struct InsertAssignmentConverter { - query_ctx: QueryContextRef, -} - -impl InsertAssignmentConverter { - fn rewrite_output_column( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - if output_idx >= plan.schema().fields().len() { - return Ok(None); - } - - match plan { - // DataFusion pushes INSERT assignment casts into Values, so inspect - // them even though the Values schema already has the target type. - LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type), - // Once the source query has produced the target type, its casts and - // coercions belong to the source query rather than INSERT assignment. - _ if plan.schema().field(output_idx).data_type() == target_type => Ok(None), - LogicalPlan::Projection(projection) => { - self.rewrite_projection(projection, output_idx, target_type) - } - LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type), - // Filter and Sort are excluded here: their predicates and keys read - // the retyped column, which would change the source query. Constants - // still reach the assignment through `lineage_literal`. Distinct::All - // holds no expressions; allowing it shifts dedup keys from raw - // strings to parsed instants. - LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - self.rewrite_passthrough(plan, output_idx, target_type) - } - _ => Ok(None), - } - } - - fn rewrite_projection( - &self, - projection: &Projection, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => { - let Some(converted) = self.convert_literal(value, target_type) else { - return Ok(None); - }; - let (qualifier, field) = projection.schema.qualified_field(output_idx); - let mut exprs = projection.expr.clone(); - exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name()); - - Projection::try_new(exprs, projection.input.clone()) - .map(LogicalPlan::Projection) - .map(Some) - } - Expr::Column(column) => { - // Retyping the input column affects every output column reading - // it, so only follow lineage with a single consumer. - if projection - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - return Ok(None); - } - - let input = projection.input.as_ref(); - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - return Ok(None); - }; - let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)? - else { - return Ok(None); - }; - - Projection::try_new(projection.expr.clone(), Arc::new(rewritten)) - .map(LogicalPlan::Projection) - .map(Some) - } - _ => Ok(None), - } - } - - fn rewrite_values( - &self, - values: &Values, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut rewritten = values.clone(); - let mut changed = false; - for row in &mut rewritten.values { - let Some(expr) = row.get_mut(output_idx) else { - return Ok(None); - }; - - if let Expr::Cast(Cast { - expr: inner, - data_type, - }) = expr - && data_type == target_type - && let Expr::Literal(value, _) = inner.as_ref() - && let Some(value) = self.convert_literal(value, target_type) - { - *expr = value; - changed = true; - continue; - } - - if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) { - return Ok(None); - } - } - - Ok(changed.then_some(LogicalPlan::Values(rewritten))) - } - - fn rewrite_union( - &self, - union: &Union, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut inputs = Vec::with_capacity(union.inputs.len()); - // A union has one shared schema, so rewrite it only when the same output - // column is an INSERT-assigned literal in every branch. - for input in &union.inputs { - let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - inputs.push(Arc::new(rewritten)); - } - - // Untouched columns may still have mismatched branch types before - // TypeCoercion runs, so rebuild loosely like the SQL planner does. - Union::try_new_with_loose_types(inputs) - .map(LogicalPlan::Union) - .map(Some) - } - - fn rewrite_passthrough( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return Ok(None); - }; - let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - - plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input]) - .map(Some) - } - - fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option { - let ScalarValue::Utf8(Some(value)) = value else { - return None; - }; - cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone())) - .ok() - .filter(|value| !value.is_null()) - .map(|value| Expr::Literal(value, None)) - } -} - -fn assignment_input_column(expr: &Expr) -> Option<&Column> { - let expr = match unalias(expr) { - Expr::Cast(Cast { expr, .. }) => expr.as_ref(), - expr => expr, - }; - let Expr::Column(column) = expr else { - return None; - }; - Some(column) -} - -fn unalias(expr: &Expr) -> &Expr { - match expr { - Expr::Alias(Alias { expr, .. }) => unalias(expr), - expr => expr, - } -} - -#[cfg(test)] -mod tests { - use datafusion_common::ScalarValue; - use datafusion_common::arrow::datatypes::TimeUnit; - use session::context::QueryContext; - - use super::*; - - #[test] - fn test_convert_literal_falls_back_for_unsupported_literal() { - let converter = InsertAssignmentConverter { - query_ctx: QueryContext::arc(), - }; - let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None); - - for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { - assert_eq!( - converter - .convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,), - None - ); - } - } -} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index 09172340af..ed55b3908e 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -852,6 +852,30 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + /// Plans `sql` and runs the DataFusion analyzer, which is where + /// `InsertAssignmentRule` sits. Planning alone stops short of it, so these + /// assertions would not see the assignment rewrite at all. + async fn analyze_insert( + engine: &QueryEngineRef, + sql: &str, + query_ctx: &QueryContextRef, + ) -> String { + let stmt = QueryLanguageParser::parse_sql(sql, query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap(); + let context = engine.engine_context(query_ctx.clone()); + let state = context.state(); + state + .analyzer() + .execute_and_check(plan, state.config_options(), |_, _| {}) + .unwrap() + .display_indent_schema() + .to_string() + } + #[tokio::test] async fn test_insert_timestamp_literals_use_query_timezone() { let query_ctx = Arc::new( @@ -884,20 +908,6 @@ mod tests { ) AS source", &[1_785_902_400_001_i64][..], ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', now()", - &[1_785_988_800_001_i64, 1_786_075_200_001_i64][..], - ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-16 12:00:00.001', now() \ - UNION \ - SELECT '2026-08-17 12:00:00.001', now()", - &[1_786_852_800_001_i64, 1_786_939_200_001_i64][..], - ), ( "INSERT INTO timestamps (ts, st) \ SELECT '2026-08-18 12:00:00.001', max(st) \ @@ -926,14 +936,7 @@ mod tests { &[1_786_766_400_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(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected_timestamp in expected_timestamps { assert!( @@ -954,19 +957,19 @@ mod tests { let engine = create_timestamp_test_engine().await; let sql = "INSERT INTO timestamps (ts, st) \ VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + // An explicit cast reaches the analyzer as an `arrow_cast` call rather + // than an `Expr::Cast`, which is how it stays out of the rewrite. assert!( plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"), "{plan}" ); + // 12:00:00.001 read as Shanghai local time; the source query keeps UTC. + assert!( + !plan.contains("TimestampMillisecond(1786104000001, None)"), + "{plan}" + ); } #[tokio::test] @@ -1003,14 +1006,7 @@ mod tests { ][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected in expected { assert!(plan.contains(expected), "{plan}"); @@ -1019,30 +1015,37 @@ mod tests { } #[tokio::test] - async fn test_insert_union_tolerates_uncoerced_untouched_column() { + async fn test_insert_union_converts_via_assignment_cast() { let query_ctx = Arc::new( QueryContextBuilder::default() .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) .build(), ); let engine = create_timestamp_test_engine().await; - // `st` stays Timestamp vs Null across branches until TypeCoercion runs. - let sql = "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', NULL"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + // Branches disagree, so the conversion stays as a cast on the + // assignment instead of folding. One cast covers every branch, which is + // why a NULL branch no longer cancels the conversion for the column and + // why UNION's dedup keys stay on the original strings. + for sql in [ + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', NULL", + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-16 12:00:00.001', now() \ + UNION \ + SELECT '2026-08-17 12:00:00.001', now()", + ] { + let plan = analyze_insert(&engine, sql, &query_ctx).await; - for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] { assert!( - plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + plan.contains("AS Timestamp(ms, \"Asia/Shanghai\")"), + "{plan}" + ); + // The branches themselves are untouched. + assert!( + plan.contains("Utf8(\"2026-08-07 12:00:00.001\")") + || plan.contains("Utf8(\"2026-08-17 12:00:00.001\")"), "{plan}" ); } @@ -1060,14 +1063,8 @@ mod tests { SELECT '2026-08-10 12:00:00.001', now() \ UNION ALL \ SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + assert!( !plan.contains("TimestampMillisecond(1786334400001, None)"), "{plan}" @@ -1076,6 +1073,11 @@ mod tests { plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"), "{plan}" ); + // TypeCoercion has already settled this union to timestamp, so the + // assignment has nothing left to reinterpret. Retargeting the cast here + // would leave a Timestamp(None) -> Timestamp(Some(tz)) step behind, + // which shifts the value instead of relabelling it. + assert!(!plan.contains("Asia/Shanghai"), "{plan}"); } #[tokio::test] diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index 413eb5d46a..41d8ca6cba 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -114,6 +114,27 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( Affected Rows: 1 +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +Affected Rows: 2 + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +Affected Rows: 2 + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + +Affected Rows: 1 + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-------------------------+-------------------------------+ @@ -130,10 +151,43 @@ SELECT ts, ts_ns FROM test3 ORDER BY ts; | 2026-08-10T12:00:00.001 | | | 2026-08-11T12:00:00.001 | | | 2026-08-12T04:00:00.123 | 2026-08-12T04:00:00.123456789 | +| 2026-08-13T04:00:00.001 | 2026-08-13T04:00:00.123456789 | +| 2026-08-14T04:00:00.001 | | +| 2026-08-15T04:00:00.001 | | | 2026-08-16T04:00:00.001 | | | 2026-08-17T04:00:00.001 | | +| 2026-08-19T04:00:00.001 | 2026-08-19T04:00:00.123456789 | +| 2026-08-20T04:00:00.001 | | +-------------------------+-------------------------------+ +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +Affected Rows: 2 + +SELECT count(*) FROM test4; + ++----------+ +| count(*) | ++----------+ +| 2 | ++----------+ + +SELECT ts FROM test4 ORDER BY ts; + ++---------------------+ +| ts | ++---------------------+ +| 2026-08-05T20:00:00 | +| 2026-08-05T20:00:00 | ++---------------------+ + SET time_zone = 'UTC'; Affected Rows: 0 @@ -150,3 +204,7 @@ DROP TABLE test3; Affected Rows: 0 +DROP TABLE test4; + +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 7cbef15d6d..f7c196699b 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -56,8 +56,34 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( '2026-08-09 12:00:00.123456789' ); +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +SELECT count(*) FROM test4; + +SELECT ts FROM test4 ORDER BY ts; + SET time_zone = 'UTC'; DROP TABLE test1; @@ -65,3 +91,5 @@ DROP TABLE test1; DROP TABLE test2; DROP TABLE test3; + +DROP TABLE test4;