Skip to content

Preserve string slice function return types #23330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
xudong963 wants to merge 1 commit into
base: main
Choose a base branch
from
+189 −291
Draft

Preserve string slice function return types #23330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 9 additions & 68 deletions datafusion/functions/src/unicode/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,79 +151,20 @@ pub(crate) fn general_left_right<F: LeftRightSlicer>(
}
}

/// Returns true if all offsets in the array fit in i32, meaning the values
/// buffer can be referenced by StringView's offset field.
fn values_fit_in_i32<T: OffsetSizeTrait>(string_array: &GenericStringArray<T>) -> bool {
string_array
.offsets()
.last()
.map(|offset| offset.as_usize() <= i32::MAX as usize)
.unwrap_or(true)
}

/// `left`/`right` for Utf8/LargeUtf8 input.
///
/// When offsets fit in i32, produces a zero-copy `StringViewArray` with views
/// pointing into the input values buffer. Otherwise falls back to building a
/// `StringViewArray` by copying.
fn general_left_right_array<T: OffsetSizeTrait, F: LeftRightSlicer>(
string_array: &GenericStringArray<T>,
n_array: &Int64Array,
) -> Result<ArrayRef> {
if !values_fit_in_i32(string_array) {
let result = string_array
.iter()
.zip(n_array.iter())
.map(|(string, n)| match (string, n) {
(Some(string), Some(n)) => Some(&string[F::slice(string, n)]),
_ => None,
})
.collect::<StringViewArray>();
return Ok(Arc::new(result) as ArrayRef);
}

let len = string_array.len();
let offsets = string_array.value_offsets();
let nulls = NullBuffer::union(string_array.nulls(), n_array.nulls());

let mut views_buf = Vec::with_capacity(len);
let mut has_out_of_line = false;

for (i, offset) in offsets.iter().enumerate().take(len) {
if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
views_buf.push(0);
continue;
}

// SAFETY: we just checked validity above
let string = unsafe { string_array.value_unchecked(i) };
let n = n_array.value(i);
let range = F::slice(string, n);
let result_bytes = &string.as_bytes()[range.clone()];
if result_bytes.len() > 12 {
has_out_of_line = true;
}

let buf_offset = offset.as_usize() as u32 + range.start as u32;
views_buf.push(make_view(result_bytes, 0, buf_offset));
}

let views = ScalarBuffer::from(views_buf);
let data_buffers = if has_out_of_line {
vec![string_array.values().clone()]
} else {
vec![]
};

// SAFETY:
// - Each view is produced by `make_view` with correct bytes and offset
// - Out-of-line views reference buffer index 0, which is the original
// values buffer included in data_buffers when has_out_of_line is true
// - values_fit_in_i32 guarantees all offsets fit in i32
unsafe {
let array = StringViewArray::new_unchecked(views, data_buffers, nulls);
Ok(Arc::new(array) as ArrayRef)
}
let result = string_array
.iter()
.zip(n_array.iter())
.map(|(string, n)| match (string, n) {
(Some(string), Some(n)) => Some(&string[F::slice(string, n)]),
_ => None,
})
.collect::<GenericStringArray<T>>();
Ok(Arc::new(result) as ArrayRef)
}

/// `general_left_right` for StringViewArray input.
Expand Down
67 changes: 39 additions & 28 deletions datafusion/functions/src/unicode/left.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,8 +79,8 @@ impl ScalarUDFImpl for LeftFunc {
&self.signature
}

fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Utf8View)
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
Ok(arg_types[0].clone())
}

/// Returns first n characters in the string, or when n is negative, returns all but last |n| characters.
Expand Down Expand Up @@ -108,8 +108,8 @@ impl ScalarUDFImpl for LeftFunc {

#[cfg(test)]
mod tests {
use arrow::array::{Array, StringViewArray};
use arrow::datatypes::DataType::Utf8View;
use arrow::array::{Array, LargeStringArray, StringArray, StringViewArray};
use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View};

use datafusion_common::{Result, ScalarValue};
use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
Expand All @@ -127,8 +127,19 @@ mod tests {
],
Ok(Some("ab")),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
vec![
ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some("abcde".to_string()))),
ColumnarValue::Scalar(ScalarValue::from(2i64)),
],
Ok(Some("ab")),
&str,
LargeUtf8,
LargeStringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -138,8 +149,8 @@ mod tests {
],
Ok(Some("abcde")),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -149,8 +160,8 @@ mod tests {
],
Ok(Some("abc")),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -160,8 +171,8 @@ mod tests {
],
Ok(Some("")),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -171,8 +182,8 @@ mod tests {
],
Ok(Some("")),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -182,8 +193,8 @@ mod tests {
],
Ok(Some("")),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -193,8 +204,8 @@ mod tests {
],
Ok(None),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -204,8 +215,8 @@ mod tests {
],
Ok(None),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -215,8 +226,8 @@ mod tests {
],
Ok(Some("joséé")),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
test_function!(
LeftFunc::new(),
Expand All @@ -226,8 +237,8 @@ mod tests {
],
Ok(Some("joséé")),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
#[cfg(not(feature = "unicode_expressions"))]
test_function!(
Expand All @@ -240,8 +251,8 @@ mod tests {
"function left requires compilation with feature flag: unicode_expressions."
),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);

// StringView cases
Expand Down Expand Up @@ -307,8 +318,8 @@ mod tests {
],
Ok(Some(expected.as_str())),
&str,
Utf8View,
StringViewArray
Utf8,
StringArray
);
}

Expand Down
Loading
Loading