sunchao commented on code in PR #25979:
URL: https://github.com/apache/datafusion/pull/25979#discussion_r4240758616
##########
datafusion/spark/src/function/datetime/from_utc_timestamp.rs:
##########
@@ -188,3 +214,529 @@ where
builder = builder.with_timezone_opt(return_tz_opt);
Ok(Arc::new(builder.finish()))
}
+
+fn parse_timezone(timezone: &str) -> Result<SparkTimeZone> {
+ SparkTimeZone::parse(timezone).map_err(|e| {
+ exec_datafusion_err!("`from_utc_timestamp`: invalid timezone
'{timezone}': {e}")
+ })
+}
+
+/// Physical `from_utc_timestamp` expression with conditional timezone
evaluation.
+///
+/// The timestamp child is evaluated once for each non-empty batch. If every
timestamp is null,
+/// the timezone child is skipped. Otherwise, literal and column timezones are
read directly, and
+/// their values are ignored on null timestamp rows. Computed timezone
expressions are evaluated
+/// only for rows with non-null timestamps. Empty batches evaluate neither
child.
+///
+/// This intentionally differs from Spark's generated evaluation of foldable
timezones, which
+/// evaluates the timezone and validates a non-null value before processing
timestamp rows. This
+/// expression can therefore skip an invalid timezone or an error in the
timezone expression when
+/// no timestamp is non-null. It also evaluates the timestamp first when the
timezone is null, so a
+/// timestamp error can be raised even though Spark's generated path would
return null. Planners
+/// using this expression should expose these differences as an
incompatibility.
+#[derive(Debug, Clone)]
+pub struct SparkFromUtcTimestampExpr {
+ timestamp: Arc<dyn PhysicalExpr>,
+ timezone: Arc<dyn PhysicalExpr>,
+}
+
+impl SparkFromUtcTimestampExpr {
+ /// Construct an expression that skips computed timezone evaluation for
null timestamps.
+ pub fn new(
+ timestamp: Arc<dyn PhysicalExpr>,
+ timezone: Arc<dyn PhysicalExpr>,
+ ) -> Self {
+ Self {
+ timestamp,
+ timezone,
+ }
+ }
+}
+
+impl Display for SparkFromUtcTimestampExpr {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "from_utc_timestamp({}, {})",
+ self.timestamp, self.timezone
+ )
+ }
+}
+
+impl PartialEq for SparkFromUtcTimestampExpr {
+ fn eq(&self, other: &Self) -> bool {
+ self.timestamp.eq(&other.timestamp) &&
self.timezone.eq(&other.timezone)
+ }
+}
+
+impl Eq for SparkFromUtcTimestampExpr {}
+
+impl Hash for SparkFromUtcTimestampExpr {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.timestamp.hash(state);
+ self.timezone.hash(state);
+ }
+}
+
+impl PhysicalExpr for SparkFromUtcTimestampExpr {
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ Display::fmt(self, f)
Review Comment:
Nit: could `fmt_sql` format each child through its `fmt_sql` implementation
rather than delegate to `Display`? For example, column `Display` produces
`timestamp@0`, whereas its SQL formatter produces `timestamp`. This would keep
SQL-formatted output consistent with other physical expressions.
##########
datafusion/spark/src/function/datetime/from_utc_timestamp.rs:
##########
@@ -188,3 +214,529 @@ where
builder = builder.with_timezone_opt(return_tz_opt);
Ok(Arc::new(builder.finish()))
}
+
+fn parse_timezone(timezone: &str) -> Result<SparkTimeZone> {
+ SparkTimeZone::parse(timezone).map_err(|e| {
+ exec_datafusion_err!("`from_utc_timestamp`: invalid timezone
'{timezone}': {e}")
+ })
+}
+
+/// Physical `from_utc_timestamp` expression with conditional timezone
evaluation.
+///
+/// The timestamp child is evaluated once for each non-empty batch. If every
timestamp is null,
+/// the timezone child is skipped. Otherwise, literal and column timezones are
read directly, and
+/// their values are ignored on null timestamp rows. Computed timezone
expressions are evaluated
+/// only for rows with non-null timestamps. Empty batches evaluate neither
child.
+///
+/// This intentionally differs from Spark's generated evaluation of foldable
timezones, which
+/// evaluates the timezone and validates a non-null value before processing
timestamp rows. This
+/// expression can therefore skip an invalid timezone or an error in the
timezone expression when
+/// no timestamp is non-null. It also evaluates the timestamp first when the
timezone is null, so a
+/// timestamp error can be raised even though Spark's generated path would
return null. Planners
+/// using this expression should expose these differences as an
incompatibility.
+#[derive(Debug, Clone)]
+pub struct SparkFromUtcTimestampExpr {
+ timestamp: Arc<dyn PhysicalExpr>,
+ timezone: Arc<dyn PhysicalExpr>,
+}
+
+impl SparkFromUtcTimestampExpr {
+ /// Construct an expression that skips computed timezone evaluation for
null timestamps.
Review Comment:
Nonblocking documentation suggestion: a short planner/consumer example would
help show where to construct this expression instead of the registered UDF, and
that its children need to be coerced before construction. Ordinary SQL still
uses the eager UDF, while this constructor simply stores the supplied physical
children. An integration example would make the boundary of the
conditional-evaluation support clearer.
##########
datafusion/spark/benches/from_utc_timestamp.rs:
##########
@@ -0,0 +1,94 @@
+// 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.
+
+//! Measure the physical evaluator with a timezone column. The narrow and wide
batches use
+//! identical timestamp and timezone columns; only the number of unused Int64
columns differs.
+//! Mixed-null timestamps expose any cost from filtering those unrelated
columns.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array, StringArray,
TimestampMicrosecondArray};
+use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use arrow::record_batch::RecordBatch;
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group,
criterion_main};
+use datafusion::physical_expr::expressions::Column;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use
datafusion_spark::function::datetime::from_utc_timestamp::SparkFromUtcTimestampExpr;
+
+fn batch(rows: usize, columns: usize, half_null: bool) -> RecordBatch {
+ let mut fields = vec![
+ Field::new(
+ "timestamp",
+ DataType::Timestamp(TimeUnit::Microsecond, None),
+ true,
+ ),
+ Field::new("timezone", DataType::Utf8, false),
+ ];
+ let mut arrays: Vec<ArrayRef> = vec![
+ Arc::new(TimestampMicrosecondArray::from_iter((0..rows).map(|row| {
+ (!half_null || row % 2 != 0).then_some(row as i64 * 86_400_000_000)
+ }))),
+ Arc::new(StringArray::from_iter_values(
+ (0..rows).map(|row| if (row / 2) % 2 == 0 { "PST" } else { "EST"
}),
+ )),
+ ];
+
+ for column in 2..columns {
+ fields.push(Field::new(
+ format!("unused_{column}"),
+ DataType::Int64,
+ false,
+ ));
+ arrays.push(Arc::new(Int64Array::from_iter_values(
+ (0..rows).map(|row| (column * rows + row) as i64),
+ )));
+ }
+
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let rows = 8192;
+ let expr = SparkFromUtcTimestampExpr::new(
+ Arc::new(Column::new("timestamp", 0)),
+ Arc::new(Column::new("timezone", 1)),
Review Comment:
Nonblocking: could we also benchmark a computed timezone expression, with
sparse and mixed null timestamps, across these batch widths? This `Column`
takes the fast path and bypasses `evaluate_selection`, so it does not exercise
the branch that filters every input column. A computed-expression case would
measure the remaining allocation/width tradeoff; a literal timezone case would
also cover scalar broadcasting.
--
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]
