This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new beb3d5a735 Improve split_part udf by using a GenericStringBuilder
(#12093)
beb3d5a735 is described below
commit beb3d5a735e9fe3036e65bfea3424750498d27ea
Author: Lordworms <[email protected]>
AuthorDate: Thu Aug 22 03:15:14 2024 -0700
Improve split_part udf by using a GenericStringBuilder (#12093)
* optimize code
* optimize code
---
datafusion/functions/src/string/split_part.rs | 206 ++++++++++++++++----------
1 file changed, 126 insertions(+), 80 deletions(-)
diff --git a/datafusion/functions/src/string/split_part.rs
b/datafusion/functions/src/string/split_part.rs
index 19721f0fad..8d292315a3 100644
--- a/datafusion/functions/src/string/split_part.rs
+++ b/datafusion/functions/src/string/split_part.rs
@@ -15,21 +15,23 @@
// specific language governing permissions and limitations
// under the License.
-use std::any::Any;
-use std::sync::Arc;
-
-use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait};
-use arrow::datatypes::DataType;
-
-use datafusion_common::cast::{
- as_generic_string_array, as_int64_array, as_string_view_array,
+use arrow::array::{
+ ArrayRef, GenericStringArray, Int64Array, OffsetSizeTrait, StringViewArray,
};
-use datafusion_common::{exec_err, Result};
+use arrow::array::{AsArray, GenericStringBuilder};
+use arrow::datatypes::DataType;
+use datafusion_common::cast::as_int64_array;
+use datafusion_common::ScalarValue;
+use datafusion_common::{exec_err, DataFusionError, Result};
use datafusion_expr::TypeSignature::*;
use datafusion_expr::{ColumnarValue, Volatility};
use datafusion_expr::{ScalarUDFImpl, Signature};
+use std::any::Any;
+use std::sync::Arc;
+
+use crate::utils::utf8_to_str_type;
-use crate::utils::{make_scalar_function, utf8_to_str_type};
+use super::common::StringArrayType;
#[derive(Debug)]
pub struct SplitPartFunc {
@@ -82,36 +84,121 @@ impl ScalarUDFImpl for SplitPartFunc {
}
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- match (args[0].data_type(), args[1].data_type()) {
- (
- DataType::Utf8 | DataType::Utf8View,
- DataType::Utf8 | DataType::Utf8View,
- ) => make_scalar_function(split_part::<i32, i32>, vec![])(args),
+ // First, determine if any of the arguments is an Array
+ let len = args.iter().find_map(|arg| match arg {
+ ColumnarValue::Array(a) => Some(a.len()),
+ _ => None,
+ });
+
+ let inferred_length = len.unwrap_or(1);
+ let is_scalar = len.is_none();
+
+ // Convert all ColumnarValues to ArrayRefs
+ let args = args
+ .iter()
+ .map(|arg| match arg {
+ ColumnarValue::Scalar(scalar) =>
scalar.to_array_of_size(inferred_length),
+ ColumnarValue::Array(array) => Ok(Arc::clone(array)),
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ // Unpack the ArrayRefs from the arguments
+ let n_array = as_int64_array(&args[2])?;
+ let result = match (args[0].data_type(), args[1].data_type()) {
+ (DataType::Utf8View, DataType::Utf8View) => {
+ split_part_impl::<&StringViewArray, &StringViewArray, i32>(
+ args[0].as_string_view(),
+ args[1].as_string_view(),
+ n_array,
+ )
+ }
+ (DataType::Utf8View, DataType::Utf8) => {
+ split_part_impl::<&StringViewArray, &GenericStringArray<i32>,
i32>(
+ args[0].as_string_view(),
+ args[1].as_string::<i32>(),
+ n_array,
+ )
+ }
+ (DataType::Utf8View, DataType::LargeUtf8) => {
+ split_part_impl::<&StringViewArray, &GenericStringArray<i64>,
i32>(
+ args[0].as_string_view(),
+ args[1].as_string::<i64>(),
+ n_array,
+ )
+ }
+ (DataType::Utf8, DataType::Utf8View) => {
+ split_part_impl::<&GenericStringArray<i32>, &StringViewArray,
i32>(
+ args[0].as_string::<i32>(),
+ args[1].as_string_view(),
+ n_array,
+ )
+ }
+ (DataType::LargeUtf8, DataType::Utf8View) => {
+ split_part_impl::<&GenericStringArray<i64>, &StringViewArray,
i64>(
+ args[0].as_string::<i64>(),
+ args[1].as_string_view(),
+ n_array,
+ )
+ }
+ (DataType::Utf8, DataType::Utf8) => {
+ split_part_impl::<&GenericStringArray<i32>,
&GenericStringArray<i32>, i32>(
+ args[0].as_string::<i32>(),
+ args[1].as_string::<i32>(),
+ n_array,
+ )
+ }
(DataType::LargeUtf8, DataType::LargeUtf8) => {
- make_scalar_function(split_part::<i64, i64>, vec![])(args)
+ split_part_impl::<&GenericStringArray<i64>,
&GenericStringArray<i64>, i64>(
+ args[0].as_string::<i64>(),
+ args[1].as_string::<i64>(),
+ n_array,
+ )
}
- (_, DataType::LargeUtf8) => {
- make_scalar_function(split_part::<i32, i64>, vec![])(args)
+ (DataType::Utf8, DataType::LargeUtf8) => {
+ split_part_impl::<&GenericStringArray<i32>,
&GenericStringArray<i64>, i32>(
+ args[0].as_string::<i32>(),
+ args[1].as_string::<i64>(),
+ n_array,
+ )
}
- (DataType::LargeUtf8, _) => {
- make_scalar_function(split_part::<i64, i32>, vec![])(args)
+ (DataType::LargeUtf8, DataType::Utf8) => {
+ split_part_impl::<&GenericStringArray<i64>,
&GenericStringArray<i32>, i64>(
+ args[0].as_string::<i64>(),
+ args[1].as_string::<i32>(),
+ n_array,
+ )
}
- (first_type, second_type) => exec_err!(
- "unsupported first type {} and second type {} for split_part
function",
- first_type,
- second_type
- ),
+ _ => exec_err!("Unsupported combination of argument types for
split_part"),
+ };
+ if is_scalar {
+ // If all inputs are scalar, keep the output as scalar
+ let result = result.and_then(|arr|
ScalarValue::try_from_array(&arr, 0));
+ result.map(ColumnarValue::Scalar)
+ } else {
+ result.map(ColumnarValue::Array)
}
}
}
-macro_rules! process_split_part {
- ($string_array: expr, $delimiter_array: expr, $n_array: expr) => {{
- let result = $string_array
- .iter()
- .zip($delimiter_array.iter())
- .zip($n_array.iter())
- .map(|((string, delimiter), n)| match (string, delimiter, n) {
+/// impl
+pub fn split_part_impl<'a, StringArrType, DelimiterArrType, StringArrayLen>(
+ string_array: StringArrType,
+ delimiter_array: DelimiterArrType,
+ n_array: &Int64Array,
+) -> Result<ArrayRef>
+where
+ StringArrType: StringArrayType<'a>,
+ DelimiterArrType: StringArrayType<'a>,
+ StringArrayLen: OffsetSizeTrait,
+{
+ let mut builder: GenericStringBuilder<StringArrayLen> =
GenericStringBuilder::new();
+
+ string_array
+ .iter()
+ .zip(delimiter_array.iter())
+ .zip(n_array.iter())
+ .try_for_each(|((string, delimiter), n)| -> Result<(),
DataFusionError> {
+ match (string, delimiter, n) {
(Some(string), Some(delimiter), Some(n)) => {
let split_string: Vec<&str> =
string.split(delimiter).collect();
let len = split_string.len();
@@ -125,58 +212,17 @@ macro_rules! process_split_part {
} as usize;
if index < len {
- Ok(Some(split_string[index]))
+ builder.append_value(split_string[index]);
} else {
- Ok(Some(""))
+ builder.append_value("");
}
}
- _ => Ok(None),
- })
- .collect::<Result<GenericStringArray<StringLen>>>()?;
- Ok(Arc::new(result) as ArrayRef)
- }};
-}
-
-/// Splits string at occurrences of delimiter and returns the n'th field
(counting from one).
-/// split_part('abc~@~def~@~ghi', '~@~', 2) = 'def'
-fn split_part<StringLen: OffsetSizeTrait, DelimiterLen: OffsetSizeTrait>(
- args: &[ArrayRef],
-) -> Result<ArrayRef> {
- let n_array = as_int64_array(&args[2])?;
- match (args[0].data_type(), args[1].data_type()) {
- (DataType::Utf8View, _) => {
- let string_array = as_string_view_array(&args[0])?;
- match args[1].data_type() {
- DataType::Utf8View => {
- let delimiter_array = as_string_view_array(&args[1])?;
- process_split_part!(string_array, delimiter_array, n_array)
- }
- _ => {
- let delimiter_array =
- as_generic_string_array::<DelimiterLen>(&args[1])?;
- process_split_part!(string_array, delimiter_array, n_array)
- }
- }
- }
- (_, DataType::Utf8View) => {
- let delimiter_array = as_string_view_array(&args[1])?;
- match args[0].data_type() {
- DataType::Utf8View => {
- let string_array = as_string_view_array(&args[0])?;
- process_split_part!(string_array, delimiter_array, n_array)
- }
- _ => {
- let string_array =
as_generic_string_array::<StringLen>(&args[0])?;
- process_split_part!(string_array, delimiter_array, n_array)
- }
+ _ => builder.append_null(),
}
- }
- (_, _) => {
- let string_array = as_generic_string_array::<StringLen>(&args[0])?;
- let delimiter_array =
as_generic_string_array::<DelimiterLen>(&args[1])?;
- process_split_part!(string_array, delimiter_array, n_array)
- }
- }
+ Ok(())
+ })?;
+
+ Ok(Arc::new(builder.finish()) as ArrayRef)
}
#[cfg(test)]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
