Skip to content

Commit

Permalink
Add support
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed Jan 30, 2023
1 parent 220f7b0 commit f06ad3c
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 7 deletions.
11 changes: 8 additions & 3 deletions datafusion/expr/src/type_coercion/functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,10 @@ use arrow::{
};
use datafusion_common::{DataFusionError, Result};

/// Performs type coercion for functions Returns the data types that
/// each argument must be coerced to match `signature`.
/// Performs type coercion for function arguments.
///
/// Returns the data types to which each argument must be coerced to
/// match `signature`.
///
/// For more details on coercion in general, please see the
/// [`type_coercion`](crate::type_coercion) module.
Expand Down Expand Up @@ -175,7 +177,10 @@ pub fn can_coerce_from(type_into: &DataType, type_from: &DataType) -> bool {
| Decimal128(_, _)
),
Timestamp(TimeUnit::Nanosecond, None) => {
matches!(type_from, Null | Timestamp(_, None))
matches!(type_from, Null | Timestamp(_, None) | Utf8 | LargeUtf8)
}
Interval(_) => {
matches!(type_from, Utf8 | LargeUtf8)
}
Utf8 | LargeUtf8 => true,
Null => can_cast_types(type_from, type_into),
Expand Down
31 changes: 27 additions & 4 deletions datafusion/optimizer/src/type_coercion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ use datafusion_expr::{
WindowFrame, WindowFrameBound, WindowFrameUnits,
};
use datafusion_expr::{ExprSchemable, Signature};
use log::debug;

use crate::utils::rewrite_preserving_name;
use crate::{OptimizerConfig, OptimizerRule};
Expand Down Expand Up @@ -550,13 +551,35 @@ fn coerce_arguments_for_signature(

let new_types = data_types(&current_types, signature)?;

expressions
.iter()
debug!("AAL current types: {current_types:?}");
debug!("AAL new types: {new_types:?}");


let exprs = expressions
.into_iter()
.enumerate()
.map(|(i, expr)| expr.clone().cast_to(&new_types[i], schema))
.collect::<Result<Vec<_>>>()
.map(|(i, expr)| cast_expr(expr, &new_types[i], schema))
.collect::<Result<Vec<_>>>();

debug!("AAL new exprs: {new_types:?}");

exprs
}

/// Cast expr to the specified type, if possible
fn cast_expr(expr: &Expr, to_type: &DataType, schema: &DFSchema) -> Result<Expr> {
// Special case until Interval coercion is handled in arrow-rs
match (expr, to_type) {
(Expr::Literal(ScalarValue::Utf8(Some(s))), DataType::Interval(_)) => {
parse_interval("millisecond", s.as_str()).map(Expr::Literal)
},
_ => expr.clone().cast_to(to_type, schema),
}
}




/// Returns the coerced exprs for each `input_exprs`.
/// Get the coerced data type from `aggregate_rule::coerce_types` and add `try_cast` if the
/// data type of `input_exprs` need to be coerced.
Expand Down

0 comments on commit f06ad3c

Please sign in to comment.