adriangb commented on code in PR #25094:
URL: https://github.com/apache/datafusion/pull/25094#discussion_r3981582469
##########
datafusion/optimizer/src/scalar_subquery_to_join.rs:
##########
@@ -398,9 +398,7 @@ fn build_join(
// itself be NULL) otherwise.
let mut compensation_exprs = HashMap::new();
if let Some(expr_map) = collected_count_expr_map {
- let mut expr_rewrite = TypeCoercionRewriter {
- schema: new_plan.schema(),
- };
+ let mut expr_rewrite = TypeCoercionRewriter::new(new_plan.schema());
Review Comment:
Same here — the reasoning is sound, it's just invisible:
```suggestion
// No session timezone: this builds a *searched* `CASE WHEN`, so the
// `CASE expr WHEN` comparison rule never applies, and the only type
// introduced is the untyped `Null` of the HAVING arm. The
expressions
// themselves come from a plan `TypeCoercion` has already coerced.
let mut expr_rewrite = TypeCoercionRewriter::new(new_plan.schema());
```
##########
datafusion/optimizer/src/utils.rs:
##########
@@ -244,7 +244,7 @@ fn evaluate_expr_with_null_column<'a>(
}
fn coerce(expr: Expr, schema: &DFSchema) -> Result<Expr> {
Review Comment:
This answers the question I raised last round about why no session timezone
is threaded in here — worth writing down rather than leaving the omission to be
rediscovered:
```suggestion
/// No session timezone is threaded in here on purpose. This runs over a
plan the
/// `TypeCoercion` analyzer has already coerced, so every binary operand pair
/// already shares a type; the only type this helper introduces is the
`Null` of
/// the dummy column in `evaluate_expr_with_null_column`, and `Null` against
a
/// timestamp is short-circuited by `null_coercion` before the aware/naive
rule.
fn coerce(expr: Expr, schema: &DFSchema) -> Result<Expr> {
```
##########
datafusion/optimizer/src/analyzer/type_coercion.rs:
##########
@@ -778,26 +809,27 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
}) => {
let expr_type = expr.get_type(self.schema)?;
let low_type = low.get_type(self.schema)?;
- let low_coerced_type = comparison_coercion(&expr_type,
&low_type)
- .ok_or_else(|| {
- internal_datafusion_err!(
- "Failed to coerce types {expr_type} and {low_type}
in BETWEEN expression"
- )
- })?;
+ let low_coerced_type =
comparison_coercion_with_session_timezone(
+ &expr_type,
+ &low_type,
+ self.session_time_zone,
+ )
+ .ok_or_else(|| {
+ internal_datafusion_err!(
+ "Failed to coerce types {expr_type} and {low_type} in
BETWEEN expression"
+ )
+ })?;
let high_type = high.get_type(self.schema)?;
- let high_coerced_type = comparison_coercion(&expr_type,
&high_type)
+ let coercion_type = comparison_coercion_with_session_timezone(
+ &low_coerced_type,
+ &high_type,
+ self.session_time_zone,
+ )
.ok_or_else(|| {
internal_datafusion_err!(
"Failed to coerce types {expr_type} and
{high_type} in BETWEEN expression"
Review Comment:
This arm coerces `low_coerced_type` against `high_type`, but the message
reports `expr_type` — the type it didn't coerce here. With the two-step fold,
`expr_type` and `low_type` have already been reconciled, so naming all three is
what actually tells you which pair failed:
```suggestion
"Failed to coerce types {expr_type}, {low_type}
and {high_type} in BETWEEN expression"
```
##########
datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs:
##########
@@ -57,6 +57,27 @@ fn test_date_timestamp_arithmetic_error() -> Result<()> {
Ok(())
}
+#[test]
+fn test_timestamp_session_timezone_coercion() -> Result<()> {
+ let aware = DataType::Timestamp(Millisecond,
Some("America/New_York".into()));
+ let naive = DataType::Timestamp(Nanosecond, None);
+ let expected = DataType::Timestamp(Nanosecond, Some("+08:00".into()));
+
+ for op in [Operator::Minus, Operator::Eq, Operator::Gt] {
+ let (lhs, rhs) = BinaryTypeCoercer::new(&aware, &op, &naive)
+ .with_session_time_zone(Some("+08:00"))
+ .get_input_types()?;
+ assert_eq!((lhs, rhs), (expected.clone(), expected.clone()));
+
+ let (lhs, rhs) = BinaryTypeCoercer::new(&naive, &op, &aware)
+ .with_session_time_zone(Some("+08:00"))
+ .get_input_types()?;
+ assert_eq!((lhs, rhs), (expected.clone(), expected.clone()));
+ }
+
+ Ok(())
+}
Review Comment:
`test_timestamp_session_timezone_coercion` only exercises `Some(tz)`; the
`None` path — the one that has to stay exactly as it was in 55.0.0 — is covered
only end to end in the slt. A unit test alongside this one pins the part that's
easy to regress:
```rust
#[test]
fn test_timestamp_without_session_timezone_coercion() -> Result<()> {
let aware_ms = DataType::Timestamp(Millisecond,
Some("America/New_York".into()));
let aware_ns = DataType::Timestamp(Nanosecond,
Some("America/New_York".into()));
let naive = DataType::Timestamp(Nanosecond, None);
// Comparisons fall back to reading the naive side in the aware side's
// timezone, in either operand order.
for op in [Operator::Eq, Operator::Gt] {
for (lhs, rhs) in [(&aware_ms, &naive), (&naive, &aware_ms)] {
let (lhs, rhs) = BinaryTypeCoercer::new(lhs, &op, rhs)
.with_session_time_zone(None)
.get_input_types()?;
assert_eq!((lhs, rhs), (aware_ns.clone(), aware_ns.clone()));
}
}
// `Minus` only agrees with that when the units differ and force a
coercion.
// At a shared unit the operands are left alone and arrow subtracts their
// raw values — which is exactly the behaviour this PR exists to correct.
let (lhs, rhs) = BinaryTypeCoercer::new(&aware_ms, &Operator::Minus,
&naive)
.with_session_time_zone(None)
.get_input_types()?;
assert_eq!((lhs, rhs), (aware_ns.clone(), aware_ns.clone()));
let (lhs, rhs) = BinaryTypeCoercer::new(&aware_ns, &Operator::Minus,
&naive)
.with_session_time_zone(None)
.get_input_types()?;
assert_eq!((lhs, rhs), (aware_ns.clone(), naive.clone()));
// A session timezone doesn't disturb pairs that are both aware or both
// naive: those already denote the same kind of value.
let (lhs, rhs) = BinaryTypeCoercer::new(&aware_ms, &Operator::Minus,
&aware_ms)
.with_session_time_zone(Some("+08:00"))
.get_input_types()?;
assert_eq!((lhs, rhs), (aware_ms.clone(), aware_ms));
let (lhs, rhs) = BinaryTypeCoercer::new(&naive, &Operator::Minus, &naive)
.with_session_time_zone(Some("+08:00"))
.get_input_types()?;
assert_eq!((lhs, rhs), (naive.clone(), naive));
Ok(())
}
```
That second `Minus` assertion is the interesting one: `Timestamp(ns,
Some(tz)) - Timestamp(ns, None)` still coerces nothing without a session
timezone, which is the 55.0.0 behaviour left intact.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
