Skip to content
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

Jump to bottom

Update concat_ws scalar function to support Utf8View #12309

Merged
merged 21 commits into from
Sep 12, 2024
Merged

Update concat_ws scalar function to support Utf8View #12309

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
22fd1b5
Update concat_ws scalar function to support Utf8View
devanbenz Sep 4, 2024
d60a75d
fmt
devanbenz Sep 4, 2024
1383f8b
adds one_of for type sig
devanbenz Sep 6, 2024
63fd5c6
revert type sig, the problem was outside the actual signature xD
devanbenz Sep 6, 2024
92fbcfa
fmt and clippy
devanbenz Sep 6, 2024
1d50a62
add utf8view
devanbenz Sep 6, 2024
31576e3
fix match
devanbenz Sep 6, 2024
b5f64f3
fix match
devanbenz Sep 6, 2024
fec232c
:thinking: why nulls why
devanbenz Sep 6, 2024
5e0c4c6
small fix
devanbenz Sep 6, 2024
a660a31
log ms
devanbenz Sep 6, 2024
a9e3c4a
pushing up -- wip
devanbenz Sep 7, 2024
cc2e03f
make it so the return type is just Utf8
devanbenz Sep 8, 2024
0d07f25
fmt
devanbenz Sep 8, 2024
b6967f0
fmt
devanbenz Sep 8, 2024
f256b45
fmt
devanbenz Sep 8, 2024
c7d35e8
fmt
devanbenz Sep 8, 2024
2926d0b
fmt
devanbenz Sep 8, 2024
bc32c90
order matters
devanbenz Sep 8, 2024
66a576d
sum inner buffer for stringviewarray data_size
devanbenz Sep 9, 2024
cef69b7
need turbo fish
devanbenz Sep 9, 2024
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
10 changes: 3 additions & 7 deletions datafusion/functions/src/string/concat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,13 +136,9 @@ impl ScalarUDFImpl for ConcatFunc {

for arg in args {
match arg {
ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value)) => {
if let Some(s) = maybe_value {
data_size += s.len() * len;
columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
}
}
ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => {
ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value))
| ColumnarValue::Scalar(ScalarValue::LargeUtf8(maybe_value))
| ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => {
if let Some(s) = maybe_value {
data_size += s.len() * len;
columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
Expand Down
236 changes: 196 additions & 40 deletions datafusion/functions/src/string/concat_ws.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,14 @@
// specific language governing permissions and limitations
// under the License.

use arrow::array::StringArray;
use arrow::array::{as_largestring_array, Array, StringArray};
use std::any::Any;
use std::sync::Arc;

use arrow::datatypes::DataType;
use arrow::datatypes::DataType::Utf8;

use datafusion_common::cast::as_string_array;
use datafusion_common::{exec_err, internal_err, Result, ScalarValue};
use datafusion_common::cast::{as_string_array, as_string_view_array};
use datafusion_common::{exec_err, internal_err, plan_err, Result, ScalarValue};
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
use datafusion_expr::{lit, ColumnarValue, Expr, Volatility};
Expand All @@ -46,9 +45,8 @@ impl Default for ConcatWsFunc {

impl ConcatWsFunc {
pub fn new() -> Self {
use DataType::*;
Self {
signature: Signature::variadic(vec![Utf8], Volatility::Immutable),
signature: Signature::variadic_any(Volatility::Immutable),
devanbenz marked this conversation as resolved.
Show resolved Hide resolved
}
}
}
Expand All @@ -66,8 +64,19 @@ impl ScalarUDFImpl for ConcatWsFunc {
&self.signature
}

fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(Utf8)
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
use DataType::*;
let mut dt = &Utf8;
arg_types.iter().for_each(|data_type| {
if data_type == &Utf8View {
dt = data_type;
}
if data_type == &LargeUtf8 && dt != &Utf8View {
dt = data_type;
}
});

Ok(dt.to_owned())
}

/// Concatenates all but the first argument, with separators. The first argument is used as the separator string, and should not be NULL. Other NULL arguments are ignored.
Expand All @@ -89,6 +98,18 @@ impl ScalarUDFImpl for ConcatWsFunc {
})
.next();

let mut return_datatype = DataType::Utf8;
args.iter().for_each(|col| {
if col.data_type() == DataType::Utf8View {
return_datatype = col.data_type();
}
if col.data_type() == DataType::LargeUtf8
&& return_datatype != DataType::Utf8View
{
return_datatype = col.data_type();
}
});

// Scalar
if array_len.is_none() {
let sep = match &args[0] {
Expand All @@ -104,27 +125,43 @@ impl ScalarUDFImpl for ConcatWsFunc {

for arg in iter.by_ref() {
match arg {
ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => {
ColumnarValue::Scalar(ScalarValue::Utf8(Some(s)))
| ColumnarValue::Scalar(ScalarValue::Utf8View(Some(s)))
| ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(s))) => {
result.push_str(s);
break;
}
ColumnarValue::Scalar(ScalarValue::Utf8(None)) => {}
ColumnarValue::Scalar(ScalarValue::Utf8(None))
| ColumnarValue::Scalar(ScalarValue::Utf8View(None))
| ColumnarValue::Scalar(ScalarValue::LargeUtf8(None)) => {}
_ => unreachable!(),
}
}

for arg in iter.by_ref() {
match arg {
ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => {
ColumnarValue::Scalar(ScalarValue::Utf8(Some(s)))
| ColumnarValue::Scalar(ScalarValue::Utf8View(Some(s)))
| ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(s))) => {
result.push_str(sep);
result.push_str(s);
}
ColumnarValue::Scalar(ScalarValue::Utf8(None)) => {}
ColumnarValue::Scalar(ScalarValue::Utf8(None))
| ColumnarValue::Scalar(ScalarValue::Utf8View(None))
| ColumnarValue::Scalar(ScalarValue::LargeUtf8(None)) => {}
_ => unreachable!(),
}
}

return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))));
return match return_datatype {
DataType::Utf8View => {
Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))))
}
DataType::LargeUtf8 => {
Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))))
}
_ => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result)))),
};
}

