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

Jump to bottom

[arrow-cast] Support cast numeric to string view #6719

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open

[arrow-cast] Support cast numeric to string view #6719

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
205e40d
[arrow-cast] Support cast numeric to string view
tlm365 Nov 12, 2024
2a937fe
fix fmt error
tlm365 Nov 12, 2024
2f8f9c3
Implement std::fmt::Write for StringViewBuilder
tlm365 Nov 12, 2024
74de9bc
Update that current impl support only for numeric -> string view
tlm365 Nov 15, 2024
cdc16c9
Merge branch 'master' into cast_numeric_to_string_view
tlm365 Nov 18, 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
7 changes: 7 additions & 0 deletions arrow-array/src/builder/generic_bytes_view_builder.rs
View file
Open in desktop
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 {
fn write_str(&mut self, s: &str) -> std::fmt::Result {
self.append_value(s);
Ok(())
}
}

/// Array builder for [`BinaryViewArray`][crate::BinaryViewArray]
///
/// Values can be appended using [`GenericByteViewBuilder::append_value`], and nulls with
Expand Down
275 changes: 168 additions & 107 deletions arrow-cast/src/cast/mod.rs
View file
Open in desktop
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,45 @@ 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
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 +5272,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 +9173,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 +9229,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