Optimize performance of string::ascii function

Author: Tai Le Manh

datafusion/functions/Cargo.toml

@@ -90,6 +90,11 @@ criterion = { workspace = true }
 rand = { workspace = true }
 tokio = { workspace = true, features = ["macros", "rt", "sync"] }
 
+[[bench]]
+harness = false
+name = "ascii"
+required-features = ["string_expressions"]
+
 [[bench]]
 harness = false
 name = "concat"

datafusion/functions/benches/ascii.rs

@@ -0,0 +1,116 @@
+// 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.
+
+extern crate criterion;
+mod helper;
+
+use arrow::datatypes::{DataType, Field};
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use datafusion_expr::ScalarFunctionArgs;
+use helper::gen_string_array;
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let ascii = datafusion_functions::string::ascii();
+
+    // All benches are single batch run with 8192 rows
+    const N_ROWS: usize = 8192;
+    const STR_LEN: usize = 16;
+    const NULL_DENSITY: f32 = 0.1;
+    const UTF8_DENSITY_OF_ALL_ASCII: f32 = 0.0;
+    const NORMAL_UTF8_DENSITY: f32 = 0.8;
+
+    // StringArray ASCII only
+    let args_string_ascii = gen_string_array(
+        N_ROWS,
+        STR_LEN,
+        NULL_DENSITY,
+        UTF8_DENSITY_OF_ALL_ASCII,
+        false,
+    );
+    c.bench_function("ascii/string_ascii_only", |b| {
+        b.iter(|| {
+            black_box(ascii.invoke_with_args(ScalarFunctionArgs {
+                args: args_string_ascii.clone(),
+                arg_fields: vec![&Field::new(
+                    "a",
+                    args_string_ascii[0].data_type(),
+                    true,
+                )],
+                number_rows: N_ROWS,
+                return_field: &Field::new("f", DataType::Utf8, true),
+            }))
+        })
+    });
+
+    // StringArray UTF8
+    let args_string_utf8 =
+        gen_string_array(N_ROWS, STR_LEN, NULL_DENSITY, NORMAL_UTF8_DENSITY, false);
+    c.bench_function("ascii/string_utf8", |b| {
+        b.iter(|| {
+            black_box(ascii.invoke_with_args(ScalarFunctionArgs {
+                args: args_string_utf8.clone(),
+                arg_fields: vec![&Field::new("a", args_string_utf8[0].data_type(), true)],
+                number_rows: N_ROWS,
+                return_field: &Field::new("f", DataType::Utf8, true),
+            }))
+        })
+    });
+
+    // StringViewArray ASCII only
+    let args_string_view_ascii = gen_string_array(
+        N_ROWS,
+        STR_LEN,
+        NULL_DENSITY,
+        UTF8_DENSITY_OF_ALL_ASCII,
+        true,
+    );
+    c.bench_function("ascii/string_view_ascii_only", |b| {
+        b.iter(|| {
+            black_box(ascii.invoke_with_args(ScalarFunctionArgs {
+                args: args_string_view_ascii.clone(),
+                arg_fields: vec![&Field::new(
+                    "a",
+                    args_string_view_ascii[0].data_type(),
+                    true,
+                )],
+                number_rows: N_ROWS,
+                return_field: &Field::new("f", DataType::Utf8, true),
+            }))
+        })
+    });
+
+    // StringViewArray UTF8
+    let args_string_view_utf8 =
+        gen_string_array(N_ROWS, STR_LEN, NULL_DENSITY, NORMAL_UTF8_DENSITY, true);
+    c.bench_function("ascii/string_view_utf8", |b| {
+        b.iter(|| {
+            black_box(ascii.invoke_with_args(ScalarFunctionArgs {
+                args: args_string_view_utf8.clone(),
+                arg_fields: vec![&Field::new(
+                    "a",
+                    args_string_view_utf8[0].data_type(),
+                    true,
+                )],
+                number_rows: N_ROWS,
+                return_field: &Field::new("f", DataType::Utf8, true),
+            }))
+        })
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

datafusion/functions/src/string/ascii.rs

@@ -16,9 +16,12 @@
 // under the License.
 
 use crate::utils::make_scalar_function;
-use arrow::array::{ArrayAccessor, ArrayIter, ArrayRef, AsArray, Int32Array};
-use arrow::datatypes::DataType;
+use arrow::array::{
+    ArrayData, ArrayRef, ArrowPrimitiveType, AsArray, Int32Array, StringArrayType,
+};
+use arrow::datatypes::{DataType, Int32Type};
 use arrow::error::ArrowError;
+use arrow_buffer::Buffer;
 use datafusion_common::types::logical_string;
 use datafusion_common::{internal_err, Result};
 use datafusion_expr::{ColumnarValue, Documentation, TypeSignatureClass};
@@ -103,19 +106,29 @@ impl ScalarUDFImpl for AsciiFunc {
 
 fn calculate_ascii<'a, V>(array: V) -> Result<ArrayRef, ArrowError>
 where
-    V: ArrayAccessor<Item = &'a str>,
+    V: StringArrayType<'a, Item = &'a str>,
 {
-    let iter = ArrayIter::new(array);
-    let result = iter
-        .map(|string| {
-            string.map(|s| {
-                let mut chars = s.chars();
-                chars.next().map_or(0, |v| v as i32)
-            })
-        })
-        .collect::<Int32Array>();
-
-    Ok(Arc::new(result) as ArrayRef)
+    let mut values = Vec::with_capacity(array.len());
+
+    for i in 0..array.len() {
+        if array.is_null(i) {
+            values.push(0);
+        } else {
+            let s = array.value(i);
+            let code = s.chars().next().map_or(0, |c| c as i32);
+            values.push(code);
+        }
+    }
+
+    let nulls = array.nulls().cloned();
+
+    let data = ArrayData::builder(Int32Type::DATA_TYPE)
+        .len(array.len())
+        .add_buffer(Buffer::from_slice_ref(&values))
+        .null_bit_buffer(nulls.map(|n| n.buffer().clone()))
+        .build()?;
+
+    Ok(Arc::new(Int32Array::from(data)) as ArrayRef)
 }
 
 /// Returns the numeric code of the first character of the argument.
@@ -182,6 +195,7 @@ mod tests {
         test_ascii!(Some(String::from("x")), Ok(Some(120)));
         test_ascii!(Some(String::from("a")), Ok(Some(97)));
         test_ascii!(Some(String::from("")), Ok(Some(0)));
+        test_ascii!(Some(String::from("🚀")), Ok(Some(128640)));
         test_ascii!(None, Ok(None));
         Ok(())
     }