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

fix: Only delegate to DataFusion cast when we know that it is compatible with Spark #461

Merged
merged 16 commits into from
May 25, 2024
Merged

fix: Only delegate to DataFusion cast when we know that it is compatible with Spark #461

Changes from 7 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
3a1c387
only delegate to DataFusion cast when we know that it is compatible w…
andygrove May 23, 2024
60354bd
add more supported casts
andygrove May 23, 2024
35585aa
improve support for dictionary-encoded string arrays
andygrove May 23, 2024
dc4f99c
clippy
andygrove May 23, 2024
f8a6f94
Merge remote-tracking branch 'apache/main' into cast-df-compat
andygrove May 23, 2024
2b21917
fix merge conflict
andygrove May 23, 2024
c2c2546
fix a regression
andygrove May 23, 2024
7f1951a
fix a regression
andygrove May 23, 2024
c0f913b
fix a regression
andygrove May 23, 2024
34edf00
fix regression
andygrove May 23, 2024
54043ae
fix regression
andygrove May 23, 2024
2f22e06
fix regression
andygrove May 23, 2024
bc18fae
remove TODO comment now that issue has been filed
andygrove May 24, 2024
340e000
remove cast int32/int64 -> decimal from datafusion compatible list
andygrove May 25, 2024
7d804f5
Revert "remove cast int32/int64 -> decimal from datafusion compatible…
andygrove May 25, 2024
3fb2258
add comment
andygrove May 25, 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
166 changes: 99 additions & 67 deletions core/src/execution/datafusion/expressions/cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Comment on lines +513 to +522
Copy link
Member Author

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

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)?
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)?
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)?
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::Dictionary(key_type, value_type), DataType::Date32)
if key_type.as_ref() == &DataType::Int32
&& (value_type.as_ref() == &DataType::Utf8
|| value_type.as_ref() == &DataType::LargeUtf8) =>
{
match value_type.as_ref() {
DataType::Utf8 => {
let unpacked_array =
cast_with_options(&array, &DataType::Utf8, &CAST_OPTIONS)?;
Self::cast_string_to_date(&unpacked_array, to_type, self.eval_mode)?
}
DataType::LargeUtf8 => {
let unpacked_array =
cast_with_options(&array, &DataType::LargeUtf8, &CAST_OPTIONS)?;
Self::cast_string_to_date(&unpacked_array, to_type, self.eval_mode)?
}
dt => unreachable!(
"{}",
format!("invalid value type {dt} for dictionary-encoded string array")
),
}
Self::cast_string_to_date(&array, to_type, self.eval_mode)
}
(DataType::Int64, DataType::Int32)
| (DataType::Int64, DataType::Int16)
Expand All @@ -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)?
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)?,
) => 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::Dictionary(key_type, value_type),
DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64,
) if key_type.as_ref() == &DataType::Int32
&& (value_type.as_ref() == &DataType::Utf8
|| value_type.as_ref() == &DataType::LargeUtf8) =>
{
// 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
match value_type.as_ref() {
DataType::Utf8 => {
let unpacked_array =
cast_with_options(&array, &DataType::Utf8, &CAST_OPTIONS)?;
Self::cast_string_to_int::<i32>(to_type, &unpacked_array, self.eval_mode)?
}
DataType::LargeUtf8 => {
let unpacked_array =
cast_with_options(&array, &DataType::LargeUtf8, &CAST_OPTIONS)?;
Self::cast_string_to_int::<i64>(to_type, &unpacked_array, self.eval_mode)?
}
dt => unreachable!(
"{}",
format!("invalid value type {dt} for dictionary-encoded string array")
),
}
}
) => 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)?
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)?
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)?
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)?
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)?
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)?
Self::cast_float64_to_decimal128(&array, *precision, *scale, self.eval_mode)
}
(DataType::Float32, DataType::Int8)
| (DataType::Float32, DataType::Int16)
Expand All @@ -622,14 +590,78 @@ 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)?)
}
_ => {
// when we have no Spark-specific casting we delegate to DataFusion
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))
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 {
match from_type {
DataType::Boolean => matches!(
to_type,
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::Float32
| DataType::Float64
| DataType::Utf8
Copy link
Contributor

Choose a reason for hiding this comment

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

So right now, there is not Int8 to Decimal128 cast supported, looks like?

Copy link
Member Author

@andygrove andygrove May 24, 2024 •
edited
Loading

Choose a reason for hiding this comment

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

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.

Copy link
Member Author

Choose a reason for hiding this comment

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

Removing that case causes a test failure:

- scalar subquery *** FAILED *** (8 seconds, 253 milliseconds)

  Cause: java.util.concurrent.ExecutionException: org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 410.0 failed 1 times, most recent failure: Lost task 0.0 in stage 410.0 (TID 1286) (192.168.64.23 executor driver): org.apache.comet.CometNativeException: Execution error: Comet Internal Error: Native cast invoked for unsupported cast from Int32 to Decimal128(38, 10)

Copy link
Member Author

Choose a reason for hiding this comment

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

This test relies on a cast that we do not yet support and enables COMET_CAST_ALLOW_INCOMPATIBLE to allow it. I will revert the last change and add a comment about this

),
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
Copy link
Contributor

Choose a reason for hiding this comment

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

For Float32/64 to Int8/16/32/64, I saw spark_cast_nonintegral_numeric_to_integral covers them above.
Is this for the case self.eval_mode == EvalMode::Try?

Copy link
Member Author

Choose a reason for hiding this comment

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

Yes, that is correct.

),
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(_, _) => {
matches!(to_type, DataType::Int64 | DataType::Date32 | DataType::Utf8)
}
_ => false,
Copy link
Contributor

Choose a reason for hiding this comment

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

Casting to narrower type like Int64 to Int32 cases are not supported when self.eval_mode == EvalMode::Try?

Copy link
Member Author

Choose a reason for hiding this comment

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

Casting from Int64 to Int32 for Try is covered here:

           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>(
Expand Down
Loading