klion26 commented on code in PR #10114:
URL: https://github.com/apache/arrow-rs/pull/10114#discussion_r3903666439
##########
parquet-variant-compute/src/type_conversion.rs:
##########
@@ -766,3 +883,237 @@ macro_rules! primitive_conversion_single_value {
}};
}
pub(crate) use primitive_conversion_single_value;
+
+#[cfg(test)]
+mod tests {
+ use crate::type_conversion::variant_to_string;
+ use arrow::array::{
+ Array, AsArray, BooleanArray, Date32Array, Int32Builder, ListBuilder,
StringArray,
+ Time64MicrosecondArray, TimestampMicrosecondArray,
TimestampNanosecondArray,
+ };
+ use arrow::compute::cast;
+ use arrow_schema::DataType;
+ use chrono::{DateTime, NaiveDate, NaiveTime};
+ use parquet_variant::{Variant, VariantBuilder, VariantBuilderExt};
+ use std::iter::zip;
+
+ #[test]
+ fn test_compatible_cast_logic_with_cast_kernel() {
+ // boolean -> string
+ let boolean_array = BooleanArray::from(vec![Some(true), Some(false)]);
+ let cast_array = cast(&boolean_array, &DataType::Utf8).unwrap();
+ let boolean_utf8_array = cast_array.as_string::<i32>();
+ let expected_array = vec![
+ variant_to_string(&Variant::BooleanTrue),
+ variant_to_string(&Variant::BooleanFalse),
+ ];
+ for (a, b) in zip(boolean_utf8_array, expected_array) {
+ assert_eq!(a.unwrap(), b.unwrap());
+ }
+
+ // date -> string
+ let epoch_days = [-10, 0, 18628];
+ let date_array = epoch_days
+ .iter()
+ .map(|d| Variant::Date(NaiveDate::from_epoch_days(*d).unwrap()))
+ .collect::<Vec<Variant>>();
+ let variant_as_string_array = date_array
+ .iter()
+ .map(|v| variant_to_string(v))
+ .collect::<Vec<Option<String>>>();
+
+ let date32_array = Date32Array::from_iter_values(epoch_days);
+ let date32_cast_array = cast(&date32_array, &DataType::Utf8).unwrap();
+ let date32_utf8_array = date32_cast_array
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ for (a, b) in zip(variant_as_string_array, date32_utf8_array) {
+ assert_eq!(a.unwrap(), b.unwrap());
+ }
+
+ // time -> string
+ let time_tuples = [(123, 0), (123, 456789000), (12345, 456789000)];
+ let time_array = time_tuples
+ .iter()
+ .map(|tuple| {
+ Variant::Time(
+ NaiveTime::from_num_seconds_from_midnight_opt(tuple.0,
tuple.1).unwrap(),
+ )
+ })
+ .collect::<Vec<Variant>>();
+ let time_variant_as_string_array = time_array
+ .iter()
+ .map(|v| variant_to_string(v))
+ .collect::<Vec<Option<String>>>();
+
+ let time_micro_array = Time64MicrosecondArray::from_iter(
+ time_tuples
+ .iter()
+ .map(|item| Some(item.0 as i64 * 1_000_000 + item.1 as i64 /
1000)),
+ );
+
+ let time_micro_cast_array = cast(&time_micro_array,
&DataType::Utf8).unwrap();
+ let time_micro_utf8_array = time_micro_cast_array
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+
+ for (a, b) in zip(time_variant_as_string_array, time_micro_utf8_array)
{
+ assert_eq!(a.unwrap(), b.unwrap());
+ }
+
+ // timestamp(micro) -> string
+ let micros = [-123456, 123456, 45678];
+ let timestamp_micro_array = micros
+ .iter()
+ .map(|m|
Variant::TimestampMicros(DateTime::from_timestamp_micros(*m).unwrap()))
+ .collect::<Vec<Variant>>();
+ let timestamp_micro_as_string_array = timestamp_micro_array
+ .iter()
+ .map(|v| variant_to_string(v))
+ .collect::<Vec<Option<String>>>();
+
+ let timestamp_micro_arrow_array =
+
TimestampMicrosecondArray::from_iter_values(micros).with_timezone("+00:00");
+ let timestamp_micro_arrow_cast_array =
+ cast(×tamp_micro_arrow_array, &DataType::Utf8).unwrap();
+ let timestamp_micro_utf8_array = timestamp_micro_arrow_cast_array
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ for (a, b) in zip(timestamp_micro_as_string_array,
timestamp_micro_utf8_array) {
+ assert_eq!(a.unwrap(), b.unwrap());
+ }
+
+ // timestamp(micro) ntz -> string
+ let micros_ntz = [-123456, 123456, 45678];
+ let timestamp_micro_ntz_variant_array = micros_ntz
+ .iter()
+ .map(|m| {
+ Variant::TimestampNtzMicros(
+ DateTime::from_timestamp_micros(*m).unwrap().naive_utc(),
+ )
+ })
+ .collect::<Vec<Variant>>();
+ let timestamp_micro_ntz_variant_as_string_array =
timestamp_micro_ntz_variant_array
+ .iter()
+ .map(|v| variant_to_string(v))
+ .collect::<Vec<Option<String>>>();
+
+ let timestamp_micro_ntz_arrow_array =
+ TimestampMicrosecondArray::from_iter_values(micros_ntz);
+ let timestamp_micro_ntz_arrow_cast_array =
+ cast(×tamp_micro_ntz_arrow_array, &DataType::Utf8).unwrap();
+ let timestamp_micro_ntz_utf8_array =
timestamp_micro_ntz_arrow_cast_array
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+
+ for (a, b) in zip(
+ timestamp_micro_ntz_variant_as_string_array,
+ timestamp_micro_ntz_utf8_array,
+ ) {
+ assert_eq!(a.unwrap(), b.unwrap());
+ }
+
+ // timestamp(nano) -> string
+ let nanos = [-2_208_936_075_000_000_000, 0, 1_662_921_288_000_000_000];
+ let timestamp_nano_variant_array = nanos
+ .iter()
+ .map(|n|
Variant::TimestampNanos(DateTime::from_timestamp_nanos(*n)))
+ .collect::<Vec<Variant>>();
+ let timestamp_nano_as_string_array = timestamp_nano_variant_array
+ .iter()
+ .map(|v| variant_to_string(v))
+ .collect::<Vec<Option<String>>>();
+
+ let timestamp_nano_arrow_array =
+
TimestampNanosecondArray::from_iter_values(nanos).with_timezone("+00:00");
+ let timestamp_nano_arrow_cast_array =
+ cast(×tamp_nano_arrow_array, &DataType::Utf8).unwrap();
+ let timestamp_nano_cast_utf8_array = timestamp_nano_arrow_cast_array
+ .as_any()
+ .downcast_ref::<StringArray>()
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
