You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by su...@apache.org on 2026/09/02 18:59:41 UTC
(datafusion-comet) branch main updated: chore: add benches datetime (#5620)
This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 5c41c06f6a chore: add benches datetime (#5620)
5c41c06f6a is described below
commit 5c41c06f6aa2f058c811cb7f9c68bc172c0147f4
Author: Bhargava Vadlamani <11...@users.noreply.github.com>
AuthorDate: Wed Sep 2 11:59:32 2026 -0700
chore: add benches datetime (#5620)
---
native/spark-expr/Cargo.toml | 16 +++++
native/spark-expr/benches/common/mod.rs | 25 ++++++++
native/spark-expr/benches/from_utc_timestamp.rs | 79 +++++++++++++++++++++++++
native/spark-expr/benches/timestamp_trunc.rs | 65 ++++++++++++++++++++
native/spark-expr/benches/to_utc_timestamp.rs | 79 +++++++++++++++++++++++++
native/spark-expr/benches/unix_timestamp.rs | 76 ++++++++++++++++++++++++
6 files changed, 340 insertions(+)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index c9dd6cae23..5041d7fe32 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -285,4 +285,20 @@ harness = false
[[bench]]
name = "contains"
+harness = false
+
+[[bench]]
+name = "timestamp_trunc"
+harness = false
+
+[[bench]]
+name = "unix_timestamp"
+harness = false
+
+[[bench]]
+name = "from_utc_timestamp"
+harness = false
+
+[[bench]]
+name = "to_utc_timestamp"
harness = false
\ No newline at end of file
diff --git a/native/spark-expr/benches/common/mod.rs b/native/spark-expr/benches/common/mod.rs
index dc25f37385..f22e3d3825 100644
--- a/native/spark-expr/benches/common/mod.rs
+++ b/native/spark-expr/benches/common/mod.rs
@@ -22,6 +22,7 @@
use arrow::array::{
builder::StringBuilder, ArrayRef, Float64Array, Int64Array, RecordBatch, StringArray,
+ TimestampMicrosecondArray,
};
use arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;
@@ -80,6 +81,30 @@ pub fn string_array(rows: usize, null_ratio: f64, value: impl Fn(usize) -> Strin
Arc::new(arr)
}
+/// A `Timestamp(Microsecond, tz)` array of `rows` rows. Pass `tz = None` to build a
+/// TimestampNTZ array (no timezone) and `Some(name)` for a timezone-stamped array — the two
+/// take different code paths in the datetime kernels (NTZ skips timezone resolution).
+pub fn timestamp_micros_array(
+ rows: usize,
+ null_ratio: f64,
+ tz: Option<&str>,
+ value: impl Fn(usize) -> i64,
+) -> ArrayRef {
+ let arr: TimestampMicrosecondArray = (0..rows)
+ .map(|i| {
+ if is_null(i, null_ratio) {
+ None
+ } else {
+ Some(value(i))
+ }
+ })
+ .collect();
+ match tz {
+ Some(tz) => Arc::new(arr.with_timezone(tz)),
+ None => Arc::new(arr),
+ }
+}
+
/// A single-column `Utf8` batch of `rows` rows, where row `i` holds `value(i)` unless
/// `i % null_modulus == 0`, in which case it is null.
pub fn string_batch(
diff --git a/native/spark-expr/benches/from_utc_timestamp.rs b/native/spark-expr/benches/from_utc_timestamp.rs
new file mode 100644
index 0000000000..2d101c89ff
--- /dev/null
+++ b/native/spark-expr/benches/from_utc_timestamp.rs
@@ -0,0 +1,79 @@
+// 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.
+
+//! `from_utc_timestamp` is a datafusion-spark passthrough. Comet's serde wires it directly, so the
+//! kernel itself is upstream; this bench guards against regressions in the version Comet pins and
+//! is the same shape as the (composed) `convert_timezone` cost.
+
+use arrow::array::Array;
+use arrow::datatypes::{DataType, Field};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion::common::config::ConfigOptions;
+use datafusion::common::ScalarValue;
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_spark::function::datetime::from_utc_timestamp::SparkFromUtcTimestamp;
+use std::hint::black_box;
+use std::sync::Arc;
+
+#[path = "common/mod.rs"]
+mod common;
+use common::{timestamp_micros_array, NULL_RATIOS, ROW_COUNTS};
+
+const TZ: &str = "America/Los_Angeles";
+const MICROS_PER_DAY: i64 = 86_400_000_000;
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let udf = SparkFromUtcTimestamp::new();
+ let return_field = Arc::new(Field::new(
+ "from_utc_timestamp",
+ DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, None),
+ true,
+ ));
+
+ let mut group = c.benchmark_group("from_utc_timestamp");
+ for rows in ROW_COUNTS {
+ for (null_ratio, tag) in NULL_RATIOS {
+ // Input is TimestampNTZ (Timestamp(Microsecond, None)); the timezone arrives as the
+ // second, scalar argument.
+ let ts =
+ timestamp_micros_array(rows, null_ratio, None, |i| (i as i64) * MICROS_PER_DAY);
+ group.bench_with_input(
+ BenchmarkId::from_parameter(format!("{rows}/{tag}")),
+ &ts,
+ |b, ts| {
+ b.iter(|| {
+ let args = ScalarFunctionArgs {
+ args: vec![
+ ColumnarValue::Array(Arc::clone(ts)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some(TZ.to_string()))),
+ ],
+ number_rows: ts.len(),
+ return_field: Arc::clone(&return_field),
+ config_options: Arc::new(ConfigOptions::default()),
+ arg_fields: vec![],
+ };
+ black_box(udf.invoke_with_args(args).unwrap())
+ });
+ },
+ );
+ }
+ }
+ group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/benches/timestamp_trunc.rs b/native/spark-expr/benches/timestamp_trunc.rs
new file mode 100644
index 0000000000..6dd5d9fad1
--- /dev/null
+++ b/native/spark-expr/benches/timestamp_trunc.rs
@@ -0,0 +1,65 @@
+// 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::array::RecordBatch;
+use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion::physical_expr::expressions::{lit, Column};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion_comet_spark_expr::TimestampTruncExpr;
+use std::hint::black_box;
+use std::sync::Arc;
+
+#[path = "common/mod.rs"]
+mod common;
+use common::{timestamp_micros_array, NULL_RATIOS, ROW_COUNTS};
+
+// Non-UTC session timezone so the benchmark exercises the `array_with_timezone` resolution path
+// (the hot part of `TimestampTruncExpr`), not just the truncation kernel.
+const TZ: &str = "America/Los_Angeles";
+const MICROS_PER_DAY: i64 = 86_400_000_000;
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "a",
+ DataType::Timestamp(TimeUnit::Microsecond, Some(TZ.into())),
+ true,
+ )]));
+
+ let mut group = c.benchmark_group("timestamp_trunc");
+ for format in ["YEAR", "MONTH", "DAY", "HOUR"] {
+ let expr =
+ TimestampTruncExpr::new(Arc::new(Column::new("a", 0)), lit(format), TZ.to_string());
+ for rows in ROW_COUNTS {
+ for (null_ratio, tag) in NULL_RATIOS {
+ let ts = timestamp_micros_array(rows, null_ratio, Some(TZ), |i| {
+ (i as i64) * MICROS_PER_DAY
+ });
+ let batch = RecordBatch::try_new(Arc::clone(&schema), vec![ts]).unwrap();
+ group.bench_with_input(
+ BenchmarkId::from_parameter(format!("{format}/{rows}/{tag}")),
+ &batch,
+ |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())),
+ );
+ }
+ }
+ }
+ group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/benches/to_utc_timestamp.rs b/native/spark-expr/benches/to_utc_timestamp.rs
new file mode 100644
index 0000000000..b95f33200f
--- /dev/null
+++ b/native/spark-expr/benches/to_utc_timestamp.rs
@@ -0,0 +1,79 @@
+// 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.
+
+//! `to_utc_timestamp` is a datafusion-spark passthrough. Comet's serde wires it directly (and
+//! `convert_timezone` is composed of `to_utc_timestamp` + `from_utc_timestamp`), so this bench
+//! guards the pinned upstream kernel against regressions.
+
+use arrow::array::Array;
+use arrow::datatypes::{DataType, Field};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion::common::config::ConfigOptions;
+use datafusion::common::ScalarValue;
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_spark::function::datetime::to_utc_timestamp::SparkToUtcTimestamp;
+use std::hint::black_box;
+use std::sync::Arc;
+
+#[path = "common/mod.rs"]
+mod common;
+use common::{timestamp_micros_array, NULL_RATIOS, ROW_COUNTS};
+
+const TZ: &str = "America/Los_Angeles";
+const MICROS_PER_DAY: i64 = 86_400_000_000;
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let udf = SparkToUtcTimestamp::new();
+ let return_field = Arc::new(Field::new(
+ "to_utc_timestamp",
+ DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, None),
+ true,
+ ));
+
+ let mut group = c.benchmark_group("to_utc_timestamp");
+ for rows in ROW_COUNTS {
+ for (null_ratio, tag) in NULL_RATIOS {
+ // Input is TimestampNTZ (Timestamp(Microsecond, None)); the timezone arrives as the
+ // second, scalar argument.
+ let ts =
+ timestamp_micros_array(rows, null_ratio, None, |i| (i as i64) * MICROS_PER_DAY);
+ group.bench_with_input(
+ BenchmarkId::from_parameter(format!("{rows}/{tag}")),
+ &ts,
+ |b, ts| {
+ b.iter(|| {
+ let args = ScalarFunctionArgs {
+ args: vec![
+ ColumnarValue::Array(Arc::clone(ts)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some(TZ.to_string()))),
+ ],
+ number_rows: ts.len(),
+ return_field: Arc::clone(&return_field),
+ config_options: Arc::new(ConfigOptions::default()),
+ arg_fields: vec![],
+ };
+ black_box(udf.invoke_with_args(args).unwrap())
+ });
+ },
+ );
+ }
+ }
+ group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/native/spark-expr/benches/unix_timestamp.rs b/native/spark-expr/benches/unix_timestamp.rs
new file mode 100644
index 0000000000..938c59990c
--- /dev/null
+++ b/native/spark-expr/benches/unix_timestamp.rs
@@ -0,0 +1,76 @@
+// 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::array::Array;
+use arrow::datatypes::{DataType, Field};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion::common::config::ConfigOptions;
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_comet_spark_expr::SparkUnixTimestamp;
+use std::hint::black_box;
+use std::sync::Arc;
+
+#[path = "common/mod.rs"]
+mod common;
+use common::{timestamp_micros_array, NULL_RATIOS, ROW_COUNTS};
+
+const MICROS_PER_DAY: i64 = 86_400_000_000;
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let return_field = Arc::new(Field::new("unix_timestamp", DataType::Int64, true));
+
+ let mut group = c.benchmark_group("unix_timestamp");
+ // TimestampNTZ skips timezone resolution (plain integer division); a tz-aware input runs
+ // through `array_with_timezone` first. Bench both so a regression in either path shows up.
+ for (variant, tz_input, udf_tz) in [
+ ("ntz", None, "UTC"),
+ (
+ "tz_aware",
+ Some("America/Los_Angeles"),
+ "America/Los_Angeles",
+ ),
+ ] {
+ let udf = SparkUnixTimestamp::new(udf_tz.to_string());
+ for rows in ROW_COUNTS {
+ for (null_ratio, tag) in NULL_RATIOS {
+ let ts = timestamp_micros_array(rows, null_ratio, tz_input, |i| {
+ (i as i64) * MICROS_PER_DAY
+ });
+ group.bench_with_input(
+ BenchmarkId::from_parameter(format!("{variant}/{rows}/{tag}")),
+ &ts,
+ |b, ts| {
+ b.iter(|| {
+ let args = ScalarFunctionArgs {
+ args: vec![ColumnarValue::Array(Arc::clone(ts))],
+ number_rows: ts.len(),
+ return_field: Arc::clone(&return_field),
+ config_options: Arc::new(ConfigOptions::default()),
+ arg_fields: vec![],
+ };
+ black_box(udf.invoke_with_args(args).unwrap())
+ });
+ },
+ );
+ }
+ }
+ }
+ group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org