fix: date_part sub-second parts and result nullability, reusing arrow kernels

datafusion/functions/benches/datetime_expressions/date_part.rs

@@ -188,6 +188,36 @@ fn generate_duration_nanosecond_array_with_nulls(
     )
 }
 
+fn generate_timestamp_ns_array_with_nulls(rng: &mut StdRng) -> TimestampNanosecondArray {
+    TimestampNanosecondArray::from(
+        (0..BATCH_SIZE)
+            .map(|_| {
+                if rng.random_bool(NULL_DENSITY) {
+                    None
+                } else {
+                    Some(rng.random_range(0..TS_BOUND * 1_000_000_000))
+                }
+            })
+            .collect::<Vec<_>>(),
+    )
+}
+
+fn generate_time64_nanosecond_array_with_nulls(
+    rng: &mut StdRng,
+) -> Time64NanosecondArray {
+    Time64NanosecondArray::from(
+        (0..BATCH_SIZE)
+            .map(|_| {
+                if rng.random_bool(NULL_DENSITY) {
+                    None
+                } else {
+                    Some(rng.random_range(0..SEC_DAY * 1_000_000_000))
+                }
+            })
+            .collect::<Vec<_>>(),
+    )
+}
+
 fn bench_date_part(
     c: &mut Criterion,
     udf: &Arc<ScalarUDF>,
@@ -244,6 +274,14 @@ fn criterion_benchmark(c: &mut Criterion) {
     let duration_ns_nulls =
         Arc::new(generate_duration_nanosecond_array_with_nulls(&mut rng)) as ArrayRef;
 
+    let ts_ns_tz =
+        Arc::new(generate_timestamp_ns_array(&mut rng).with_timezone("America/New_York"))
+            as ArrayRef;
+    let ts_ns_nulls =
+        Arc::new(generate_timestamp_ns_array_with_nulls(&mut rng)) as ArrayRef;
+    let time64_ns_nulls =
+        Arc::new(generate_time64_nanosecond_array_with_nulls(&mut rng)) as ArrayRef;
+
     let udf = date_part();
 
     for part in ["year", "month", "week", "day", "hour", "minute"] {
@@ -365,6 +403,41 @@ fn criterion_benchmark(c: &mut Criterion) {
         );
     }
 
+    // Durations past `i32::MAX` milliseconds or microseconds overflow `Int32`, so
+    // only `nanosecond` runs on them
+    for (part, return_type) in [
+        ("millisecond", DataType::Int32),
+        ("microsecond", DataType::Int32),
+        ("nanosecond", DataType::Int64),
+    ] {
+        let mut arrays = vec![
+            ("ns_tz", &ts_ns_tz),
+            ("ns_nulls", &ts_ns_nulls),
+            ("time32_ms", &time32_ms),
+            ("time64_ns", &time64_ns),
+            ("time64_ns_nulls", &time64_ns_nulls),
+            ("interval_dt", &interval_dt),
+            ("interval_mdn", &interval_mdn),
+            ("duration_ns_small", &duration_ns_small),
+        ];
+        if part == "nanosecond" {
+            arrays.extend([
+                ("duration_ns", &duration_ns),
+                ("duration_ns_nulls", &duration_ns_nulls),
+            ]);
+        }
+        for (name, array) in arrays {
+            bench_date_part(
+                c,
+                &udf,
+                &format!("date_part_{part}_{name}_1000"),
+                part,
+                Arc::clone(array),
+                return_type.clone(),
+            );
+        }
+    }
+
     for part in ["quarter", "isoyear", "doy", "dow", "isodow"] {
         bench_date_part(
             c,