alamb commented on code in PR #18789:
URL: https://github.com/apache/datafusion/pull/18789#discussion_r2543034309
##########
datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs:
##########
@@ -1956,6 +1962,36 @@ impl<S: SimplifyInfo> TreeNodeRewriter for
Simplifier<'_, S> {
}))
}
+ // =======================================
Review Comment:
I agree moving this into `date_part`'s implementation makes sense
> @alamb This works as is right now, but I like the idea of adding a new
method ScalarUDFImpl::simplify_predicate() as a unified api for all scalar
functions. We then can move this code in the DatePartFunc implementation of
ScalarUDFImpl and create a simplification rule for scalar functions.
I didn't quite follow why this couldn't be added to the existing `simplify`
method 🤔
https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.ScalarUDFImpl.html#method.simplify
##########
datafusion/optimizer/src/simplify_expressions/unwrap_date_part.rs:
##########
@@ -0,0 +1,371 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::datatypes::{DataType, TimeUnit};
+use chrono::NaiveDate;
+use datafusion_common::{
+ internal_err, tree_node::Transformed, DataFusionError, Result, ScalarValue,
+};
+use datafusion_expr::{
+ and, expr::ScalarFunction, lit, or, simplify::SimplifyInfo, BinaryExpr,
Expr,
+ Operator,
+};
+
+pub(super) fn unwrap_date_part_in_comparison_for_binary<S: SimplifyInfo>(
+ info: &S,
+ cast_expr: Expr,
+ literal: Expr,
+ op: Operator,
+) -> Result<Transformed<Expr>> {
+ let (args, lit_value) = match (cast_expr, literal) {
+ (
+ Expr::ScalarFunction(ScalarFunction { func, args }),
+ Expr::Literal(lit_value, _),
+ ) if func.name() == "date_part" => (args, lit_value),
+ _ => return internal_err!("Expect date_part expr and literal"),
+ };
+ let expr = Box::new(args[1].clone());
+
+ let Ok(expr_type) = info.get_data_type(&expr) else {
+ return internal_err!("Can't get the data type of the expr {:?}",
&expr);
+ };
+
+ // Helper to cast literal
+ let cast_year = |updated_year: &ScalarValue| -> Result<ScalarValue,
DataFusionError> {
+ year_literal_to_type(updated_year, &expr_type).ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "Can't cast {lit_value} to type {expr_type}"
+ ))
+ })
+ };
+
+ let rewritten_expr = match op {
+ Operator::Lt | Operator::GtEq => {
+ let v = cast_year(&lit_value)?;
+ Expr::BinaryExpr(BinaryExpr {
+ left: expr,
+ op,
+ right: Box::new(lit(v)),
+ })
+ }
+ Operator::Gt => {
+ let year = match lit_value {
+ ScalarValue::Int32(Some(y)) => y + 1,
+ _ => return internal_err!("Expected Int32 Literal"),
+ };
+ let updated_year = ScalarValue::Int32(Some(year));
+ let v = cast_year(&updated_year)?;
+ Expr::BinaryExpr(BinaryExpr {
+ left: expr,
+ op: Operator::GtEq,
+ right: Box::new(lit(v)),
+ })
+ }
+ Operator::LtEq => {
+ let year = match lit_value {
+ ScalarValue::Int32(Some(y)) => y + 1,
+ _ => return internal_err!("Expected Int32 Literal"),
+ };
+ let updated_year = ScalarValue::Int32(Some(year));
+ let v = cast_year(&updated_year)?;
+ Expr::BinaryExpr(BinaryExpr {
+ left: expr,
+ op: Operator::Lt,
+ right: Box::new(lit(v)),
+ })
+ }
+ Operator::Eq => {
+ let year = match lit_value {
+ ScalarValue::Int32(Some(y)) => y + 1,
+ _ => return internal_err!("Expected Int32 Literal"),
+ };
+ let updated_year = ScalarValue::Int32(Some(year));
+ let lower = cast_year(&lit_value)?;
+ let upper = cast_year(&updated_year)?;
+ and(
+ Expr::BinaryExpr(BinaryExpr {
+ left: expr.clone(),
+ op: Operator::GtEq,
+ right: Box::new(lit(lower)),
+ }),
+ Expr::BinaryExpr(BinaryExpr {
+ left: expr,
+ op: Operator::Lt,
+ right: Box::new(lit(upper)),
+ }),
+ )
+ }
+ Operator::NotEq => {
+ let year = match lit_value {
+ ScalarValue::Int32(Some(y)) => y + 1,
+ _ => return internal_err!("Expected Int32 Literal"),
+ };
+ let updated_year = ScalarValue::Int32(Some(year));
+ let lower = cast_year(&lit_value)?;
+ let upper = cast_year(&updated_year)?;
+ or(
+ Expr::BinaryExpr(BinaryExpr {
+ left: expr.clone(),
+ op: Operator::Lt,
+ right: Box::new(lit(lower)),
+ }),
+ Expr::BinaryExpr(BinaryExpr {
+ left: expr,
+ op: Operator::GtEq,
+ right: Box::new(lit(upper)),
+ }),
+ )
+ }
+ _ => return internal_err!("Expect comparison operators"),
+ };
+ Ok(Transformed::yes(rewritten_expr))
+}
+
+pub(super) fn
is_date_part_expr_and_support_unwrap_date_part_in_comparison_for_binary<
+ S: SimplifyInfo,
+>(
+ info: &S,
+ expr: &Expr,
+ op: Operator,
+ literal: &Expr,
+) -> bool {
+ match (expr, op, literal) {
+ (
+ Expr::ScalarFunction(ScalarFunction { func, args }),
+ Operator::Eq
+ | Operator::NotEq
+ | Operator::Gt
+ | Operator::Lt
+ | Operator::GtEq
+ | Operator::LtEq,
+ Expr::Literal(lit_val, _),
+ ) if func.name() == "date_part" => {
+ let left_expr = Box::new(args[1].clone());
+ let Ok(expr_type) = info.get_data_type(&left_expr) else {
+ return false;
+ };
+ let Ok(_lit_type) = info.get_data_type(literal) else {
+ return false;
+ };
+
+ year_literal_to_type(lit_val, &expr_type).is_some()
+ }
+ _ => false,
+ }
+}
+
+/// Cast the year to the right datatype
+fn year_literal_to_type(
+ lit_value: &ScalarValue,
+ target_type: &DataType,
+) -> Option<ScalarValue> {
+ let year = match lit_value {
+ ScalarValue::Int32(Some(y)) => *y,
+ _ => return None,
+ };
+ // Can only extract year from Date32/64 and Timestamp
+ match target_type {
+ DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _) => {}
+ _ => return None,
+ }
+
+ let naive_date = NaiveDate::from_ymd_opt(year, 1, 1).expect("Invalid
year");
+
+ let casted = match target_type {
+ DataType::Date32 => {
+ let days = naive_date
+ .signed_duration_since(NaiveDate::from_ymd_opt(1970, 1, 1)?)
+ .num_days() as i32;
+ ScalarValue::Date32(Some(days))
+ }
+ DataType::Date64 => {
+ let milis = naive_date
+ .signed_duration_since(NaiveDate::from_ymd_opt(1970, 1, 1)?)
+ .num_milliseconds();
+ ScalarValue::Date64(Some(milis))
+ }
+ DataType::Timestamp(unit, tz) => {
+ let days = naive_date
+ .signed_duration_since(NaiveDate::from_ymd_opt(1970, 1, 1)?)
+ .num_days();
+ match unit {
+ TimeUnit::Second => {
+ ScalarValue::TimestampSecond(Some(days * 86_400),
tz.clone())
+ }
+ TimeUnit::Millisecond => {
+ ScalarValue::TimestampMillisecond(Some(days * 86_400_000),
tz.clone())
+ }
+ TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(
+ Some(days * 86_400_000_000),
+ tz.clone(),
+ ),
+ TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(
+ Some(days * 86_400_000_000_000),
+ tz.clone(),
+ ),
+ }
+ }
+ _ => return None,
+ };
+
+ Some(casted)
+}
+
+#[cfg(test)]
+mod tests {
Review Comment:
I think we should also have some "end to end" SQL level tests of this
transform
Namely, add sqllogictests with various date/time timestamps and then both:
1. Run queries to ensure the correct results
2. Run `EXPLAIN` to ensure the rewrites are happening as expected
--
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]
