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.

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
Prev Previous commit
Next Next commit
Refine string_view_array
Signed-off-by: Austin Liu <austin362667@gmail.com>
  • Loading branch information
austin362667 committed Nov 2, 2024
commit 2ad7170e2bfab0bd9e501727a08af1e16502434c
26 changes: 10 additions & 16 deletions arrow-string/src/length.rs
Original file line number Diff line number Diff line change
@@ -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>(
@@ -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));
@@ -138,20 +141,12 @@ pub fn bit_length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
Ok(bit_length_impl::<Int64Type>(list.offsets(), list.nulls()))
}
DataType::Utf8View => {
let string_view_array = array
.as_any()
.downcast_ref::<StringViewArray>()
.ok_or_else(|| ArrowError::ComputeError("Expected Utf8View array".to_string()))?;
let mut bit_lengths = Vec::with_capacity(array.len());
for i in 0..array.len() {
let bit_length = if string_view_array.is_valid(i) {
(string_view_array.value(i).len() * 8) as i32
} else {
0
};
bit_lengths.push(bit_length);
}

let list = array.as_string_view();
let bit_lengths = list
.views()
.iter()
.map(|view| (*view as i32) * 8)
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 length is currently treated as a u32 elsewhere in the code

However the actual Arrow spec I think says use i32 so this is probably fine

Copy link
Contributor

Choose a reason for hiding this comment

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

I also still do think you need to handle the null cases (check for which slots are null) as there might be a view of a non zero length there

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Sure! Thanks @alamb for review

.collect::<Vec<i32>>();
Ok(Arc::new(Int32Array::new(
bit_lengths.into(),
array.nulls().cloned(),
@@ -490,7 +485,6 @@ mod tests {
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::<Int32Array>().unwrap();
expected.iter().enumerate().for_each(|(i, value)| {
assert_eq!(*value, result.value(i));