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        field,
167    }) = expr
168    else {
169        return Ok(false);
170    };
171    let DataType::Timestamp(unit, None) = field.data_type() else {
172        return Ok(false);
173    };
174    let unit = *unit;
175
176    if !matches!(
177        source.get_type(schema)?,
178        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
179    ) {
180        return Ok(false);
181    }
182
183    // Fold literals with the same parser the plain `INSERT ... VALUES` path
184    // uses, so a given string means the same thing however it reaches a column.
185    // The parsers disagree on ambiguous local times: this one resolves them,
186    // arrow rejects them.
187    let folded = source_literal(source.as_ref(), source_plan)
188        .and_then(|literal| convert_literal(&literal, unit, &timezone.parsed));
189    if let Some(folded) = folded {
190        *expr = folded;
191        return Ok(true);
192    }
193
194    let source = source.as_ref().clone();
195    *expr = Expr::Cast(Cast::new(
196        Box::new(Expr::Cast(Cast::new(
197            Box::new(source),
198            DataType::Timestamp(unit, Some(timezone.name.clone())),
199        ))),
200        DataType::Timestamp(unit, None),
201    ));
202    Ok(true)
203}
204
205fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option<ScalarValue> {
206    match source {
207        Expr::Literal(value, _) => Some(value.clone()),
208        Expr::Column(column) => {
209            let plan = source_plan?;
210            let index = plan.schema().maybe_index_of_column(column)?;
211            lineage_literal(plan, index).cloned()
212        }
213        _ => None,
214    }
215}
216
217/// Resolves a literal when every row carries the same value at `output_idx`.
218///
219/// Read-only: the literal is folded into the assignment above, so nodes that
220/// drop, reorder or deduplicate rows can be traversed — none of them changes
221/// the value a surviving row carries, and folding above them leaves their keys
222/// on the original strings.
223fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> {
224    if output_idx >= plan.schema().fields().len() {
225        return None;
226    }
227
228    match plan {
229        LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) {
230            Expr::Literal(value, _) => Some(value),
231            Expr::Column(column) => {
232                let input_idx = projection.input.schema().maybe_index_of_column(column)?;
233                lineage_literal(projection.input.as_ref(), input_idx)
234            }
235            _ => None,
236        },
237        LogicalPlan::Filter(_)
238        | LogicalPlan::Sort(_)
239        | LogicalPlan::Limit(_)
240        | LogicalPlan::SubqueryAlias(_)
241        | LogicalPlan::Distinct(Distinct::All(_)) => {
242            let inputs = plan.inputs();
243            let [input] = inputs.as_slice() else {
244                return None;
245            };
246            lineage_literal(input, output_idx)
247        }
248        _ => None,
249    }
250}
251
252fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option<Expr> {
253    let ScalarValue::Utf8(Some(value)) = value else {
254        return None;
255    };
256    cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone))
257        .ok()
258        .filter(|value| !value.is_null())
259        .map(|value| Expr::Literal(value, None))
260}
261
262fn unalias(expr: &Expr) -> &Expr {
263    match expr {
264        Expr::Alias(Alias { expr, .. }) => unalias(expr),
265        expr => expr,
266    }
267}
268
269fn unalias_mut(expr: &mut Expr) -> &mut Expr {
270    match expr {
271        Expr::Alias(Alias { expr, .. }) => unalias_mut(expr),
272        expr => expr,
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use datafusion_common::DFSchema;
279    use datafusion_expr::expr::Placeholder;
280
281    use super::*;
282
283    fn shanghai() -> SessionTimezone {
284        let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap();
285        SessionTimezone {
286            name: Arc::from(parsed.to_string()),
287            parsed,
288        }
289    }
290
291    /// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped
292    /// placeholder, which must survive for parameter substitution.
293    #[test]
294    fn test_untyped_placeholder_assignment_is_left_alone() {
295        let schema = Arc::new(DFSchema::empty());
296        let mut expr = Expr::Cast(Cast::new(
297            Box::new(Expr::Placeholder(Placeholder::new_with_field(
298                "$1".to_string(),
299                None,
300            ))),
301            DataType::Timestamp(TimeUnit::Millisecond, None),
302        ));
303        let original = expr.clone();
304
305        let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap();
306
307        assert!(!changed);
308        assert_eq!(expr, original);
309    }
310}
query/optimizer/insert_assignment.rs

query/optimizer/
insert_assignment.rs
