Skip to main content

datafusion_optimizer/analyzer/
type_coercion.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Optimizer rule for type validation and coercion
19
20use arrow::compute::can_cast_types;
21use datafusion_expr::binary::BinaryTypeCoercer;
22use itertools::{Itertools as _, izip};
23use std::sync::{Arc, LazyLock};
24
25use crate::analyzer::AnalyzerRule;
26use crate::utils::NamePreserver;
27
28use arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit};
29use arrow::temporal_conversions::SECONDS_IN_DAY;
30use datafusion_common::config::ConfigOptions;
31use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
32use datafusion_common::{
33    Column, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference,
34    exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err,
35    plan_err,
36};
37use datafusion_expr::expr::{
38    self, AggregateFunctionParams, Alias, Between, BinaryExpr, Case, Exists,
39    HigherOrderFunction, InList, InSubquery, Like, ScalarFunction, SetComparison, Sort,
40    WindowFunction,
41};
42use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema;
43use datafusion_expr::expr_schema::cast_subquery;
44use datafusion_expr::logical_plan::Subquery;
45use datafusion_expr::type_coercion::binary::{
46    comparison_coercion, like_coercion, regex_coercion, type_union_coercion,
47};
48use datafusion_expr::type_coercion::functions::{
49    UDFCoercionExt, fields_with_udf, value_fields_with_higher_order_udf_and_lambdas,
50};
51use datafusion_expr::type_coercion::other::{
52    get_coerce_type_for_case_expression, get_coerce_type_for_case_when,
53    get_coerce_type_for_list,
54};
55use datafusion_expr::type_coercion::{
56    is_datetime, is_interval, is_signed_numeric, is_timestamp,
57};
58use datafusion_expr::utils::merge_schema;
59use datafusion_expr::{
60    Cast, DmlStatement, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator,
61    Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits,
62    WriteOp, is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown,
63    lit, not,
64};
65
66/// Performs type coercion by determining the schema
67/// and performing the expression rewrites.
68#[derive(Default, Debug)]
69pub struct TypeCoercion {}
70
71impl TypeCoercion {
72    pub fn new() -> Self {
73        Self {}
74    }
75}
76
77/// Coerce output schema based upon optimizer config.
78fn coerce_output(plan: LogicalPlan, config: &ConfigOptions) -> Result<LogicalPlan> {
79    if !config.optimizer.expand_views_at_output {
80        return Ok(plan);
81    }
82
83    let outer_refs = plan.expressions();
84    if outer_refs.is_empty() {
85        return Ok(plan);
86    }
87
88    if let Some(dfschema) = transform_schema_to_nonview(plan.schema()) {
89        coerce_plan_expr_for_schema(plan, &dfschema?)
90    } else {
91        Ok(plan)
92    }
93}
94
95impl AnalyzerRule for TypeCoercion {
96    fn name(&self) -> &str {
97        "type_coercion"
98    }
99
100    fn analyze(&self, plan: LogicalPlan, config: &ConfigOptions) -> Result<LogicalPlan> {
101        static EMPTY_SCHEMA: LazyLock<DFSchema> = LazyLock::new(DFSchema::empty);
102
103        // recurse
104        let transformed_plan = plan
105            .transform_up_with_subqueries(|plan| analyze_internal(&EMPTY_SCHEMA, plan))?
106            .data;
107
108        // finish
109        coerce_output(transformed_plan, config)
110    }
111}
112
113/// use the external schema to handle the correlated subqueries case
114///
115/// Assumes that children have already been optimized
116fn analyze_internal(
117    external_schema: &DFSchema,
118    plan: LogicalPlan,
119) -> Result<Transformed<LogicalPlan>> {
120    // get schema representing all available input fields. This is used for data type
121    // resolution only, so order does not matter here
122    let mut schema = merge_schema(&plan.inputs());
123
124    if let LogicalPlan::TableScan(ts) = &plan {
125        let source_schema = DFSchema::try_from_qualified_schema(
126            ts.table_name.clone(),
127            &ts.source.schema(),
128        )?;
129        schema.merge(&source_schema);
130    }
131
132    // MERGE expressions (ON / WHEN clauses) reference the target table, which
133    // is not one of `plan.inputs()`. Rebuild the target schema from the DML's
134    // `table_name` and `target` so those columns resolve during coercion.
135    if let LogicalPlan::Dml(DmlStatement {
136        op: WriteOp::MergeInto(_),
137        table_name,
138        target,
139        ..
140    }) = &plan
141    {
142        let target_schema =
143            DFSchema::try_from_qualified_schema(table_name.clone(), &target.schema())?;
144        schema.merge(&target_schema);
145    }
146
147    // merge the outer schema for correlated subqueries
148    // like case:
149    // select t2.c2 from t1 where t1.c1 in (select t2.c1 from t2 where t2.c2=t1.c3)
150    schema.merge(external_schema);
151
152    // Coerce filter predicates to boolean (handles `WHERE NULL`)
153    let plan = if let LogicalPlan::Filter(mut filter) = plan {
154        filter.predicate = filter.predicate.cast_to(&DataType::Boolean, &schema)?;
155        LogicalPlan::Filter(filter)
156    } else {
157        plan
158    };
159
160    let mut expr_rewrite = TypeCoercionRewriter::new(&schema);
161
162    let name_preserver = NamePreserver::new(&plan);
163    // apply coercion rewrite all expressions in the plan individually
164    plan.map_expressions(|expr| {
165        let original_name = name_preserver.save(&expr);
166        expr.rewrite(&mut expr_rewrite)
167            .map(|transformed| transformed.update_data(|e| original_name.restore(e)))
168    })?
169    // some plans need extra coercion after their expressions are coerced
170    .map_data(|plan| expr_rewrite.coerce_plan(plan))?
171    // recompute the schema after the expressions have been rewritten as the types may have changed
172    .map_data(|plan| plan.recompute_schema())
173}
174
175/// Rewrite expressions to apply type coercion.
176pub struct TypeCoercionRewriter<'a> {
177    pub(crate) schema: &'a DFSchema,
178}
179
180impl<'a> TypeCoercionRewriter<'a> {
181    /// Create a new [`TypeCoercionRewriter`] with a provided schema
182    /// representing both the inputs and output of the [`LogicalPlan`] node.
183    pub fn new(schema: &'a DFSchema) -> Self {
184        Self { schema }
185    }
186
187    /// Coerce the [`LogicalPlan`].
188    ///
189    /// Refer to [`TypeCoercionRewriter::coerce_join`] and [`TypeCoercionRewriter::coerce_union`]
190    /// for type-coercion approach.
191    pub fn coerce_plan(&mut self, plan: LogicalPlan) -> Result<LogicalPlan> {
192        match plan {
193            LogicalPlan::Join(join) => self.coerce_join(join),
194            LogicalPlan::Union(union) => Self::coerce_union(union),
195            LogicalPlan::Limit(limit) => Self::coerce_limit(limit),
196            LogicalPlan::Dml(dml) => self.coerce_dml(dml),
197            _ => Ok(plan),
198        }
199    }
200
201    fn coerce_dml(&self, mut dml: DmlStatement) -> Result<LogicalPlan> {
202        let WriteOp::MergeInto(merge_op) = &dml.op else {
203            return Ok(LogicalPlan::Dml(dml));
204        };
205
206        let target_schema = DFSchema::try_from_qualified_schema(
207            dml.table_name.clone(),
208            &dml.target.schema(),
209        )?;
210        let mut merge_op = (**merge_op).clone();
211        merge_op.on = self.coerce_predicate(merge_op.on, "MERGE ON condition")?;
212        for clause in &mut merge_op.clauses {
213            clause.predicate = clause
214                .predicate
215                .take()
216                .map(|expr| self.coerce_predicate(expr, "MERGE WHEN condition"))
217                .transpose()?;
218
219            match &mut clause.action {
220                datafusion_expr::dml::MergeIntoAction::Update(assignments) => {
221                    for (column, value) in assignments {
222                        let field = target_schema.field_with_unqualified_name(column)?;
223                        *value = value.clone().cast_to(field.data_type(), self.schema)?;
224                    }
225                }
226                datafusion_expr::dml::MergeIntoAction::Insert { columns, values } => {
227                    if columns.is_empty() {
228                        for (value, field) in
229                            values.iter_mut().zip(target_schema.fields())
230                        {
231                            *value =
232                                value.clone().cast_to(field.data_type(), self.schema)?;
233                        }
234                    } else {
235                        for (column, value) in columns.iter().zip(values) {
236                            let field =
237                                target_schema.field_with_unqualified_name(column)?;
238                            *value =
239                                value.clone().cast_to(field.data_type(), self.schema)?;
240                        }
241                    }
242                }
243                datafusion_expr::dml::MergeIntoAction::Delete => {}
244            }
245        }
246        dml.op = WriteOp::MergeInto(Box::new(merge_op));
247        Ok(LogicalPlan::Dml(dml))
248    }
249
250    /// Coerce join equality expressions and join filter
251    ///
252    /// Joins must be treated specially as their equality expressions are stored
253    /// as a parallel list of left and right expressions, rather than a single
254    /// equality expression
255    ///
256    /// For example, on_exprs like `t1.a = t2.b AND t1.x = t2.y` will be stored
257    /// as a list of `(t1.a, t2.b), (t1.x, t2.y)`
258    pub fn coerce_join(&mut self, mut join: Join) -> Result<LogicalPlan> {
259        join.on = join
260            .on
261            .into_iter()
262            .map(|(lhs, rhs)| {
263                // coerce the arguments as though they were a single binary equality
264                // expression
265                let left_schema = join.left.schema();
266                let right_schema = join.right.schema();
267                let (lhs, rhs) = self.coerce_binary_op(
268                    lhs,
269                    left_schema,
270                    Operator::Eq,
271                    rhs,
272                    right_schema,
273                )?;
274                Ok((lhs, rhs))
275            })
276            .collect::<Result<Vec<_>>>()?;
277
278        // Join filter must be boolean
279        join.filter = join
280            .filter
281            .map(|expr| self.coerce_predicate(expr, "Join condition"))
282            .transpose()?;
283
284        Ok(LogicalPlan::Join(join))
285    }
286
287    /// Coerce the union’s inputs to a common schema compatible with all inputs.
288    /// This occurs after wildcard expansion and the coercion of the input expressions.
289    pub fn coerce_union(union_plan: Union) -> Result<LogicalPlan> {
290        let union_schema = Arc::new(coerce_union_schema_with_schema(
291            &union_plan.inputs,
292            &union_plan.schema,
293        )?);
294        let new_inputs = union_plan
295            .inputs
296            .into_iter()
297            .map(|p| {
298                let plan =
299                    coerce_plan_expr_for_schema(Arc::unwrap_or_clone(p), &union_schema)?;
300                match plan {
301                    LogicalPlan::Projection(Projection { expr, input, .. }) => {
302                        Ok(Arc::new(project_with_column_index(
303                            expr,
304                            input,
305                            Arc::clone(&union_schema),
306                        )?))
307                    }
308                    other_plan => Ok(Arc::new(other_plan)),
309                }
310            })
311            .collect::<Result<Vec<_>>>()?;
312        Ok(LogicalPlan::Union(Union {
313            inputs: new_inputs,
314            schema: union_schema,
315        }))
316    }
317
318    /// Coerce the fetch and skip expression to Int64 type.
319    fn coerce_limit(limit: Limit) -> Result<LogicalPlan> {
320        fn coerce_limit_expr(
321            expr: Expr,
322            schema: &DFSchema,
323            expr_name: &str,
324        ) -> Result<Expr> {
325            let dt = expr.get_type(schema)?;
326            if dt.is_integer() || dt.is_null() {
327                expr.cast_to(&DataType::Int64, schema)
328            } else {
329                plan_err!("Expected {expr_name} to be an integer or null, but got {dt}")
330            }
331        }
332
333        let empty_schema = DFSchema::empty();
334        let new_fetch = limit
335            .fetch
336            .map(|expr| coerce_limit_expr(*expr, &empty_schema, "LIMIT"))
337            .transpose()?;
338        let new_skip = limit
339            .skip
340            .map(|expr| coerce_limit_expr(*expr, &empty_schema, "OFFSET"))
341            .transpose()?;
342        Ok(LogicalPlan::Limit(Limit {
343            input: limit.input,
344            fetch: new_fetch.map(Box::new),
345            skip: new_skip.map(Box::new),
346        }))
347    }
348
349    fn coerce_predicate(&self, expr: Expr, description: &str) -> Result<Expr> {
350        let expr_type = expr.get_type(self.schema)?;
351        match expr_type {
352            DataType::Boolean => Ok(expr),
353            DataType::Null => expr.cast_to(&DataType::Boolean, self.schema),
354            other => {
355                plan_err!("{description} must be boolean type, but got {other:?}")
356            }
357        }
358    }
359
360    fn coerce_binary_op(
361        &self,
362        left: Expr,
363        left_schema: &DFSchema,
364        op: Operator,
365        right: Expr,
366        right_schema: &DFSchema,
367    ) -> Result<(Expr, Expr)> {
368        let left_data_type = left.get_type(left_schema)?;
369        let right_data_type = right.get_type(right_schema)?;
370        let (left_type, right_type) =
371            BinaryTypeCoercer::new(&left_data_type, &op, &right_data_type)
372                .get_input_types()?;
373        let left_cast_ok = can_cast_types(&left_data_type, &left_type);
374        let right_cast_ok = can_cast_types(&right_data_type, &right_type);
375
376        // handle special cases for
377        // * Date +/- int => Date
378        // * Date + time => Timestamp
379        let left_expr = if !left_cast_ok {
380            Self::coerce_date_time_math_op(
381                left,
382                &op,
383                &left_data_type,
384                &left_type,
385                &right_type,
386            )?
387        } else {
388            left.cast_to(&left_type, left_schema)?
389        };
390
391        let right_expr = if !right_cast_ok {
392            Self::coerce_date_time_math_op(
393                right,
394                &op,
395                &right_data_type,
396                &right_type,
397                &left_type,
398            )?
399        } else {
400            right.cast_to(&right_type, right_schema)?
401        };
402
403        Ok((left_expr, right_expr))
404    }
405
406    fn coerce_date_time_math_op(
407        expr: Expr,
408        op: &Operator,
409        left_current_type: &DataType,
410        left_target_type: &DataType,
411        right_target_type: &DataType,
412    ) -> Result<Expr, DataFusionError> {
413        use DataType::*;
414
415        fn cast(expr: Expr, target_type: DataType) -> Expr {
416            Expr::Cast(Cast::new(Box::new(expr), target_type))
417        }
418
419        fn time_to_nanos(
420            expr: Expr,
421            expr_type: &DataType,
422        ) -> Result<Expr, DataFusionError> {
423            let expr = match expr_type {
424                Time32(TimeUnit::Second) => {
425                    cast(cast(expr, Int32), Int64)
426                        * lit(ScalarValue::Int64(Some(1_000_000_000)))
427                }
428                Time32(TimeUnit::Millisecond) => {
429                    cast(cast(expr, Int32), Int64)
430                        * lit(ScalarValue::Int64(Some(1_000_000)))
431                }
432                Time64(TimeUnit::Microsecond) => {
433                    cast(expr, Int64) * lit(ScalarValue::Int64(Some(1_000)))
434                }
435                Time64(TimeUnit::Nanosecond) => cast(expr, Int64),
436                t => return internal_err!("Unexpected time data type {t}"),
437            };
438
439            Ok(expr)
440        }
441
442        let e = match (
443            &op,
444            &left_current_type,
445            &left_target_type,
446            &right_target_type,
447        ) {
448            // int +/- date => date
449            (
450                Operator::Plus | Operator::Minus,
451                Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64,
452                Interval(IntervalUnit::MonthDayNano),
453                Date32 | Date64,
454            ) => {
455                // cast to i64 first
456                let expr = match *left_current_type {
457                    Int64 => expr,
458                    _ => cast(expr, Int64),
459                };
460                // next, multiply by 86400 to get seconds
461                let expr = expr * lit(ScalarValue::from(SECONDS_IN_DAY));
462                // cast to duration
463                let expr = cast(expr, Duration(TimeUnit::Second));
464                // finally cast to interval
465                cast(expr, Interval(IntervalUnit::MonthDayNano))
466            }
467            // These might seem to be a bit convoluted, however for arrow to do date + time arithmetic
468            // date must be cast to Timestamp(Nanosecond) and time cast to Duration(Nanosecond)
469            // (they must be the same timeunit).
470            //
471            // For Time32/64 we first need to cast to an Int64, convert that to nanoseconds based
472            // on the time unit, then cast that to duration.
473            //
474            // Time + date -> timestamp or
475            (
476                Operator::Plus | Operator::Minus,
477                Time32(_) | Time64(_),
478                Duration(TimeUnit::Nanosecond),
479                Timestamp(TimeUnit::Nanosecond, None),
480            ) => {
481                // cast to int64, convert to nanoseconds
482                let expr = time_to_nanos(expr, left_current_type)?;
483                // cast to duration
484                cast(expr, Duration(TimeUnit::Nanosecond))
485            }
486            // Similar to above, for arrow to do time - time we need to convert to an interval.
487            // To do that we first need to cast to an Int64, convert that to nanoseconds based
488            // on the time unit, then cast that to duration, then finally cast to an interval.
489            //
490            // Time - time -> timestamp
491            (
492                Operator::Plus | Operator::Minus,
493                Time32(_) | Time64(_),
494                Interval(IntervalUnit::MonthDayNano),
495                Interval(IntervalUnit::MonthDayNano),
496            ) => {
497                // cast to int64, convert to nanoseconds
498                let expr = time_to_nanos(expr, left_current_type)?;
499                // cast to duration
500                let expr = cast(expr, Duration(TimeUnit::Nanosecond));
501                // finally cast to interval
502                cast(expr, Interval(IntervalUnit::MonthDayNano))
503            }
504            _ => {
505                return plan_err!(
506                    "Cannot automatically convert {left_current_type} to {left_target_type}"
507                );
508            }
509        };
510
511        Ok(e)
512    }
513
514    /// Coerce the value and pattern expressions of a string pattern matching
515    /// expression (`LIKE`, `ILIKE` or `SIMILAR TO`) to a common type using
516    /// the provided coercion rules. `LIKE` can preserve a dictionary-encoded
517    /// value expression, while regex array kernels require both operands to
518    /// have the same physical string type.
519    fn coerce_like_operands(
520        &self,
521        expr: Expr,
522        pattern: Expr,
523        coercion: fn(&DataType, &DataType) -> Option<DataType>,
524        op_name: &str,
525        preserve_utf8_dictionary: bool,
526    ) -> Result<(Box<Expr>, Box<Expr>)> {
527        let left_type = expr.get_type(self.schema)?;
528        let right_type = pattern.get_type(self.schema)?;
529        let coerced_type = coercion(&left_type, &right_type).ok_or_else(|| {
530            plan_datafusion_err!(
531                "There isn't a common type to coerce {left_type} and {right_type} in {op_name} expression"
532            )
533        })?;
534        let expr = match left_type {
535            DataType::Dictionary(_, inner)
536                if preserve_utf8_dictionary && *inner == DataType::Utf8 =>
537            {
538                Box::new(expr)
539            }
540            _ => Box::new(expr.cast_to(&coerced_type, self.schema)?),
541        };
542        let pattern = Box::new(pattern.cast_to(&coerced_type, self.schema)?);
543        Ok((expr, pattern))
544    }
545}
546
547impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
548    type Node = Expr;
549
550    fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
551        match expr {
552            Expr::Unnest(_) => not_impl_err!(
553                "Unnest should be rewritten to LogicalPlan::Unnest before type coercion"
554            ),
555            Expr::ScalarSubquery(Subquery {
556                subquery,
557                outer_ref_columns,
558                spans,
559            }) => {
560                let new_plan =
561                    analyze_internal(self.schema, Arc::unwrap_or_clone(subquery))?.data;
562                Ok(Transformed::yes(Expr::ScalarSubquery(Subquery {
563                    subquery: Arc::new(new_plan),
564                    outer_ref_columns,
565                    spans,
566                })))
567            }
568            Expr::Exists(Exists { subquery, negated }) => {
569                let new_plan = analyze_internal(
570                    self.schema,
571                    Arc::unwrap_or_clone(subquery.subquery),
572                )?
573                .data;
574                Ok(Transformed::yes(Expr::Exists(Exists {
575                    subquery: Subquery {
576                        subquery: Arc::new(new_plan),
577                        outer_ref_columns: subquery.outer_ref_columns,
578                        spans: subquery.spans,
579                    },
580                    negated,
581                })))
582            }
583            Expr::InSubquery(InSubquery {
584                expr,
585                subquery,
586                negated,
587            }) => {
588                let new_plan = analyze_internal(
589                    self.schema,
590                    Arc::unwrap_or_clone(subquery.subquery),
591                )?
592                .data;
593                let expr_type = expr.get_type(self.schema)?;
594                let subquery_type = new_plan.schema().field(0).data_type();
595                let common_type = comparison_coercion(&expr_type, subquery_type).ok_or(
596                    plan_datafusion_err!(
597                    "expr type {expr_type} can't cast to {subquery_type} in InSubquery"
598                ),
599                )?;
600                let new_subquery = Subquery {
601                    subquery: Arc::new(new_plan),
602                    outer_ref_columns: subquery.outer_ref_columns,
603                    spans: subquery.spans,
604                };
605                Ok(Transformed::yes(Expr::InSubquery(InSubquery::new(
606                    Box::new(expr.cast_to(&common_type, self.schema)?),
607                    cast_subquery(new_subquery, &common_type)?,
608                    negated,
609                ))))
610            }
611            Expr::SetComparison(SetComparison {
612                expr,
613                subquery,
614                op,
615                quantifier,
616            }) => {
617                let new_plan = analyze_internal(
618                    self.schema,
619                    Arc::unwrap_or_clone(subquery.subquery),
620                )?
621                .data;
622                let expr_type = expr.get_type(self.schema)?;
623                let subquery_type = new_plan.schema().field(0).data_type();
624                if (expr_type.is_numeric() && subquery_type.is_string())
625                    || (subquery_type.is_numeric() && expr_type.is_string())
626                {
627                    return plan_err!(
628                        "expr type {expr_type} can't cast to {subquery_type} in SetComparison"
629                    );
630                }
631                let common_type = comparison_coercion(&expr_type, subquery_type).ok_or(
632                    plan_datafusion_err!(
633                        "expr type {expr_type} can't cast to {subquery_type} in SetComparison"
634                    ),
635                )?;
636                let new_subquery = Subquery {
637                    subquery: Arc::new(new_plan),
638                    outer_ref_columns: subquery.outer_ref_columns,
639                    spans: subquery.spans,
640                };
641                Ok(Transformed::yes(Expr::SetComparison(SetComparison::new(
642                    Box::new(expr.cast_to(&common_type, self.schema)?),
643                    cast_subquery(new_subquery, &common_type)?,
644                    op,
645                    quantifier,
646                ))))
647            }
648            Expr::Not(expr) => Ok(Transformed::yes(not(get_casted_expr_for_bool_op(
649                *expr,
650                self.schema,
651            )?))),
652            Expr::IsTrue(expr) => Ok(Transformed::yes(is_true(
653                get_casted_expr_for_bool_op(*expr, self.schema)?,
654            ))),
655            Expr::IsNotTrue(expr) => Ok(Transformed::yes(is_not_true(
656                get_casted_expr_for_bool_op(*expr, self.schema)?,
657            ))),
658            Expr::IsFalse(expr) => Ok(Transformed::yes(is_false(
659                get_casted_expr_for_bool_op(*expr, self.schema)?,
660            ))),
661            Expr::IsNotFalse(expr) => Ok(Transformed::yes(is_not_false(
662                get_casted_expr_for_bool_op(*expr, self.schema)?,
663            ))),
664            Expr::IsUnknown(expr) => Ok(Transformed::yes(is_unknown(
665                get_casted_expr_for_bool_op(*expr, self.schema)?,
666            ))),
667            Expr::IsNotUnknown(expr) => Ok(Transformed::yes(is_not_unknown(
668                get_casted_expr_for_bool_op(*expr, self.schema)?,
669            ))),
670            Expr::Negative(expr) => {
671                let data_type = expr.get_type(self.schema)?;
672                if data_type.is_null()
673                    || is_signed_numeric(&data_type)
674                    || is_interval(&data_type)
675                    || is_timestamp(&data_type)
676                {
677                    Ok(Transformed::no(Expr::Negative(expr)))
678                } else {
679                    plan_err!(
680                        "Negation only supports numeric, interval and timestamp types"
681                    )
682                }
683            }
684            Expr::Like(Like {
685                negated,
686                expr,
687                pattern,
688                escape_char,
689                case_insensitive,
690            }) => {
691                let op_name = if case_insensitive { "ILIKE" } else { "LIKE" };
692                let (expr, pattern) = self.coerce_like_operands(
693                    *expr,
694                    *pattern,
695                    like_coercion,
696                    op_name,
697                    true,
698                )?;
699                Ok(Transformed::yes(Expr::Like(Like::new(
700                    negated,
701                    expr,
702                    pattern,
703                    escape_char,
704                    case_insensitive,
705                ))))
706            }
707            Expr::SimilarTo(Like {
708                negated,
709                expr,
710                pattern,
711                escape_char,
712                case_insensitive,
713            }) => {
714                // `SIMILAR TO` is planned as a regex operator, so its operands
715                // must be coerced to a common string type using the same
716                // coercion rules as the physical regex operators. Otherwise
717                // mismatched operand types panic during execution.
718                let (expr, pattern) = self.coerce_like_operands(
719                    *expr,
720                    *pattern,
721                    regex_coercion,
722                    "SIMILAR TO",
723                    false,
724                )?;
725                Ok(Transformed::yes(Expr::SimilarTo(Like::new(
726                    negated,
727                    expr,
728                    pattern,
729                    escape_char,
730                    case_insensitive,
731                ))))
732            }
733            Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
734                let (left, right) =
735                    self.coerce_binary_op(*left, self.schema, op, *right, self.schema)?;
736                Ok(Transformed::yes(Expr::BinaryExpr(BinaryExpr::new(
737                    Box::new(left),
738                    op,
739                    Box::new(right),
740                ))))
741            }
742            Expr::Between(Between {
743                expr,
744                negated,
745                low,
746                high,
747            }) => {
748                let expr_type = expr.get_type(self.schema)?;
749                let low_type = low.get_type(self.schema)?;
750                let low_coerced_type = comparison_coercion(&expr_type, &low_type)
751                    .ok_or_else(|| {
752                        internal_datafusion_err!(
753                            "Failed to coerce types {expr_type} and {low_type} in BETWEEN expression"
754                        )
755                    })?;
756                let high_type = high.get_type(self.schema)?;
757                let high_coerced_type = comparison_coercion(&expr_type, &high_type)
758                    .ok_or_else(|| {
759                        internal_datafusion_err!(
760                            "Failed to coerce types {expr_type} and {high_type} in BETWEEN expression"
761                        )
762                    })?;
763                let coercion_type =
764                    comparison_coercion(&low_coerced_type, &high_coerced_type)
765                        .ok_or_else(|| {
766                            internal_datafusion_err!(
767                                "Failed to coerce types {expr_type} and {high_type} in BETWEEN expression"
768                            )
769                        })?;
770                Ok(Transformed::yes(Expr::Between(Between::new(
771                    Box::new(expr.cast_to(&coercion_type, self.schema)?),
772                    negated,
773                    Box::new(low.cast_to(&coercion_type, self.schema)?),
774                    Box::new(high.cast_to(&coercion_type, self.schema)?),
775                ))))
776            }
777            Expr::InList(InList {
778                expr,
779                list,
780                negated,
781            }) => {
782                let expr_data_type = expr.get_type(self.schema)?;
783                let list_data_types = list
784                    .iter()
785                    .map(|list_expr| list_expr.get_type(self.schema))
786                    .collect::<Result<Vec<_>>>()?;
787                let result_type =
788                    get_coerce_type_for_list(&expr_data_type, &list_data_types);
789                match result_type {
790                    None => plan_err!(
791                        "Can not find compatible types to compare {expr_data_type} with [{}]",
792                        list_data_types.iter().join(", ")
793                    ),
794                    Some(coerced_type) => {
795                        // find the coerced type
796                        let cast_expr = expr.cast_to(&coerced_type, self.schema)?;
797                        let cast_list_expr = list
798                            .into_iter()
799                            .map(|list_expr| {
800                                list_expr.cast_to(&coerced_type, self.schema)
801                            })
802                            .collect::<Result<Vec<_>>>()?;
803                        Ok(Transformed::yes(Expr::InList(InList::new(
804                            Box::new(cast_expr),
805                            cast_list_expr,
806                            negated,
807                        ))))
808                    }
809                }
810            }
811            Expr::Case(case) => {
812                let case = coerce_case_expression(case, self.schema)?;
813                Ok(Transformed::yes(Expr::Case(case)))
814            }
815            Expr::ScalarFunction(ScalarFunction { func, args }) => {
816                let new_expr =
817                    coerce_arguments_for_signature(args, self.schema, func.as_ref())?;
818                Ok(Transformed::yes(Expr::ScalarFunction(
819                    ScalarFunction::new_udf(func, new_expr),
820                )))
821            }
822            Expr::AggregateFunction(expr::AggregateFunction {
823                func,
824                params:
825                    AggregateFunctionParams {
826                        args,
827                        distinct,
828                        filter,
829                        order_by,
830                        null_treatment,
831                    },
832            }) => {
833                let new_expr =
834                    coerce_arguments_for_signature(args, self.schema, func.as_ref())?;
835
836                let filter = filter
837                    .map(|filter| filter.cast_to(&DataType::Boolean, self.schema))
838                    .transpose()?
839                    .map(Box::new);
840
841                Ok(Transformed::yes(Expr::AggregateFunction(
842                    expr::AggregateFunction::new_udf(
843                        func,
844                        new_expr,
845                        distinct,
846                        filter,
847                        order_by,
848                        null_treatment,
849                    ),
850                )))
851            }
852            Expr::WindowFunction(window_fun) => {
853                let WindowFunction {
854                    fun,
855                    params:
856                        expr::WindowFunctionParams {
857                            args,
858                            partition_by,
859                            order_by,
860                            window_frame,
861                            filter,
862                            null_treatment,
863                            distinct,
864                        },
865                } = *window_fun;
866                let window_frame =
867                    coerce_window_frame(window_frame, self.schema, &order_by)?;
868
869                let args = match &fun {
870                    expr::WindowFunctionDefinition::AggregateUDF(udf) => {
871                        coerce_arguments_for_signature(args, self.schema, udf.as_ref())?
872                    }
873                    expr::WindowFunctionDefinition::WindowUDF(udf) => {
874                        coerce_arguments_for_signature(args, self.schema, udf.as_ref())?
875                    }
876                };
877
878                let filter = filter
879                    .map(|filter| filter.cast_to(&DataType::Boolean, self.schema))
880                    .transpose()?
881                    .map(Box::new);
882
883                let new_expr = Expr::from(WindowFunction {
884                    fun,
885                    params: expr::WindowFunctionParams {
886                        args,
887                        partition_by,
888                        order_by,
889                        window_frame,
890                        filter,
891                        null_treatment,
892                        distinct,
893                    },
894                });
895                Ok(Transformed::yes(new_expr))
896            }
897            Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => {
898                let current_fields = args
899                    .iter()
900                    .map(|arg| match arg {
901                        Expr::Lambda(lambda) => Ok(ValueOrLambda::Lambda(
902                            lambda.body.to_field(self.schema)?.1,
903                        )),
904                        _ => Ok(ValueOrLambda::Value(arg.to_field(self.schema)?.1)),
905                    })
906                    .collect::<Result<Vec<_>>>()?;
907
908                let new_fields = value_fields_with_higher_order_udf_and_lambdas(
909                    &current_fields,
910                    func.as_ref(),
911                )?;
912
913                let new_args = std::iter::zip(args, new_fields)
914                    .map(|(arg, new_field)| match (&arg, new_field) {
915                        (Expr::Lambda(_lambda), ValueOrLambda::Lambda(_)) => Ok(arg),
916                        (Expr::Lambda(_lambda), ValueOrLambda::Value(_)) => internal_err!("value_fields_with_higher_order_udf returned a value for a lambda argument"),
917                        (_, ValueOrLambda::Value(new_field)) => arg.cast_to(new_field.data_type(), self.schema),
918                        (_, ValueOrLambda::Lambda(_)) => internal_err!("value_fields_with_higher_order_udf returned a lambda for a value argument"),
919                    })
920                    .collect::<Result<_>>()?;
921
922                Ok(Transformed::yes(Expr::HigherOrderFunction(
923                    HigherOrderFunction::new(func, new_args),
924                )))
925            }
926            // TODO: remove the next line after `Expr::Wildcard` is removed
927            #[expect(deprecated)]
928            Expr::Alias(_)
929            | Expr::Column(_)
930            | Expr::ScalarVariable(_, _)
931            | Expr::Literal(_, _)
932            | Expr::IsNotNull(_)
933            | Expr::IsNull(_)
934            | Expr::Cast(_)
935            | Expr::TryCast(_)
936            | Expr::Wildcard { .. }
937            | Expr::GroupingSet(_)
938            | Expr::Placeholder(_)
939            | Expr::OuterReferenceColumn(_, _)
940            | Expr::Lambda(_)
941            | Expr::LambdaVariable(_) => Ok(Transformed::no(expr)),
942        }
943    }
944}
945
946/// Transform a schema to use non-view types for Utf8View and BinaryView
947fn transform_schema_to_nonview(dfschema: &DFSchemaRef) -> Option<Result<DFSchema>> {
948    let metadata = dfschema.as_arrow().metadata.clone();
949    let mut transformed = false;
950
951    let (qualifiers, transformed_fields): (Vec<Option<TableReference>>, Vec<Arc<Field>>) =
952        dfschema
953            .iter()
954            .map(|(qualifier, field)| match field.data_type() {
955                DataType::Utf8View => {
956                    transformed = true;
957                    (
958                        qualifier.cloned() as Option<TableReference>,
959                        Arc::new(Field::new(
960                            field.name(),
961                            DataType::LargeUtf8,
962                            field.is_nullable(),
963                        )),
964                    )
965                }
966                DataType::BinaryView => {
967                    transformed = true;
968                    (
969                        qualifier.cloned() as Option<TableReference>,
970                        Arc::new(Field::new(
971                            field.name(),
972                            DataType::LargeBinary,
973                            field.is_nullable(),
974                        )),
975                    )
976                }
977                _ => (
978                    qualifier.cloned() as Option<TableReference>,
979                    Arc::clone(field),
980                ),
981            })
982            .unzip();
983
984    if !transformed {
985        return None;
986    }
987
988    let schema = Schema::new_with_metadata(transformed_fields, metadata);
989    Some(DFSchema::from_field_specific_qualified_schema(
990        qualifiers,
991        &Arc::new(schema),
992    ))
993}
994
995/// Casts the given `value` to `target_type`. Note that this function
996/// only considers `Null` or `Utf8` values.
997fn coerce_scalar(target_type: &DataType, value: &ScalarValue) -> Result<ScalarValue> {
998    match value {
999        // Coerce Utf8 values:
1000        ScalarValue::Utf8(Some(val)) => {
1001            ScalarValue::try_from_string(val.clone(), target_type)
1002        }
1003        s => {
1004            if s.is_null() {
1005                // Coerce `Null` values:
1006                ScalarValue::try_from(target_type)
1007            } else {
1008                // Values except `Utf8`/`Null` variants already have the right type
1009                // (casted before) since we convert `sqlparser` outputs to `Utf8`
1010                // for all possible cases. Therefore, we return a clone here.
1011                Ok(s.clone())
1012            }
1013        }
1014    }
1015}
1016
1017/// This function coerces `value` to `target_type` in a range-aware fashion.
1018/// If the coercion is successful, we return an `Ok` value with the result.
1019/// If the coercion fails because `target_type` is not wide enough (i.e. we
1020/// can not coerce to `target_type`, but we can to a wider type in the same
1021/// family), we return a `Null` value of this type to signal this situation.
1022/// Downstream code uses this signal to treat these values as *unbounded*.
1023fn coerce_scalar_range_aware(
1024    target_type: &DataType,
1025    value: &ScalarValue,
1026) -> Result<ScalarValue> {
1027    coerce_scalar(target_type, value).or_else(|err| {
1028        // If type coercion fails, check if the largest type in family works:
1029        if let Some(largest_type) = get_widest_type_in_family(target_type) {
1030            coerce_scalar(largest_type, value).map_or_else(
1031                |_| exec_err!("Cannot cast {value} to {target_type}"),
1032                |_| ScalarValue::try_from(target_type),
1033            )
1034        } else {
1035            Err(err)
1036        }
1037    })
1038}
1039
1040/// This function returns the widest type in the family of `given_type`.
1041/// If the given type is already the widest type, it returns `None`.
1042/// For example, if `given_type` is `Int8`, it returns `Int64`.
1043fn get_widest_type_in_family(given_type: &DataType) -> Option<&DataType> {
1044    match given_type {
1045        DataType::UInt8 | DataType::UInt16 | DataType::UInt32 => Some(&DataType::UInt64),
1046        DataType::Int8 | DataType::Int16 | DataType::Int32 => Some(&DataType::Int64),
1047        DataType::Float16 | DataType::Float32 => Some(&DataType::Float64),
1048        _ => None,
1049    }
1050}
1051
1052/// Coerces the given (window frame) `bound` to `target_type`.
1053fn coerce_frame_bound(
1054    target_type: &DataType,
1055    bound: WindowFrameBound,
1056) -> Result<WindowFrameBound> {
1057    match bound {
1058        WindowFrameBound::Preceding(v) => {
1059            coerce_scalar_range_aware(target_type, &v).map(WindowFrameBound::Preceding)
1060        }
1061        WindowFrameBound::CurrentRow => Ok(WindowFrameBound::CurrentRow),
1062        WindowFrameBound::Following(v) => {
1063            coerce_scalar_range_aware(target_type, &v).map(WindowFrameBound::Following)
1064        }
1065    }
1066}
1067
1068fn extract_window_frame_target_type(col_type: &DataType) -> Result<DataType> {
1069    if col_type.is_numeric()
1070        || col_type.is_string()
1071        || col_type.is_null()
1072        || matches!(
1073            col_type,
1074            DataType::List(_)
1075                | DataType::LargeList(_)
1076                | DataType::FixedSizeList(_, _)
1077                | DataType::Boolean
1078        )
1079    {
1080        Ok(col_type.clone())
1081    } else if is_datetime(col_type) {
1082        Ok(DataType::Interval(IntervalUnit::MonthDayNano))
1083    } else if let DataType::Dictionary(_, value_type) = col_type {
1084        extract_window_frame_target_type(value_type)
1085    } else {
1086        internal_err!("Cannot run range queries on datatype: {col_type}")
1087    }
1088}
1089
1090// Coerces the given `window_frame` to use appropriate natural types.
1091// For example, ROWS and GROUPS frames use `UInt64` during calculations.
1092fn coerce_window_frame(
1093    window_frame: WindowFrame,
1094    schema: &DFSchema,
1095    expressions: &[Sort],
1096) -> Result<WindowFrame> {
1097    let mut window_frame = window_frame;
1098    let target_type = match window_frame.units {
1099        WindowFrameUnits::Range => {
1100            let current_types = expressions
1101                .first()
1102                .map(|s| s.expr.get_type(schema))
1103                .transpose()?;
1104            if let Some(col_type) = current_types {
1105                extract_window_frame_target_type(&col_type)?
1106            } else {
1107                return internal_err!("ORDER BY column cannot be empty");
1108            }
1109        }
1110        WindowFrameUnits::Rows | WindowFrameUnits::Groups => DataType::UInt64,
1111    };
1112    window_frame.start_bound =
1113        coerce_frame_bound(&target_type, window_frame.start_bound)?;
1114    window_frame.end_bound = coerce_frame_bound(&target_type, window_frame.end_bound)?;
1115    Ok(window_frame)
1116}
1117
1118// Support the `IsTrue` `IsNotTrue` `IsFalse` `IsNotFalse` type coercion.
1119// The above op will be rewrite to the binary op when creating the physical op.
1120fn get_casted_expr_for_bool_op(expr: Expr, schema: &DFSchema) -> Result<Expr> {
1121    let left_type = expr.get_type(schema)?;
1122    BinaryTypeCoercer::new(&left_type, &Operator::IsDistinctFrom, &DataType::Boolean)
1123        .get_input_types()?;
1124    expr.cast_to(&DataType::Boolean, schema)
1125}
1126
1127/// Returns `expressions` coerced to types compatible with
1128/// `signature`, if possible.
1129///
1130/// See the module level documentation for more detail on coercion.
1131fn coerce_arguments_for_signature<F: UDFCoercionExt>(
1132    expressions: Vec<Expr>,
1133    schema: &DFSchema,
1134    func: &F,
1135) -> Result<Vec<Expr>> {
1136    let current_fields = expressions
1137        .iter()
1138        .map(|e| e.to_field(schema).map(|(_, f)| f))
1139        .collect::<Result<Vec<_>>>()?;
1140
1141    let coerced_types = fields_with_udf(&current_fields, func)?
1142        .into_iter()
1143        .map(|f| f.data_type().clone())
1144        .collect::<Vec<_>>();
1145
1146    expressions
1147        .into_iter()
1148        .enumerate()
1149        .map(|(i, expr)| expr.cast_to(&coerced_types[i], schema))
1150        .collect()
1151}
1152
1153fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result<Case> {
1154    // Given expressions like:
1155    //
1156    // CASE a1
1157    //   WHEN a2 THEN b1
1158    //   WHEN a3 THEN b2
1159    //   ELSE b3
1160    // END
1161    //
1162    // or:
1163    //
1164    // CASE
1165    //   WHEN x1 THEN b1
1166    //   WHEN x2 THEN b2
1167    //   ELSE b3
1168    // END
1169    //
1170    // Then all aN (a1, a2, a3) must be converted to a common data type in the first example
1171    // (case-when expression coercion)
1172    //
1173    // All xN (x1, x2) must be converted to a boolean data type in the second example
1174    // (when-boolean expression coercion)
1175    //
1176    // And all bN (b1, b2, b3) must be converted to a common data type in both examples
1177    // (then-else expression coercion)
1178    //
1179    // If any fail to find and cast to a common/specific data type, will return error
1180    //
1181    // Note that case-when and when-boolean expression coercions are mutually exclusive
1182    // Only one or the other can occur for a case expression, whilst then-else expression coercion will always occur
1183
1184    // prepare types
1185    let case_type = case
1186        .expr
1187        .as_ref()
1188        .map(|expr| expr.get_type(schema))
1189        .transpose()?;
1190    let then_types = case
1191        .when_then_expr
1192        .iter()
1193        .map(|(_when, then)| then.get_type(schema))
1194        .collect::<Result<Vec<_>>>()?;
1195    let else_type = case
1196        .else_expr
1197        .as_ref()
1198        .map(|expr| expr.get_type(schema))
1199        .transpose()?;
1200
1201    // find common coercible types
1202    let case_when_coerce_type = case_type
1203        .as_ref()
1204        .map(|case_type| {
1205            let when_types = case
1206                .when_then_expr
1207                .iter()
1208                .map(|(when, _then)| when.get_type(schema))
1209                .collect::<Result<Vec<_>>>()?;
1210            let coerced_type = get_coerce_type_for_case_when(&when_types, case_type);
1211            coerced_type.ok_or_else(|| {
1212                plan_datafusion_err!(
1213                    "Failed to coerce case ({case_type}) and when ({}) \
1214                     to common types in CASE WHEN expression",
1215                    when_types.iter().join(", ")
1216                )
1217            })
1218        })
1219        .transpose()?;
1220    let then_else_coerce_type =
1221        get_coerce_type_for_case_expression(&then_types, else_type.as_ref()).ok_or_else(
1222            || {
1223                if let Some(else_type) = else_type {
1224                    plan_datafusion_err!(
1225                        "Failed to coerce then ({}) and else ({else_type}) \
1226                         to common types in CASE WHEN expression",
1227                        then_types.iter().join(", ")
1228                    )
1229                } else {
1230                    plan_datafusion_err!(
1231                        "Failed to coerce then ({}) and else (None) \
1232                         to common types in CASE WHEN expression",
1233                        then_types.iter().join(", ")
1234                    )
1235                }
1236            },
1237        )?;
1238
1239    // do cast if found common coercible types
1240    let case_expr = case
1241        .expr
1242        .zip(case_when_coerce_type.as_ref())
1243        .map(|(case_expr, coercible_type)| case_expr.cast_to(coercible_type, schema))
1244        .transpose()?
1245        .map(Box::new);
1246    let when_then = case
1247        .when_then_expr
1248        .into_iter()
1249        .map(|(when, then)| {
1250            let when_type = case_when_coerce_type.as_ref().unwrap_or(&DataType::Boolean);
1251            let when = when.cast_to(when_type, schema).map_err(|e| {
1252                DataFusionError::Context(
1253                    format!(
1254                        "WHEN expressions in CASE couldn't be \
1255                         converted to common type ({when_type})"
1256                    ),
1257                    Box::new(e),
1258                )
1259            })?;
1260            let then = then.cast_to(&then_else_coerce_type, schema)?;
1261            Ok((Box::new(when), Box::new(then)))
1262        })
1263        .collect::<Result<Vec<_>>>()?;
1264    let else_expr = case
1265        .else_expr
1266        .map(|expr| expr.cast_to(&then_else_coerce_type, schema))
1267        .transpose()?
1268        .map(Box::new);
1269
1270    Ok(Case::new(case_expr, when_then, else_expr))
1271}
1272
1273/// Get a common schema that is compatible with all inputs of UNION.
1274///
1275/// This method presumes that the wildcard expansion is unneeded, or has already
1276/// been applied.
1277///
1278/// ## Schema and Field Handling in Union Coercion
1279///
1280/// **Processing order**: The function starts with the base schema (first input) and then
1281/// processes remaining inputs sequentially, with later inputs taking precedence in merging.
1282///
1283/// **Schema-level metadata merging**: Later schemas take precedence for duplicate keys.
1284///
1285/// **Field-level metadata merging**: Later fields take precedence for duplicate metadata keys.
1286///
1287/// **Type coercion precedence**: The coerced type is determined by iteratively applying
1288/// `type_union_coercion()` between the accumulated type and each new input's type. The
1289/// result depends on type coercion rules, not input order.
1290///
1291/// **Nullability merging**: Nullability is accumulated using logical OR (`||`).
1292/// Once any input field is nullable, the result field becomes nullable permanently.
1293/// Later inputs can make a field nullable but cannot make it non-nullable.
1294///
1295/// **Field precedence**: Field names come from the first (base) schema, but the field properties
1296/// (nullability and field-level metadata) have later schemas taking precedence.
1297///
1298/// **Example**:
1299/// ```sql
1300/// SELECT a, b FROM table1  -- a: Int32, metadata {"source": "t1"}, nullable=false
1301/// UNION
1302/// SELECT a, b FROM table2  -- a: Int64, metadata {"source": "t2"}, nullable=true
1303/// UNION
1304/// SELECT a, b FROM table3  -- a: Int32, metadata {"encoding": "utf8"}, nullable=false
1305/// -- Result:
1306/// -- a: Int64 (from type coercion), nullable=true (from table2),
1307/// -- metadata: {"source": "t2", "encoding": "utf8"} (later inputs take precedence)
1308/// ```
1309///
1310/// **Precedence Summary**:
1311/// - **Datatypes**: Determined by `type_union_coercion()` rules, not input order
1312/// - **Nullability**: Later inputs can add nullability but cannot remove it (logical OR)
1313/// - **Metadata**: Later inputs take precedence for same keys (HashMap::extend semantics)
1314pub fn coerce_union_schema(inputs: &[Arc<LogicalPlan>]) -> Result<DFSchema> {
1315    coerce_union_schema_with_schema(&inputs[1..], inputs[0].schema())
1316}
1317fn coerce_union_schema_with_schema(
1318    inputs: &[Arc<LogicalPlan>],
1319    base_schema: &DFSchemaRef,
1320) -> Result<DFSchema> {
1321    let mut union_datatypes = base_schema
1322        .fields()
1323        .iter()
1324        .map(|f| f.data_type().clone())
1325        .collect::<Vec<_>>();
1326    let mut union_nullabilities = base_schema
1327        .fields()
1328        .iter()
1329        .map(|f| f.is_nullable())
1330        .collect::<Vec<_>>();
1331    let mut union_field_meta = base_schema
1332        .fields()
1333        .iter()
1334        .map(|f| f.metadata().clone())
1335        .collect::<Vec<_>>();
1336
1337    let mut metadata = base_schema.metadata().clone();
1338
1339    for (i, plan) in inputs.iter().enumerate() {
1340        let plan_schema = plan.schema();
1341        metadata.extend(plan_schema.metadata().clone());
1342
1343        if plan_schema.fields().len() != base_schema.fields().len() {
1344            return plan_err!(
1345                "Union schemas have different number of fields: \
1346                query 1 has {} fields whereas query {} has {} fields",
1347                base_schema.fields().len(),
1348                i + 1,
1349                plan_schema.fields().len()
1350            );
1351        }
1352
1353        // coerce data type and nullability for each field
1354        for (union_datatype, union_nullable, union_field_map, plan_field) in izip!(
1355            union_datatypes.iter_mut(),
1356            union_nullabilities.iter_mut(),
1357            union_field_meta.iter_mut(),
1358            plan_schema.fields().iter()
1359        ) {
1360            let coerced_type =
1361                type_union_coercion(union_datatype, plan_field.data_type()).ok_or_else(
1362                    || {
1363                        plan_datafusion_err!(
1364                            "Incompatible inputs for Union: Previous inputs were \
1365                            of type {}, but got incompatible type {} on column '{}'",
1366                            union_datatype,
1367                            plan_field.data_type(),
1368                            plan_field.name()
1369                        )
1370                    },
1371                )?;
1372
1373            *union_datatype = coerced_type;
1374            *union_nullable = *union_nullable || plan_field.is_nullable();
1375            union_field_map.extend(plan_field.metadata().clone());
1376        }
1377    }
1378    let union_qualified_fields = izip!(
1379        base_schema.fields(),
1380        union_datatypes.into_iter(),
1381        union_nullabilities,
1382        union_field_meta.into_iter()
1383    )
1384    .map(|(field, datatype, nullable, metadata)| {
1385        let mut field = Field::new(field.name().clone(), datatype, nullable);
1386        field.set_metadata(metadata);
1387        (None, field.into())
1388    })
1389    .collect::<Vec<_>>();
1390
1391    DFSchema::new_with_metadata(union_qualified_fields, metadata)
1392}
1393
1394/// See `<https://github.com/apache/datafusion/pull/2108>`
1395fn project_with_column_index(
1396    expr: Vec<Expr>,
1397    input: Arc<LogicalPlan>,
1398    schema: DFSchemaRef,
1399) -> Result<LogicalPlan> {
1400    let alias_expr = expr
1401        .into_iter()
1402        .enumerate()
1403        .map(|(i, e)| match e {
1404            Expr::Alias(Alias { ref name, .. }) if name != schema.field(i).name() => {
1405                Ok(e.unalias().alias(schema.field(i).name()))
1406            }
1407            Expr::Column(Column {
1408                relation: _,
1409                ref name,
1410                spans: _,
1411            }) if name != schema.field(i).name() => Ok(e.alias(schema.field(i).name())),
1412            Expr::Alias { .. } | Expr::Column { .. } => Ok(e),
1413            #[expect(deprecated)]
1414            Expr::Wildcard { .. } => {
1415                plan_err!("Wildcard should be expanded before type coercion")
1416            }
1417            _ => Ok(e.alias(schema.field(i).name())),
1418        })
1419        .collect::<Result<Vec<_>>>()?;
1420
1421    Projection::try_new_with_schema(alias_expr, input, schema)
1422        .map(LogicalPlan::Projection)
1423}
1424
1425#[cfg(test)]
1426mod test {
1427
1428    use std::sync::Arc;
1429
1430    use arrow::datatypes::DataType::Utf8;
1431    use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, TimeUnit};
1432    use insta::assert_snapshot;
1433
1434    use crate::analyzer::Analyzer;
1435    use crate::analyzer::type_coercion::{
1436        TypeCoercion, TypeCoercionRewriter, coerce_case_expression,
1437    };
1438    use crate::assert_analyzed_plan_with_config_eq_snapshot;
1439    use datafusion_common::config::ConfigOptions;
1440    use datafusion_common::tree_node::{TransformedResult, TreeNode};
1441    use datafusion_common::{
1442        DFSchema, DFSchemaRef, Result, ScalarValue, Spans, TableReference,
1443    };
1444    use datafusion_expr::expr::{self, InSubquery, Like, ScalarFunction};
1445    use datafusion_expr::logical_plan::{EmptyRelation, Projection, Sort};
1446    use datafusion_expr::test::function_stub::avg_udaf;
1447    use datafusion_expr::{
1448        AccumulatorFactoryFunction, AggregateUDF, BinaryExpr, Case, ColumnarValue, Expr,
1449        ExprSchemable, Filter, LogicalPlan, Operator, ScalarFunctionArgs, ScalarUDF,
1450        ScalarUDFImpl, Signature, SimpleAggregateUDF, Subquery, Union, Volatility, cast,
1451        col, create_udaf, is_true, lit,
1452    };
1453    use datafusion_functions_aggregate::average::AvgAccumulator;
1454
1455    fn empty() -> Arc<LogicalPlan> {
1456        Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
1457            produce_one_row: false,
1458            schema: Arc::new(DFSchema::empty()),
1459        }))
1460    }
1461
1462    fn empty_with_type(data_type: DataType) -> Arc<LogicalPlan> {
1463        Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
1464            produce_one_row: false,
1465            schema: Arc::new(
1466                DFSchema::from_unqualified_fields(
1467                    vec![Field::new("a", data_type, true)].into(),
1468                    std::collections::HashMap::new(),
1469                )
1470                .unwrap(),
1471            ),
1472        }))
1473    }
1474
1475    macro_rules! assert_analyzed_plan_eq {
1476        (
1477            $plan: expr,
1478            @ $expected: literal $(,)?
1479        ) => {{
1480            let options = ConfigOptions::default();
1481            let rule = Arc::new(TypeCoercion::new());
1482            assert_analyzed_plan_with_config_eq_snapshot!(
1483                options,
1484                rule,
1485                $plan,
1486                @ $expected,
1487            )
1488            }};
1489    }
1490
1491    macro_rules! coerce_on_output_if_viewtype {
1492        (
1493            $is_viewtype: expr,
1494            $plan: expr,
1495            @ $expected: literal $(,)?
1496        ) => {{
1497            let mut options = ConfigOptions::default();
1498            // coerce on output
1499            if $is_viewtype {options.optimizer.expand_views_at_output = true;}
1500            let rule = Arc::new(TypeCoercion::new());
1501
1502            assert_analyzed_plan_with_config_eq_snapshot!(
1503                options,
1504                rule,
1505                $plan,
1506                @ $expected,
1507            )
1508        }};
1509    }
1510
1511    fn assert_type_coercion_error(
1512        plan: LogicalPlan,
1513        expected_substr: &str,
1514    ) -> Result<()> {
1515        let options = ConfigOptions::default();
1516        let analyzer = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())]);
1517
1518        match analyzer.execute_and_check(plan, &options, |_, _| {}) {
1519            Ok(succeeded_plan) => {
1520                panic!(
1521                    "Expected a type coercion error, but analysis succeeded: \n{succeeded_plan:#?}"
1522                );
1523            }
1524            Err(e) => {
1525                let msg = e.to_string();
1526                assert!(
1527                    msg.contains(expected_substr),
1528                    "Error did not contain expected substring.\n  expected to find: `{expected_substr}`\n  actual error: `{msg}`"
1529                );
1530            }
1531        }
1532
1533        Ok(())
1534    }
1535
1536    #[test]
1537    fn simple_case() -> Result<()> {
1538        let expr = col("a").lt(lit(2_u32));
1539        let empty = empty_with_type(DataType::Float64);
1540        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
1541
1542        assert_analyzed_plan_eq!(
1543            plan,
1544            @r"
1545        Projection: a < CAST(UInt32(2) AS Float64)
1546          EmptyRelation: rows=0
1547        "
1548        )
1549    }
1550
1551    #[test]
1552    fn negative_expr_wrapped_by_is_null_errors() -> Result<()> {
1553        let predicate = Expr::IsNull(Box::new(Expr::Negative(Box::new(lit("a")))));
1554        let plan = LogicalPlan::Filter(Filter::try_new(predicate, empty())?);
1555
1556        assert_type_coercion_error(
1557            plan,
1558            "Negation only supports numeric, interval and timestamp types",
1559        )
1560    }
1561
1562    #[test]
1563    fn test_coerce_union() -> Result<()> {
1564        let left_plan = Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
1565            produce_one_row: false,
1566            schema: Arc::new(
1567                DFSchema::try_from_qualified_schema(
1568                    TableReference::full("datafusion", "test", "foo"),
1569                    &Schema::new(vec![Field::new("a", DataType::Int32, false)]),
1570                )
1571                .unwrap(),
1572            ),
1573        }));
1574        let right_plan = Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
1575            produce_one_row: false,
1576            schema: Arc::new(
1577                DFSchema::try_from_qualified_schema(
1578                    TableReference::full("datafusion", "test", "foo"),
1579                    &Schema::new(vec![Field::new("a", DataType::Int64, false)]),
1580                )
1581                .unwrap(),
1582            ),
1583        }));
1584        let union = LogicalPlan::Union(Union::try_new_with_loose_types(vec![
1585            left_plan, right_plan,
1586        ])?);
1587        let analyzed_union = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())])
1588            .execute_and_check(union, &ConfigOptions::default(), |_, _| {})?;
1589        let top_level_plan = LogicalPlan::Projection(Projection::try_new(
1590            vec![col("a")],
1591            Arc::new(analyzed_union),
1592        )?);
1593
1594        assert_analyzed_plan_eq!(
1595            top_level_plan,
1596            @r"
1597        Projection: a
1598          Union
1599            Projection: CAST(datafusion.test.foo.a AS Int64) AS a
1600              EmptyRelation: rows=0
1601            EmptyRelation: rows=0
1602        "
1603        )
1604    }
1605
1606    #[test]
1607    fn merge_into_resolves_and_coerces_target_and_source_columns() -> Result<()> {
1608        use datafusion_expr::dml::{
1609            MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp,
1610        };
1611        use datafusion_expr::logical_plan::table_scan;
1612        use datafusion_expr::{DmlStatement, WriteOp};
1613
1614        // Target table `target(id: UInt32)`.
1615        let target_table_name = TableReference::bare("target");
1616        let target_arrow_schema =
1617            Schema::new(vec![Field::new("id", DataType::UInt32, false)]);
1618        let target_plan =
1619            table_scan(Some(target_table_name.clone()), &target_arrow_schema, None)?
1620                .build()?;
1621        let target_source = match &target_plan {
1622            LogicalPlan::TableScan(ts) => Arc::clone(&ts.source),
1623            _ => unreachable!("table_scan() always builds a TableScan"),
1624        };
1625
1626        // Source plan `source(id: Int64)` — deliberately a different numeric
1627        // type than `target.id` so the `ON` comparison needs a CAST.
1628        let source_arrow_schema =
1629            Schema::new(vec![Field::new("id", DataType::Int64, false)]);
1630        let source_plan =
1631            table_scan(Some("source"), &source_arrow_schema, None)?.build()?;
1632
1633        // `ON target.id = source.id`. Resolving `target.id` requires the
1634        // target schema to be visible to the analyzer, which only sees
1635        // `plan.inputs()` (the source plan) by default.
1636        let on = col("target.id").eq(col("source.id"));
1637        let merge_op = MergeIntoOp {
1638            on,
1639            clauses: vec![
1640                MergeIntoClause {
1641                    kind: MergeIntoClauseKind::Matched,
1642                    predicate: None,
1643                    action: MergeIntoAction::Update(vec![(
1644                        "id".to_string(),
1645                        col("source.id"),
1646                    )]),
1647                },
1648                MergeIntoClause {
1649                    kind: MergeIntoClauseKind::NotMatched,
1650                    predicate: None,
1651                    action: MergeIntoAction::Insert {
1652                        columns: vec!["id".to_string()],
1653                        values: vec![col("source.id")],
1654                    },
1655                },
1656            ],
1657        };
1658        let plan = LogicalPlan::Dml(DmlStatement::new(
1659            target_table_name,
1660            target_source,
1661            WriteOp::MergeInto(Box::new(merge_op)),
1662            Arc::new(source_plan),
1663        ));
1664
1665        let analyzed = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())])
1666            .execute_and_check(plan, &ConfigOptions::default(), |_, _| {})?;
1667        let LogicalPlan::Dml(dml) = analyzed else {
1668            panic!("expected Dml");
1669        };
1670        let WriteOp::MergeInto(merge_op) = dml.op else {
1671            panic!("expected MergeInto");
1672        };
1673        assert_eq!(
1674            merge_op.on.to_string(),
1675            "CAST(target.id AS Int64) = source.id"
1676        );
1677        let MergeIntoAction::Update(assignments) = &merge_op.clauses[0].action else {
1678            panic!("expected UPDATE");
1679        };
1680        assert_eq!(assignments[0].1.to_string(), "CAST(source.id AS UInt32)");
1681        let MergeIntoAction::Insert { values, .. } = &merge_op.clauses[1].action else {
1682            panic!("expected INSERT");
1683        };
1684        assert_eq!(values[0].to_string(), "CAST(source.id AS UInt32)");
1685        Ok(())
1686    }
1687
1688    #[test]
1689    fn coerce_utf8view_output() -> Result<()> {
1690        // Plan A
1691        // scenario: outermost utf8view projection
1692        let expr = col("a");
1693        let empty = empty_with_type(DataType::Utf8View);
1694        let plan = LogicalPlan::Projection(Projection::try_new(
1695            vec![expr.clone()],
1696            Arc::clone(&empty),
1697        )?);
1698
1699        // Plan A: no coerce
1700        coerce_on_output_if_viewtype!(
1701            false,
1702            plan.clone(),
1703            @r"
1704        Projection: a
1705          EmptyRelation: rows=0
1706        "
1707        )?;
1708
1709        // Plan A: coerce requested: Utf8View => LargeUtf8
1710        coerce_on_output_if_viewtype!(
1711            true,
1712            plan.clone(),
1713            @r"
1714        Projection: CAST(a AS LargeUtf8) AS a
1715          EmptyRelation: rows=0
1716        "
1717        )?;
1718
1719        // Plan B
1720        // scenario: outermost bool projection
1721        let bool_expr = col("a").lt(lit("foo"));
1722        let bool_plan = LogicalPlan::Projection(Projection::try_new(
1723            vec![bool_expr],
1724            Arc::clone(&empty),
1725        )?);
1726        // Plan B: no coerce
1727        coerce_on_output_if_viewtype!(
1728            false,
1729            bool_plan.clone(),
1730            @r#"
1731        Projection: a < CAST(Utf8("foo") AS Utf8View)
1732          EmptyRelation: rows=0
1733        "#
1734        )?;
1735
1736        coerce_on_output_if_viewtype!(
1737            false,
1738            plan.clone(),
1739            @r"
1740        Projection: a
1741          EmptyRelation: rows=0
1742        "
1743        )?;
1744
1745        // Plan B: coerce requested: no coercion applied
1746        coerce_on_output_if_viewtype!(
1747            true,
1748            plan.clone(),
1749            @r"
1750        Projection: CAST(a AS LargeUtf8) AS a
1751          EmptyRelation: rows=0
1752        "
1753        )?;
1754
1755        // Plan C
1756        // scenario: with a non-projection root logical plan node
1757        let sort_expr = expr.sort(true, true);
1758        let sort_plan = LogicalPlan::Sort(Sort {
1759            expr: vec![sort_expr],
1760            input: Arc::new(plan),
1761            fetch: None,
1762        });
1763
1764        // Plan C: no coerce
1765        coerce_on_output_if_viewtype!(
1766            false,
1767            sort_plan.clone(),
1768            @r"
1769        Sort: a ASC NULLS FIRST
1770          Projection: a
1771            EmptyRelation: rows=0
1772        "
1773        )?;
1774
1775        // Plan C: coerce requested: Utf8View => LargeUtf8
1776        coerce_on_output_if_viewtype!(
1777            true,
1778            sort_plan.clone(),
1779            @r"
1780        Projection: CAST(a AS LargeUtf8) AS a
1781          Sort: a ASC NULLS FIRST
1782            Projection: a
1783              EmptyRelation: rows=0
1784        "
1785        )?;
1786
1787        // Plan D
1788        // scenario: two layers of projections with view types
1789        let plan = LogicalPlan::Projection(Projection::try_new(
1790            vec![col("a")],
1791            Arc::new(sort_plan),
1792        )?);
1793        // Plan D: no coerce
1794        coerce_on_output_if_viewtype!(
1795            false,
1796            plan.clone(),
1797            @r"
1798        Projection: a
1799          Sort: a ASC NULLS FIRST
1800            Projection: a
1801              EmptyRelation: rows=0
1802        "
1803        )?;
1804        // Plan B: coerce requested: Utf8View => LargeUtf8 only on outermost
1805        coerce_on_output_if_viewtype!(
1806            true,
1807            plan.clone(),
1808            @r"
1809        Projection: CAST(a AS LargeUtf8) AS a
1810          Sort: a ASC NULLS FIRST
1811            Projection: a
1812              EmptyRelation: rows=0
1813        "
1814        )?;
1815
1816        Ok(())
1817    }
1818
1819    #[test]
1820    fn coerce_binaryview_output() -> Result<()> {
1821        // Plan A
1822        // scenario: outermost binaryview projection
1823        let expr = col("a");
1824        let empty = empty_with_type(DataType::BinaryView);
1825        let plan = LogicalPlan::Projection(Projection::try_new(
1826            vec![expr.clone()],
1827            Arc::clone(&empty),
1828        )?);
1829
1830        // Plan A: no coerce
1831        coerce_on_output_if_viewtype!(
1832            false,
1833            plan.clone(),
1834            @r"
1835        Projection: a
1836          EmptyRelation: rows=0
1837        "
1838        )?;
1839
1840        // Plan A: coerce requested: BinaryView => LargeBinary
1841        coerce_on_output_if_viewtype!(
1842            true,
1843            plan.clone(),
1844            @r"
1845        Projection: CAST(a AS LargeBinary) AS a
1846          EmptyRelation: rows=0
1847        "
1848        )?;
1849
1850        // Plan B
1851        // scenario: outermost bool projection
1852        let bool_expr = col("a").lt(lit(vec![8, 1, 8, 1]));
1853        let bool_plan = LogicalPlan::Projection(Projection::try_new(
1854            vec![bool_expr],
1855            Arc::clone(&empty),
1856        )?);
1857
1858        // Plan B: no coerce
1859        coerce_on_output_if_viewtype!(
1860            false,
1861            bool_plan.clone(),
1862            @r#"
1863        Projection: a < CAST(Binary("8,1,8,1") AS BinaryView)
1864          EmptyRelation: rows=0
1865        "#
1866        )?;
1867
1868        // Plan B: coerce requested: no coercion applied
1869        coerce_on_output_if_viewtype!(
1870            true,
1871            bool_plan.clone(),
1872            @r#"
1873        Projection: a < CAST(Binary("8,1,8,1") AS BinaryView)
1874          EmptyRelation: rows=0
1875        "#
1876        )?;
1877
1878        // Plan C
1879        // scenario: with a non-projection root logical plan node
1880        let sort_expr = expr.sort(true, true);
1881        let sort_plan = LogicalPlan::Sort(Sort {
1882            expr: vec![sort_expr],
1883            input: Arc::new(plan),
1884            fetch: None,
1885        });
1886
1887        // Plan C: no coerce
1888        coerce_on_output_if_viewtype!(
1889            false,
1890            sort_plan.clone(),
1891            @r"
1892        Sort: a ASC NULLS FIRST
1893          Projection: a
1894            EmptyRelation: rows=0
1895        "
1896        )?;
1897        // Plan C: coerce requested: BinaryView => LargeBinary
1898        coerce_on_output_if_viewtype!(
1899            true,
1900            sort_plan.clone(),
1901            @r"
1902        Projection: CAST(a AS LargeBinary) AS a
1903          Sort: a ASC NULLS FIRST
1904            Projection: a
1905              EmptyRelation: rows=0
1906        "
1907        )?;
1908
1909        // Plan D
1910        // scenario: two layers of projections with view types
1911        let plan = LogicalPlan::Projection(Projection::try_new(
1912            vec![col("a")],
1913            Arc::new(sort_plan),
1914        )?);
1915
1916        // Plan D: no coerce
1917        coerce_on_output_if_viewtype!(
1918            false,
1919            plan.clone(),
1920            @r"
1921        Projection: a
1922          Sort: a ASC NULLS FIRST
1923            Projection: a
1924              EmptyRelation: rows=0
1925        "
1926        )?;
1927
1928        // Plan B: coerce requested: BinaryView => LargeBinary only on outermost
1929        coerce_on_output_if_viewtype!(
1930            true,
1931            plan.clone(),
1932            @r"
1933        Projection: CAST(a AS LargeBinary) AS a
1934          Sort: a ASC NULLS FIRST
1935            Projection: a
1936              EmptyRelation: rows=0
1937        "
1938        )?;
1939
1940        Ok(())
1941    }
1942
1943    #[test]
1944    fn nested_case() -> Result<()> {
1945        let expr = col("a").lt(lit(2_u32));
1946        let empty = empty_with_type(DataType::Float64);
1947
1948        let plan = LogicalPlan::Projection(Projection::try_new(
1949            vec![expr.clone().or(expr)],
1950            empty,
1951        )?);
1952
1953        assert_analyzed_plan_eq!(
1954            plan,
1955            @r"
1956        Projection: a < CAST(UInt32(2) AS Float64) OR a < CAST(UInt32(2) AS Float64)
1957          EmptyRelation: rows=0
1958        "
1959        )
1960    }
1961
1962    #[derive(Debug, PartialEq, Eq, Hash)]
1963    struct TestScalarUDF {
1964        signature: Signature,
1965    }
1966
1967    impl ScalarUDFImpl for TestScalarUDF {
1968        fn name(&self) -> &str {
1969            "TestScalarUDF"
1970        }
1971
1972        fn signature(&self) -> &Signature {
1973            &self.signature
1974        }
1975
1976        fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
1977            Ok(Utf8)
1978        }
1979
1980        fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
1981            Ok(ColumnarValue::Scalar(ScalarValue::from("a")))
1982        }
1983    }
1984
1985    #[derive(Debug, Hash, PartialEq, Eq)]
1986    struct TestArrayElementUDF;
1987
1988    impl ScalarUDFImpl for TestArrayElementUDF {
1989        fn name(&self) -> &str {
1990            "TestArrayElementUDF"
1991        }
1992
1993        fn signature(&self) -> &Signature {
1994            static SIGNATURE: std::sync::LazyLock<Signature> =
1995                std::sync::LazyLock::new(|| {
1996                    Signature::array_and_index(Volatility::Immutable)
1997                });
1998            &SIGNATURE
1999        }
2000
2001        fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
2002            Ok(Utf8)
2003        }
2004
2005        fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
2006            Ok(ColumnarValue::Scalar(ScalarValue::from("a")))
2007        }
2008    }
2009
2010    #[test]
2011    fn scalar_udf() -> Result<()> {
2012        let empty = empty();
2013
2014        let udf = ScalarUDF::from(TestScalarUDF {
2015            signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable),
2016        })
2017        .call(vec![lit(123_i32)]);
2018        let plan = LogicalPlan::Projection(Projection::try_new(vec![udf], empty)?);
2019
2020        assert_analyzed_plan_eq!(
2021            plan,
2022            @r"
2023        Projection: TestScalarUDF(CAST(Int32(123) AS Float32))
2024          EmptyRelation: rows=0
2025        "
2026        )
2027    }
2028
2029    #[test]
2030    fn scalar_udf_invalid_input() -> Result<()> {
2031        let empty = empty();
2032        let udf = ScalarUDF::from(TestScalarUDF {
2033            signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable),
2034        })
2035        .call(vec![lit("Apple")]);
2036        Projection::try_new(vec![udf], empty)
2037            .expect_err("Expected an error due to incorrect function input");
2038
2039        Ok(())
2040    }
2041
2042    #[test]
2043    fn scalar_function() -> Result<()> {
2044        // test that automatic argument type coercion for scalar functions work
2045        let empty = empty();
2046        let lit_expr = lit(10i64);
2047        let fun = ScalarUDF::new_from_impl(TestScalarUDF {
2048            signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable),
2049        });
2050        let scalar_function_expr =
2051            Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(fun), vec![lit_expr]));
2052        let plan = LogicalPlan::Projection(Projection::try_new(
2053            vec![scalar_function_expr],
2054            empty,
2055        )?);
2056
2057        assert_analyzed_plan_eq!(
2058            plan,
2059            @r"
2060        Projection: TestScalarUDF(CAST(Int64(10) AS Float32))
2061          EmptyRelation: rows=0
2062        "
2063        )
2064    }
2065
2066    #[test]
2067    fn agg_udaf() -> Result<()> {
2068        let empty = empty();
2069        let my_avg = create_udaf(
2070            "MY_AVG",
2071            vec![DataType::Float64],
2072            Arc::new(DataType::Float64),
2073            Volatility::Immutable,
2074            Arc::new(|_| Ok(Box::<AvgAccumulator>::default())),
2075            Arc::new(vec![DataType::UInt64, DataType::Float64]),
2076        );
2077        let udaf = Expr::AggregateFunction(expr::AggregateFunction::new_udf(
2078            Arc::new(my_avg),
2079            vec![lit(10i64)],
2080            false,
2081            None,
2082            vec![],
2083            None,
2084        ));
2085        let plan = LogicalPlan::Projection(Projection::try_new(vec![udaf], empty)?);
2086
2087        assert_analyzed_plan_eq!(
2088            plan,
2089            @r"
2090        Projection: MY_AVG(CAST(Int64(10) AS Float64))
2091          EmptyRelation: rows=0
2092        "
2093        )
2094    }
2095
2096    #[test]
2097    fn agg_udaf_invalid_input() -> Result<()> {
2098        let empty = empty();
2099        let return_type = DataType::Float64;
2100        let accumulator: AccumulatorFactoryFunction =
2101            Arc::new(|_| Ok(Box::<AvgAccumulator>::default()));
2102        let my_avg = AggregateUDF::from(SimpleAggregateUDF::new_with_signature(
2103            "MY_AVG",
2104            Signature::uniform(1, vec![DataType::Float64], Volatility::Immutable),
2105            return_type,
2106            accumulator,
2107            vec![
2108                Field::new("count", DataType::UInt64, true).into(),
2109                Field::new("avg", DataType::Float64, true).into(),
2110            ],
2111        ));
2112        let udaf = Expr::AggregateFunction(expr::AggregateFunction::new_udf(
2113            Arc::new(my_avg),
2114            vec![lit("10")],
2115            false,
2116            None,
2117            vec![],
2118            None,
2119        ));
2120
2121        let err = Projection::try_new(vec![udaf], empty).err().unwrap();
2122        assert!(
2123            err.strip_backtrace().starts_with("Error during planning: Failed to coerce arguments to satisfy a call to 'MY_AVG' function: coercion from Utf8 to the signature Uniform(1, [Float64]) failed")
2124        );
2125        Ok(())
2126    }
2127
2128    #[test]
2129    fn agg_function_case() -> Result<()> {
2130        let empty = empty();
2131        let agg_expr = Expr::AggregateFunction(expr::AggregateFunction::new_udf(
2132            avg_udaf(),
2133            vec![lit(12f64)],
2134            false,
2135            None,
2136            vec![],
2137            None,
2138        ));
2139        let plan = LogicalPlan::Projection(Projection::try_new(vec![agg_expr], empty)?);
2140
2141        assert_analyzed_plan_eq!(
2142            plan,
2143            @r"
2144        Projection: avg(Float64(12))
2145          EmptyRelation: rows=0
2146        "
2147        )?;
2148
2149        let empty = empty_with_type(DataType::Int32);
2150        let agg_expr = Expr::AggregateFunction(expr::AggregateFunction::new_udf(
2151            avg_udaf(),
2152            vec![cast(col("a"), DataType::Float64)],
2153            false,
2154            None,
2155            vec![],
2156            None,
2157        ));
2158        let plan = LogicalPlan::Projection(Projection::try_new(vec![agg_expr], empty)?);
2159
2160        assert_analyzed_plan_eq!(
2161            plan,
2162            @r"
2163        Projection: avg(CAST(a AS Float64))
2164          EmptyRelation: rows=0
2165        "
2166        )
2167    }
2168
2169    #[test]
2170    fn agg_function_invalid_input_avg() -> Result<()> {
2171        let empty = empty();
2172        let agg_expr = Expr::AggregateFunction(expr::AggregateFunction::new_udf(
2173            avg_udaf(),
2174            vec![lit("1")],
2175            false,
2176            None,
2177            vec![],
2178            None,
2179        ));
2180        let err = Projection::try_new(vec![agg_expr], empty)
2181            .err()
2182            .unwrap()
2183            .strip_backtrace();
2184        assert!(
2185            err.contains("Function 'avg' failed to match any signature"),
2186            "Err: {err:?}"
2187        );
2188        Ok(())
2189    }
2190
2191    #[test]
2192    fn binary_op_date32_op_interval() -> Result<()> {
2193        // CAST(Utf8("1998-03-18") AS Date32) + IntervalDayTime("...")
2194        let expr = cast(lit("1998-03-18"), DataType::Date32)
2195            + lit(ScalarValue::new_interval_dt(123, 456));
2196        let empty = empty();
2197        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2198
2199        assert_analyzed_plan_eq!(
2200            plan,
2201            @r#"
2202        Projection: CAST(Utf8("1998-03-18") AS Date32) + IntervalDayTime("IntervalDayTime { days: 123, milliseconds: 456 }")
2203          EmptyRelation: rows=0
2204        "#
2205        )
2206    }
2207
2208    #[test]
2209    fn inlist_case() -> Result<()> {
2210        // a in (1,4,8), a is int64
2211        let expr = col("a").in_list(vec![lit(1_i32), lit(4_i8), lit(8_i64)], false);
2212        let empty = empty_with_type(DataType::Int64);
2213        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2214        assert_analyzed_plan_eq!(
2215            plan,
2216            @r"
2217        Projection: a IN ([CAST(Int32(1) AS Int64), CAST(Int8(4) AS Int64), Int64(8)])
2218          EmptyRelation: rows=0
2219        ")?;
2220
2221        // a in (1,4,8), a is decimal
2222        let expr = col("a").in_list(vec![lit(1_i32), lit(4_i8), lit(8_i64)], false);
2223        let empty = Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
2224            produce_one_row: false,
2225            schema: Arc::new(DFSchema::from_unqualified_fields(
2226                vec![Field::new("a", DataType::Decimal128(12, 4), true)].into(),
2227                std::collections::HashMap::new(),
2228            )?),
2229        }));
2230        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2231        assert_analyzed_plan_eq!(
2232            plan,
2233            @r"
2234        Projection: CAST(a AS Decimal128(24, 4)) IN ([CAST(Int32(1) AS Decimal128(24, 4)), CAST(Int8(4) AS Decimal128(24, 4)), CAST(Int64(8) AS Decimal128(24, 4))])
2235          EmptyRelation: rows=0
2236        ")
2237    }
2238
2239    #[test]
2240    fn between_case() -> Result<()> {
2241        let expr = col("a").between(
2242            lit("2002-05-08"),
2243            // (cast('2002-05-08' as date) + interval '1 months')
2244            cast(lit("2002-05-08"), DataType::Date32)
2245                + lit(ScalarValue::new_interval_ym(0, 1)),
2246        );
2247        let empty = empty_with_type(Utf8);
2248        let plan = LogicalPlan::Filter(Filter::try_new(expr, empty)?);
2249
2250        assert_analyzed_plan_eq!(
2251            plan,
2252            @r#"
2253        Filter: CAST(a AS Date32) BETWEEN CAST(Utf8("2002-05-08") AS Date32) AND CAST(Utf8("2002-05-08") AS Date32) + IntervalYearMonth("1")
2254          EmptyRelation: rows=0
2255        "#
2256        )
2257    }
2258
2259    #[test]
2260    fn between_infer_cheap_type() -> Result<()> {
2261        let expr = col("a").between(
2262            // (cast('2002-05-08' as date) + interval '1 months')
2263            cast(lit("2002-05-08"), DataType::Date32)
2264                + lit(ScalarValue::new_interval_ym(0, 1)),
2265            lit("2002-12-08"),
2266        );
2267        let empty = empty_with_type(Utf8);
2268        let plan = LogicalPlan::Filter(Filter::try_new(expr, empty)?);
2269
2270        // TODO: we should cast col(a).
2271        assert_analyzed_plan_eq!(
2272            plan,
2273            @r#"
2274        Filter: CAST(a AS Date32) BETWEEN CAST(Utf8("2002-05-08") AS Date32) + IntervalYearMonth("1") AND CAST(Utf8("2002-12-08") AS Date32)
2275          EmptyRelation: rows=0
2276        "#
2277        )
2278    }
2279
2280    #[test]
2281    fn between_null() -> Result<()> {
2282        let expr = lit(ScalarValue::Null).between(lit(ScalarValue::Null), lit(2i64));
2283        let empty = empty();
2284        let plan = LogicalPlan::Filter(Filter::try_new(expr, empty)?);
2285
2286        assert_analyzed_plan_eq!(
2287            plan,
2288            @r"
2289        Filter: CAST(NULL AS Int64) BETWEEN CAST(NULL AS Int64) AND Int64(2)
2290          EmptyRelation: rows=0
2291        "
2292        )
2293    }
2294
2295    #[test]
2296    fn is_bool_for_type_coercion() -> Result<()> {
2297        // is true
2298        let expr = col("a").is_true();
2299        let empty = empty_with_type(DataType::Boolean);
2300        let plan =
2301            LogicalPlan::Projection(Projection::try_new(vec![expr.clone()], empty)?);
2302
2303        assert_analyzed_plan_eq!(
2304            plan,
2305            @r"
2306        Projection: a IS TRUE
2307          EmptyRelation: rows=0
2308        "
2309        )?;
2310
2311        let empty = empty_with_type(DataType::Int64);
2312        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2313        assert_type_coercion_error(
2314            plan,
2315            "Cannot infer common argument type for comparison operation Int64 IS DISTINCT FROM Boolean",
2316        )?;
2317
2318        // is not true
2319        let expr = col("a").is_not_true();
2320        let empty = empty_with_type(DataType::Boolean);
2321        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2322
2323        assert_analyzed_plan_eq!(
2324            plan,
2325            @r"
2326        Projection: a IS NOT TRUE
2327          EmptyRelation: rows=0
2328        "
2329        )?;
2330
2331        // is false
2332        let expr = col("a").is_false();
2333        let empty = empty_with_type(DataType::Boolean);
2334        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2335
2336        assert_analyzed_plan_eq!(
2337            plan,
2338            @r"
2339        Projection: a IS FALSE
2340          EmptyRelation: rows=0
2341        "
2342        )?;
2343
2344        // is not false
2345        let expr = col("a").is_not_false();
2346        let empty = empty_with_type(DataType::Boolean);
2347        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2348
2349        assert_analyzed_plan_eq!(
2350            plan,
2351            @r"
2352        Projection: a IS NOT FALSE
2353          EmptyRelation: rows=0
2354        "
2355        )
2356    }
2357
2358    #[test]
2359    fn like_for_type_coercion() -> Result<()> {
2360        // like : utf8 like "abc"
2361        let expr = Box::new(col("a"));
2362        let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
2363        let like_expr = Expr::Like(Like::new(false, expr, pattern, None, false));
2364        let empty = empty_with_type(Utf8);
2365        let plan = LogicalPlan::Projection(Projection::try_new(vec![like_expr], empty)?);
2366
2367        assert_analyzed_plan_eq!(
2368            plan,
2369            @r#"
2370        Projection: a LIKE Utf8("abc")
2371          EmptyRelation: rows=0
2372        "#
2373        )?;
2374
2375        let expr = Box::new(col("a"));
2376        let pattern = Box::new(lit(ScalarValue::Null));
2377        let like_expr = Expr::Like(Like::new(false, expr, pattern, None, false));
2378        let empty = empty_with_type(Utf8);
2379        let plan = LogicalPlan::Projection(Projection::try_new(vec![like_expr], empty)?);
2380
2381        assert_analyzed_plan_eq!(
2382            plan,
2383            @r"
2384        Projection: a LIKE CAST(NULL AS Utf8)
2385          EmptyRelation: rows=0
2386        "
2387        )?;
2388
2389        let expr = Box::new(col("a"));
2390        let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
2391        let like_expr = Expr::Like(Like::new(false, expr, pattern, None, false));
2392        let empty = empty_with_type(DataType::Int64);
2393        let plan = LogicalPlan::Projection(Projection::try_new(vec![like_expr], empty)?);
2394        assert_type_coercion_error(
2395            plan,
2396            "There isn't a common type to coerce Int64 and Utf8 in LIKE expression",
2397        )?;
2398
2399        // ilike
2400        let expr = Box::new(col("a"));
2401        let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
2402        let ilike_expr = Expr::Like(Like::new(false, expr, pattern, None, true));
2403        let empty = empty_with_type(Utf8);
2404        let plan = LogicalPlan::Projection(Projection::try_new(vec![ilike_expr], empty)?);
2405
2406        assert_analyzed_plan_eq!(
2407            plan,
2408            @r#"
2409        Projection: a ILIKE Utf8("abc")
2410          EmptyRelation: rows=0
2411        "#
2412        )?;
2413
2414        let expr = Box::new(col("a"));
2415        let pattern = Box::new(lit(ScalarValue::Null));
2416        let ilike_expr = Expr::Like(Like::new(false, expr, pattern, None, true));
2417        let empty = empty_with_type(Utf8);
2418        let plan = LogicalPlan::Projection(Projection::try_new(vec![ilike_expr], empty)?);
2419
2420        assert_analyzed_plan_eq!(
2421            plan,
2422            @r"
2423        Projection: a ILIKE CAST(NULL AS Utf8)
2424          EmptyRelation: rows=0
2425        "
2426        )?;
2427
2428        let expr = Box::new(col("a"));
2429        let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
2430        let ilike_expr = Expr::Like(Like::new(false, expr, pattern, None, true));
2431        let empty = empty_with_type(DataType::Int64);
2432        let plan = LogicalPlan::Projection(Projection::try_new(vec![ilike_expr], empty)?);
2433        assert_type_coercion_error(
2434            plan,
2435            "There isn't a common type to coerce Int64 and Utf8 in ILIKE expression",
2436        )?;
2437
2438        Ok(())
2439    }
2440
2441    #[test]
2442    fn similar_to_for_type_coercion() -> Result<()> {
2443        // similar to : utf8 similar to "abc"
2444        let expr = Box::new(col("a"));
2445        let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
2446        let similar_to_expr =
2447            Expr::SimilarTo(Like::new(false, expr, pattern, None, false));
2448        let empty = empty_with_type(Utf8);
2449        let plan =
2450            LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?);
2451
2452        assert_analyzed_plan_eq!(
2453            plan,
2454            @r#"
2455        Projection: a SIMILAR TO Utf8("abc")
2456          EmptyRelation: rows=0
2457        "#
2458        )?;
2459
2460        // NULL pattern is coerced to a typed NULL instead of panicking
2461        // (https://github.com/apache/datafusion/issues/22886)
2462        let expr = Box::new(col("a"));
2463        let pattern = Box::new(lit(ScalarValue::Null));
2464        let similar_to_expr =
2465            Expr::SimilarTo(Like::new(false, expr, pattern, None, false));
2466        let empty = empty_with_type(Utf8);
2467        let plan =
2468            LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?);
2469
2470        assert_analyzed_plan_eq!(
2471            plan,
2472            @r"
2473        Projection: a SIMILAR TO CAST(NULL AS Utf8)
2474          EmptyRelation: rows=0
2475        "
2476        )?;
2477
2478        // Utf8View value and Utf8 pattern are coerced to Utf8View
2479        let expr = Box::new(col("a"));
2480        let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
2481        let similar_to_expr =
2482            Expr::SimilarTo(Like::new(false, expr, pattern, None, false));
2483        let empty = empty_with_type(DataType::Utf8View);
2484        let plan =
2485            LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?);
2486
2487        assert_analyzed_plan_eq!(
2488            plan,
2489            @r#"
2490        Projection: a SIMILAR TO CAST(Utf8("abc") AS Utf8View)
2491          EmptyRelation: rows=0
2492        "#
2493        )?;
2494
2495        // Utf8 value and Utf8View pattern are coerced to Utf8View
2496        let expr = Box::new(col("a"));
2497        let pattern = Box::new(lit(ScalarValue::Utf8View(Some("abc".to_string()))));
2498        let similar_to_expr =
2499            Expr::SimilarTo(Like::new(false, expr, pattern, None, false));
2500        let empty = empty_with_type(Utf8);
2501        let plan =
2502            LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?);
2503
2504        assert_analyzed_plan_eq!(
2505            plan,
2506            @r#"
2507        Projection: CAST(a AS Utf8View) SIMILAR TO Utf8View("abc")
2508          EmptyRelation: rows=0
2509        "#
2510        )?;
2511
2512        // Dictionary values are coerced to the common regex operand type
2513        let expr = Box::new(col("a"));
2514        let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
2515        let similar_to_expr =
2516            Expr::SimilarTo(Like::new(false, expr, pattern, None, false));
2517        let empty = empty_with_type(DataType::Dictionary(
2518            Box::new(DataType::Int32),
2519            Box::new(Utf8),
2520        ));
2521        let plan =
2522            LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?);
2523
2524        assert_analyzed_plan_eq!(
2525            plan,
2526            @r#"
2527        Projection: CAST(a AS Utf8) SIMILAR TO Utf8("abc")
2528          EmptyRelation: rows=0
2529        "#
2530        )?;
2531
2532        // incompatible types are a planning error, not a panic
2533        let expr = Box::new(col("a"));
2534        let pattern = Box::new(lit(ScalarValue::new_utf8("abc")));
2535        let similar_to_expr =
2536            Expr::SimilarTo(Like::new(false, expr, pattern, None, false));
2537        let empty = empty_with_type(DataType::Int64);
2538        let plan =
2539            LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?);
2540        assert_type_coercion_error(
2541            plan,
2542            "There isn't a common type to coerce Int64 and Utf8 in SIMILAR TO expression",
2543        )?;
2544
2545        Ok(())
2546    }
2547
2548    #[test]
2549    fn unknown_for_type_coercion() -> Result<()> {
2550        // unknown
2551        let expr = col("a").is_unknown();
2552        let empty = empty_with_type(DataType::Boolean);
2553        let plan =
2554            LogicalPlan::Projection(Projection::try_new(vec![expr.clone()], empty)?);
2555
2556        assert_analyzed_plan_eq!(
2557            plan,
2558            @r"
2559        Projection: a IS UNKNOWN
2560          EmptyRelation: rows=0
2561        "
2562        )?;
2563
2564        let empty = empty_with_type(Utf8);
2565        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2566        assert_type_coercion_error(
2567            plan,
2568            "Cannot infer common argument type for comparison operation Utf8 IS DISTINCT FROM Boolean",
2569        )?;
2570
2571        // is not unknown
2572        let expr = col("a").is_not_unknown();
2573        let empty = empty_with_type(DataType::Boolean);
2574        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2575
2576        assert_analyzed_plan_eq!(
2577            plan,
2578            @r"
2579        Projection: a IS NOT UNKNOWN
2580          EmptyRelation: rows=0
2581        "
2582        )
2583    }
2584
2585    #[test]
2586    fn concat_for_type_coercion() -> Result<()> {
2587        let empty = empty_with_type(Utf8);
2588        let args = [col("a"), lit("b"), lit(true), lit(false), lit(13)];
2589
2590        // concat-type signature
2591        let expr = ScalarUDF::new_from_impl(TestScalarUDF {
2592            signature: Signature::variadic(vec![Utf8], Volatility::Immutable),
2593        })
2594        .call(args.to_vec());
2595        let plan =
2596            LogicalPlan::Projection(Projection::try_new(vec![expr], Arc::clone(&empty))?);
2597        assert_analyzed_plan_eq!(
2598            plan,
2599            @r#"
2600        Projection: TestScalarUDF(a, Utf8("b"), CAST(Boolean(true) AS Utf8), CAST(Boolean(false) AS Utf8), CAST(Int32(13) AS Utf8))
2601          EmptyRelation: rows=0
2602        "#
2603        )
2604    }
2605
2606    #[test]
2607    fn test_type_coercion_rewrite() -> Result<()> {
2608        // gt
2609        let schema = Arc::new(DFSchema::from_unqualified_fields(
2610            vec![Field::new("a", DataType::Int64, true)].into(),
2611            std::collections::HashMap::new(),
2612        )?);
2613        let mut rewriter = TypeCoercionRewriter { schema: &schema };
2614        let expr = is_true(lit(12i32).gt(lit(13i64)));
2615        let expected = is_true(cast(lit(12i32), DataType::Int64).gt(lit(13i64)));
2616        let result = expr.rewrite(&mut rewriter).data()?;
2617        assert_eq!(expected, result);
2618
2619        // eq
2620        let schema = Arc::new(DFSchema::from_unqualified_fields(
2621            vec![Field::new("a", DataType::Int64, true)].into(),
2622            std::collections::HashMap::new(),
2623        )?);
2624        let mut rewriter = TypeCoercionRewriter { schema: &schema };
2625        let expr = is_true(lit(12i32).eq(lit(13i64)));
2626        let expected = is_true(cast(lit(12i32), DataType::Int64).eq(lit(13i64)));
2627        let result = expr.rewrite(&mut rewriter).data()?;
2628        assert_eq!(expected, result);
2629
2630        // lt
2631        let schema = Arc::new(DFSchema::from_unqualified_fields(
2632            vec![Field::new("a", DataType::Int64, true)].into(),
2633            std::collections::HashMap::new(),
2634        )?);
2635        let mut rewriter = TypeCoercionRewriter { schema: &schema };
2636        let expr = is_true(lit(12i32).lt(lit(13i64)));
2637        let expected = is_true(cast(lit(12i32), DataType::Int64).lt(lit(13i64)));
2638        let result = expr.rewrite(&mut rewriter).data()?;
2639        assert_eq!(expected, result);
2640
2641        Ok(())
2642    }
2643
2644    #[test]
2645    fn binary_op_date32_eq_ts() -> Result<()> {
2646        let expr = cast(
2647            lit("1998-03-18"),
2648            DataType::Timestamp(TimeUnit::Nanosecond, None),
2649        )
2650        .eq(cast(lit("1998-03-18"), DataType::Date32));
2651        let empty = empty();
2652        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2653
2654        assert_analyzed_plan_eq!(
2655            plan,
2656            @r#"
2657        Projection: CAST(Utf8("1998-03-18") AS Timestamp(ns)) = CAST(CAST(Utf8("1998-03-18") AS Date32) AS Timestamp(ns))
2658          EmptyRelation: rows=0
2659        "#
2660        )
2661    }
2662
2663    fn cast_if_not_same_type(
2664        expr: Box<Expr>,
2665        data_type: &DataType,
2666        schema: &DFSchemaRef,
2667    ) -> Box<Expr> {
2668        if &expr.get_type(schema).unwrap() != data_type {
2669            Box::new(cast(*expr, data_type.clone()))
2670        } else {
2671            expr
2672        }
2673    }
2674
2675    fn cast_helper(
2676        case: Case,
2677        case_when_type: &DataType,
2678        then_else_type: &DataType,
2679        schema: &DFSchemaRef,
2680    ) -> Case {
2681        let expr = case
2682            .expr
2683            .map(|e| cast_if_not_same_type(e, case_when_type, schema));
2684        let when_then_expr = case
2685            .when_then_expr
2686            .into_iter()
2687            .map(|(when, then)| {
2688                (
2689                    cast_if_not_same_type(when, case_when_type, schema),
2690                    cast_if_not_same_type(then, then_else_type, schema),
2691                )
2692            })
2693            .collect::<Vec<_>>();
2694        let else_expr = case
2695            .else_expr
2696            .map(|e| cast_if_not_same_type(e, then_else_type, schema));
2697
2698        Case {
2699            expr,
2700            when_then_expr,
2701            else_expr,
2702        }
2703    }
2704
2705    #[test]
2706    fn test_case_expression_coercion() -> Result<()> {
2707        let schema = Arc::new(DFSchema::from_unqualified_fields(
2708            vec![
2709                Field::new("boolean", DataType::Boolean, true),
2710                Field::new("integer", DataType::Int32, true),
2711                Field::new("float", DataType::Float32, true),
2712                Field::new(
2713                    "timestamp",
2714                    DataType::Timestamp(TimeUnit::Nanosecond, None),
2715                    true,
2716                ),
2717                Field::new("date", DataType::Date32, true),
2718                Field::new(
2719                    "interval",
2720                    DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano),
2721                    true,
2722                ),
2723                Field::new("binary", DataType::Binary, true),
2724                Field::new("string", Utf8, true),
2725                Field::new("decimal", DataType::Decimal128(10, 10), true),
2726            ]
2727            .into(),
2728            std::collections::HashMap::new(),
2729        )?);
2730
2731        let case = Case {
2732            expr: None,
2733            when_then_expr: vec![
2734                (Box::new(col("boolean")), Box::new(col("integer"))),
2735                (Box::new(col("integer")), Box::new(col("float"))),
2736                (Box::new(col("string")), Box::new(col("string"))),
2737            ],
2738            else_expr: None,
2739        };
2740        let case_when_common_type = DataType::Boolean;
2741        let then_else_common_type = Utf8;
2742        let expected = cast_helper(
2743            case.clone(),
2744            &case_when_common_type,
2745            &then_else_common_type,
2746            &schema,
2747        );
2748        let actual = coerce_case_expression(case, &schema)?;
2749        assert_eq!(expected, actual);
2750
2751        // CASE string WHEN float/integer/string: comparison coercion
2752        // prefers numeric, so the common type for the CASE expr and
2753        // WHEN values is Float32.
2754        let case = Case {
2755            expr: Some(Box::new(col("string"))),
2756            when_then_expr: vec![
2757                (Box::new(col("float")), Box::new(col("integer"))),
2758                (Box::new(col("integer")), Box::new(col("float"))),
2759                (Box::new(col("string")), Box::new(col("string"))),
2760            ],
2761            else_expr: Some(Box::new(col("string"))),
2762        };
2763        let case_when_common_type = DataType::Float32;
2764        let then_else_common_type = Utf8;
2765        let expected = cast_helper(
2766            case.clone(),
2767            &case_when_common_type,
2768            &then_else_common_type,
2769            &schema,
2770        );
2771        let actual = coerce_case_expression(case, &schema)?;
2772        assert_eq!(expected, actual);
2773
2774        let case = Case {
2775            expr: Some(Box::new(col("interval"))),
2776            when_then_expr: vec![
2777                (Box::new(col("float")), Box::new(col("integer"))),
2778                (Box::new(col("binary")), Box::new(col("float"))),
2779                (Box::new(col("string")), Box::new(col("string"))),
2780            ],
2781            else_expr: Some(Box::new(col("string"))),
2782        };
2783        let err = coerce_case_expression(case, &schema).unwrap_err();
2784        assert_snapshot!(
2785            err.strip_backtrace(),
2786            @"Error during planning: Failed to coerce case (Interval(MonthDayNano)) and when (Float32, Binary, Utf8) to common types in CASE WHEN expression"
2787        );
2788
2789        let case = Case {
2790            expr: Some(Box::new(col("string"))),
2791            when_then_expr: vec![
2792                (Box::new(col("float")), Box::new(col("date"))),
2793                (Box::new(col("string")), Box::new(col("float"))),
2794                (Box::new(col("string")), Box::new(col("binary"))),
2795            ],
2796            else_expr: Some(Box::new(col("timestamp"))),
2797        };
2798        let err = coerce_case_expression(case, &schema).unwrap_err();
2799        assert_snapshot!(
2800            err.strip_backtrace(),
2801            @"Error during planning: Failed to coerce then (Date32, Float32, Binary) and else (Timestamp(ns)) to common types in CASE WHEN expression"
2802        );
2803
2804        Ok(())
2805    }
2806
2807    macro_rules! test_case_expression {
2808        ($expr:expr, $when_then:expr, $case_when_type:expr, $then_else_type:expr, $schema:expr) => {
2809            let case = Case {
2810                expr: $expr.map(|e| Box::new(col(e))),
2811                when_then_expr: $when_then,
2812                else_expr: None,
2813            };
2814
2815            let expected =
2816                cast_helper(case.clone(), &$case_when_type, &$then_else_type, &$schema);
2817
2818            let actual = coerce_case_expression(case, &$schema)?;
2819            assert_eq!(expected, actual);
2820        };
2821    }
2822
2823    #[test]
2824    fn tes_case_when_list() -> Result<()> {
2825        let inner_field = Arc::new(Field::new_list_field(DataType::Int64, true));
2826        let schema = Arc::new(DFSchema::from_unqualified_fields(
2827            vec![
2828                Field::new(
2829                    "large_list",
2830                    DataType::LargeList(Arc::clone(&inner_field)),
2831                    true,
2832                ),
2833                Field::new(
2834                    "fixed_list",
2835                    DataType::FixedSizeList(Arc::clone(&inner_field), 3),
2836                    true,
2837                ),
2838                Field::new("list", DataType::List(inner_field), true),
2839            ]
2840            .into(),
2841            std::collections::HashMap::new(),
2842        )?);
2843
2844        test_case_expression!(
2845            Some("list"),
2846            vec![(Box::new(col("large_list")), Box::new(lit("1")))],
2847            DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int64, true))),
2848            Utf8,
2849            schema
2850        );
2851
2852        test_case_expression!(
2853            Some("large_list"),
2854            vec![(Box::new(col("list")), Box::new(lit("1")))],
2855            DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int64, true))),
2856            Utf8,
2857            schema
2858        );
2859
2860        test_case_expression!(
2861            Some("list"),
2862            vec![(Box::new(col("fixed_list")), Box::new(lit("1")))],
2863            DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
2864            Utf8,
2865            schema
2866        );
2867
2868        test_case_expression!(
2869            Some("fixed_list"),
2870            vec![(Box::new(col("list")), Box::new(lit("1")))],
2871            DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
2872            Utf8,
2873            schema
2874        );
2875
2876        test_case_expression!(
2877            Some("fixed_list"),
2878            vec![(Box::new(col("large_list")), Box::new(lit("1")))],
2879            DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int64, true))),
2880            Utf8,
2881            schema
2882        );
2883
2884        test_case_expression!(
2885            Some("large_list"),
2886            vec![(Box::new(col("fixed_list")), Box::new(lit("1")))],
2887            DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int64, true))),
2888            Utf8,
2889            schema
2890        );
2891        Ok(())
2892    }
2893
2894    #[test]
2895    fn test_then_else_list() -> Result<()> {
2896        let inner_field = Arc::new(Field::new_list_field(DataType::Int64, true));
2897        let schema = Arc::new(DFSchema::from_unqualified_fields(
2898            vec![
2899                Field::new("boolean", DataType::Boolean, true),
2900                Field::new(
2901                    "large_list",
2902                    DataType::LargeList(Arc::clone(&inner_field)),
2903                    true,
2904                ),
2905                Field::new(
2906                    "fixed_list",
2907                    DataType::FixedSizeList(Arc::clone(&inner_field), 3),
2908                    true,
2909                ),
2910                Field::new("list", DataType::List(inner_field), true),
2911            ]
2912            .into(),
2913            std::collections::HashMap::new(),
2914        )?);
2915
2916        // large list and list
2917        test_case_expression!(
2918            None::<String>,
2919            vec![
2920                (Box::new(col("boolean")), Box::new(col("large_list"))),
2921                (Box::new(col("boolean")), Box::new(col("list")))
2922            ],
2923            DataType::Boolean,
2924            DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int64, true))),
2925            schema
2926        );
2927
2928        test_case_expression!(
2929            None::<String>,
2930            vec![
2931                (Box::new(col("boolean")), Box::new(col("list"))),
2932                (Box::new(col("boolean")), Box::new(col("large_list")))
2933            ],
2934            DataType::Boolean,
2935            DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int64, true))),
2936            schema
2937        );
2938
2939        // fixed list and list
2940        test_case_expression!(
2941            None::<String>,
2942            vec![
2943                (Box::new(col("boolean")), Box::new(col("fixed_list"))),
2944                (Box::new(col("boolean")), Box::new(col("list")))
2945            ],
2946            DataType::Boolean,
2947            DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
2948            schema
2949        );
2950
2951        test_case_expression!(
2952            None::<String>,
2953            vec![
2954                (Box::new(col("boolean")), Box::new(col("list"))),
2955                (Box::new(col("boolean")), Box::new(col("fixed_list")))
2956            ],
2957            DataType::Boolean,
2958            DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
2959            schema
2960        );
2961
2962        // fixed list and large list
2963        test_case_expression!(
2964            None::<String>,
2965            vec![
2966                (Box::new(col("boolean")), Box::new(col("fixed_list"))),
2967                (Box::new(col("boolean")), Box::new(col("large_list")))
2968            ],
2969            DataType::Boolean,
2970            DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int64, true))),
2971            schema
2972        );
2973
2974        test_case_expression!(
2975            None::<String>,
2976            vec![
2977                (Box::new(col("boolean")), Box::new(col("large_list"))),
2978                (Box::new(col("boolean")), Box::new(col("fixed_list")))
2979            ],
2980            DataType::Boolean,
2981            DataType::LargeList(Arc::new(Field::new_list_field(DataType::Int64, true))),
2982            schema
2983        );
2984        Ok(())
2985    }
2986
2987    #[test]
2988    fn test_map_with_diff_name() -> Result<()> {
2989        let mut builder = SchemaBuilder::new();
2990        builder.push(Field::new("key", Utf8, false));
2991        builder.push(Field::new("value", DataType::Float64, true));
2992        let struct_fields = builder.finish().fields;
2993
2994        let fields =
2995            Field::new("entries", DataType::Struct(struct_fields.clone()), false);
2996        let map_type_entries = DataType::Map(Arc::new(fields), false);
2997
2998        let fields = Field::new("key_value", DataType::Struct(struct_fields), false);
2999        let may_type_custom = DataType::Map(Arc::new(fields), false);
3000
3001        let expr = col("a").eq(cast(col("a"), may_type_custom));
3002        let empty = empty_with_type(map_type_entries);
3003        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
3004
3005        assert_analyzed_plan_eq!(
3006            plan,
3007            @r#"
3008        Projection: a = CAST(CAST(a AS Map("key_value": non-null Struct("key": non-null Utf8, "value": Float64), unsorted)) AS Map("entries": non-null Struct("key": non-null Utf8, "value": Float64), unsorted))
3009          EmptyRelation: rows=0
3010        "#
3011        )
3012    }
3013
3014    #[test]
3015    fn array_element_preserves_parquet_list_field_name() -> Result<()> {
3016        let list_type = DataType::List(Arc::new(Field::new(
3017            "element",
3018            DataType::Struct(
3019                vec![
3020                    Field::new("id", Utf8, true),
3021                    Field::new("prim", DataType::Boolean, true),
3022                ]
3023                .into(),
3024            ),
3025            true,
3026        )));
3027
3028        let expr = ScalarUDF::from(TestArrayElementUDF).call(vec![col("a"), lit(1_i64)]);
3029        let empty = empty_with_type(list_type);
3030        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
3031
3032        assert_analyzed_plan_eq!(
3033            plan,
3034            @r#"
3035        Projection: TestArrayElementUDF(a, Int64(1))
3036          EmptyRelation: rows=0
3037        "#
3038        )
3039    }
3040
3041    #[test]
3042    fn interval_plus_timestamp() -> Result<()> {
3043        // SELECT INTERVAL '1' YEAR + '2000-01-01T00:00:00'::timestamp;
3044        let expr = Expr::BinaryExpr(BinaryExpr::new(
3045            Box::new(lit(ScalarValue::IntervalYearMonth(Some(12)))),
3046            Operator::Plus,
3047            Box::new(cast(
3048                lit("2000-01-01T00:00:00"),
3049                DataType::Timestamp(TimeUnit::Nanosecond, None),
3050            )),
3051        ));
3052        let empty = empty();
3053        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
3054
3055        assert_analyzed_plan_eq!(
3056            plan,
3057            @r#"
3058        Projection: IntervalYearMonth("12") + CAST(Utf8("2000-01-01T00:00:00") AS Timestamp(ns))
3059          EmptyRelation: rows=0
3060        "#
3061        )
3062    }
3063
3064    #[test]
3065    fn timestamp_subtract_timestamp() -> Result<()> {
3066        let expr = Expr::BinaryExpr(BinaryExpr::new(
3067            Box::new(cast(
3068                lit("1998-03-18"),
3069                DataType::Timestamp(TimeUnit::Nanosecond, None),
3070            )),
3071            Operator::Minus,
3072            Box::new(cast(
3073                lit("1998-03-18"),
3074                DataType::Timestamp(TimeUnit::Nanosecond, None),
3075            )),
3076        ));
3077        let empty = empty();
3078        let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
3079
3080        assert_analyzed_plan_eq!(
3081            plan,
3082            @r#"
3083        Projection: CAST(Utf8("1998-03-18") AS Timestamp(ns)) - CAST(Utf8("1998-03-18") AS Timestamp(ns))
3084          EmptyRelation: rows=0
3085        "#
3086        )
3087    }
3088
3089    #[test]
3090    fn in_subquery_cast_subquery() -> Result<()> {
3091        let empty_int32 = empty_with_type(DataType::Int32);
3092        let empty_int64 = empty_with_type(DataType::Int64);
3093
3094        let in_subquery_expr = Expr::InSubquery(InSubquery::new(
3095            Box::new(col("a")),
3096            Subquery {
3097                subquery: empty_int32,
3098                outer_ref_columns: vec![],
3099                spans: Spans::new(),
3100            },
3101            false,
3102        ));
3103        let plan = LogicalPlan::Filter(Filter::try_new(in_subquery_expr, empty_int64)?);
3104        // add cast for subquery
3105
3106        assert_analyzed_plan_eq!(
3107            plan,
3108            @r"
3109        Filter: a IN (<subquery>)
3110          Subquery:
3111            Projection: CAST(a AS Int64)
3112              EmptyRelation: rows=0
3113          EmptyRelation: rows=0
3114        "
3115        )
3116    }
3117
3118    #[test]
3119    fn in_subquery_cast_expr() -> Result<()> {
3120        let empty_int32 = empty_with_type(DataType::Int32);
3121        let empty_int64 = empty_with_type(DataType::Int64);
3122
3123        let in_subquery_expr = Expr::InSubquery(InSubquery::new(
3124            Box::new(col("a")),
3125            Subquery {
3126                subquery: empty_int64,
3127                outer_ref_columns: vec![],
3128                spans: Spans::new(),
3129            },
3130            false,
3131        ));
3132        let plan = LogicalPlan::Filter(Filter::try_new(in_subquery_expr, empty_int32)?);
3133
3134        // add cast for subquery
3135        assert_analyzed_plan_eq!(
3136            plan,
3137            @r"
3138        Filter: CAST(a AS Int64) IN (<subquery>)
3139          Subquery:
3140            EmptyRelation: rows=0
3141          EmptyRelation: rows=0
3142        "
3143        )
3144    }
3145
3146    #[test]
3147    fn in_subquery_cast_all() -> Result<()> {
3148        let empty_inside = empty_with_type(DataType::Decimal128(10, 5));
3149        let empty_outside = empty_with_type(DataType::Decimal128(8, 8));
3150
3151        let in_subquery_expr = Expr::InSubquery(InSubquery::new(
3152            Box::new(col("a")),
3153            Subquery {
3154                subquery: empty_inside,
3155                outer_ref_columns: vec![],
3156                spans: Spans::new(),
3157            },
3158            false,
3159        ));
3160        let plan = LogicalPlan::Filter(Filter::try_new(in_subquery_expr, empty_outside)?);
3161
3162        // add cast for subquery
3163        assert_analyzed_plan_eq!(
3164            plan,
3165            @r"
3166        Filter: CAST(a AS Decimal128(13, 8)) IN (<subquery>)
3167          Subquery:
3168            Projection: CAST(a AS Decimal128(13, 8))
3169              EmptyRelation: rows=0
3170          EmptyRelation: rows=0
3171        "
3172        )
3173    }
3174}