perf: optimize spark_cast_utf8_to_boolean (>2x faster)

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

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);

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::<F>().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<bool> {
+    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<OffsetSize>(
     from: &dyn Array,
     eval_mode: EvalMode,
@@ -266,22 +288,26 @@ where
         .downcast_ref::<GenericStringArray<OffsetSize>>()
         .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::<Result<BooleanArray, _>>()?;
+        .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() {