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

Support casting between BinaryView <--> Utf8 and LargeUtf8 #6180

Merged
merged 5 commits into from
Aug 8, 2024
Merged

Support casting between BinaryView <--> Utf8 and LargeUtf8 #6180

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5477e80
support cast between binaryview and string
xinlifoobar Aug 2, 2024
31b1fe4
update impl. and add bench mark
xinlifoobar Aug 4, 2024
9b3f00f
Add ut for views
xinlifoobar Aug 5, 2024
efba2fe
Merge branch 'master' of github.com:apache/arrow-rs into dev/xinli/bi…
xinlifoobar Aug 8, 2024
79c47ac
Apply coments
xinlifoobar Aug 8, 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
89 changes: 78 additions & 11 deletions arrow-cast/src/cast/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,8 +227,8 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
| Timestamp(Nanosecond, _)
| Interval(_),
) => true,
(Utf8 | LargeUtf8, Utf8View) => true,
(BinaryView, Binary | LargeBinary) => true,
(Utf8 | LargeUtf8, Utf8View | BinaryView) => true,
(BinaryView, Binary | LargeBinary | Utf8 | LargeUtf8 ) => true,
(Utf8 | LargeUtf8, _) => to_type.is_numeric() && to_type != &Float16,
(_, Utf8 | LargeUtf8) => from_type.is_primitive(),

Expand Down Expand Up @@ -1229,6 +1229,13 @@ pub fn cast_with_options(
cast_byte_container::<BinaryType, LargeBinaryType>(&binary)
}
Utf8View => Ok(Arc::new(StringViewArray::from(array.as_string::<i32>()))),
BinaryView => Ok(Arc::new(BinaryViewArray::from(
array
.as_string::<i32>()
.into_iter()
.map(|x| x.map(|x| x.as_bytes()))
.collect::<Vec<_>>(),
))),
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
LargeUtf8 => cast_byte_container::<Utf8Type, LargeUtf8Type>(array),
Time32(TimeUnit::Second) => parse_string::<Time32SecondType, i32>(array, cast_options),
Time32(TimeUnit::Millisecond) => {
Expand Down Expand Up @@ -1339,6 +1346,13 @@ pub fn cast_with_options(
array.as_string::<i64>().clone(),
))),
Utf8View => Ok(Arc::new(StringViewArray::from(array.as_string::<i64>()))),
BinaryView => Ok(Arc::new(BinaryViewArray::from(
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
array
.as_string::<i64>()
.into_iter()
.map(|x| x.map(|x| x.as_bytes()))
.collect::<Vec<_>>(),
))),
Time32(TimeUnit::Second) => parse_string::<Time32SecondType, i64>(array, cast_options),
Time32(TimeUnit::Millisecond) => {
parse_string::<Time32MillisecondType, i64>(array, cast_options)
Expand Down Expand Up @@ -1413,10 +1427,15 @@ pub fn cast_with_options(
"Casting from {from_type:?} to {to_type:?} not supported",
))),
},
(BinaryView, Binary) => cast_view_to_byte::<BinaryViewType, GenericBinaryType<i32>>(array),
(BinaryView, LargeBinary) => {
cast_view_to_byte::<BinaryViewType, GenericBinaryType<i64>>(array)
}
(BinaryView, _) => match to_type {
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
Binary => cast_view_to_byte::<BinaryViewType, GenericBinaryType<i32>>(array),
LargeBinary => cast_view_to_byte::<BinaryViewType, GenericBinaryType<i64>>(array),
Utf8 => cast_slice_view_to_byte::<BinaryViewType, GenericStringType<i32>>(array),
LargeUtf8 => cast_slice_view_to_byte::<BinaryViewType, GenericStringType<i64>>(array),
_ => Err(ArrowError::CastError(format!(
"Casting from {from_type:?} to {to_type:?} not supported",
))),
},
(from_type, LargeUtf8) if from_type.is_primitive() => {
value_to_string::<i64>(array, cast_options)
}
Expand Down Expand Up @@ -2417,6 +2436,37 @@ where
Ok(Arc::new(byte_array_builder.finish()))
}

/// Specialized function to cast from one `ByteViewType` array to `ByteArrayType` array.
/// Equvilent to [`cast_view_to_byte`] but with additional constraint on the `FROM::Native` type.
fn cast_slice_view_to_byte<FROM, TO>(array: &dyn Array) -> Result<ArrayRef, ArrowError>
where
FROM: ByteViewType,
TO: ByteArrayType,
FROM::Native: AsRef<[u8]>,
str: AsRef<TO::Native>,
{
let data = array.to_data();
let view_array = GenericByteViewArray::<FROM>::from(data);

let len = view_array.len();
let bytes = view_array
.views()
.iter()
.map(|v| ByteView::from(*v).length as usize)
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
.sum::<usize>();

let mut byte_array_builder = GenericByteBuilder::<TO>::with_capacity(len, bytes);

for val in view_array.iter() {
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
let val = val
.map(|val| std::str::from_utf8(val.as_ref()))
.transpose()?;
byte_array_builder.append_option(val);
}

Ok(Arc::new(byte_array_builder.finish()))
}

#[cfg(test)]
mod tests {
use arrow_buffer::{Buffer, IntervalDayTime, NullBuffer};
Expand Down Expand Up @@ -5281,11 +5331,22 @@ mod tests {
&DataType::Utf8View
));

assert!(can_cast_types(
string_array.data_type(),
&DataType::BinaryView
));

let string_view_array = cast(&string_array, &DataType::Utf8View).unwrap();
assert_eq!(string_view_array.data_type(), &DataType::Utf8View);

let binary_view_array = cast(&string_array, &DataType::BinaryView).unwrap();
assert_eq!(binary_view_array.data_type(), &DataType::BinaryView);

let expect_string_view_array = StringViewArray::from_iter(VIEW_TEST_DATA);
assert_eq!(string_view_array.as_ref(), &expect_string_view_array);

let expect_binary_view_array = BinaryViewArray::from_iter(VIEW_TEST_DATA);
assert_eq!(binary_view_array.as_ref(), &expect_binary_view_array);
}

#[test]
Expand Down Expand Up @@ -5380,23 +5441,29 @@ mod tests {
where
O: OffsetSizeTrait,
{
let view_array = {
let string_view_array = {
let mut builder = StringViewBuilder::new().with_fixed_block_size(8); // multiple buffers.
for s in VIEW_TEST_DATA.iter() {
builder.append_option(*s);
}
builder.finish()
};

let binary_view_array = BinaryViewArray::from_iter(VIEW_TEST_DATA);

let expected_string_array = GenericStringArray::<O>::from_iter(VIEW_TEST_DATA);
let expected_type = expected_string_array.data_type();

assert!(can_cast_types(view_array.data_type(), expected_type));
assert!(can_cast_types(string_view_array.data_type(), expected_type));
assert!(can_cast_types(binary_view_array.data_type(), expected_type));

let string_array = cast(&view_array, expected_type).unwrap();
assert_eq!(string_array.data_type(), expected_type);
let string_view_casted_array = cast(&string_view_array, expected_type).unwrap();
assert_eq!(string_view_casted_array.data_type(), expected_type);
assert_eq!(string_view_casted_array.as_ref(), &expected_string_array);

assert_eq!(string_array.as_ref(), &expected_string_array);
let binary_view_casted_array = cast(&binary_view_array, expected_type).unwrap();
assert_eq!(binary_view_casted_array.data_type(), expected_type);
assert_eq!(binary_view_casted_array.as_ref(), &expected_string_array);
}

#[test]
Expand Down
Loading