Optimize the evaluation of date_part(<col>) == <constant> when pushed down

[sdf-jkl]

Which issue does this PR close?

• closes Optimize the evaluation of date_part(<col>) == <constant> when pushed down #19889.

Rationale for this change

Check issue.

What changes are included in this PR?

Added preimage impl for date_part udf.

Added sqllogictests for the impl.

Are these changes tested?

Yes, sqllogictests.

Are there any user-facing changes?

No

[kosiew]

LGTM

Left some minor suggestions.

[kosiew]

Consider adding a section to docs/source/library-user-guide/functions/adding-udfs.md explaining:

• What preimage is and when to implement it
• How it enables predicate pushdown
• Example implementation (perhaps referencing date_part)

[sdf-jkl]

There is a PR for preimage doc improvement here #20008.

I am however, not sure that this doc needs to explain preimage. I think the doc's goal is to be a very minimal guide on adding and registering a function. There is also no mention of simplify too.

[alamb]

I agree that the API docs is probably adequate. We could potentially add a note to adding-udfs.md that says something generic like "The ScalarUDFImpl has additional methods that support specialized optimizations such as preimage -- see the API documentation for additional details"

[sdf-jkl]

A separate PR should do.

[sdf-jkl]

Thanks for your review @kosiew. I've addressed your comments and also added proper handling for timestamps with timezones, please check it out too when you have time.

[sdf-jkl]

@alamb I'm still thinking about is_literal_or_literal_cast check in expr_simplifier: #19722 (comment)

[alamb]

Thank you @sdf-jkl -- this is looking very close

I have some suggestions on avoiding adding new date/time arithmetic here.

I also recommend an additional test for non-int32 arguments (coercion works fine, but it is good to show that):

diff --git a/datafusion/sqllogictest/test_files/udf_preimage.slt b/datafusion/sqllogictest/test_files/udf_preimage.slt
index 12f3a9a45b..c9d408227a 100644
--- a/datafusion/sqllogictest/test_files/udf_preimage.slt
+++ b/datafusion/sqllogictest/test_files/udf_preimage.slt
@@ -30,6 +30,11 @@ select c from t1 where extract(year from c) = 2024;
 ----
 2024-01-01

+query D
+select c from t1 where extract(year from c) = cast(2024 as bigint);
+----
+2024-01-01
+
 query D
 select c from t1 where extract(year from c) <> 2024;
 ----
@@ -83,6 +88,16 @@ physical_plan
 01)FilterExec: c@0 >= 2024-01-01 AND c@0 < 2025-01-01
 02)--DataSourceExec: partitions=1, partition_sizes=[1]

+query TT
+explain select c from t1 where extract (year from c) = cast(2024 as bigint)
+----
+logical_plan
+01)Filter: t1.c >= Date32("2024-01-01") AND t1.c < Date32("2025-01-01")
+02)--TableScan: t1 projection=[c]
+physical_plan
+01)FilterExec: c@0 >= 2024-01-01 AND c@0 < 2025-01-01
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+
 query TT
 explain select c from t1 where extract (year from c) <> 2024
 ----

[alamb]

I agree that the API docs is probably adequate. We could potentially add a note to adding-udfs.md that says something generic like "The ScalarUDFImpl has additional methods that support specialized optimizations such as preimage -- see the API documentation for additional details"

[alamb]

It feels to me like this code should be able to re-use the conversion functions in arrow rather than re-implementing them here

For example

        Date32 => ScalarValue::Date32(Some(Date32Type::from_naive_date(date))),
        Date64 => ScalarValue::Date64(Some(Date64Type::from_naive_date(date))),
...

I didn't have a chance to figure out how to do it for Timestamp, but it seems like there should be a function like this for the timestamps too -- for example

https://docs.rs/arrow/latest/arrow/array/types/struct.TimestampSecondType.html and https://docs.rs/arrow/latest/arrow/datatypes/trait.ArrowTimestampType.html

Maybe something like

TimestampSecondType::make_value(date)

(We will have to figure out timestamps)

[sdf-jkl]

Fixed date32/64 62b0841

As for timestamp: dyn TimestampType::make_value() is using a NaiveDate, not DateTime<Tz>. We'd still have to do some tz math to create an offset for NaiveDate. (if only there was an existing API to help...)