andygrove commented on code in PR #5150: URL: https://github.com/apache/datafusion-comet/pull/5150#discussion_r3693364347
########## native/spark-expr/benches/cast_string_to_timestamp.rs: ########## @@ -0,0 +1,210 @@ +// 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::{builder::StringBuilder, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_expr::{expressions::Column, PhysicalExpr}; +use datafusion_comet_spark_expr::{Cast, EvalMode, SparkCastOptions}; +use std::sync::Arc; + +const BATCH_SIZE: usize = 8192; + +fn criterion_benchmark(c: &mut Criterion) { + let expr = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>; + + // Input shapes, chosen to cover each branch `timestamp_parser` can take: the canonical + // form, the fractional-second form, an offset suffix (which takes the extract-offset + // path), a date-only string, whitespace padding (the trim), and a mix that includes + // invalid values so the null path is measured too. + let batches = [ + ( + "canonical", + create_batch(|i| { + format!( + "{:04}-{:02}-{:02} {:02}:{:02}:{:02}", + 1970 + i % 60, + i % 12 + 1, + i % 28 + 1, + i % 24, + i % 60, + i % 60 + ) + }), + ), + ( + "microseconds", + create_batch(|i| { + format!( + "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:06}", + 1970 + i % 60, + i % 12 + 1, + i % 28 + 1, + i % 24, + i % 60, + i % 60, + i % 1_000_000 + ) + }), + ), + ( + "offset_suffix", + create_batch(|i| { + format!( + "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}+05:30", + 1970 + i % 60, + i % 12 + 1, + i % 28 + 1, + i % 24, + i % 60, + i % 60 + ) + }), + ), + ( + "date_only", + create_batch(|i| format!("{:04}-{:02}-{:02}", 1970 + i % 60, i % 12 + 1, i % 28 + 1)), + ), + ( + "padded", + create_batch(|i| { + format!( + " {:04}-{:02}-{:02} {:02}:00:00 ", + 1970 + i % 60, + i % 12 + 1, + i % 28 + 1, + i % 24 + ) + }), + ), + ( + "mixed", + create_batch(|i| match i % 5 { + 0 => format!( + "{:04}-{:02}-{:02} 12:34:56", + 1970 + i % 60, + i % 12 + 1, + i % 28 + 1 + ), + 1 => format!( + "{:04}-{:02}-{:02}T12:34:56.123456Z", + 1970 + i % 60, + i % 12 + 1, + i % 28 + 1 + ), + 2 => format!( + " {:04}-{:02}-{:02} ", + 1900 + i % 200, + i % 12 + 1, + i % 28 + 1 + ), + 3 => "T12:34:56".to_string(), + _ => "not a timestamp".to_string(), + }), + ), + ]; + + // Timezone-aware and NTZ go through different parsers (`timestamp_parser` vs + // `timestamp_ntz_parser`), and a non-UTC session timezone exercises the offset lookup that + // UTC short-circuits, so all three are measured. + for (target_name, to_type, timezone) in [ + ( + "timestamp", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + "UTC", + ), + ( + "timestamp_non_utc", + DataType::Timestamp(TimeUnit::Microsecond, Some("America/Los_Angeles".into())), + "America/Los_Angeles", + ), + ( + "timestamp_ntz", + DataType::Timestamp(TimeUnit::Microsecond, None), + "UTC", + ), + ] { + for (mode, mode_name) in [ + (EvalMode::Legacy, "legacy"), + (EvalMode::Ansi, "ansi"), + (EvalMode::Try, "try"), + ] { + let mut group = + c.benchmark_group(format!("cast_string_to_{}/{}", target_name, mode_name)); + for (name, batch) in &batches { + // ANSI raises on the first invalid value, so timing it against a batch that is + // mostly invalid would measure the error path rather than the parser. + if mode == EvalMode::Ansi && *name == "mixed" { + continue; + } + let cast = Cast::new( + Arc::clone(&expr), + to_type.clone(), + SparkCastOptions::new(mode, timezone, false), + None, + None, + ); + group.bench_function(*name, |b| { + b.iter(|| cast.evaluate(batch).unwrap()); + }); + } + group.finish(); + } + } + + // The Spark 4 path adds a leading-whitespace check for T-prefixed time-only strings, so it + // is measured separately on the inputs where that check can fire. + let mut group = c.benchmark_group("cast_string_to_timestamp/spark4_legacy"); + for name in ["padded", "mixed"] { + let batch = &batches.iter().find(|(n, _)| *n == name).unwrap().1; + let cast = Cast::new( + Arc::clone(&expr), + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + SparkCastOptions::new_with_version(EvalMode::Legacy, "UTC", false, true), + None, + None, + ); + group.bench_function(name, |b| { + b.iter(|| cast.evaluate(batch).unwrap()); + }); + } + group.finish(); +} + +fn create_batch(f: impl Fn(usize) -> String) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)])); + let mut builder = StringBuilder::new(); + for i in 0..BATCH_SIZE { + if i % 17 == 0 { + builder.append_null(); + } else { + builder.append_value(f(i)); + } + } + RecordBatch::try_new(schema, vec![Arc::new(builder.finish())]).unwrap() +} Review Comment: Fixed in 434efec. There is a new `benches/common/mod.rs` with one parameterized `string_batch(rows, null_modulus, value_fn)`, pulled into all three cast-from-string benches with `#[path]`. `create_batch` here is now a one-line delegate: ```rust fn create_batch(f: impl Fn(usize) -> String) -> RecordBatch { common::string_batch(BATCH_SIZE, 17, f) } ``` and `cast_string_to_date.rs` uses the same helper, so the byte-for-byte copy is gone. The helper lives in a subdirectory deliberately: Cargo's bench auto-discovery only scans `benches/*.rs`, so a top-level `common.rs` would be built as its own bench target. ########## native/spark-expr/benches/cast_from_string.rs: ########## Review Comment: Fixed in 434efec: `EVAL_MODES` is now a single const used at all three sites. `grep -c` for the inline array in that file is down from 3 to 1 (the const itself). ########## native/spark-expr/benches/cast_from_string.rs: ########## @@ -133,6 +133,80 @@ fn criterion_benchmark(c: &mut Criterion) { } group.finish(); } + + // str -> boolean and str -> float benchmarks, with and without the leading/trailing + // whitespace that exercises the trim helpers in `conversion_funcs::trim` + let bool_batch = create_boolean_string_batch(false); + let bool_padded_batch = create_boolean_string_batch(true); + let float_batch = create_float_string_batch(false); + let float_padded_batch = create_float_string_batch(true); + for (mode, mode_name) in [ + (EvalMode::Legacy, "legacy"), + (EvalMode::Ansi, "ansi"), + (EvalMode::Try, "try"), + ] { + let spark_cast_options = SparkCastOptions::new(mode, "", false); + let mut group = c.benchmark_group(format!("cast_string_to_bool_and_float/{}", mode_name)); + for (data_type, name, batch) in [ + (DataType::Boolean, "boolean", &bool_batch), + (DataType::Boolean, "boolean_padded", &bool_padded_batch), + (DataType::Float32, "float", &float_batch), + (DataType::Float64, "double", &float_batch), + (DataType::Float64, "double_padded", &float_padded_batch), + ] { + let cast = Cast::new( + expr.clone(), + data_type, + spark_cast_options.clone(), + None, + None, + ); + group.bench_function(name, |b| { + b.iter(|| cast.evaluate(batch).unwrap()); + }); + } + group.finish(); + } +} + +/// Create batch with the boolean spellings Spark accepts, optionally space-padded +fn create_boolean_string_batch(padded: bool) -> RecordBatch { + let words = ["true", "FALSE", "t", "n", "yes", "0", "TRUE", "no"]; + create_string_batch(|i| { + let word = words[i % words.len()]; + if padded { + format!(" {} ", word) + } else { + word.to_string() + } + }) +} + +/// Create batch with floating point strings, optionally space-padded +fn create_float_string_batch(padded: bool) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(42); + create_string_batch(move |_| { + let value = rng.random_range(-1_000_000.0..1_000_000.0f64); + if padded { + format!(" {} ", value) + } else { + format!("{}", value) + } + }) +} + +/// Create a single-column Utf8 batch of 8192 rows, every tenth one null +fn create_string_batch(mut value: impl FnMut(usize) -> String) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)])); + let mut b = StringBuilder::new(); + for i in 0..8192 { + if i % 10 == 0 { + b.append_null(); + } else { + b.append_value(value(i)); + } + } + RecordBatch::try_new(schema, vec![Arc::new(b.finish())]).unwrap() } Review Comment: Fixed in 434efec: all five near-copies (`create_small_int_string_batch`, `create_int_string_batch`, `create_string_batch`, and the other two) collapse into `common::string_batch`. `cast_from_string.rs` is 185 lines lighter as a result. All three benches smoke-tested with `--test`. -- 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]
