Directly write to MutableBuffer in substring (#1423)

arrow/Cargo.toml

@@ -167,3 +167,7 @@ harness = false
 [[bench]]
 name = "buffer_create"
 harness = false
+
+[[bench]]
+name = "string_kernels"
+harness = false

arrow/benches/string_kernels.rs

@@ -0,0 +1,45 @@
+// 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.
+
+#[macro_use]
+extern crate criterion;
+use criterion::Criterion;
+
+extern crate arrow;
+
+use arrow::array::*;
+use arrow::compute::kernels::substring::substring;
+use arrow::util::bench_util::*;
+
+fn bench_substring(arr: &StringArray, start: i64, length: usize) {
+    substring(criterion::black_box(arr), start, &Some(length as u64)).unwrap();
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    let size = 65536;
+    let str_len = 1000;
+
+    let arr_string = create_string_array_with_len::<i32>(size, 0.0, str_len);
+    let start = 0;
+
+    c.bench_function("substring", |b| {
+        b.iter(|| bench_substring(&arr_string, start, str_len))
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

arrow/src/compute/kernels/substring.rs

@@ -17,6 +17,7 @@
 
 //! Defines kernel to extract a substring of a \[Large\]StringArray
 
+use crate::buffer::MutableBuffer;
 use crate::{array::*, buffer::Buffer};
 use crate::{
     datatypes::DataType,
@@ -40,7 +41,7 @@ fn generic_substring<OffsetSize: StringOffsetSizeTrait>(
     let values = &array.data_ref().buffers()[1];
     let data = values.as_slice();
 
-    let mut new_values = Vec::new(); // we have no way to estimate how much this will be.
+    let mut new_values = MutableBuffer::new(0); // we have no way to estimate how much this will be.
     let mut new_offsets: Vec<OffsetSize> = Vec::with_capacity(array.len() + 1);
 
     let mut length_so_far = OffsetSize::zero();
@@ -81,10 +82,7 @@ fn generic_substring<OffsetSize: StringOffsetSizeTrait>(
             None,
             null_bit_buffer,
             0,
-            vec![
-                Buffer::from_slice_ref(&new_offsets),
-                Buffer::from_slice_ref(&new_values),
-            ],
+            vec![Buffer::from_slice_ref(&new_offsets), new_values.into()],
             vec![],
         )
     };

arrow/src/util/bench_util.rs

@@ -94,6 +94,15 @@ where
 pub fn create_string_array<Offset: StringOffsetSizeTrait>(
     size: usize,
     null_density: f32,
+) -> GenericStringArray<Offset> {
+    create_string_array_with_len(size, null_density, 4)
+}
+
+/// Creates a random (but fixed-seeded) array of a given size, null density and length
+pub fn create_string_array_with_len<Offset: StringOffsetSizeTrait>(
+    size: usize,
+    null_density: f32,
+    str_len: usize,
 ) -> GenericStringArray<Offset> {
     let rng = &mut seedable_rng();
 
@@ -102,7 +111,7 @@ pub fn create_string_array<Offset: StringOffsetSizeTrait>(
             if rng.gen::<f32>() < null_density {
                 None
             } else {
-                let value = rng.sample_iter(&Alphanumeric).take(4).collect();
+                let value = rng.sample_iter(&Alphanumeric).take(str_len).collect();
                 let value = String::from_utf8(value).unwrap();
                 Some(value)
             }