// Array
Expand Down Expand Up @@ -155,52 +192,145 @@ impl ScalarUDFImpl for ConcatWsFunc {
let mut columns = Vec::with_capacity(args.len() - 1);
for arg in &args[1..] {
match arg {
ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value)) => {
ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value))
| ColumnarValue::Scalar(ScalarValue::LargeUtf8(maybe_value))
| ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => {
if let Some(s) = maybe_value {
data_size += s.len() * len;
columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
}
}
ColumnarValue::Array(array) => {
let string_array = as_string_array(array)?;
data_size += string_array.values().len();
let column = if array.is_nullable() {
ColumnarValueRef::NullableArray(string_array)
} else {
ColumnarValueRef::NonNullableArray(string_array)
match array.data_type() {
DataType::Utf8 => {
let string_array = as_string_array(array)?;

data_size += string_array.values().len();
let column = if array.is_nullable() {
ColumnarValueRef::NullableArray(string_array)
} else {
ColumnarValueRef::NonNullableArray(string_array)
};
columns.push(column);
},
DataType::LargeUtf8 => {
let string_array = as_largestring_array(array);

data_size += string_array.values().len();
let column = if array.is_nullable() {
ColumnarValueRef::NullableLargeStringArray(string_array)
} else {
ColumnarValueRef::NonNullableLargeStringArray(string_array)
};
columns.push(column);
},
DataType::Utf8View => {
let string_array = as_string_view_array(array)?;

data_size += string_array.len();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

data_size is the estimated string array size after concatenation, I think here it should increment sum of inner buffers' size instead

(And everything else LGTM)

Copy link
Contributor Author

@devanbenz devanbenz Sep 9, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using data_size += string_array.data_buffers().len(); ?

disregard I think I get what you mean I modified it to sum up the inner buffer vec on each iteration and use that as the size

let column = if array.is_nullable() {
ColumnarValueRef::NullableStringViewArray(string_array)
} else {
ColumnarValueRef::NonNullableStringViewArray(string_array)
};
columns.push(column);
},
other => {
return plan_err!("Input was {other} which is not a supported datatype for concat function")
}
};
columns.push(column);
}
_ => unreachable!(),
}
}

