query/optimizer/
insert_assignment.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use common_time::Timezone;
18use datafusion::config::ConfigOptions;
19use datafusion_common::{DFSchemaRef, Result, ScalarValue};
20use datafusion_expr::expr::{Alias, Cast};
21use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Values};
22use datafusion_optimizer::analyzer::AnalyzerRule;
23use datafusion_optimizer::analyzer::type_coercion::TypeCoercion;
24use datatypes::arrow::datatypes::{DataType, TimeUnit};
25use session::context::QueryContextRef;
26
27use crate::optimizer::type_conversion::cast_string_to_timestamp;
28
29/// Interprets strings assigned to timestamp columns at an `INSERT` boundary
30/// using the session timezone. `plan` is the assignment projection under a
31/// `WriteOp::Insert`.
32///
33/// DataFusion plans `INSERT` as a projection casting each source column to its
34/// target column type, and that cast reads a naive string as UTC. Arrow does
35/// apply a timezone when the cast target carries one, so the assignment is
36/// routed through `Timestamp(unit, Some(tz))` and back. Stripping the timezone
37/// afterwards is value-preserving — arrow only shifts values in the opposite
38/// direction.
39///
40/// The source query is left untouched. Reinterpreting a value where it is
41/// *produced* would change what the source query means: pushing the conversion
42/// below a `UNION`'s `DISTINCT`, for instance, moves the dedup key from the raw
43/// strings to parsed instants and silently drops rows.
44///
45/// # Why `TypeCoercion` runs here
46///
47/// The rewrite reads source types, and those are only settled once a `UNION`'s
48/// branch types have been reconciled: before coercion a union carries its loose
49/// schema (the first branch's types), so a mixed
50/// `SELECT 'string' UNION ALL SELECT CAST(.. AS TIMESTAMP)` still looks like a
51/// string. Retargeting that cast would leave `Timestamp(None) ->
52/// Timestamp(Some(tz))` behind once coercion retypes the union — the one
53/// direction in which arrow shifts the value instead of relabelling it.
54///
55/// Coercing here rather than deferring to the analyzer is forced by where an
56/// INSERT is still identifiable: `exec_dml_statement` strips the `Dml` node and
57/// executes its input, so by the time the analyzer runs, an assignment
58/// projection is indistinguishable from any other projection.
59///
60/// Explicit casts stay out of this: the SQL layer turns a user's
61/// `CAST(x AS TIMESTAMP)` into an `arrow_cast` call, which only becomes an
62/// `Expr::Cast` in the optimizer's `SimplifyExpressions`. Assignment casts are
63/// therefore the only `Expr::Cast` reaching a timestamp column here.
64///
65/// # Reach
66///
67/// The emitted cast only carries its timezone where the expression is evaluated
68/// on this node. Substrait drops the timezone name when a plan is pushed down —
69/// it encodes any zoned timestamp as `PrecisionTimestampTz` and decodes it back
70/// as UTC — so a source reading from a table falls back to UTC, the behaviour it
71/// had before this rule existed. Sources that never leave this node (literals,
72/// `VALUES`, and `UNION`s of them) keep the session timezone, and those are what
73/// an INSERT's timestamp assignment is in practice.
74pub(crate) fn rewrite_insert_assignments(
75    plan: LogicalPlan,
76    query_ctx: &QueryContextRef,
77    config: &ConfigOptions,
78) -> Result<LogicalPlan> {
79    let Some(timezone) = session_timezone(query_ctx) else {
80        return Ok(plan);
81    };
82
83    let plan = TypeCoercion::new().analyze(plan, config)?;
84    rewrite_assignment(plan, &timezone)
85}
86
87/// Session timezone, in both forms the rewrite needs.
88struct SessionTimezone {
89    /// Parses literals, matching the plain `INSERT ... VALUES` path.
90    parsed: Timezone,
91    /// Names the intermediate arrow cast target.
92    name: Arc<str>,
93}
94
95fn session_timezone(query_ctx: &QueryContextRef) -> Option<SessionTimezone> {
96    let parsed = query_ctx.timezone();
97
98    // A UTC session already gets UTC semantics from the plain assignment cast.
99    if parsed.is_utc() {
100        return None;
101    }
102
103    Some(SessionTimezone {
104        name: Arc::from(parsed.to_string()),
105        parsed,
106    })
107}
108
109fn rewrite_assignment(plan: LogicalPlan, timezone: &SessionTimezone) -> Result<LogicalPlan> {
110    let LogicalPlan::Projection(assignment) = plan else {
111        return Ok(plan);
112    };
113
114    let mut exprs = assignment.expr.clone();
115    let mut changed = false;
116    for expr in &mut exprs {
117        changed |= retarget_assignment_cast(
118            expr,
119            assignment.input.schema(),
120            Some(assignment.input.as_ref()),
121            timezone,
122        )?;
123    }
124
125    // The planner types `VALUES` against the target table, so the assignment
126    // cast lands inside the `Values` rows instead of on the projection above.
127    let mut input = assignment.input.clone();
128    if let LogicalPlan::Values(values) = assignment.input.as_ref()
129        && let Some(rewritten) = rewrite_values(values, timezone)?
130    {
131        input = Arc::new(LogicalPlan::Values(rewritten));
132        changed = true;
133    }
134
135    if !changed {
136        return Ok(LogicalPlan::Projection(assignment));
137    }
138    Projection::try_new(exprs, input).map(LogicalPlan::Projection)
139}
140
141fn rewrite_values(values: &Values, timezone: &SessionTimezone) -> Result<Option<Values>> {
142    let mut rewritten = values.clone();
143    let mut changed = false;
144    for row in &mut rewritten.values {
145        for expr in row.iter_mut() {
146            changed |= retarget_assignment_cast(expr, &values.schema, None, timezone)?;
147        }
148    }
149
150    Ok(changed.then_some(rewritten))
151}
152
153/// Reinterprets one assignment cast, returning whether it was rewritten.
154///
155/// `source_plan` is the projection's input, used to resolve a literal behind a
156/// column reference; `Values` rows carry their expression inline and pass `None`.
157fn retarget_assignment_cast(
158    expr: &mut Expr,
159    schema: &DFSchemaRef,
160    source_plan: Option<&LogicalPlan>,
161    timezone: &SessionTimezone,
162) -> Result<bool> {
163    let expr = unalias_mut(expr);
164    let Expr::Cast(Cast {
165        expr: source,
166        data_type: DataType::Timestamp(unit, None),
167    }) = expr
168    else {
169        return Ok(false);
170    };
171    let unit = *unit;
172
173    if !matches!(
174        source.get_type(schema)?,
175        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
176    ) {
177        return Ok(false);
178    }
179
180    // Fold literals with the same parser the plain `INSERT ... VALUES` path
181    // uses, so a given string means the same thing however it reaches a column.
182    // The parsers disagree on ambiguous local times: this one resolves them,
183    // arrow rejects them.
184    let folded = source_literal(source.as_ref(), source_plan)
185        .and_then(|literal| convert_literal(&literal, unit, &timezone.parsed));
186    if let Some(folded) = folded {
187        *expr = folded;
188        return Ok(true);
189    }
190
191    let source = source.as_ref().clone();
192    *expr = Expr::Cast(Cast::new(
193        Box::new(Expr::Cast(Cast::new(
194            Box::new(source),
195            DataType::Timestamp(unit, Some(timezone.name.clone())),
196        ))),
197        DataType::Timestamp(unit, None),
198    ));
199    Ok(true)
200}
201
202fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option<ScalarValue> {
203    match source {
204        Expr::Literal(value, _) => Some(value.clone()),
205        Expr::Column(column) => {
206            let plan = source_plan?;
207            let index = plan.schema().maybe_index_of_column(column)?;
208            lineage_literal(plan, index).cloned()
209        }
210        _ => None,
211    }
212}
213
214/// Resolves a literal when every row carries the same value at `output_idx`.
215///
216/// Read-only: the literal is folded into the assignment above, so nodes that
217/// drop, reorder or deduplicate rows can be traversed — none of them changes
218/// the value a surviving row carries, and folding above them leaves their keys
219/// on the original strings.
220fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> {
221    if output_idx >= plan.schema().fields().len() {
222        return None;
223    }
224
225    match plan {
226        LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) {
227            Expr::Literal(value, _) => Some(value),
228            Expr::Column(column) => {
229                let input_idx = projection.input.schema().maybe_index_of_column(column)?;
230                lineage_literal(projection.input.as_ref(), input_idx)
231            }
232            _ => None,
233        },
234        LogicalPlan::Filter(_)
235        | LogicalPlan::Sort(_)
236        | LogicalPlan::Limit(_)
237        | LogicalPlan::SubqueryAlias(_)
238        | LogicalPlan::Distinct(Distinct::All(_)) => {
239            let inputs = plan.inputs();
240            let [input] = inputs.as_slice() else {
241                return None;
242            };
243            lineage_literal(input, output_idx)
244        }
245        _ => None,
246    }
247}
248
249fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option<Expr> {
250    let ScalarValue::Utf8(Some(value)) = value else {
251        return None;
252    };
253    cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone))
254        .ok()
255        .filter(|value| !value.is_null())
256        .map(|value| Expr::Literal(value, None))
257}
258
259fn unalias(expr: &Expr) -> &Expr {
260    match expr {
261        Expr::Alias(Alias { expr, .. }) => unalias(expr),
262        expr => expr,
263    }
264}
265
266fn unalias_mut(expr: &mut Expr) -> &mut Expr {
267    match expr {
268        Expr::Alias(Alias { expr, .. }) => unalias_mut(expr),
269        expr => expr,
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use datafusion_common::DFSchema;
276    use datafusion_expr::expr::Placeholder;
277
278    use super::*;
279
280    fn shanghai() -> SessionTimezone {
281        let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap();
282        SessionTimezone {
283            name: Arc::from(parsed.to_string()),
284            parsed,
285        }
286    }
287
288    /// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped
289    /// placeholder, which must survive for parameter substitution.
290    #[test]
291    fn test_untyped_placeholder_assignment_is_left_alone() {
292        let schema = Arc::new(DFSchema::empty());
293        let mut expr = Expr::Cast(Cast::new(
294            Box::new(Expr::Placeholder(Placeholder::new_with_field(
295                "$1".to_string(),
296                None,
297            ))),
298            DataType::Timestamp(TimeUnit::Millisecond, None),
299        ));
300        let original = expr.clone();
301
302        let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap();
303
304        assert!(!changed);
305        assert_eq!(expr, original);
306    }
307}
query/optimizer/insert_assignment.rs

query/optimizer/
insert_assignment.rs
