diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..be14acdad8 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -139,3 +139,7 @@ harness = false [[bench]] name = "to_json" harness = false + +[[bench]] +name = "cast_string_to_boolean" +harness = false diff --git a/native/spark-expr/benches/cast_string_to_boolean.rs b/native/spark-expr/benches/cast_string_to_boolean.rs new file mode 100644 index 0000000000..ce6f5063eb --- /dev/null +++ b/native/spark-expr/benches/cast_string_to_boolean.rs @@ -0,0 +1,95 @@ +// 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, ArrayRef}; +use arrow::datatypes::DataType; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::{spark_cast, EvalMode, SparkCastOptions}; +use std::sync::Arc; + +const NUM_ROWS: usize = 8192; + +/// Strings that Spark accepts as booleans, in the mixed casing and padding that shows +/// up in real data. +const VALID: [&str; 12] = [ + "true", "TRUE", "True", "false", "FALSE", " t ", "y", "YES", "no", "N", "1", "0", +]; + +/// A mix of valid boolean spellings and values that cast to NULL. +fn mixed_batch() -> ArrayRef { + let mut b = StringBuilder::new(); + for i in 0..NUM_ROWS { + match i % 10 { + 0 => b.append_null(), + 1 => b.append_value("maybe"), + 2 => b.append_value("2"), + _ => b.append_value(VALID[i % VALID.len()]), + } + } + Arc::new(b.finish()) +} + +/// Only valid boolean spellings, so every row takes the parsing path. +fn valid_batch() -> ArrayRef { + let mut b = StringBuilder::new(); + for i in 0..NUM_ROWS { + b.append_value(VALID[i % VALID.len()]); + } + Arc::new(b.finish()) +} + +fn criterion_benchmark(c: &mut Criterion) { + let mixed = mixed_batch(); + let valid = valid_batch(); + + let mut group = c.benchmark_group("cast_string_to_boolean"); + for (mode, mode_name) in [ + (EvalMode::Legacy, "legacy"), + (EvalMode::Try, "try"), + (EvalMode::Ansi, "ansi"), + ] { + let options = SparkCastOptions::new(mode, "", false); + // ANSI mode errors on any unparseable value, so it only gets the valid batch. + if mode != EvalMode::Ansi { + group.bench_function(format!("{mode_name}/mixed"), |b| { + b.iter(|| { + spark_cast( + ColumnarValue::Array(Arc::clone(&mixed)), + &DataType::Boolean, + &options, + ) + .unwrap() + }); + }); + } + group.bench_function(format!("{mode_name}/valid"), |b| { + b.iter(|| { + spark_cast( + ColumnarValue::Array(Arc::clone(&valid)), + &DataType::Boolean, + &options, + ) + .unwrap() + }); + }); + } + group.finish(); +} + +criterion_group!(benches, 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..1a234ae0be 100644 --- a/native/spark-expr/src/conversion_funcs/string.rs +++ b/native/spark-expr/src/conversion_funcs/string.rs @@ -17,8 +17,8 @@ use crate::{timezone, EvalMode, SparkError, SparkResult}; use arrow::array::{ - Array, ArrayRef, ArrowPrimitiveType, BooleanArray, Decimal128Builder, GenericStringArray, - OffsetSizeTrait, PrimitiveArray, PrimitiveBuilder, StringArray, + Array, ArrayRef, ArrowPrimitiveType, BooleanArray, BooleanBuilder, Decimal128Builder, + GenericStringArray, OffsetSizeTrait, PrimitiveArray, PrimitiveBuilder, StringArray, }; use arrow::compute::DecimalCast; use arrow::datatypes::{ @@ -254,6 +254,28 @@ where pruned_float_str.parse::().ok() } +/// Parse the boolean literals Spark accepts in a string-to-boolean cast, returning +/// `None` for anything else. +/// +/// Trimming before comparing is equivalent to trimming afterwards: ASCII case folding +/// never turns a whitespace character into a non-whitespace one, or vice versa. +#[inline] +fn parse_utf8_to_boolean(value: &str) -> Option { + let trimmed = value.trim(); + match trimmed.len() { + 1 => match trimmed.as_bytes()[0] { + b't' | b'T' | b'y' | b'Y' | b'1' => Some(true), + b'f' | b'F' | b'n' | b'N' | b'0' => Some(false), + _ => None, + }, + 2 if trimmed.eq_ignore_ascii_case("no") => Some(false), + 3 if trimmed.eq_ignore_ascii_case("yes") => Some(true), + 4 if trimmed.eq_ignore_ascii_case("true") => Some(true), + 5 if trimmed.eq_ignore_ascii_case("false") => Some(false), + _ => None, + } +} + pub(crate) fn spark_cast_utf8_to_boolean( from: &dyn Array, eval_mode: EvalMode, @@ -266,22 +288,26 @@ where .downcast_ref::>() .unwrap(); - let output_array = array + // Only ANSI mode can fail, so the other modes avoid the per-row `Result` entirely. + if eval_mode == EvalMode::Ansi { + let mut builder = BooleanBuilder::with_capacity(array.len()); + for value in array.iter() { + match value { + Some(value) => match parse_utf8_to_boolean(value) { + Some(parsed) => builder.append_value(parsed), + None => return Err(invalid_value(value, "STRING", "BOOLEAN")), + }, + None => builder.append_null(), + } + } + + return Ok(Arc::new(builder.finish())); + } + + let output_array: BooleanArray = array .iter() - .map(|value| match value { - Some(value) => match value.to_ascii_lowercase().trim() { - "t" | "true" | "y" | "yes" | "1" => Ok(Some(true)), - "f" | "false" | "n" | "no" | "0" => Ok(Some(false)), - _ if eval_mode == EvalMode::Ansi => Err(SparkError::CastInvalidValue { - value: value.to_string(), - from_type: "STRING".to_string(), - to_type: "BOOLEAN".to_string(), - }), - _ => Ok(None), - }, - _ => Ok(None), - }) - .collect::>()?; + .map(|value| value.and_then(parse_utf8_to_boolean)) + .collect(); Ok(Arc::new(output_array)) } @@ -1989,6 +2015,38 @@ mod tests { } } + #[test] + fn test_parse_utf8_to_boolean() { + for truthy in [ + "t", "T", "true", "TRUE", "tRuE", "y", "Y", "yes", "YES", "1", + ] { + assert_eq!(parse_utf8_to_boolean(truthy), Some(true), "{truthy}"); + } + for falsy in [ + "f", "F", "false", "FALSE", "fAlSe", "n", "N", "no", "NO", "0", + ] { + assert_eq!(parse_utf8_to_boolean(falsy), Some(false), "{falsy}"); + } + // Surrounding whitespace is ignored, as in Spark. + assert_eq!(parse_utf8_to_boolean(" \ttrue\n "), Some(true)); + assert_eq!(parse_utf8_to_boolean(" no "), Some(false)); + // Anything else casts to NULL. + for invalid in [ + "", + " ", + "tru", + "truely", + "2", + "-1", + "on", + "off", + "úñî", + "ｔｒｕｅ", + ] { + assert_eq!(parse_utf8_to_boolean(invalid), None, "{invalid}"); + } + } + #[test] #[cfg_attr(miri, ignore)] // test takes too long with miri fn test_cast_string_to_timestamp() {