let mut builder = StringArrayBuilder::with_capacity(len, data_size);
for i in 0..len {
if !sep.is_valid(i) {
builder.append_offset();
continue;
}
match return_datatype {
DataType::Utf8 => {
let mut builder = StringArrayBuilder::with_capacity(len, data_size);
for i in 0..len {
if !sep.is_valid(i) {
builder.append_offset();
continue;
}

let mut iter = columns.iter();
for column in iter.by_ref() {
if column.is_valid(i) {
builder.write::<false>(column, i);
break;
let mut iter = columns.iter();
for column in iter.by_ref() {
if column.is_valid(i) {
builder.write::<false>(column, i);
break;
}
}

for column in iter {
if column.is_valid(i) {
builder.write::<false>(&sep, i);
builder.write::<false>(column, i);
}
}

builder.append_offset();
}

Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls()))))
}
DataType::Utf8View => {
let mut builder = StringViewArrayBuilder::with_capacity(len, data_size);
for i in 0..len {
if !sep.is_valid(i) {
builder.append_offset();
continue;
}

for column in iter {
if column.is_valid(i) {
builder.write::<false>(&sep, i);
builder.write::<false>(column, i);
let mut iter = columns.iter();
for column in iter.by_ref() {
if column.is_valid(i) {
builder.write::<false>(column, i);
break;
}
}

for column in iter {
if column.is_valid(i) {
builder.write::<false>(&sep, i);
builder.write::<false>(column, i);
}
}

builder.append_offset();
}

Ok(ColumnarValue::Array(Arc::new(builder.finish())))
}
DataType::LargeUtf8 => {
let mut builder = LargeStringArrayBuilder::with_capacity(len, data_size);
for i in 0..len {
if !sep.is_valid(i) {
builder.append_offset();
continue;
}

builder.append_offset();
}
let mut iter = columns.iter();
for column in iter.by_ref() {
if column.is_valid(i) {
builder.write::<false>(column, i);
break;
}
}

Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls()))))
for column in iter {
if column.is_valid(i) {
builder.write::<false>(&sep, i);
builder.write::<false>(column, i);
}
}

builder.append_offset();
}

Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls()))))
}
_ => unreachable!(),
}
}

/// Simply the `concat_ws` function by
Expand Down Expand Up @@ -304,7 +434,7 @@ mod tests {
use std::sync::Arc;

use arrow::array::{Array, ArrayRef, StringArray};
use arrow::datatypes::DataType::Utf8;
use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View};

use crate::string::concat_ws::ConcatWsFunc;
use datafusion_common::Result;
Expand Down Expand Up @@ -365,6 +495,32 @@ mod tests {
Utf8,
StringArray
);
test_function!(
ConcatWsFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::from("|")),
ColumnarValue::Scalar(ScalarValue::from("aa")),
ColumnarValue::Scalar(ScalarValue::LargeUtf8(None)),
ColumnarValue::Scalar(ScalarValue::from("cc")),
],
Ok(Some("aa|cc")),
&str,
LargeUtf8,
StringArray
);
test_function!(
ConcatWsFunc::new(),
&[
ColumnarValue::Scalar(ScalarValue::from("|")),
ColumnarValue::Scalar(ScalarValue::from("aa")),
ColumnarValue::Scalar(ScalarValue::Utf8View(None)),
ColumnarValue::Scalar(ScalarValue::from("cc")),
],
Ok(Some("aa|cc")),
&str,
Utf8View,
StringArray
);

Ok(())
}
Expand Down
Loading
Loading