perf: use bulk-NULL semantics in split and substring, skip Vec allocation in split (#4403)

mbutrovich · web-flow · commit 2a266f694013 · 2026-05-22T11:15:43.000-04:00
diff --git a/native/spark-expr/src/string_funcs/split.rs b/native/spark-expr/src/string_funcs/split.rs
@@ -15,7 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{Array, ArrayRef, GenericStringArray, ListArray};
+use arrow::array::{
+    Array, ArrayBuilder, ArrayRef, GenericListArray, GenericStringArray, GenericStringBuilder,
+    ListArray, OffsetSizeTrait,
+};
+use arrow::buffer::OffsetBuffer;
 use arrow::datatypes::{DataType, Field};
 use datafusion::common::{
     cast::as_generic_string_array, exec_err, DataFusionError, Result as DataFusionResult,
@@ -115,89 +119,99 @@ fn split_array(
         DataFusionError::Execution(format!("Invalid regex pattern '{}': {}", pattern, e))
     })?;
 
-    let string_array = match string_array.data_type() {
-        DataType::Utf8 => as_generic_string_array::<i32>(string_array)?,
-        DataType::LargeUtf8 => {
-            // Convert LargeUtf8 to Utf8 for processing
-            let large_array = as_generic_string_array::<i64>(string_array)?;
-            return split_large_string_array(large_array, &regex, limit);
+    match string_array.data_type() {
+        DataType::Utf8 => {
+            split_generic::<i32>(as_generic_string_array::<i32>(string_array)?, &regex, limit)
         }
-        _ => {
-            return exec_err!(
-                "split expects Utf8 or LargeUtf8 string array, got {:?}",
-                string_array.data_type()
-            );
+        DataType::LargeUtf8 => {
+            split_generic::<i64>(as_generic_string_array::<i64>(string_array)?, &regex, limit)
         }
-    };
+        _ => exec_err!(
+            "split expects Utf8 or LargeUtf8 string array, got {:?}",
+            string_array.data_type()
+        ),
+    }
+}
 
-    // Build the result ListArray
-    let mut offsets: Vec<i32> = Vec::with_capacity(string_array.len() + 1);
-    let mut values: Vec<String> = Vec::new();
-    let mut null_buffer_builder = arrow::array::BooleanBufferBuilder::new(string_array.len());
-    offsets.push(0);
-
-    for i in 0..string_array.len() {
-        if string_array.is_null(i) {
-            // NULL input produces NULL in result (Spark behavior)
-            offsets.push(offsets[i]);
-            null_buffer_builder.append(false); // false = NULL
-        } else {
-            let string_val = string_array.value(i);
-            let parts = split_with_regex(string_val, &regex, limit);
-            values.extend(parts);
-            offsets.push(values.len() as i32);
-            null_buffer_builder.append(true); // true = valid
+fn split_generic<O: OffsetSizeTrait>(
+    string_array: &GenericStringArray<O>,
+    regex: &Regex,
+    limit: i32,
+) -> DataFusionResult<ColumnarValue> {
+    let len = string_array.len();
+    let mut offsets: Vec<O> = Vec::with_capacity(len + 1);
+    let mut values_builder = GenericStringBuilder::<O>::new();
+    offsets.push(O::usize_as(0));
+
+    // Bulk-NULL: output null mask equals input's, so reuse it instead of
+    // tracking per-row in a NullBufferBuilder. Null rows contribute no parts
+    // (offset does not advance) and the cloned NullBuffer marks them.
+    for i in 0..len {
+        if !string_array.is_null(i) {
+            let s = string_array.value(i);
+            push_split_parts(s, regex, limit, &mut values_builder);
         }
+        offsets.push(O::usize_as(values_builder.len()));
     }
 
-    let values_array = Arc::new(GenericStringArray::<i32>::from(values)) as ArrayRef;
-    let field = Arc::new(Field::new("item", DataType::Utf8, false));
-    let nulls = arrow::buffer::NullBuffer::new(null_buffer_builder.finish());
-    let list_array = ListArray::new(
+    let values_array = Arc::new(values_builder.finish()) as ArrayRef;
+    let item_type = if O::IS_LARGE {
+        DataType::LargeUtf8
+    } else {
+        DataType::Utf8
+    };
+    let field = Arc::new(Field::new("item", item_type, false));
+    let list_array = GenericListArray::<O>::new(
         field,
-        arrow::buffer::OffsetBuffer::new(offsets.into()),
+        OffsetBuffer::new(offsets.into()),
         values_array,
-        Some(nulls),
+        string_array.nulls().cloned(),
     );
 
     Ok(ColumnarValue::Array(Arc::new(list_array)))
 }
 
-fn split_large_string_array(
-    string_array: &GenericStringArray<i64>,
+/// Push the splits of `string` into `builder`. Avoids materializing an
+/// intermediate `Vec<String>` — appends each `&str` slice from the regex
+/// iterator directly (the builder copies into its own buffer).
+fn push_split_parts<O: OffsetSizeTrait>(
+    string: &str,
     regex: &Regex,
     limit: i32,
-) -> DataFusionResult<ColumnarValue> {
-    let mut offsets: Vec<i32> = Vec::with_capacity(string_array.len() + 1);
-    let mut values: Vec<String> = Vec::new();
-    let mut null_buffer_builder = arrow::array::BooleanBufferBuilder::new(string_array.len());
-    offsets.push(0);
-
-    for i in 0..string_array.len() {
-        if string_array.is_null(i) {
-            // NULL input produces NULL in result (Spark behavior)
-            offsets.push(offsets[i]);
-            null_buffer_builder.append(false); // false = NULL
+    builder: &mut GenericStringBuilder<O>,
+) {
+    if limit == 0 {
+        // limit = 0: split all, drop trailing empties. Need to know the end
+        // before pushing, so collect borrowed slices first (no string copies).
+        let mut parts: Vec<&str> = regex.split(string).collect();
+        while parts.last().is_some_and(|s| s.is_empty()) {
+            parts.pop();
+        }
+        if parts.is_empty() {
+            builder.append_value("");
         } else {
-            let string_val = string_array.value(i);
-            let parts = split_with_regex(string_val, regex, limit);
-            values.extend(parts);
-            offsets.push(values.len() as i32);
-            null_buffer_builder.append(true); // true = valid
+            for p in parts {
+                builder.append_value(p);
+            }
+        }
+    } else if limit > 0 {
+        // limit > 0: at most limit-1 splits.
+        let mut last_end = 0;
+        let cap = (limit - 1) as usize;
+        for (count, mat) in regex.find_iter(string).enumerate() {
+            if count >= cap {
+                break;
+            }
+            builder.append_value(&string[last_end..mat.start()]);
+            last_end = mat.end();
+        }
+        builder.append_value(&string[last_end..]);
+    } else {
+        // limit < 0: split all, keep trailing empties.
+        for p in regex.split(string) {
+            builder.append_value(p);
         }
     }
-
-    let values_array = Arc::new(GenericStringArray::<i32>::from(values)) as ArrayRef;
-    let field = Arc::new(Field::new("item", DataType::Utf8, false));
-    let nulls = arrow::buffer::NullBuffer::new(null_buffer_builder.finish());
-    let list_array = ListArray::new(
-        field,
-        arrow::buffer::OffsetBuffer::new(offsets.into()),
-        values_array,
-        Some(nulls),
-    );
-
-    Ok(ColumnarValue::Array(Arc::new(list_array)))
 }
 
 fn split_string(string: &str, pattern: &str, limit: i32) -> DataFusionResult<Vec<String>> {
diff --git a/native/spark-expr/src/string_funcs/substring.rs b/native/spark-expr/src/string_funcs/substring.rs
@@ -18,7 +18,9 @@
 #![allow(deprecated)]
 
 use crate::kernels::strings::substring;
-use arrow::array::{as_dictionary_array, as_largestring_array, as_string_array, Array, ArrayRef};
+use arrow::array::{
+    as_dictionary_array, as_largestring_array, as_string_array, Array, ArrayRef, GenericStringArray,
+};
 use arrow::datatypes::{DataType, Int32Type, Schema};
 use arrow::record_batch::RecordBatch;
 use datafusion::logical_expr::ColumnarValue;
@@ -128,67 +130,69 @@ fn spark_substring_negative_start(
     start: i64,
     len: u64,
 ) -> datafusion::common::Result<ArrayRef> {
-    use arrow::array::{
-        BinaryArray, DictionaryArray, GenericBinaryBuilder, GenericStringBuilder, LargeBinaryArray,
-    };
-
-    match array.data_type() {
-        DataType::Utf8 => {
-            let str_array = as_string_array(array);
-            let mut builder = GenericStringBuilder::<i32>::new();
-            for i in 0..str_array.len() {
-                if str_array.is_null(i) {
-                    builder.append_null();
-                } else {
-                    builder.append_value(spark_substr_negative(str_array.value(i), start, len));
-                }
-            }
-            Ok(Arc::new(builder.finish()) as ArrayRef)
-        }
-        DataType::LargeUtf8 => {
-            let str_array = as_largestring_array(array);
-            let mut builder = GenericStringBuilder::<i64>::new();
-            for i in 0..str_array.len() {
-                if str_array.is_null(i) {
-                    builder.append_null();
-                } else {
-                    builder.append_value(spark_substr_negative(str_array.value(i), start, len));
-                }
-            }
-            Ok(Arc::new(builder.finish()) as ArrayRef)
+    use arrow::array::{DictionaryArray, GenericBinaryArray, OffsetSizeTrait};
+
+    fn substr_str<O: OffsetSizeTrait>(
+        str_array: &GenericStringArray<O>,
+        start: i64,
+        len: u64,
+    ) -> ArrayRef {
+        use arrow::array::GenericStringBuilder;
+        let mut builder = GenericStringBuilder::<O>::with_capacity(str_array.len(), 0);
+        for i in 0..str_array.len() {
+            // Always append; nulls are reattached in bulk below. This avoids
+            // per-row NullBufferBuilder maintenance.
+            let s = if str_array.is_null(i) {
+                ""
+            } else {
+                spark_substr_negative(str_array.value(i), start, len)
+            };
+            builder.append_value(s);
         }
-        DataType::Binary => {
-            let bin_array = array.as_any().downcast_ref::<BinaryArray>().unwrap();
-            let mut builder = GenericBinaryBuilder::<i32>::new();
-            for i in 0..bin_array.len() {
-                if bin_array.is_null(i) {
-                    builder.append_null();
-                } else {
-                    builder.append_value(spark_binary_substr_negative(
-                        bin_array.value(i),
-                        start,
-                        len,
-                    ));
-                }
-            }
-            Ok(Arc::new(builder.finish()) as ArrayRef)
-        }
-        DataType::LargeBinary => {
-            let bin_array = array.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
-            let mut builder = GenericBinaryBuilder::<i64>::new();
-            for i in 0..bin_array.len() {
-                if bin_array.is_null(i) {
-                    builder.append_null();
-                } else {
-                    builder.append_value(spark_binary_substr_negative(
-                        bin_array.value(i),
-                        start,
-                        len,
-                    ));
-                }
-            }
-            Ok(Arc::new(builder.finish()) as ArrayRef)
+        let (offsets, values, _) = builder.finish().into_parts();
+        Arc::new(GenericStringArray::<O>::new(
+            offsets,
+            values,
+            str_array.nulls().cloned(),
+        ))
+    }
+
+    fn substr_bin<O: OffsetSizeTrait>(
+        bin_array: &GenericBinaryArray<O>,
+        start: i64,
+        len: u64,
+    ) -> ArrayRef {
+        use arrow::array::GenericBinaryBuilder;
+        let mut builder = GenericBinaryBuilder::<O>::with_capacity(bin_array.len(), 0);
+        for i in 0..bin_array.len() {
+            let b: &[u8] = if bin_array.is_null(i) {
+                &[]
+            } else {
+                spark_binary_substr_negative(bin_array.value(i), start, len)
+            };
+            builder.append_value(b);
         }
+        let (offsets, values, _) = builder.finish().into_parts();
+        Arc::new(GenericBinaryArray::<O>::new(
+            offsets,
+            values,
+            bin_array.nulls().cloned(),
+        ))
+    }
+
+    match array.data_type() {
+        DataType::Utf8 => Ok(substr_str::<i32>(as_string_array(array), start, len)),
+        DataType::LargeUtf8 => Ok(substr_str::<i64>(as_largestring_array(array), start, len)),
+        DataType::Binary => Ok(substr_bin::<i32>(
+            array.as_any().downcast_ref().unwrap(),
+            start,
+            len,
+        )),
+        DataType::LargeBinary => Ok(substr_bin::<i64>(
+            array.as_any().downcast_ref().unwrap(),
+            start,
+            len,
+        )),
         DataType::Dictionary(_, _) => {
             let dict = as_dictionary_array::<Int32Type>(array);
             let values = spark_substring_negative_start(dict.values(), start, len)?;
