Skip to content

Commit

Permalink
Lint and format
Browse files Browse the repository at this point in the history
  • Loading branch information
@HawaiianSpork
HawaiianSpork committed May 2, 2024
1 parent 8649d25 commit 164ae42
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 57 deletions.
21 changes: 15 additions & 6 deletions arrow-cast/src/cast/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@

use crate::cast::*;


/// Helper function that takes a map container and casts the inner datatype.
pub(crate) fn cast_map_values(
from: &MapArray,
Expand All @@ -28,19 +27,29 @@ pub(crate) fn cast_map_values(
let entries_field = if let DataType::Map(entries_field, _) = to_data_type {
entries_field
} else {
return Err(ArrowError::CastError("Internal Error: to_data_type is not a map type.".to_string()));
return Err(ArrowError::CastError(
"Internal Error: to_data_type is not a map type.".to_string(),
));
};

let key_field = key_field(entries_field).ok_or(ArrowError::CastError("map is missing key field".to_string()))?;
let value_field = value_field(entries_field).ok_or(ArrowError::CastError("map is missing value field".to_string()))?;
let key_field = key_field(entries_field).ok_or(ArrowError::CastError(
"map is missing key field".to_string(),
))?;
let value_field = value_field(entries_field).ok_or(ArrowError::CastError(
"map is missing value field".to_string(),
))?;

let key_array = cast_with_options(from.keys(), key_field.data_type(), cast_options)?;
let value_array = cast_with_options(from.values(), value_field.data_type(), cast_options)?;

Ok(Arc::new(MapArray::new(
entries_field.clone(),
from.offsets().clone(),
StructArray::new(Fields::from(vec![key_field, value_field]), vec![key_array, value_array], from.nulls().cloned()),
StructArray::new(
Fields::from(vec![key_field, value_field]),
vec![key_array, value_array],
from.nulls().cloned(),
),
from.nulls().cloned(),
to_ordered,
)))
Expand All @@ -49,7 +58,7 @@ pub(crate) fn cast_map_values(
/// Gets the key field from the entries of a map. For all other types returns None.
pub(crate) fn key_field(entries_field: &FieldRef) -> Option<FieldRef> {
if let DataType::Struct(fields) = entries_field.data_type() {
fields.get(0).cloned()
fields.first().cloned()
} else {
None
}
Expand Down
123 changes: 72 additions & 51 deletions arrow-cast/src/cast/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,6 @@ use arrow_select::take::take;
use num::cast::AsPrimitive;
use num::{NumCast, ToPrimitive};


/// CastOptions provides a way to override the default cast behaviors
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CastOptions<'a> {
Expand Down Expand Up @@ -163,7 +162,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
(Map(from_entries,ordered_from), Map(to_entries, ordered_to)) if ordered_from == ordered_to =>
match (key_field(from_entries), key_field(to_entries), value_field(from_entries), value_field(to_entries)) {
(Some(from_key), Some(to_key), Some(from_value), Some(to_value)) =>
can_cast_types(from_key.data_type(), to_key.data_type()) && can_cast_types(&from_value.data_type(), &to_value.data_type()),
can_cast_types(from_key.data_type(), to_key.data_type()) && can_cast_types(from_value.data_type(), to_value.data_type()),
_ => false
},
// cast one decimal type to another decimal type
Expand Down Expand Up @@ -719,13 +718,9 @@ pub fn cast_with_options(
(_, FixedSizeList(ref to, size)) if *size == 1 => {
cast_values_to_fixed_size_list(array, to, *size, cast_options)
}
(Map(_, ordered1), Map(_, ordered2)) if ordered1 == ordered2 =>
cast_map_values(
array.as_map(),
&to_type,
cast_options,
ordered1.to_owned()
),
(Map(_, ordered1), Map(_, ordered2)) if ordered1 == ordered2 => {
cast_map_values(array.as_map(), to_type, cast_options, ordered1.to_owned())
}
(Decimal128(_, s1), Decimal128(p2, s2)) => {
cast_decimal_to_decimal_same_type::<Decimal128Type>(
array.as_primitive(),
Expand Down Expand Up @@ -7141,18 +7136,26 @@ mod tests {
let array = builder.finish();

let new_ordered = true;
let new_type = DataType::Map(Arc::new(Field::new("entries", DataType::Struct(
vec![
Field::new("key", DataType::Utf8, false),
Field::new("value", DataType::Utf8, false),
].into()
), false)), new_ordered);
let new_type = DataType::Map(
Arc::new(Field::new(
"entries",
DataType::Struct(
vec![
Field::new("key", DataType::Utf8, false),
Field::new("value", DataType::Utf8, false),
]
.into(),
),
false,
)),
new_ordered,
);

let new_array_result = cast(
&array,
&new_type.clone());
let new_array_result = cast(&array, &new_type.clone());
assert!(!can_cast_types(array.data_type(), &new_type));
assert!(matches!(new_array_result, Err(ArrowError::CastError(t)) if t == r#"Casting from Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Utf8, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, false) to Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, true) not supported"#));
assert!(
matches!(new_array_result, Err(ArrowError::CastError(t)) if t == r#"Casting from Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Utf8, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, false) to Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, true) not supported"#)
);
}

#[test]
Expand Down Expand Up @@ -7180,18 +7183,26 @@ mod tests {
let array = builder.finish();

let new_ordered = true;
let new_type = DataType::Map(Arc::new(Field::new("entries", DataType::Struct(
vec![
Field::new("key", DataType::Utf8, false),
Field::new("value", DataType::Duration(TimeUnit::Second), false),
].into()
), false)), new_ordered);
let new_type = DataType::Map(
Arc::new(Field::new(
"entries",
DataType::Struct(
vec![
Field::new("key", DataType::Utf8, false),
Field::new("value", DataType::Duration(TimeUnit::Second), false),
]
.into(),
),
false,
)),
new_ordered,
);

let new_array_result = cast(
&array,
&new_type.clone());
let new_array_result = cast(&array, &new_type.clone());
assert!(!can_cast_types(array.data_type(), &new_type));
assert!(matches!(new_array_result, Err(ArrowError::CastError(t)) if t == r#"Casting from Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Interval(DayTime), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, false) to Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Duration(Second), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, true) not supported"#));
assert!(
matches!(new_array_result, Err(ArrowError::CastError(t)) if t == r#"Casting from Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Interval(DayTime), nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, false) to Map(Field { name: "entries", data_type: Struct([Field { name: "key", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "value", data_type: Duration(Second), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }]), nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, true) not supported"#)
);
}

#[test]
Expand All @@ -7217,18 +7228,24 @@ mod tests {

let array = builder.finish();

let new_type = DataType::Map(Arc::new(Field::new("entries_new", DataType::Struct(
vec![
Field::new("key_new", DataType::Utf8, false),
Field::new("value_values", DataType::Utf8, false),
].into()
), false)), false);
let new_type = DataType::Map(
Arc::new(Field::new(
"entries_new",
DataType::Struct(
vec![
Field::new("key_new", DataType::Utf8, false),
Field::new("value_values", DataType::Utf8, false),
]
.into(),
),
false,
)),
false,
);

assert_ne!(new_type, array.data_type().clone());

let new_array = cast(
&array,
&new_type.clone()).unwrap();
let new_array = cast(&array, &new_type.clone()).unwrap();
assert_eq!(new_type, new_array.data_type().clone());
let map_array = new_array.as_map();

Expand All @@ -7245,8 +7262,7 @@ mod tests {
.collect::<Vec<_>>();
assert_eq!(&key_string, &vec!["0", "1"]);

let values_string_array = cast(map_array
.values(), &DataType::Utf8).unwrap();
let values_string_array = cast(map_array.values(), &DataType::Utf8).unwrap();
let values_string = values_string_array
.as_any()
.downcast_ref::<StringArray>()
Expand Down Expand Up @@ -7280,16 +7296,22 @@ mod tests {

let array = builder.finish();

let new_type = DataType::Map(Arc::new(Field::new("entries", DataType::Struct(
vec![
Field::new("key", DataType::Utf8, false),
Field::new("value", DataType::Utf8, false),
].into()
), false)), false);
let new_type = DataType::Map(
Arc::new(Field::new(
"entries",
DataType::Struct(
vec![
Field::new("key", DataType::Utf8, false),
Field::new("value", DataType::Utf8, false),
]
.into(),
),
false,
)),
false,
);

let new_array = cast(
&array,
&new_type.clone()).unwrap();
let new_array = cast(&array, &new_type.clone()).unwrap();
assert_eq!(new_type, new_array.data_type().clone());
let map_array = new_array.as_map();

Expand All @@ -7306,8 +7328,7 @@ mod tests {
.collect::<Vec<_>>();
assert_eq!(&key_string, &vec!["0", "1"]);

let values_string_array = cast(map_array
.values(), &DataType::Utf8).unwrap();
let values_string_array = cast(map_array.values(), &DataType::Utf8).unwrap();
let values_string = values_string_array
.as_any()
.downcast_ref::<StringArray>()
Expand Down

0 comments on commit 164ae42

Please sign in to comment.