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

[arrow-cast] Support cast from Numeric (Int, UInt, etc) to Utf8View #6719

Closed
wants to merge 9 commits into from
7 changes: 7 additions & 0 deletions arrow-array/src/builder/generic_bytes_view_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,13 @@ impl<T: ByteViewType + ?Sized, V: AsRef<T::Native>> Extend<Option<V>>
/// ```
pub type StringViewBuilder = GenericByteViewBuilder<StringViewType>;

impl std::fmt::Write for StringViewBuilder {
Copy link
Contributor

Choose a reason for hiding this comment

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

this is also what is contemplated by #6373 (aka I think this PR fixes that ticket as well)

fn write_str(&mut self, s: &str) -> std::fmt::Result {
self.append_value(s);
Copy link
Contributor

@alamb alamb Nov 22, 2024

Choose a reason for hiding this comment

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

I was writing some tests for this, and it turns out this is different behavior than StringViewBuilder

https://docs.rs/arrow/latest/arrow/array/builder/type.GenericStringBuilder.html#example-incrementally-writing-strings-with-stdfmtwrite

Specifically, calling write_str doesn't compete the row 🤔

I made a PR showing the problem: tlm365#1

Copy link
Contributor

Choose a reason for hiding this comment

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

I am working on a potential solution so we can unblock this PR

Copy link
Contributor

Choose a reason for hiding this comment

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

Ok(())
}
}

/// Array builder for [`BinaryViewArray`][crate::BinaryViewArray]
///
/// Values can be appended using [`GenericByteViewBuilder::append_value`], and nulls with
Expand Down
276 changes: 169 additions & 107 deletions arrow-cast/src/cast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -182,10 +182,11 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
(Decimal128(_, _) | Decimal256(_, _), UInt8 | UInt16 | UInt32 | UInt64) |
// decimal to signed numeric
(Decimal128(_, _) | Decimal256(_, _), Null | Int8 | Int16 | Int32 | Int64 | Float32 | Float64) => true,
// decimal to Utf8
(Decimal128(_, _) | Decimal256(_, _), Utf8 | LargeUtf8) => true,
// decimal to string
(Decimal128(_, _) | Decimal256(_, _), Utf8View | Utf8 | LargeUtf8) => true,
// string to decimal
(Utf8View | Utf8 | LargeUtf8, Decimal128(_, _) | Decimal256(_, _)) => true,

