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

Update the CONCAT scalar function to support Utf8View #12224

Merged
merged 18 commits into from
Sep 3, 2024
Merged

Update the CONCAT scalar function to support Utf8View #12224

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
723ceb7
wip
devanbenz Aug 28, 2024
f7abdd5
feat: Update the CONCAT scalar function to support Utf8View
devanbenz Aug 29, 2024
503d5b9
fmt
devanbenz Aug 29, 2024
b30330d
fmt and add default return type for concat
devanbenz Aug 29, 2024
76d6b5f
fix clippy lint
devanbenz Aug 29, 2024
9798ed3
fmt
devanbenz Aug 29, 2024
91d04ff
add more tests for sqllogic
devanbenz Aug 29, 2024
6d28927
make sure no casting with LargeUtf8
devanbenz Aug 29, 2024
769d99d
fixing utf8large
devanbenz Aug 29, 2024
ac30a83
fix large utf8
devanbenz Aug 30, 2024
7ea6e0a
fix large utf8
devanbenz Aug 30, 2024
dd3ad39
add test
devanbenz Aug 30, 2024
504459c
fmt
devanbenz Aug 30, 2024
e081934
make it so Utf8View just returns Utf8
devanbenz Aug 30, 2024
0069c1a
wip -- trying to build a stringview with columnar refs
devanbenz Sep 1, 2024
0929a4b
built stringview builder but it does allocate a new String each iter :(
devanbenz Sep 1, 2024
f16de44
add some testing
devanbenz Sep 1, 2024
d0bf3ba
clippy
devanbenz Sep 1, 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
195 changes: 190 additions & 5 deletions datafusion/functions/src/string/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,11 @@ use std::sync::Arc;

use arrow::array::{
new_null_array, Array, ArrayAccessor, ArrayDataBuilder, ArrayIter, ArrayRef,
GenericStringArray, GenericStringBuilder, OffsetSizeTrait, StringArray,
StringBuilder, StringViewArray,
GenericStringArray, GenericStringBuilder, LargeStringArray, OffsetSizeTrait,
StringArray, StringBuilder, StringViewArray, StringViewBuilder,
};
use arrow::buffer::{Buffer, MutableBuffer, NullBuffer};
use arrow::datatypes::DataType;

use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
use datafusion_common::Result;
use datafusion_common::{exec_err, ScalarValue};
Expand Down Expand Up @@ -251,26 +250,41 @@ where
}
}

