mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-13 07:22:32 +00:00
fix(query): preserve timestamp literal semantics in inserts (#8889)
* fix(query): follow timestamp insert assignment lineage Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix(query): fold constant insert timestamp literals at the assignment Following lineage by retyping the source column changed every output column that reads it: a string column sharing the literal was silently rewritten to a formatted timestamp, and a nanosecond column was truncated to the precision of whichever column was converted first. Resolve the constant read-only and fold it into the assignment expression instead, which leaves the source query untouched and also covers literals behind WHERE, ORDER BY and DISTINCT. VALUES rows and UNION branches carry per-row values, so they keep the in-place rewrite, now guarded against columns with more than one consumer. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test(query): strengthen insert lineage regression Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test(query): trim redundant insert coverage Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix(query): rebuild insert unions loosely and cover UNION distinct Per review: rebuilding a rewritten union with the strict constructor rejected legal pre-coercion plans whose untouched columns still differ across branches. Use try_new_with_loose_types, matching the SQL planner. Distinct::All joins the rewrite passthrough so UNION (distinct) literals get session-timezone parsing like UNION ALL; deduplication then keys on parsed instants instead of raw strings. The top-level rewrite path gains the same single-consumer guard as rewrite_projection for hand-built DML plans that share a source column between targets. Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * style(query): tighten insert assignment comments Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com> (cherry picked from commit b2fc4b9cba4c0e137cc6f704d88dfe7f187b57aa) (cherry picked from commit 1775e31a7916cf09a7da8238cebea5e3d74817f7)
This commit is contained in:
dennis zhuang
discord9
@@ -12,6 +12,8 @@
|
||||
// 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;
|
||||
@@ -19,7 +21,7 @@ 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::{Cast, InList};
|
||||
use datafusion_expr::expr::InList;
|
||||
use datafusion_expr::{
|
||||
Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan, WriteOp,
|
||||
};
|
||||
@@ -29,6 +31,7 @@ use session::context::QueryContextRef;
|
||||
|
||||
use crate::QueryEngineContext;
|
||||
use crate::optimizer::ExtensionAnalyzerRule;
|
||||
use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments;
|
||||
use crate::plan::ExtractExpr;
|
||||
|
||||
/// TypeConversionRule converts some literal values in logical plan to other types according
|
||||
@@ -307,120 +310,6 @@ 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,
|
||||
@@ -472,7 +361,7 @@ mod tests {
|
||||
|
||||
use datafusion_common::arrow::datatypes::Field;
|
||||
use datafusion_common::{Column, DFSchema, NullEquality};
|
||||
use datafusion_expr::expr::Exists;
|
||||
use datafusion_expr::expr::{Cast, Exists};
|
||||
use datafusion_expr::{Join, JoinConstraint, JoinType, Literal, LogicalPlanBuilder, Subquery};
|
||||
use datafusion_sql::TableReference;
|
||||
use session::context::QueryContext;
|
||||
@@ -594,45 +483,6 @@ 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;
|
||||
|
||||
@@ -0,0 +1,336 @@
|
||||
// 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<LogicalPlan> {
|
||||
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<Option<LogicalPlan>> {
|
||||
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<Option<LogicalPlan>> {
|
||||
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<Option<LogicalPlan>> {
|
||||
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<Option<LogicalPlan>> {
|
||||
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<Option<LogicalPlan>> {
|
||||
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<Expr> {
|
||||
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
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
+199
-4
@@ -697,12 +697,18 @@ mod tests {
|
||||
ConcreteDataType::timestamp_millisecond_datatype(),
|
||||
false,
|
||||
),
|
||||
ColumnSchema::new("note", ConcreteDataType::string_datatype(), true),
|
||||
ColumnSchema::new(
|
||||
"ts_ns",
|
||||
ConcreteDataType::timestamp_nanosecond_datatype(),
|
||||
true,
|
||||
),
|
||||
];
|
||||
let schema = Arc::new(Schema::new(columns));
|
||||
let table_meta = TableMetaBuilder::empty()
|
||||
.schema(schema)
|
||||
.primary_key_indices(vec![])
|
||||
.value_indices(vec![0, 1])
|
||||
.value_indices(vec![0, 1, 2, 3])
|
||||
.next_column_id(1024)
|
||||
.build()
|
||||
.unwrap();
|
||||
@@ -855,16 +861,69 @@ mod tests {
|
||||
);
|
||||
let engine = create_timestamp_test_engine().await;
|
||||
|
||||
for (sql, expected_timestamp) in [
|
||||
for (sql, expected_timestamps) in [
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
VALUES ('2026-08-02 12:00:00.001', now())",
|
||||
1_785_643_200_001_i64,
|
||||
&[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,
|
||||
&[1_785_729_600_001_i64][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT '2026-08-04 12:00:00.001', now() LIMIT 1",
|
||||
&[1_785_816_000_001_i64][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT * FROM (\
|
||||
SELECT '2026-08-05 12:00:00.001', now()\
|
||||
) 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) \
|
||||
FROM timestamps GROUP BY note",
|
||||
&[1_787_025_600_001_i64][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT c, s FROM (\
|
||||
SELECT '2026-08-13 12:00:00.001' AS c, now() AS s\
|
||||
) AS t WHERE c > '2026-01-01'",
|
||||
&[1_786_593_600_001_i64][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT c, s FROM (\
|
||||
SELECT '2026-08-14 12:00:00.001' AS c, now() AS s\
|
||||
) AS t ORDER BY c",
|
||||
&[1_786_680_000_001_i64][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, st) \
|
||||
SELECT DISTINCT c, s FROM (\
|
||||
SELECT '2026-08-15 12:00:00.001' AS c, now() AS s\
|
||||
) AS t",
|
||||
&[1_786_766_400_001_i64][..],
|
||||
),
|
||||
] {
|
||||
let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap();
|
||||
@@ -876,6 +935,112 @@ mod tests {
|
||||
.display_indent()
|
||||
.to_string();
|
||||
|
||||
for expected_timestamp in expected_timestamps {
|
||||
assert!(
|
||||
plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")),
|
||||
"{plan}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_insert_explicit_timestamp_cast_keeps_datafusion_semantics() {
|
||||
let query_ctx = Arc::new(
|
||||
QueryContextBuilder::default()
|
||||
.timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
|
||||
.build(),
|
||||
);
|
||||
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();
|
||||
|
||||
assert!(
|
||||
plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"),
|
||||
"{plan}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_insert_converts_source_literal_shared_by_several_columns() {
|
||||
let query_ctx = Arc::new(
|
||||
QueryContextBuilder::default()
|
||||
.timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
|
||||
.build(),
|
||||
);
|
||||
let engine = create_timestamp_test_engine().await;
|
||||
|
||||
// Rewriting `c` in place would also retype `note` and truncate `ts_ns`.
|
||||
for (sql, expected) in [
|
||||
(
|
||||
"INSERT INTO timestamps (ts, note) SELECT a, b FROM (\
|
||||
SELECT c AS a, c AS b FROM (\
|
||||
SELECT '2026-08-12 12:00:00.001' AS c\
|
||||
) AS t1\
|
||||
) AS t2",
|
||||
&[
|
||||
"TimestampMillisecond(1786507200001, None)",
|
||||
"[a:Utf8, b:Utf8]",
|
||||
][..],
|
||||
),
|
||||
(
|
||||
"INSERT INTO timestamps (ts, ts_ns) SELECT a, b FROM (\
|
||||
SELECT c AS a, c AS b FROM (\
|
||||
SELECT '2026-08-12 12:00:00.123456789' AS c\
|
||||
) AS t1\
|
||||
) AS t2",
|
||||
&[
|
||||
"TimestampMillisecond(1786507200123, None)",
|
||||
"TimestampNanosecond(1786507200123456789, None)",
|
||||
][..],
|
||||
),
|
||||
] {
|
||||
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();
|
||||
|
||||
for expected in expected {
|
||||
assert!(plan.contains(expected), "{plan}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_insert_union_tolerates_uncoerced_untouched_column() {
|
||||
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();
|
||||
|
||||
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}"
|
||||
@@ -883,6 +1048,36 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_insert_mixed_union_keeps_source_coercion() {
|
||||
let query_ctx = Arc::new(
|
||||
QueryContextBuilder::default()
|
||||
.timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap())
|
||||
.build(),
|
||||
);
|
||||
let engine = create_timestamp_test_engine().await;
|
||||
let sql = "INSERT INTO timestamps (ts, st) \
|
||||
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();
|
||||
assert!(
|
||||
!plan.contains("TimestampMillisecond(1786334400001, None)"),
|
||||
"{plan}"
|
||||
);
|
||||
assert!(
|
||||
plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"),
|
||||
"{plan}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_extract_placeholder_cast_types_multiple() {
|
||||
let plan = parse_sql_to_plan(
|
||||
|
||||
@@ -68,10 +68,48 @@ INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now();
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
INSERT INTO test3 (ts, st) SELECT '2026-08-04 12:00:00.001', now() LIMIT 1;
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
INSERT INTO test3 (ts, st)
|
||||
SELECT '2026-08-06 12:00:00.001', now()
|
||||
UNION ALL
|
||||
SELECT '2026-08-07 12:00:00.001', now();
|
||||
|
||||
Affected Rows: 2
|
||||
|
||||
INSERT INTO test3 (ts, st) VALUES (
|
||||
CAST('2026-08-08 12:00:00.001' AS TIMESTAMP),
|
||||
now()
|
||||
);
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
INSERT INTO test3 (ts, st)
|
||||
SELECT '2026-08-10 12:00:00.001', now()
|
||||
UNION ALL
|
||||
SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now();
|
||||
|
||||
Affected Rows: 2
|
||||
|
||||
INSERT INTO test3 (ts, st)
|
||||
SELECT '2026-08-16 12:00:00.001', now()
|
||||
UNION
|
||||
SELECT '2026-08-17 12:00:00.001', now();
|
||||
|
||||
Affected Rows: 2
|
||||
|
||||
INSERT INTO test3 (ts, ts_ns) SELECT a, b FROM (
|
||||
SELECT c AS a, c AS b FROM (SELECT '2026-08-12 12:00:00.123456789' AS c) AS t1
|
||||
) AS t2;
|
||||
|
||||
Affected Rows: 1
|
||||
|
||||
INSERT INTO test3 (ts, st, ts_ns) VALUES (
|
||||
'2026-08-04 12:00:00.001',
|
||||
'2026-08-09 12:00:00.001',
|
||||
now(),
|
||||
'2026-08-04 12:00:00.123456789'
|
||||
'2026-08-09 12:00:00.123456789'
|
||||
);
|
||||
|
||||
Affected Rows: 1
|
||||
@@ -84,7 +122,16 @@ SELECT ts, ts_ns FROM test3 ORDER BY ts;
|
||||
| 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 |
|
||||
| 2026-08-04T04:00:00.001 | |
|
||||
| 2026-08-06T04:00:00.001 | |
|
||||
| 2026-08-07T04:00:00.001 | |
|
||||
| 2026-08-08T12:00:00.001 | |
|
||||
| 2026-08-09T04:00:00.001 | 2026-08-09T04:00:00.123456789 |
|
||||
| 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-16T04:00:00.001 | |
|
||||
| 2026-08-17T04:00:00.001 | |
|
||||
+-------------------------+-------------------------------+
|
||||
|
||||
SET time_zone = 'UTC';
|
||||
|
||||
@@ -24,10 +24,36 @@ 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) SELECT '2026-08-04 12:00:00.001', now() LIMIT 1;
|
||||
|
||||
INSERT INTO test3 (ts, st)
|
||||
SELECT '2026-08-06 12:00:00.001', now()
|
||||
UNION ALL
|
||||
SELECT '2026-08-07 12:00:00.001', now();
|
||||
|
||||
INSERT INTO test3 (ts, st) VALUES (
|
||||
CAST('2026-08-08 12:00:00.001' AS TIMESTAMP),
|
||||
now()
|
||||
);
|
||||
|
||||
INSERT INTO test3 (ts, st)
|
||||
SELECT '2026-08-10 12:00:00.001', now()
|
||||
UNION ALL
|
||||
SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now();
|
||||
|
||||
INSERT INTO test3 (ts, st)
|
||||
SELECT '2026-08-16 12:00:00.001', now()
|
||||
UNION
|
||||
SELECT '2026-08-17 12:00:00.001', now();
|
||||
|
||||
INSERT INTO test3 (ts, ts_ns) SELECT a, b FROM (
|
||||
SELECT c AS a, c AS b FROM (SELECT '2026-08-12 12:00:00.123456789' AS c) AS t1
|
||||
) AS t2;
|
||||
|
||||
INSERT INTO test3 (ts, st, ts_ns) VALUES (
|
||||
'2026-08-04 12:00:00.001',
|
||||
'2026-08-09 12:00:00.001',
|
||||
now(),
|
||||
'2026-08-04 12:00:00.123456789'
|
||||
'2026-08-09 12:00:00.123456789'
|
||||
);
|
||||
|
||||
SELECT ts, ts_ns FROM test3 ORDER BY ts;
|
||||
|
||||
Reference in New Issue
Block a user