Skip to content

Commit 01d00fd

Browse files
authored
Implement cast between Date and Timestamp (#4726)
* Date to timestamp conversion * bit of cleaning * fmt
1 parent a8f1f8a commit 01d00fd

4 files changed

Lines changed: 105 additions & 3 deletions

File tree

‎datafusion/core/tests/sql/timestamp.rs‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1679,3 +1679,68 @@ async fn test_current_time() -> Result<()> {
16791679
assert_batches_eq!(expected, &results);
16801680
Ok(())
16811681
}
1682+
1683+
#[tokio::test]
1684+
async fn test_ts_dt_binary_ops() -> Result<()> {
1685+
let ctx = SessionContext::new();
1686+
1687+
// test cast in where clause
1688+
let sql =
1689+
"select count(1) result from (select now() as n) a where n = '2000-01-01'::date";
1690+
let results = execute_to_batches(&ctx, sql).await;
1691+
1692+
let expected = vec![
1693+
"+--------+",
1694+
"| result |",
1695+
"+--------+",
1696+
"| 0 |",
1697+
"+--------+",
1698+
];
1699+
1700+
assert_batches_eq!(expected, &results);
1701+
1702+
// test cast in where ge clause
1703+
let sql =
1704+
"select count(1) result from (select now() as n) a where n >= '2000-01-01'::date";
1705+
let results = execute_to_batches(&ctx, sql).await;
1706+
1707+
let expected = vec![
1708+
"+--------+",
1709+
"| result |",
1710+
"+--------+",
1711+
"| 1 |",
1712+
"+--------+",
1713+
];
1714+
1715+
assert_batches_eq!(expected, &results);
1716+
1717+
// test cast in equal select
1718+
let sql = "select now() = '2000-01-01'::date as result";
1719+
let results = execute_to_batches(&ctx, sql).await;
1720+
1721+
let expected = vec![
1722+
"+--------+",
1723+
"| result |",
1724+
"+--------+",
1725+
"| false |",
1726+
"+--------+",
1727+
];
1728+
1729+
assert_batches_eq!(expected, &results);
1730+
1731+
// test cast in gt select
1732+
let sql = "select now() >= '2000-01-01'::date as result";
1733+
let results = execute_to_batches(&ctx, sql).await;
1734+
1735+
let expected = vec![
1736+
"+--------+",
1737+
"| result |",
1738+
"+--------+",
1739+
"| true |",
1740+
"+--------+",
1741+
];
1742+
1743+
assert_batches_eq!(expected, &results);
1744+
1745+
Ok(())
1746+
}

‎datafusion/expr/src/type_coercion/binary.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -575,6 +575,8 @@ fn temporal_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataTyp
575575
},
576576
(Timestamp(_, tz), Utf8) => Some(Timestamp(TimeUnit::Nanosecond, tz.clone())),
577577
(Utf8, Timestamp(_, tz)) => Some(Timestamp(TimeUnit::Nanosecond, tz.clone())),
578+
(Timestamp(_, _), Date32) => Some(Date32),
579+
(Timestamp(_, _), Date64) => Some(Date64),
578580
(Timestamp(lhs_unit, lhs_tz), Timestamp(rhs_unit, rhs_tz)) => {
579581
let tz = match (lhs_tz, rhs_tz) {
580582
// can't cast across timezones

‎datafusion/expr/src/type_coercion/functions.rs‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,10 @@ fn maybe_data_types(
131131
/// See the module level documentation for more detail on coercion.
132132
pub fn can_coerce_from(type_into: &DataType, type_from: &DataType) -> bool {
133133
use self::DataType::*;
134+
135+
if type_into == type_from {
136+
return true;
137+
}
134138
// Null can convert to most of types
135139
match type_into {
136140
Int8 => matches!(type_from, Null | Int8),
@@ -173,12 +177,35 @@ pub fn can_coerce_from(type_into: &DataType, type_from: &DataType) -> bool {
173177
Timestamp(TimeUnit::Nanosecond, None) => {
174178
matches!(type_from, Null | Timestamp(_, None))
175179
}
180+
Date32 => {
181+
matches!(type_from, Null | Timestamp(_, None))
182+
}
176183
Utf8 | LargeUtf8 => true,
177184
Null => can_cast_types(type_from, type_into),
178185
_ => false,
179186
}
180187
}
181188

189+
/// Returns a common coerced datatype between 2 given datatypes
190+
///
191+
/// See the module level documentation for more detail on coercion.
192+
pub fn get_common_coerced_type(
193+
left_datatype: DataType,
194+
right_datatype: DataType,
195+
) -> Result<DataType> {
196+
if left_datatype == right_datatype || can_coerce_from(&left_datatype, &right_datatype)
197+
{
198+
Ok(left_datatype)
199+
} else if can_coerce_from(&right_datatype, &left_datatype) {
200+
Ok(right_datatype)
201+
} else {
202+
Err(DataFusionError::Plan(format!(
203+
"Datatypes cannot be casted into common type {:?} <-> {:?}",
204+
left_datatype, right_datatype
205+
)))
206+
}
207+
}
208+
182209
#[cfg(test)]
183210
mod tests {
184211
use super::*;

‎datafusion/physical-expr/src/planner.rs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ use crate::{
2828
use arrow::datatypes::{DataType, Schema};
2929
use datafusion_common::{DFSchema, DataFusionError, Result, ScalarValue};
3030
use datafusion_expr::expr::Cast;
31+
use datafusion_expr::type_coercion::functions::get_common_coerced_type;
3132
use datafusion_expr::{
3233
binary_expr, Between, BinaryExpr, Expr, GetIndexedField, Like, Operator, TryCast,
3334
};
@@ -198,9 +199,16 @@ pub fn create_physical_expr(
198199
input_schema,
199200
)?)),
200201
_ => {
201-
// assume that we can coerce both sides into a common type
202-
// and then perform a binary operation
203-
binary(lhs, *op, rhs, input_schema)
202+
let target_datatype = get_common_coerced_type(
203+
lhs.data_type(input_schema)?,
204+
rhs.data_type(input_schema)?,
205+
)?;
206+
binary(
207+
expressions::cast(lhs, input_schema, target_datatype.clone())?,
208+
*op,
209+
expressions::cast(rhs, input_schema, target_datatype)?,
210+
input_schema,
211+
)
204212
}
205213
}
206214
}

0 commit comments

Comments
 (0)