From 6487f91d6e00a64535abcf98f3b44c8901e3c886 Mon Sep 17 00:00:00 2001 From: theirix Date: Sun, 5 Jul 2026 20:52:43 +0100 Subject: [PATCH 1/4] Add date_part bench --- datafusion/functions/Cargo.toml | 5 + datafusion/functions/benches/date_part.rs | 345 ++++++++++++++++++++++ 2 files changed, 350 insertions(+) create mode 100644 datafusion/functions/benches/date_part.rs diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 4eca16961fa8c..94830ee360585 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -182,6 +182,11 @@ harness = false name = "date_trunc" required-features = ["datetime_expressions"] +[[bench]] +harness = false +name = "date_part" +required-features = ["datetime_expressions"] + [[bench]] harness = false name = "to_char" diff --git a/datafusion/functions/benches/date_part.rs b/datafusion/functions/benches/date_part.rs new file mode 100644 index 0000000000000..1b0751255a780 --- /dev/null +++ b/datafusion/functions/benches/date_part.rs @@ -0,0 +1,345 @@ +// 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 std::hint::black_box; +use std::sync::Arc; + +use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; +use arrow::array::{ + Array, ArrayRef, Date32Array, Date64Array, DurationNanosecondArray, + IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, + Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, + Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, +}; +use arrow::datatypes::{DataType, Field}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; +use datafusion_functions::datetime::date_part; +use rand::Rng; +use rand::rngs::ThreadRng; + +const BATCH_SIZE: usize = 1000; +const TS_BOUND: i64 = 2_006_463_600; +const SEC_DAY: i64 = 86_400; + +fn generate_timestamp_ns_array(rng: &mut ThreadRng) -> TimestampNanosecondArray { + TimestampNanosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND * 1_000_000_000)) + .collect::>(), + ) +} + +fn generate_timestamp_us_array(rng: &mut ThreadRng) -> TimestampMicrosecondArray { + TimestampMicrosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND * 1_000_000)) + .collect::>(), + ) +} + +fn generate_timestamp_ms_array(rng: &mut ThreadRng) -> TimestampMillisecondArray { + TimestampMillisecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND * 1_000)) + .collect::>(), + ) +} + +fn generate_timestamp_s_array(rng: &mut ThreadRng) -> TimestampSecondArray { + TimestampSecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND)) + .collect::>(), + ) +} + +fn generate_date32_array(rng: &mut ThreadRng) -> Date32Array { + Date32Array::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..30_000)) + .collect::>(), + ) +} + +fn generate_date64_array(rng: &mut ThreadRng) -> Date64Array { + Date64Array::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0i64..30_000)) + .collect::>(), + ) +} + +fn generate_time32_second_array(rng: &mut ThreadRng) -> Time32SecondArray { + Time32SecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..SEC_DAY as i32)) + .collect::>(), + ) +} + +fn generate_time32_millisecond_array(rng: &mut ThreadRng) -> Time32MillisecondArray { + Time32MillisecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..(SEC_DAY * 1_000) as i32)) + .collect::>(), + ) +} + +fn generate_time64_microsecond_array(rng: &mut ThreadRng) -> Time64MicrosecondArray { + Time64MicrosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..SEC_DAY * 1_000_000)) + .collect::>(), + ) +} + +fn generate_time64_nanosecond_array(rng: &mut ThreadRng) -> Time64NanosecondArray { + Time64NanosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..SEC_DAY * 1_000_000_000)) + .collect::>(), + ) +} + +fn generate_interval_year_month_array(rng: &mut ThreadRng) -> IntervalYearMonthArray { + IntervalYearMonthArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..1_200)) + .collect::>(), + ) +} + +fn generate_interval_day_time_array(rng: &mut ThreadRng) -> IntervalDayTimeArray { + IntervalDayTimeArray::from( + (0..BATCH_SIZE) + .map(|_| IntervalDayTime { + days: rng.random_range(0..365), + milliseconds: rng.random_range(0..(SEC_DAY * 1_000) as i32), + }) + .collect::>(), + ) +} + +fn generate_interval_mdn_array(rng: &mut ThreadRng) -> IntervalMonthDayNanoArray { + IntervalMonthDayNanoArray::from( + (0..BATCH_SIZE) + .map(|_| IntervalMonthDayNano { + months: rng.random_range(0..120), + days: rng.random_range(0..365), + nanoseconds: rng.random_range(0..SEC_DAY * 1_000_000_000), + }) + .collect::>(), + ) +} + +fn generate_duration_nanosecond_array(rng: &mut ThreadRng) -> DurationNanosecondArray { + DurationNanosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND * 1_000_000_000)) + .collect::>(), + ) +} + +fn bench_date_part( + c: &mut Criterion, + udf: &Arc, + bench_name: &str, + part: &str, + array: ArrayRef, + return_type: DataType, +) { + let batch_len = array.len(); + let part_cv = ColumnarValue::Scalar(ScalarValue::Utf8(Some(part.to_string()))); + let array_cv = ColumnarValue::Array(array); + let return_field = Arc::new(Field::new("date_part", return_type, true)); + let arg_fields = vec![ + Field::new("a", part_cv.data_type(), true).into(), + Field::new("b", array_cv.data_type(), true).into(), + ]; + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_function(bench_name, |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![part_cv.clone(), array_cv.clone()], + arg_fields: arg_fields.clone(), + number_rows: batch_len, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .expect("date_part should work on valid values"), + ) + }) + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = rand::rng(); + + let ts_s = Arc::new(generate_timestamp_s_array(&mut rng)) as ArrayRef; + let ts_ms = Arc::new(generate_timestamp_ms_array(&mut rng)) as ArrayRef; + let ts_us = Arc::new(generate_timestamp_us_array(&mut rng)) as ArrayRef; + let ts_ns = Arc::new(generate_timestamp_ns_array(&mut rng)) as ArrayRef; + let time32_s = Arc::new(generate_time32_second_array(&mut rng)) as ArrayRef; + let time32_ms = Arc::new(generate_time32_millisecond_array(&mut rng)) as ArrayRef; + let time64_us = Arc::new(generate_time64_microsecond_array(&mut rng)) as ArrayRef; + let time64_ns = Arc::new(generate_time64_nanosecond_array(&mut rng)) as ArrayRef; + let interval_ym = Arc::new(generate_interval_year_month_array(&mut rng)) as ArrayRef; + let interval_dt = Arc::new(generate_interval_day_time_array(&mut rng)) as ArrayRef; + let interval_mdn = Arc::new(generate_interval_mdn_array(&mut rng)) as ArrayRef; + let duration_ns = Arc::new(generate_duration_nanosecond_array(&mut rng)) as ArrayRef; + let date32 = Arc::new(generate_date32_array(&mut rng)) as ArrayRef; + let date64 = Arc::new(generate_date64_array(&mut rng)) as ArrayRef; + + let udf = date_part(); + + for part in ["year", "month", "week", "day", "hour", "minute"] { + for (name, array) in + [("s", &ts_s), ("ms", &ts_ms), ("us", &ts_us), ("ns", &ts_ns)] + { + bench_date_part( + c, + &udf, + &format!("date_part_{part}_{name}_1000"), + part, + Arc::clone(array), + DataType::Int32, + ); + } + } + for part in ["year", "month", "week", "day"] { + bench_date_part( + c, + &udf, + &format!("date_part_{part}_date32_1000"), + part, + Arc::clone(&date32), + DataType::Int32, + ); + bench_date_part( + c, + &udf, + &format!("date_part_{part}_date64_1000"), + part, + Arc::clone(&date64), + DataType::Int32, + ); + } + + for part in ["second", "millisecond", "microsecond"] { + for (name, array) in + [("s", &ts_s), ("ms", &ts_ms), ("us", &ts_us), ("ns", &ts_ns)] + { + bench_date_part( + c, + &udf, + &format!("date_part_{part}_{name}_1000"), + part, + Arc::clone(array), + DataType::Int32, + ); + } + bench_date_part( + c, + &udf, + &format!("date_part_{part}_date32_1000"), + part, + Arc::clone(&date32), + DataType::Int32, + ); + bench_date_part( + c, + &udf, + &format!("date_part_{part}_date64_1000"), + part, + Arc::clone(&date64), + DataType::Int32, + ); + } + + for (name, array) in [("s", &ts_s), ("ms", &ts_ms), ("us", &ts_us), ("ns", &ts_ns)] { + bench_date_part( + c, + &udf, + &format!("date_part_nanosecond_{name}_1000"), + "nanosecond", + Arc::clone(array), + DataType::Int64, + ); + } + bench_date_part( + c, + &udf, + "date_part_nanosecond_date32_1000", + "nanosecond", + Arc::clone(&date32), + DataType::Int64, + ); + bench_date_part( + c, + &udf, + "date_part_nanosecond_date64_1000", + "nanosecond", + Arc::clone(&date64), + DataType::Int64, + ); + + for (name, array) in [ + ("s", &ts_s), + ("ms", &ts_ms), + ("us", &ts_us), + ("ns", &ts_ns), + ("date32", &date32), + ("date64", &date64), + ("time32_s", &time32_s), + ("time32_ms", &time32_ms), + ("time64_us", &time64_us), + ("time64_ns", &time64_ns), + ("interval_ym", &interval_ym), + ("interval_dt", &interval_dt), + ("interval_mdn", &interval_mdn), + ("duration_ns", &duration_ns), + ] { + bench_date_part( + c, + &udf, + &format!("date_part_epoch_{name}_1000"), + "epoch", + Arc::clone(array), + DataType::Float64, + ); + } + + for part in ["quarter", "isoyear", "doy", "dow", "isodow"] { + bench_date_part( + c, + &udf, + &format!("date_part_{part}_timestamp_ns_1000"), + part, + Arc::clone(&ts_ns), + DataType::Int32, + ); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); From 0d339180284bc539f27be3372dc06b6077a20428 Mon Sep 17 00:00:00 2001 From: theirix Date: Thu, 9 Jul 2026 19:44:39 +0100 Subject: [PATCH 2/4] Better ranges for datetime arrays --- datafusion/functions/benches/date_part.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/datafusion/functions/benches/date_part.rs b/datafusion/functions/benches/date_part.rs index 1b0751255a780..9e32a06bd94d6 100644 --- a/datafusion/functions/benches/date_part.rs +++ b/datafusion/functions/benches/date_part.rs @@ -71,10 +71,11 @@ fn generate_timestamp_s_array(rng: &mut ThreadRng) -> TimestampSecondArray { ) } -fn generate_date32_array(rng: &mut ThreadRng) -> Date32Array { +fn generate_date32_array(rng: &mut StdRng) -> Date32Array { + let days_since_epoch = (TS_BOUND as i32) / SEC_DAY as i32; Date32Array::from( (0..BATCH_SIZE) - .map(|_| rng.random_range(0..30_000)) + .map(|_| rng.random_range(0..days_since_epoch)) .collect::>(), ) } @@ -82,7 +83,7 @@ fn generate_date32_array(rng: &mut ThreadRng) -> Date32Array { fn generate_date64_array(rng: &mut ThreadRng) -> Date64Array { Date64Array::from( (0..BATCH_SIZE) - .map(|_| rng.random_range(0i64..30_000)) + .map(|_| rng.random_range(0..TS_BOUND * 1_000)) .collect::>(), ) } @@ -119,10 +120,11 @@ fn generate_time64_nanosecond_array(rng: &mut ThreadRng) -> Time64NanosecondArra ) } -fn generate_interval_year_month_array(rng: &mut ThreadRng) -> IntervalYearMonthArray { +fn generate_interval_year_month_array(rng: &mut StdRng) -> IntervalYearMonthArray { + let years = 10; IntervalYearMonthArray::from( (0..BATCH_SIZE) - .map(|_| rng.random_range(0..1_200)) + .map(|_| rng.random_range(0..12 * years)) .collect::>(), ) } @@ -142,7 +144,7 @@ fn generate_interval_mdn_array(rng: &mut ThreadRng) -> IntervalMonthDayNanoArray IntervalMonthDayNanoArray::from( (0..BATCH_SIZE) .map(|_| IntervalMonthDayNano { - months: rng.random_range(0..120), + months: rng.random_range(0..12), days: rng.random_range(0..365), nanoseconds: rng.random_range(0..SEC_DAY * 1_000_000_000), }) From e4158160d7ca8dd0d978e5153fd9ed24833e6cff Mon Sep 17 00:00:00 2001 From: theirix Date: Thu, 9 Jul 2026 19:45:09 +0100 Subject: [PATCH 3/4] Use seedable rng --- datafusion/functions/benches/date_part.rs | 30 +++++++++++------------ 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/datafusion/functions/benches/date_part.rs b/datafusion/functions/benches/date_part.rs index 9e32a06bd94d6..7bd69bf392e20 100644 --- a/datafusion/functions/benches/date_part.rs +++ b/datafusion/functions/benches/date_part.rs @@ -32,14 +32,14 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; use datafusion_functions::datetime::date_part; -use rand::Rng; -use rand::rngs::ThreadRng; +use rand::prelude::StdRng; +use rand::{Rng, SeedableRng}; const BATCH_SIZE: usize = 1000; const TS_BOUND: i64 = 2_006_463_600; const SEC_DAY: i64 = 86_400; -fn generate_timestamp_ns_array(rng: &mut ThreadRng) -> TimestampNanosecondArray { +fn generate_timestamp_ns_array(rng: &mut StdRng) -> TimestampNanosecondArray { TimestampNanosecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..TS_BOUND * 1_000_000_000)) @@ -47,7 +47,7 @@ fn generate_timestamp_ns_array(rng: &mut ThreadRng) -> TimestampNanosecondArray ) } -fn generate_timestamp_us_array(rng: &mut ThreadRng) -> TimestampMicrosecondArray { +fn generate_timestamp_us_array(rng: &mut StdRng) -> TimestampMicrosecondArray { TimestampMicrosecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..TS_BOUND * 1_000_000)) @@ -55,7 +55,7 @@ fn generate_timestamp_us_array(rng: &mut ThreadRng) -> TimestampMicrosecondArray ) } -fn generate_timestamp_ms_array(rng: &mut ThreadRng) -> TimestampMillisecondArray { +fn generate_timestamp_ms_array(rng: &mut StdRng) -> TimestampMillisecondArray { TimestampMillisecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..TS_BOUND * 1_000)) @@ -63,7 +63,7 @@ fn generate_timestamp_ms_array(rng: &mut ThreadRng) -> TimestampMillisecondArray ) } -fn generate_timestamp_s_array(rng: &mut ThreadRng) -> TimestampSecondArray { +fn generate_timestamp_s_array(rng: &mut StdRng) -> TimestampSecondArray { TimestampSecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..TS_BOUND)) @@ -80,7 +80,7 @@ fn generate_date32_array(rng: &mut StdRng) -> Date32Array { ) } -fn generate_date64_array(rng: &mut ThreadRng) -> Date64Array { +fn generate_date64_array(rng: &mut StdRng) -> Date64Array { Date64Array::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..TS_BOUND * 1_000)) @@ -88,7 +88,7 @@ fn generate_date64_array(rng: &mut ThreadRng) -> Date64Array { ) } -fn generate_time32_second_array(rng: &mut ThreadRng) -> Time32SecondArray { +fn generate_time32_second_array(rng: &mut StdRng) -> Time32SecondArray { Time32SecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..SEC_DAY as i32)) @@ -96,7 +96,7 @@ fn generate_time32_second_array(rng: &mut ThreadRng) -> Time32SecondArray { ) } -fn generate_time32_millisecond_array(rng: &mut ThreadRng) -> Time32MillisecondArray { +fn generate_time32_millisecond_array(rng: &mut StdRng) -> Time32MillisecondArray { Time32MillisecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..(SEC_DAY * 1_000) as i32)) @@ -104,7 +104,7 @@ fn generate_time32_millisecond_array(rng: &mut ThreadRng) -> Time32MillisecondAr ) } -fn generate_time64_microsecond_array(rng: &mut ThreadRng) -> Time64MicrosecondArray { +fn generate_time64_microsecond_array(rng: &mut StdRng) -> Time64MicrosecondArray { Time64MicrosecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..SEC_DAY * 1_000_000)) @@ -112,7 +112,7 @@ fn generate_time64_microsecond_array(rng: &mut ThreadRng) -> Time64MicrosecondAr ) } -fn generate_time64_nanosecond_array(rng: &mut ThreadRng) -> Time64NanosecondArray { +fn generate_time64_nanosecond_array(rng: &mut StdRng) -> Time64NanosecondArray { Time64NanosecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..SEC_DAY * 1_000_000_000)) @@ -129,7 +129,7 @@ fn generate_interval_year_month_array(rng: &mut StdRng) -> IntervalYearMonthArra ) } -fn generate_interval_day_time_array(rng: &mut ThreadRng) -> IntervalDayTimeArray { +fn generate_interval_day_time_array(rng: &mut StdRng) -> IntervalDayTimeArray { IntervalDayTimeArray::from( (0..BATCH_SIZE) .map(|_| IntervalDayTime { @@ -140,7 +140,7 @@ fn generate_interval_day_time_array(rng: &mut ThreadRng) -> IntervalDayTimeArray ) } -fn generate_interval_mdn_array(rng: &mut ThreadRng) -> IntervalMonthDayNanoArray { +fn generate_interval_mdn_array(rng: &mut StdRng) -> IntervalMonthDayNanoArray { IntervalMonthDayNanoArray::from( (0..BATCH_SIZE) .map(|_| IntervalMonthDayNano { @@ -152,7 +152,7 @@ fn generate_interval_mdn_array(rng: &mut ThreadRng) -> IntervalMonthDayNanoArray ) } -fn generate_duration_nanosecond_array(rng: &mut ThreadRng) -> DurationNanosecondArray { +fn generate_duration_nanosecond_array(rng: &mut StdRng) -> DurationNanosecondArray { DurationNanosecondArray::from( (0..BATCH_SIZE) .map(|_| rng.random_range(0..TS_BOUND * 1_000_000_000)) @@ -195,7 +195,7 @@ fn bench_date_part( } fn criterion_benchmark(c: &mut Criterion) { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(42); let ts_s = Arc::new(generate_timestamp_s_array(&mut rng)) as ArrayRef; let ts_ms = Arc::new(generate_timestamp_ms_array(&mut rng)) as ArrayRef; From 1c5749199dfd7bbccedd9812f8e64cd01c2ea99d Mon Sep 17 00:00:00 2001 From: theirix Date: Thu, 9 Jul 2026 20:08:35 +0100 Subject: [PATCH 4/4] Align Date arrays to days --- datafusion/functions/benches/date_part.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/datafusion/functions/benches/date_part.rs b/datafusion/functions/benches/date_part.rs index 7bd69bf392e20..fb93ebd03b4e2 100644 --- a/datafusion/functions/benches/date_part.rs +++ b/datafusion/functions/benches/date_part.rs @@ -38,6 +38,7 @@ use rand::{Rng, SeedableRng}; const BATCH_SIZE: usize = 1000; const TS_BOUND: i64 = 2_006_463_600; const SEC_DAY: i64 = 86_400; +const DAYS_SINCE_EPOCH: i64 = TS_BOUND / SEC_DAY; fn generate_timestamp_ns_array(rng: &mut StdRng) -> TimestampNanosecondArray { TimestampNanosecondArray::from( @@ -72,18 +73,19 @@ fn generate_timestamp_s_array(rng: &mut StdRng) -> TimestampSecondArray { } fn generate_date32_array(rng: &mut StdRng) -> Date32Array { - let days_since_epoch = (TS_BOUND as i32) / SEC_DAY as i32; + // Provide days since epoch Date32Array::from( (0..BATCH_SIZE) - .map(|_| rng.random_range(0..days_since_epoch)) + .map(|_| rng.random_range(0..DAYS_SINCE_EPOCH as i32)) .collect::>(), ) } fn generate_date64_array(rng: &mut StdRng) -> Date64Array { + // Provide milliseconds since epoch aligned to day boundaries Date64Array::from( (0..BATCH_SIZE) - .map(|_| rng.random_range(0..TS_BOUND * 1_000)) + .map(|_| rng.random_range(0..DAYS_SINCE_EPOCH) * SEC_DAY * 1_000) .collect::>(), ) }