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
+127 −14
Merged

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

Show file tree
Hide file tree
Changes from 3 commits
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
136 changes: 118 additions & 18 deletions arrow-cast/src/cast/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -225,10 +225,11 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
| Timestamp(Millisecond, _)
| Timestamp(Microsecond, _)
| Timestamp(Nanosecond, _)
| Interval(_),
| Interval(_)
| BinaryView,
) => true,
(Utf8 | LargeUtf8, Utf8View) => true,
(BinaryView, Binary | LargeBinary) => true,
(BinaryView, Binary | LargeBinary | Utf8 | LargeUtf8 | Utf8View ) => true,
(Utf8 | LargeUtf8, _) => to_type.is_numeric() && to_type != &Float16,
(_, Utf8 | LargeUtf8) => from_type.is_primitive(),

Expand Down Expand Up @@ -1229,6 +1230,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 @@ -1282,6 +1290,15 @@ pub fn cast_with_options(
Date64 => parse_string_view::<Date64Type>(array, cast_options),
Binary => cast_view_to_byte::<StringViewType, GenericBinaryType<i32>>(array),
LargeBinary => cast_view_to_byte::<StringViewType, GenericBinaryType<i64>>(array),
BinaryView => {
if let Some(arr) = array.as_any().downcast_ref::<StringViewArray>() {
Ok(Arc::new(arr.clone().to_binary_view()))
} else {
Err(ArrowError::CastError(
"Cannot cast StringView to BinaryView".to_string(),
))
}
}
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
Utf8 => cast_view_to_byte::<StringViewType, GenericStringType<i32>>(array),
LargeUtf8 => cast_view_to_byte::<StringViewType, GenericStringType<i64>>(array),
Time32(TimeUnit::Second) => parse_string_view::<Time32SecondType>(array, cast_options),
Expand Down Expand Up @@ -1339,6 +1356,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 +1437,34 @@ 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 => {
let binary_arr =
cast_view_to_byte::<BinaryViewType, GenericBinaryType<i32>>(array)?;
cast_binary_to_string::<i32>(&binary_arr, cast_options)
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
}
LargeUtf8 => {
let binary_arr =
cast_view_to_byte::<BinaryViewType, GenericBinaryType<i64>>(array)?;
cast_binary_to_string::<i64>(&binary_arr, cast_options)
}
Utf8View => {
if let Some(arr) = array.as_any().downcast_ref::<BinaryViewArray>() {
arr.clone()
.to_string_view()
.map(|x| Arc::new(x) as ArrayRef)
} else {
Err(ArrowError::CastError(
"Cannot cast BinaryView to StringView".to_string(),
))
}
}
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
_ => 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 @@ -2008,7 +2056,6 @@ pub fn cast_with_options(
})?,
))
}

(Date64, Timestamp(TimeUnit::Second, None)) => Ok(Arc::new(
array
.as_primitive::<Date64Type>()
Expand Down Expand Up @@ -5256,12 +5303,6 @@ mod tests {
}
}

#[test]
fn test_string_to_view() {
_test_string_to_view::<i32>();
_test_string_to_view::<i64>();
}

const VIEW_TEST_DATA: [Option<&str>; 5] = [
Some("hello"),
Some("repeated"),
Expand All @@ -5270,6 +5311,48 @@ mod tests {
Some("repeated"),
];

#[test]
xinlifoobar marked this conversation as resolved.
Show resolved Hide resolved
fn test_between_views() {
_test_string_view_to_binary_view();
_test_binary_view_to_string_view();
}

fn _test_string_view_to_binary_view() {
let string_view_array = StringViewArray::from_iter(VIEW_TEST_DATA);

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

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

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

fn _test_binary_view_to_string_view() {
let binary_view_array = BinaryViewArray::from_iter(VIEW_TEST_DATA);

assert!(can_cast_types(
binary_view_array.data_type(),
&DataType::Utf8View
));

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

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

#[test]
fn test_string_to_view() {
_test_string_to_view::<i32>();
_test_string_to_view::<i64>();
}

fn _test_string_to_view<O>()
where
O: OffsetSizeTrait,
Expand All @@ -5281,11 +5364,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 +5474,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
39 changes: 39 additions & 0 deletions arrow/benches/cast_kernels.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,18 @@ fn build_decimal256_array(size: usize, precision: u8, scale: i8) -> ArrayRef {
)
}

fn build_string_array(size: usize) -> ArrayRef {
let mut builder = StringBuilder::new();
for v in 0..size {
match v % 3 {
0 => builder.append_value("small"),
1 => builder.append_value("larger string more than 12 bytes"),
_ => builder.append_null(),
}
}
Arc::new(builder.finish())
}

fn build_dict_array(size: usize) -> ArrayRef {
let values = StringArray::from_iter([
Some("small"),
Expand Down Expand Up @@ -148,9 +160,12 @@ fn add_benchmark(c: &mut Criterion) {

let decimal128_array = build_decimal128_array(512, 10, 3);
let decimal256_array = build_decimal256_array(512, 50, 3);
let string_array = build_string_array(512);
let wide_string_array = cast(&string_array, &DataType::LargeUtf8).unwrap();

let dict_array = build_dict_array(10_000);
let string_view_array = cast(&dict_array, &DataType::Utf8View).unwrap();
let binary_view_array = cast(&string_view_array, &DataType::BinaryView).unwrap();

c.bench_function("cast int32 to int32 512", |b| {
b.iter(|| cast_array(&i32_array, DataType::Int32))
Expand Down Expand Up @@ -262,6 +277,30 @@ fn add_benchmark(c: &mut Criterion) {
)
})
});
c.bench_function("cast string view to string", |b| {
b.iter(|| cast_array(&string_view_array, DataType::Utf8))
});
c.bench_function("cast string view to wide string", |b| {
b.iter(|| cast_array(&string_view_array, DataType::LargeUtf8))
});
c.bench_function("cast binary view to string", |b| {
b.iter(|| cast_array(&binary_view_array, DataType::Utf8))
});
c.bench_function("cast binary view to wide string", |b| {
b.iter(|| cast_array(&binary_view_array, DataType::LargeUtf8))
});
c.bench_function("cast string to binary view 512", |b| {
b.iter(|| cast_array(&string_array, DataType::BinaryView))
});
c.bench_function("cast wide string to binary view 512", |b| {
b.iter(|| cast_array(&wide_string_array, DataType::BinaryView))
});
c.bench_function("cast string view to binary view", |b| {
b.iter(|| cast_array(&string_view_array, DataType::BinaryView))
});
c.bench_function("cast binary view to string view", |b| {
b.iter(|| cast_array(&binary_view_array, DataType::Utf8View))
});
}

criterion_group!(benches, add_benchmark);
Expand Down
Loading