#[derive(Debug)]
pub(crate) enum ColumnarValueRef<'a> {
Scalar(&'a [u8]),
NullableArray(&'a StringArray),
NonNullableArray(&'a StringArray),
NullableLargeStringArray(&'a LargeStringArray),
NonNullableLargeStringArray(&'a LargeStringArray),
NullableStringViewArray(&'a StringViewArray),
NonNullableStringViewArray(&'a StringViewArray),
}

impl<'a> ColumnarValueRef<'a> {
#[inline]
pub fn is_valid(&self, i: usize) -> bool {
match &self {
Self::Scalar(_) | Self::NonNullableArray(_) => true,
Self::Scalar(_)
| Self::NonNullableArray(_)
| Self::NonNullableLargeStringArray(_)
| Self::NonNullableStringViewArray(_) => true,
Self::NullableArray(array) => array.is_valid(i),
Self::NullableStringViewArray(array) => array.is_valid(i),
Self::NullableLargeStringArray(array) => array.is_valid(i),
}
}

#[inline]
pub fn nulls(&self) -> Option<NullBuffer> {
match &self {
Self::Scalar(_) | Self::NonNullableArray(_) => None,
Self::Scalar(_)
| Self::NonNullableArray(_)
| Self::NonNullableStringViewArray(_)
| Self::NonNullableLargeStringArray(_) => None,
Self::NullableArray(array) => array.nulls().cloned(),
Self::NullableStringViewArray(array) => array.nulls().cloned(),
Self::NullableLargeStringArray(array) => array.nulls().cloned(),
}
}
}
Expand Down Expand Up @@ -389,10 +403,30 @@ impl StringArrayBuilder {
.extend_from_slice(array.value(i).as_bytes());
}
}
ColumnarValueRef::NullableLargeStringArray(array) => {
if !CHECK_VALID || array.is_valid(i) {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
}
ColumnarValueRef::NullableStringViewArray(array) => {
if !CHECK_VALID || array.is_valid(i) {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
}
ColumnarValueRef::NonNullableArray(array) => {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
ColumnarValueRef::NonNullableLargeStringArray(array) => {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
ColumnarValueRef::NonNullableStringViewArray(array) => {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
}
}

Expand All @@ -418,6 +452,157 @@ impl StringArrayBuilder {
}
}

pub(crate) struct StringViewArrayBuilder {
builder: StringViewBuilder,
block: String,
}

impl StringViewArrayBuilder {
Copy link
Contributor Author

Choose a reason for hiding this comment

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

StringViewArrayBuilder works but it almost feels "hacky"

Copy link
Contributor

Choose a reason for hiding this comment

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

I think the original intent of StringViewBuilder when @JasonLi-cn added it was to have a way to append (but not complete) strings to avoid a copy and not incrementally update the null mask.

Specifically, StringBuilder in arrow https://docs.rs/arrow/latest/arrow/array/type.StringBuilder.html#method.append_value lets you append an entire string but it requires the input be contiguous bytes

StringArrayBuilder is that you can push bytes (via write) wthout copying the data.

The reason I think it feels hacky is that the API for building up single values in a StringArray or StringViewArray without a copy is not present

As we see we want to do the same thing with Utf8View I wonder if we should spend some time making the API eaiser to use 🤔

Copy link
Contributor

Choose a reason for hiding this comment

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

I filed apache/arrow-rs#6347 to discuss adding an API like this in arrow as well

pub fn with_capacity(_item_capacity: usize, data_capacity: usize) -> Self {
let builder = StringViewBuilder::with_capacity(data_capacity);
Self {
builder,
block: String::new(),
}
}

pub fn write<const CHECK_VALID: bool>(
&mut self,
column: &ColumnarValueRef,
i: usize,
) {
match column {
ColumnarValueRef::Scalar(s) => {
self.block.push_str(std::str::from_utf8(s).unwrap());
}
ColumnarValueRef::NullableArray(array) => {
if !CHECK_VALID || array.is_valid(i) {
self.block.push_str(
std::str::from_utf8(array.value(i).as_bytes()).unwrap(),
);
}
}
ColumnarValueRef::NullableLargeStringArray(array) => {
if !CHECK_VALID || array.is_valid(i) {
self.block.push_str(
std::str::from_utf8(array.value(i).as_bytes()).unwrap(),
);
}
}
ColumnarValueRef::NullableStringViewArray(array) => {
if !CHECK_VALID || array.is_valid(i) {
self.block.push_str(
std::str::from_utf8(array.value(i).as_bytes()).unwrap(),
);
}
}
ColumnarValueRef::NonNullableArray(array) => {
self.block
.push_str(std::str::from_utf8(array.value(i).as_bytes()).unwrap());
}
ColumnarValueRef::NonNullableLargeStringArray(array) => {
self.block
.push_str(std::str::from_utf8(array.value(i).as_bytes()).unwrap());
}
ColumnarValueRef::NonNullableStringViewArray(array) => {
self.block
.push_str(std::str::from_utf8(array.value(i).as_bytes()).unwrap());
}
}
}

pub fn append_offset(&mut self) {
self.builder.append_value(&self.block);
self.block = String::new();
}

pub fn finish(mut self) -> StringViewArray {
self.builder.finish()
}
}

pub(crate) struct LargeStringArrayBuilder {
Copy link
Contributor Author

Choose a reason for hiding this comment

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

I wanted to make the already existing StringArrayBuilder generic but was having issues 😢

offsets_buffer: MutableBuffer,
value_buffer: MutableBuffer,
}

impl LargeStringArrayBuilder {
pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
let mut offsets_buffer = MutableBuffer::with_capacity(
(item_capacity + 1) * std::mem::size_of::<i64>(),
);
// SAFETY: the first offset value is definitely not going to exceed the bounds.
unsafe { offsets_buffer.push_unchecked(0_i64) };
Self {
offsets_buffer,
value_buffer: MutableBuffer::with_capacity(data_capacity),
}
}

pub fn write<const CHECK_VALID: bool>(
&mut self,
column: &ColumnarValueRef,
i: usize,
) {
match column {
ColumnarValueRef::Scalar(s) => {
self.value_buffer.extend_from_slice(s);
}
ColumnarValueRef::NullableArray(array) => {
if !CHECK_VALID || array.is_valid(i) {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
}
ColumnarValueRef::NullableLargeStringArray(array) => {
if !CHECK_VALID || array.is_valid(i) {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
}
ColumnarValueRef::NullableStringViewArray(array) => {
if !CHECK_VALID || array.is_valid(i) {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
}
ColumnarValueRef::NonNullableArray(array) => {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
ColumnarValueRef::NonNullableLargeStringArray(array) => {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
ColumnarValueRef::NonNullableStringViewArray(array) => {
self.value_buffer
.extend_from_slice(array.value(i).as_bytes());
}
}
}

pub fn append_offset(&mut self) {
let next_offset: i64 = self
.value_buffer
.len()
.try_into()
.expect("byte array offset overflow");
unsafe { self.offsets_buffer.push_unchecked(next_offset) };
}

pub fn finish(self, null_buffer: Option<NullBuffer>) -> LargeStringArray {
let array_builder = ArrayDataBuilder::new(DataType::LargeUtf8)
.len(self.offsets_buffer.len() / std::mem::size_of::<i64>() - 1)
.add_buffer(self.offsets_buffer.into())
.add_buffer(self.value_buffer.into())
.nulls(null_buffer);
// SAFETY: all data that was appended was valid Large UTF8 and the values
// and offsets were created correctly
let array_data = unsafe { array_builder.build_unchecked() };
LargeStringArray::from(array_data)
}
}

fn case_conversion_array<'a, O, F>(array: &'a ArrayRef, op: F) -> Result<ArrayRef>
where
O: OffsetSizeTrait,
Expand Down
Loading