1use 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#[derive(Default, Debug)]
69pub struct TypeCoercion {}
70
71impl TypeCoercion {
72 pub fn new() -> Self {
73 Self {}
74 }
75}
76
77fn 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 let transformed_plan = plan
105 .transform_up_with_subqueries(|plan| analyze_internal(&EMPTY_SCHEMA, plan))?
106 .data;
107
108 coerce_output(transformed_plan, config)
110 }
111}
112
113fn analyze_internal(
117 external_schema: &DFSchema,
118 plan: LogicalPlan,
119) -> Result<Transformed<LogicalPlan>> {
120 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 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 schema.merge(external_schema);
151
152 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 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 .map_data(|plan| expr_rewrite.coerce_plan(plan))?
171 .map_data(|plan| plan.recompute_schema())
173}
174
175pub struct TypeCoercionRewriter<'a> {
177 pub(crate) schema: &'a DFSchema,
178}
179
180impl<'a> TypeCoercionRewriter<'a> {
181 pub fn new(schema: &'a DFSchema) -> Self {
184 Self { schema }
185 }
186
187 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 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 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 = join
280 .filter
281 .map(|expr| self.coerce_predicate(expr, "Join condition"))
282 .transpose()?;
283
284 Ok(LogicalPlan::Join(join))
285 }
286
287 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 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 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 (
450 Operator::Plus | Operator::Minus,
451 Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64,
452 Interval(IntervalUnit::MonthDayNano),
453 Date32 | Date64,
454 ) => {
455 let expr = match *left_current_type {
457 Int64 => expr,
458 _ => cast(expr, Int64),
459 };
460 let expr = expr * lit(ScalarValue::from(SECONDS_IN_DAY));
462 let expr = cast(expr, Duration(TimeUnit::Second));
464 cast(expr, Interval(IntervalUnit::MonthDayNano))
466 }
467 (
476 Operator::Plus | Operator::Minus,
477 Time32(_) | Time64(_),
478 Duration(TimeUnit::Nanosecond),
479 Timestamp(TimeUnit::Nanosecond, None),
480 ) => {
481 let expr = time_to_nanos(expr, left_current_type)?;
483 cast(expr, Duration(TimeUnit::Nanosecond))
485 }
486 (
492 Operator::Plus | Operator::Minus,
493 Time32(_) | Time64(_),
494 Interval(IntervalUnit::MonthDayNano),
495 Interval(IntervalUnit::MonthDayNano),
496 ) => {
497 let expr = time_to_nanos(expr, left_current_type)?;
499 let expr = cast(expr, Duration(TimeUnit::Nanosecond));
501 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 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 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 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 ¤t_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 #[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
946fn 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
995fn coerce_scalar(target_type: &DataType, value: &ScalarValue) -> Result<ScalarValue> {
998 match value {
999 ScalarValue::Utf8(Some(val)) => {
1001 ScalarValue::try_from_string(val.clone(), target_type)
1002 }
1003 s => {
1004 if s.is_null() {
1005 ScalarValue::try_from(target_type)
1007 } else {
1008 Ok(s.clone())
1012 }
1013 }
1014 }
1015}
1016
1017fn 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 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
1040fn 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
1052fn 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
1090fn 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
1118fn 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
1127fn 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(¤t_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 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 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 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
1273pub 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 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
1394fn 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 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 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 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 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 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 coerce_on_output_if_viewtype!(
1701 false,
1702 plan.clone(),
1703 @r"
1704 Projection: a
1705 EmptyRelation: rows=0
1706 "
1707 )?;
1708
1709 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 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 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 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 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 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 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 let plan = LogicalPlan::Projection(Projection::try_new(
1790 vec![col("a")],
1791 Arc::new(sort_plan),
1792 )?);
1793 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 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 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 coerce_on_output_if_viewtype!(
1832 false,
1833 plan.clone(),
1834 @r"
1835 Projection: a
1836 EmptyRelation: rows=0
1837 "
1838 )?;
1839
1840 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 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 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 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 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 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 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 let plan = LogicalPlan::Projection(Projection::try_new(
1912 vec![col("a")],
1913 Arc::new(sort_plan),
1914 )?);
1915
1916 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 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 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 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 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 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(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(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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}