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 Utf8View for bit_length kernel #6671

Merged
merged 9 commits into from
Nov 5, 2024
Merged

Support Utf8View for bit_length kernel #6671

Changes from 4 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
507dcfa
Support `Utf8View` for string function `bit_length()`
austin362667 Nov 1, 2024
67201b7
Add test & handle view bytes length counting
austin362667 Nov 1, 2024
2ad7170
Refine `string_view_array`
austin362667 Nov 2, 2024
de5e33a
Make length from `i32` to `u32` & check nullity
austin362667 Nov 2, 2024
4ab6cfd
Clean up
austin362667 Nov 2, 2024
cdfc465
Refine
austin362667 Nov 4, 2024
a3650ef
Use `from_unary` instead
austin362667 Nov 4, 2024
098cf3e
Prevent inspect the string data
austin362667 Nov 4, 2024
2a0c8b2
Clean up
austin362667 Nov 4, 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
46 changes: 43 additions & 3 deletions arrow-string/src/length.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,9 @@

use arrow_array::*;
use arrow_array::{cast::AsArray, types::*};
use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer};
use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer, ScalarBuffer};
use arrow_schema::{ArrowError, DataType};
use cast::as_string_array;
use std::sync::Arc;

fn length_impl<P: ArrowPrimitiveType>(
Expand Down Expand Up @@ -115,6 +116,8 @@ pub fn length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
/// * bit_length of null is null.
/// * bit_length is in number of bits
pub fn bit_length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
println!("In Array bit_length()");
austin362667 marked this conversation as resolved.
Show resolved Hide resolved

if let Some(d) = array.as_any_dictionary_opt() {
let lengths = bit_length(d.values().as_ref())?;
return Ok(d.with_values(lengths));
Expand All @@ -137,6 +140,26 @@ pub fn bit_length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
let list = array.as_string::<i64>();
Ok(bit_length_impl::<Int64Type>(list.offsets(), list.nulls()))
}
DataType::Utf8View => {
let list = array.as_string_view();
let values = list
.views()
.iter()
.enumerate()
.map(|(i, _)| {
austin362667 marked this conversation as resolved.
Show resolved Hide resolved
if list.is_valid(i) {
list.views()[i] as u32 * 8
} else {
0
Copy link
Member

Choose a reason for hiding this comment

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

In #6662 (comment) it was suggested to use from_unary.
It this applicable here?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@findepi Thanks for bringing this up. I tried it and it works trivially well.

Copy link
Contributor

@tustvold tustvold Nov 4, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

Using from_unary reverts the optimisation of only inspecting the views, and not the string data.

Perhaps you might use the example I wrote in #6671 (comment)

}
})
.collect::<Vec<u32>>();

Ok(Arc::new(UInt32Array::new(
values.into(),
array.nulls().cloned(),
)))
}
DataType::Binary => {
let list = array.as_binary::<i32>();
Ok(bit_length_impl::<Int32Type>(list.offsets(), list.nulls()))
Expand All @@ -157,6 +180,8 @@ pub fn bit_length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {

#[cfg(test)]
mod tests {
use std::ptr::null;

use super::*;
use arrow_buffer::Buffer;
use arrow_data::ArrayData;
Expand Down Expand Up @@ -418,7 +443,7 @@ mod tests {
assert_eq!(&expected, result);
}

fn bit_length_cases() -> Vec<(Vec<&'static str>, usize, Vec<i32>)> {
fn bit_length_cases() -> Vec<(Vec<&'static str>, usize, Vec<u32>)> {
// a large array
let values = ["one", "on", "o", ""];
let values = values.into_iter().cycle().take(4096).collect();
Expand All @@ -440,7 +465,7 @@ mod tests {
let array = StringArray::from(input);
let result = bit_length(&array).unwrap();
assert_eq!(len, result.len());
let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
let result = result.as_any().downcast_ref::<UInt32Array>().unwrap();
Copy link
Member

Choose a reason for hiding this comment

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

I agree that bit_length should not beed to return signed integer, but it seems better not to change this within this PR.

expected.iter().enumerate().for_each(|(i, value)| {
assert_eq!(*value, result.value(i));
});
Expand All @@ -462,6 +487,21 @@ mod tests {
})
}

#[test]
fn bit_length_test_utf8view() {
bit_length_cases()
austin362667 marked this conversation as resolved.
Show resolved Hide resolved
.into_iter()
.for_each(|(input, len, expected)| {
let string_array = StringViewArray::from(input);
let result = bit_length(&string_array).unwrap();
assert_eq!(len, result.len());
let result = result.as_any().downcast_ref::<UInt32Array>().unwrap();
expected.iter().enumerate().for_each(|(i, value)| {
assert_eq!(*value, result.value(i));
});
})
}

#[test]
fn bit_length_binary() {
let value: Vec<&[u8]> = vec![b"one", &[0xff, 0xf8], b"three"];
Expand Down