diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..fbad7d20ce 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -60,6 +60,10 @@ path = "src/lib.rs" name = "cast_from_string" harness = false +[[bench]] +name = "cast_string_to_date" +harness = false + [[bench]] name = "cast_numeric" harness = false diff --git a/native/spark-expr/benches/cast_string_to_date.rs b/native/spark-expr/benches/cast_string_to_date.rs new file mode 100644 index 0000000000..9bfb33470d --- /dev/null +++ b/native/spark-expr/benches/cast_string_to_date.rs @@ -0,0 +1,84 @@ +// 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}; +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)); + let cast_to_date = Cast::new( + expr, + DataType::Date32, + SparkCastOptions::new(EvalMode::Legacy, "UTC", false), + None, + None, + ); + + let canonical = + create_batch(|i| format!("{:04}-{:02}-{:02}", 1970 + i % 60, i % 12 + 1, i % 28 + 1)); + let mixed = create_batch(|i| match i % 5 { + 0 => format!("{:04}-{:02}-{:02}", 1970 + i % 60, i % 12 + 1, i % 28 + 1), + 1 => format!("{}-{}-{} 12:34:56", 1970 + i % 60, i % 12 + 1, i % 28 + 1), + 2 => format!( + " {:04}-{:02}-{:02} ", + 1900 + i % 200, + i % 12 + 1, + i % 28 + 1 + ), + 3 => format!("{:04}", 1970 + i % 60), + _ => "not a date".to_string(), + }); + + let mut group = c.benchmark_group("cast_string_to_date"); + group.bench_function("canonical", |b| { + b.iter(|| cast_to_date.evaluate(&canonical).unwrap()); + }); + group.bench_function("mixed", |b| { + b.iter(|| cast_to_date.evaluate(&mixed).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() +} + +fn config() -> Criterion { + Criterion::default() +} + +criterion_group! { + name = benches; + config = config(); + targets = criterion_benchmark +} +criterion_main!(benches); diff --git a/native/spark-expr/src/conversion_funcs/string.rs b/native/spark-expr/src/conversion_funcs/string.rs index adfd6e2390..b4e03cdf6f 100644 --- a/native/spark-expr/src/conversion_funcs/string.rs +++ b/native/spark-expr/src/conversion_funcs/string.rs @@ -20,12 +20,11 @@ use arrow::array::{ Array, ArrayRef, ArrowPrimitiveType, BooleanArray, Decimal128Builder, GenericStringArray, OffsetSizeTrait, PrimitiveArray, PrimitiveBuilder, StringArray, }; -use arrow::compute::DecimalCast; use arrow::datatypes::{ i256, is_validate_decimal_precision, DataType, Date32Type, Decimal256Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, TimestampMicrosecondType, }; -use chrono::{DateTime, LocalResult, NaiveDate, NaiveTime, Offset, TimeZone, Timelike}; +use chrono::{LocalResult, NaiveDate, NaiveTime, Offset, TimeZone, Timelike}; use num::traits::CheckedNeg; use num::{CheckedSub, Integer}; use regex::Regex; @@ -1161,6 +1160,25 @@ fn days_from_civil(y: i64, m: i64, d: i64) -> i64 { era * 146097 + doe - 719468 } +fn is_leap_year(year: i64) -> bool { + year % 4 == 0 && (year % 100 != 0 || year % 400 == 0) +} + +/// Days since 1970-01-01 for a proleptic Gregorian year/month/day, or `None` when the +/// combination is not a real calendar date. Unlike `NaiveDate::from_ymd_opt`, this accepts +/// any year that fits in `i64`. +fn ymd_to_epoch_day(year: i64, month: i64, day: i64) -> Option { + const DAYS_IN_MONTH: [i64; 12] = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]; + let mut max_day = *DAYS_IN_MONTH.get(usize::try_from(month.checked_sub(1)?).ok()?)?; + if month == 2 && is_leap_year(year) { + max_day = 29; + } + if day < 1 || day > max_day { + return None; + } + Some(days_from_civil(year, month, day)) +} + fn parse_timestamp_to_micros( timestamp_info: &TimeStampInfo, tz: &T, @@ -1226,27 +1244,11 @@ fn parse_timestamp_to_micros( return Ok(None); } // Validate month and day manually for extreme years. - let m = timestamp_info.month; - let d = timestamp_info.day; - if !(1..=12).contains(&m) { - return Ok(None); - } - let max_day = match m { - 1 | 3 | 5 | 7 | 8 | 10 | 12 => 31, - 4 | 6 | 9 | 11 => 30, - 2 => { - let leap = year % 4 == 0 && (year % 100 != 0 || year % 400 == 0); - if leap { - 29 - } else { - 28 - } - } - _ => return Ok(None), - }; - if d < 1 || d > max_day { - return Ok(None); - } + let days = + match ymd_to_epoch_day(year, timestamp_info.month as i64, timestamp_info.day as i64) { + Some(days) => days, + None => return Ok(None), + }; // Compute the timezone offset using epoch as a surrogate probe point. // Extreme-year timestamps are only valid with a UTC-like fixed offset (any DST // zone would overflow). Using epoch gives us the standard offset. @@ -1263,7 +1265,6 @@ fn parse_timestamp_to_micros( // Use i128 for the intermediate multiply-by-1_000_000 step: the seconds value can be // just outside the i64 range while the final microseconds result is still within range // (e.g., Long.MinValue boundary: seconds = -9_223_372_036_855, result = i64::MIN). - let days = days_from_civil(year, m as i64, d as i64); let time_secs = timestamp_info.hour as i64 * 3600 + timestamp_info.minute as i64 * 60 + timestamp_info.second as i64; @@ -1281,33 +1282,16 @@ fn parse_timestamp_to_micros( fn local_datetime_to_micros(timestamp_info: &TimeStampInfo) -> SparkResult> { let year = timestamp_info.year as i64; - let m = timestamp_info.month; - let d = timestamp_info.day; - if !(1..=12).contains(&m) { - return Ok(None); - } - let max_day = match m { - 1 | 3 | 5 | 7 | 8 | 10 | 12 => 31u32, - 4 | 6 | 9 | 11 => 30, - 2 => { - let leap = year % 4 == 0 && (year % 100 != 0 || year % 400 == 0); - if leap { - 29 - } else { - 28 - } - } - _ => return Ok(None), + let days = match ymd_to_epoch_day(year, timestamp_info.month as i64, timestamp_info.day as i64) + { + Some(days) => days, + None => return Ok(None), }; - if d < 1 || d > max_day { - return Ok(None); - } if timestamp_info.hour >= 24 || timestamp_info.minute >= 60 || timestamp_info.second >= 60 { return Ok(None); } - let days = days_from_civil(year, m as i64, d as i64); let time_secs = timestamp_info.hour as i64 * 3600 + timestamp_info.minute as i64 * 60 + timestamp_info.second as i64; @@ -1864,11 +1848,17 @@ fn date_parser(date_str: &str, eval_mode: EvalMode) -> SparkResult> byte.is_ascii_whitespace() || byte.is_ascii_control() } + /// Decodes a run of ASCII digits, or `None` if any byte is not a digit. + fn decode_digits(bytes: &[u8]) -> Option { + bytes.iter().try_fold(0i64, |acc, b| { + b.is_ascii_digit().then(|| acc * 10 + (b - b'0') as i64) + }) + } + fn is_valid_digits(segment: i32, digits: usize) -> bool { - // NaiveDate is bounded to [-262142, 262142] (6 digits). We allow up to 7 digits to support - // leading-zero year strings like "0002020" (= year 2020), matching Spark's - // isValidDigits. Values outside the bounds are caught by an explicit bounds - // check below. + // Years are bounded to [-262143, 262142] by the explicit range check below. We allow up + // to 7 digits to support leading-zero year strings like "0002020" (= year 2020), + // matching Spark's isValidDigits. let max_digits_year = 7; // year (segment 0) can be between 4 to 7 digits, // month and day (segment 1 and 2) can be between 1 to 2 digits @@ -1893,12 +1883,6 @@ fn date_parser(date_str: &str, eval_mode: EvalMode) -> SparkResult> return return_result(date_str, eval_mode); } - //values of date segments year, month and day defaulting to 1 - let mut date_segments = [1, 1, 1]; - let mut sign = 1; - let mut current_segment = 0; - let mut current_segment_value = Wrapping(0); - let mut current_segment_digits = 0; let bytes = date_str.as_bytes(); let mut j = get_trimmed_start(bytes); @@ -1908,6 +1892,29 @@ fn date_parser(date_str: &str, eval_mode: EvalMode) -> SparkResult> return return_result(date_str, eval_mode); } + // Fast path for the canonical `yyyy-mm-dd` form. A 4-digit year is always inside the + // range checked below, so the only way this can fail is an invalid calendar date, which + // is a null in every eval mode rather than an ANSI error. Any other shape (including a + // leading sign, which makes the first byte a non-digit) falls through to the general + // parser below. + let trimmed = &bytes[j..str_end_trimmed]; + if trimmed.len() == 10 && trimmed[4] == b'-' && trimmed[7] == b'-' { + if let (Some(year), Some(month), Some(day)) = ( + decode_digits(&trimmed[..4]), + decode_digits(&trimmed[5..7]), + decode_digits(&trimmed[8..10]), + ) { + return Ok(ymd_to_epoch_day(year, month, day).map(|days| days as i32)); + } + } + + //values of date segments year, month and day defaulting to 1 + let mut date_segments = [1, 1, 1]; + let mut sign = 1; + let mut current_segment = 0; + let mut current_segment_value = Wrapping(0); + let mut current_segment_digits = 0; + // assign a sign to the date; both '-' and '+' are accepted (Spark stringToDate line 357-360) if bytes[j] == b'-' { sign = -1; @@ -1960,15 +1967,12 @@ fn date_parser(date_str: &str, eval_mode: EvalMode) -> SparkResult> return Ok(None); } - match NaiveDate::from_ymd_opt(year, date_segments[1] as u32, date_segments[2] as u32) { - Some(date) => { - let duration_since_epoch = date - .signed_duration_since(DateTime::UNIX_EPOCH.naive_utc().date()) - .num_days(); - Ok(Some(duration_since_epoch.to_i32().unwrap())) - } - None => Ok(None), - } + Ok(ymd_to_epoch_day( + year as i64, + date_segments[1] as i64, + date_segments[2] as i64, + ) + .map(|days| days as i32)) } #[cfg(test)]