fix: Only delegate to DataFusion cast when we know that it is compatible with Spark #461
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -503,41 +503,37 @@ impl Cast { | |
| fn cast_array(&self, array: ArrayRef) -> DataFusionResult<ArrayRef> { | ||
| let to_type = &self.data_type; | ||
| let array = array_with_timezone(array, self.timezone.clone(), Some(to_type)); | ||
| let from_type = array.data_type().clone(); | ||
|
|
||
| // unpack dictionary string arrays first | ||
| // TODO: we are unpacking a dictionary-encoded array and then performing | ||
| // the cast. We could potentially improve performance here by casting the | ||
| // dictionary values directly without unpacking the array first, although this | ||
| // would add more complexity to the code | ||
| let array = match &from_type { | ||
| DataType::Dictionary(key_type, value_type) | ||
| if key_type.as_ref() == &DataType::Int32 | ||
| && (value_type.as_ref() == &DataType::Utf8 | ||
| || value_type.as_ref() == &DataType::LargeUtf8) => | ||
| { | ||
| cast_with_options(&array, value_type.as_ref(), &CAST_OPTIONS)? | ||
| } | ||
| _ => array, | ||
| }; | ||
| let from_type = array.data_type(); | ||
|
|
||
| let cast_result = match (from_type, to_type) { | ||
| (DataType::Utf8, DataType::Boolean) => { | ||
| Self::spark_cast_utf8_to_boolean::<i32>(&array, self.eval_mode) | ||
| } | ||
| (DataType::LargeUtf8, DataType::Boolean) => { | ||
| Self::spark_cast_utf8_to_boolean::<i64>(&array, self.eval_mode) | ||
| } | ||
| (DataType::Utf8, DataType::Timestamp(_, _)) => { | ||
| Self::cast_string_to_timestamp(&array, to_type, self.eval_mode) | ||
| } | ||
| (DataType::Utf8, DataType::Date32) => { | ||
| Self::cast_string_to_date(&array, to_type, self.eval_mode) | ||
| } | ||
| (DataType::Int64, DataType::Int32) | ||
| | (DataType::Int64, DataType::Int16) | ||
|
|
@@ -547,61 +543,33 @@ impl Cast { | |
| | (DataType::Int16, DataType::Int8) | ||
| if self.eval_mode != EvalMode::Try => | ||
| { | ||
| Self::spark_cast_int_to_int(&array, self.eval_mode, from_type, to_type) | ||
| } | ||
| ( | ||
| DataType::Utf8, | ||
| DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64, | ||
| ) => Self::cast_string_to_int::<i32>(to_type, &array, self.eval_mode), | ||
| ( | ||
| DataType::LargeUtf8, | ||
| DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64, | ||
| ) => Self::cast_string_to_int::<i64>(to_type, &array, self.eval_mode), | ||
| (DataType::Float64, DataType::Utf8) => { | ||
| Self::spark_cast_float64_to_utf8::<i32>(&array, self.eval_mode) | ||
| } | ||
| (DataType::Float64, DataType::LargeUtf8) => { | ||
| Self::spark_cast_float64_to_utf8::<i64>(&array, self.eval_mode) | ||
| } | ||
| (DataType::Float32, DataType::Utf8) => { | ||
| Self::spark_cast_float32_to_utf8::<i32>(&array, self.eval_mode) | ||
| } | ||
| (DataType::Float32, DataType::LargeUtf8) => { | ||
| Self::spark_cast_float32_to_utf8::<i64>(&array, self.eval_mode) | ||
| } | ||
| (DataType::Float32, DataType::Decimal128(precision, scale)) => { | ||
| Self::cast_float32_to_decimal128(&array, *precision, *scale, self.eval_mode) | ||
| } | ||
| (DataType::Float64, DataType::Decimal128(precision, scale)) => { | ||
| Self::cast_float64_to_decimal128(&array, *precision, *scale, self.eval_mode) | ||
| } | ||
| (DataType::Float32, DataType::Int8) | ||
| | (DataType::Float32, DataType::Int16) | ||
|
|
@@ -622,14 +590,91 @@ impl Cast { | |
| self.eval_mode, | ||
| from_type, | ||
| to_type, | ||
| ) | ||
| } | ||
| _ if Self::is_datafusion_spark_compatible(from_type, to_type) => { | ||
| // use DataFusion cast only when we know that it is compatible with Spark | ||
| Ok(cast_with_options(&array, to_type, &CAST_OPTIONS)?) | ||
| } | ||
| _ => { | ||
| // we should never reach this code because the Scala code should be checking | ||
| // for supported cast operations and falling back to Spark for anything that | ||
| // is not yet supported | ||
| Err(CometError::Internal(format!( | ||
| "Native cast invoked for unsupported cast from {from_type:?} to {to_type:?}" | ||
| ))) | ||
| } | ||
| }; | ||
| Ok(spark_cast(cast_result?, from_type, to_type)) | ||
| } | ||
|
|
||
| /// Determines if DataFusion supports the given cast in a way that is | ||
| /// compatible with Spark | ||
| fn is_datafusion_spark_compatible(from_type: &DataType, to_type: &DataType) -> bool { | ||
| if from_type == to_type { | ||
| return true; | ||
| } | ||
| match from_type { | ||
| DataType::Boolean => matches!( | ||
| to_type, | ||
| DataType::Int8 | ||
| | DataType::Int16 | ||
| | DataType::Int32 | ||
| | DataType::Int64 | ||
| | DataType::Float32 | ||
| | DataType::Float64 | ||
| | DataType::Utf8 | ||
|
There was a problem hiding this comment. So right now, there is not There was a problem hiding this comment. The current code says that datafusion is compatible with Spark for all int types -> decimal: DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => matches!(
to_type,
DataType::Boolean
...
| DataType::Decimal128(_, _)However, this is actually not correct since DataFusion does not have overflow checks for int32 and int64 -> decimal and is not compatible with Spark. I will look at removing those. There was a problem hiding this comment. Removing that case causes a test failure: There was a problem hiding this comment. This test relies on a cast that we do not yet support and enables |
||
| ), | ||
| DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => matches!( | ||
| to_type, | ||
| DataType::Boolean | ||
| | DataType::Int8 | ||
| | DataType::Int16 | ||
| | DataType::Int32 | ||
| | DataType::Int64 | ||
| | DataType::Float32 | ||
| | DataType::Float64 | ||
| | DataType::Decimal128(_, _) | ||
| | DataType::Utf8 | ||
| ), | ||
| DataType::Float32 | DataType::Float64 => matches!( | ||
| to_type, | ||
| DataType::Boolean | ||
| | DataType::Int8 | ||
| | DataType::Int16 | ||
| | DataType::Int32 | ||
| | DataType::Int64 | ||
| | DataType::Float32 | ||
| | DataType::Float64 | ||
|
Comment on lines
+645
to
+653
There was a problem hiding this comment. For |
||
| ), | ||
| DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => matches!( | ||
viirya marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| to_type, | ||
| DataType::Int8 | ||
| | DataType::Int16 | ||
| | DataType::Int32 | ||
| | DataType::Int64 | ||
| | DataType::Float32 | ||
| | DataType::Float64 | ||
| | DataType::Decimal128(_, _) | ||
| | DataType::Decimal256(_, _) | ||
| ), | ||
| DataType::Utf8 => matches!(to_type, DataType::Binary), | ||
| DataType::Date32 => matches!(to_type, DataType::Utf8), | ||
| DataType::Timestamp(_, _) => { | ||
| // TODO need to add tests to see if we really do support all | ||
| // timestamp to timestamp conversions | ||
|
There was a problem hiding this comment. I filed #467 for adding timestamp to timestamp tests |
||
| matches!( | ||
| to_type, | ||
| DataType::Int64 | DataType::Date32 | DataType::Utf8 | DataType::Timestamp(_, _) | ||
| ) | ||
| } | ||
| DataType::Binary => { | ||
| // note that this is not completely Spark compatible because | ||
| // DataFusion only supports binary data containing valid UTF-8 strings | ||
| matches!(to_type, DataType::Utf8) | ||
| } | ||
| _ => false, | ||
|
There was a problem hiding this comment. Casting to narrower type like There was a problem hiding this comment. Casting from DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => matches!(
to_type,
DataType::Boolean
| DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::Float32
| DataType::Float64
| DataType::Decimal128(_, _)
| DataType::Utf8
), |
||
| } | ||
| } | ||
|
|
||
| fn cast_string_to_int<OffsetSize: OffsetSizeTrait>( | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We were previously unpacking dictionary-encoded string arrays only for string to int and string to date. I just moved it earlier on so that we don't have to handle it specifically for certain casts from string