(Struct(from_fields), Struct(to_fields)) => {
from_fields.len() == to_fields.len() &&
from_fields.iter().zip(to_fields.iter()).all(|(f1, f2)| {
Expand Down Expand Up @@ -232,6 +233,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
(BinaryView, Binary | LargeBinary | Utf8 | LargeUtf8 | Utf8View ) => true,
(Utf8View | Utf8 | LargeUtf8, _) => to_type.is_numeric() && to_type != &Float16,
(_, Utf8 | LargeUtf8) => from_type.is_primitive(),
(_, Utf8View) => from_type.is_numeric(),

(_, Binary | LargeBinary) => from_type.is_integer(),

Expand Down Expand Up @@ -917,6 +919,7 @@ pub fn cast_with_options(
Float64 => cast_decimal_to_float::<Decimal128Type, Float64Type, _>(array, |x| {
x as f64 / 10_f64.powi(*scale as i32)
}),
Utf8View => value_to_string_view(array, cast_options),
Utf8 => value_to_string::<i32>(array, cast_options),
LargeUtf8 => value_to_string::<i64>(array, cast_options),
Null => Ok(new_null_array(to_type, array.len())),
Expand Down Expand Up @@ -982,6 +985,7 @@ pub fn cast_with_options(
Float64 => cast_decimal_to_float::<Decimal256Type, Float64Type, _>(array, |x| {
x.to_f64().unwrap() / 10_f64.powi(*scale as i32)
}),
Utf8View => value_to_string_view(array, cast_options),
Utf8 => value_to_string::<i32>(array, cast_options),
LargeUtf8 => value_to_string::<i64>(array, cast_options),
Null => Ok(new_null_array(to_type, array.len())),
Expand Down Expand Up @@ -1462,6 +1466,9 @@ pub fn cast_with_options(
(BinaryView, _) => Err(ArrowError::CastError(format!(
"Casting from {from_type:?} to {to_type:?} not supported",
))),
(from_type, Utf8View) if from_type.is_numeric() => {
value_to_string_view(array, cast_options)
}
(from_type, LargeUtf8) if from_type.is_primitive() => {
value_to_string::<i64>(array, cast_options)
}
Expand Down Expand Up @@ -3707,6 +3714,40 @@ mod tests {
assert_eq!(10.0, c.value(3));
}

#[test]
fn test_cast_int_to_utf8view() {
assert!(can_cast_types(&DataType::Int8, &DataType::Utf8View));
assert!(can_cast_types(&DataType::Int16, &DataType::Utf8View));
assert!(can_cast_types(&DataType::Int32, &DataType::Utf8View));
assert!(can_cast_types(&DataType::Int64, &DataType::Utf8View));

let array = Int32Array::from(vec![None, Some(8), Some(9), Some(10)]);
let arr = cast(&array, &DataType::Utf8View).unwrap();
assert_eq!(4, arr.len());
assert_eq!(1, arr.null_count());
let c = arr.as_string_view();
assert!(c.is_null(0));
assert_eq!("8", c.value(1));
assert_eq!("9", c.value(2));
assert_eq!("10", c.value(3));
}

#[test]
fn test_cast_float_to_utf8view() {
assert!(can_cast_types(&DataType::Float16, &DataType::Utf8View));
assert!(can_cast_types(&DataType::Float32, &DataType::Utf8View));
assert!(can_cast_types(&DataType::Float64, &DataType::Utf8View));

let array = Float32Array::from(vec![Some(8.64), Some(9.81), None]);
let arr = cast(&array, &DataType::Utf8View).unwrap();
assert_eq!(3, arr.len());
assert_eq!(1, arr.null_count());
let c = arr.as_string_view();
assert_eq!("8.64", c.value(0));
assert_eq!("9.81", c.value(1));
assert!(c.is_null(2));
}

#[test]
fn test_cast_utf8_to_i32() {
let array = StringArray::from(vec!["5", "6", "seven", "8", "9.1"]);
Expand Down Expand Up @@ -5180,41 +5221,46 @@ mod tests {
assert_eq!("2018-12-25T00:00:00", c.value(1));
}

// Cast Timestamp to Utf8View is not supported yet
// TODO: Implement casting from Timestamp to Utf8View
alamb marked this conversation as resolved.
Show resolved Hide resolved
// https://github.com/apache/arrow-rs/issues/6734
macro_rules! assert_cast_timestamp_to_string {
($array:expr, $datatype:expr, $output_array_type: ty, $expected:expr) => {{
let out = cast(&$array, &$datatype).unwrap();
let actual = out
.as_any()
.downcast_ref::<$output_array_type>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(actual, $expected);
}};
($array:expr, $datatype:expr, $output_array_type: ty, $options:expr, $expected:expr) => {{
let out = cast_with_options(&$array, &$datatype, &$options).unwrap();
let actual = out
.as_any()
.downcast_ref::<$output_array_type>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(actual, $expected);
}};
}

#[test]
fn test_cast_timestamp_to_strings() {
// "2018-12-25T00:00:02.001", "1997-05-19T00:00:03.005", None
let array =
TimestampMillisecondArray::from(vec![Some(864000003005), Some(1545696002001), None]);
let out = cast(&array, &DataType::Utf8).unwrap();
let out = out
.as_any()
.downcast_ref::<StringArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
out,
vec![
Some("1997-05-19T00:00:03.005"),
Some("2018-12-25T00:00:02.001"),
None
]
);
let out = cast(&array, &DataType::LargeUtf8).unwrap();
let out = out
.as_any()
.downcast_ref::<LargeStringArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
out,
vec![
Some("1997-05-19T00:00:03.005"),
Some("2018-12-25T00:00:02.001"),
None
]
);
let expected = vec![
Some("1997-05-19T00:00:03.005"),
Some("2018-12-25T00:00:02.001"),
None,
];

// assert_cast_timestamp_to_string!(array, DataType::Utf8View, StringViewArray, expected);
assert_cast_timestamp_to_string!(array, DataType::Utf8, StringArray, expected);
assert_cast_timestamp_to_string!(array, DataType::LargeUtf8, LargeStringArray, expected);
}

#[test]
Expand All @@ -5227,73 +5273,53 @@ mod tests {
.with_timestamp_format(Some(ts_format))
.with_timestamp_tz_format(Some(ts_format)),
};

// "2018-12-25T00:00:02.001", "1997-05-19T00:00:03.005", None
let array_without_tz =
TimestampMillisecondArray::from(vec![Some(864000003005), Some(1545696002001), None]);
let out = cast_with_options(&array_without_tz, &DataType::Utf8, &cast_options).unwrap();
let out = out
.as_any()
.downcast_ref::<StringArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
out,
vec![
Some("1997-05-19 00:00:03.005000"),
Some("2018-12-25 00:00:02.001000"),
None
]
let expected = vec![
Some("1997-05-19 00:00:03.005000"),
Some("2018-12-25 00:00:02.001000"),
None,
];
// assert_cast_timestamp_to_string!(array_without_tz, DataType::Utf8View, StringViewArray, cast_options, expected);
assert_cast_timestamp_to_string!(
array_without_tz,
DataType::Utf8,
StringArray,
cast_options,
expected
);
let out =
cast_with_options(&array_without_tz, &DataType::LargeUtf8, &cast_options).unwrap();
let out = out
.as_any()
.downcast_ref::<LargeStringArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
out,
vec![
Some("1997-05-19 00:00:03.005000"),
Some("2018-12-25 00:00:02.001000"),
None
]
assert_cast_timestamp_to_string!(
array_without_tz,
DataType::LargeUtf8,
LargeStringArray,
cast_options,
expected
);

let array_with_tz =
TimestampMillisecondArray::from(vec![Some(864000003005), Some(1545696002001), None])
.with_timezone(tz.to_string());
let out = cast_with_options(&array_with_tz, &DataType::Utf8, &cast_options).unwrap();
let out = out
.as_any()
.downcast_ref::<StringArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
out,
vec![
Some("1997-05-19 05:45:03.005000"),
Some("2018-12-25 05:45:02.001000"),
None
]
let expected = vec![
Some("1997-05-19 05:45:03.005000"),
Some("2018-12-25 05:45:02.001000"),
None,
];
// assert_cast_timestamp_to_string!(array_with_tz, DataType::Utf8View, StringViewArray, cast_options, expected);
assert_cast_timestamp_to_string!(
array_with_tz,
DataType::Utf8,
StringArray,
cast_options,
expected
);
let out = cast_with_options(&array_with_tz, &DataType::LargeUtf8, &cast_options).unwrap();
let out = out
.as_any()
.downcast_ref::<LargeStringArray>()
.unwrap()
.into_iter()
.collect::<Vec<_>>();
assert_eq!(
out,
vec![
Some("1997-05-19 05:45:03.005000"),
Some("2018-12-25 05:45:02.001000"),
None
]
assert_cast_timestamp_to_string!(
array_with_tz,
DataType::LargeUtf8,
LargeStringArray,
cast_options,
expected
);
}

Expand Down Expand Up @@ -9148,26 +9174,51 @@ mod tests {
}

#[test]
fn test_cast_decimal_to_utf8() {
fn test_cast_decimal_to_string() {
assert!(can_cast_types(
&DataType::Decimal128(10, 4),
&DataType::Utf8View
));
assert!(can_cast_types(
&DataType::Decimal256(38, 10),
&DataType::Utf8View
));

macro_rules! assert_decimal_values {
($array:expr) => {
let c = $array;
assert_eq!("1123.454", c.value(0));
assert_eq!("2123.456", c.value(1));
assert_eq!("-3123.453", c.value(2));
assert_eq!("-3123.456", c.value(3));
assert_eq!("0.000", c.value(4));
assert_eq!("0.123", c.value(5));
assert_eq!("1234.567", c.value(6));
assert_eq!("-1234.567", c.value(7));
assert!(c.is_null(8));
};
}

fn test_decimal_to_string<IN: ArrowPrimitiveType, OffsetSize: OffsetSizeTrait>(
output_type: DataType,
array: PrimitiveArray<IN>,
) {
let b = cast(&array, &output_type).unwrap();

assert_eq!(b.data_type(), &output_type);
let c = b.as_string::<OffsetSize>();

assert_eq!("1123.454", c.value(0));
assert_eq!("2123.456", c.value(1));
assert_eq!("-3123.453", c.value(2));
assert_eq!("-3123.456", c.value(3));
assert_eq!("0.000", c.value(4));
assert_eq!("0.123", c.value(5));
assert_eq!("1234.567", c.value(6));
assert_eq!("-1234.567", c.value(7));
assert!(c.is_null(8));
match b.data_type() {
DataType::Utf8View => {
let c = b.as_string_view();
assert_decimal_values!(c);
}
DataType::Utf8 | DataType::LargeUtf8 => {
let c = b.as_string::<OffsetSize>();
assert_decimal_values!(c);
}
_ => (),
}
}

let array128: Vec<Option<i128>> = vec![
Some(1123454),
Some(2123456),
Expand All @@ -9179,22 +9230,33 @@ mod tests {
Some(-123456789),
None,
];
let array256: Vec<Option<i256>> = array128
.iter()
.map(|num| num.map(i256::from_i128))
.collect();

let array256: Vec<Option<i256>> = array128.iter().map(|v| v.map(i256::from_i128)).collect();

test_decimal_to_string::<arrow_array::types::Decimal128Type, i32>(
test_decimal_to_string::<Decimal128Type, i32>(
DataType::Utf8View,
create_decimal_array(array128.clone(), 7, 3).unwrap(),
);
test_decimal_to_string::<Decimal128Type, i32>(
DataType::Utf8,
create_decimal_array(array128.clone(), 7, 3).unwrap(),
);
test_decimal_to_string::<arrow_array::types::Decimal128Type, i64>(
test_decimal_to_string::<Decimal128Type, i64>(
DataType::LargeUtf8,
create_decimal_array(array128, 7, 3).unwrap(),
);
test_decimal_to_string::<arrow_array::types::Decimal256Type, i32>(

test_decimal_to_string::<Decimal256Type, i32>(
DataType::Utf8View,
create_decimal256_array(array256.clone(), 7, 3).unwrap(),
);
test_decimal_to_string::<Decimal256Type, i32>(
DataType::Utf8,
create_decimal256_array(array256.clone(), 7, 3).unwrap(),
);
test_decimal_to_string::<arrow_array::types::Decimal256Type, i64>(
test_decimal_to_string::<Decimal256Type, i64>(
DataType::LargeUtf8,
create_decimal256_array(array256, 7, 3).unwrap(),
);
Expand Down
Loading
Loading