Skip to content

Commit 290e3bb

Browse files
hareshkhandygrovealamb
authored
[branch-50]: fix: Add overflow checks to SparkDateAdd/Sub to avoid panics (apache#18013) (apache#18131)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes apache#123` indicates that this PR will close issue apache#123. --> - Related to apache/datafusion-comet#2539 - Related to apache#18072 ## Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> Return errors rather than panicking. ## What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> Co-authored-by: Andy Grove <agrove@apache.org> Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent 5226650 commit 290e3bb

File tree

3 files changed

+54
-12
lines changed

3 files changed

+54
-12
lines changed

datafusion/spark/src/function/datetime/date_add.rs

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use std::sync::Arc;
2121
use arrow::array::ArrayRef;
2222
use arrow::compute;
2323
use arrow::datatypes::{DataType, Date32Type};
24+
use arrow::error::ArrowError;
2425
use datafusion_common::cast::{
2526
as_date32_array, as_int16_array, as_int32_array, as_int8_array,
2627
};
@@ -96,26 +97,38 @@ fn spark_date_add(args: &[ArrayRef]) -> Result<ArrayRef> {
9697
let result = match days_arg.data_type() {
9798
DataType::Int8 => {
9899
let days_array = as_int8_array(days_arg)?;
99-
compute::binary::<_, _, _, Date32Type>(
100+
compute::try_binary::<_, _, _, Date32Type>(
100101
date_array,
101102
days_array,
102-
|date, days| date + days as i32,
103+
|date, days| {
104+
date.checked_add(days as i32).ok_or_else(|| {
105+
ArrowError::ArithmeticOverflow("date_add".to_string())
106+
})
107+
},
103108
)?
104109
}
105110
DataType::Int16 => {
106111
let days_array = as_int16_array(days_arg)?;
107-
compute::binary::<_, _, _, Date32Type>(
112+
compute::try_binary::<_, _, _, Date32Type>(
108113
date_array,
109114
days_array,
110-
|date, days| date + days as i32,
115+
|date, days| {
116+
date.checked_add(days as i32).ok_or_else(|| {
117+
ArrowError::ArithmeticOverflow("date_add".to_string())
118+
})
119+
},
111120
)?
112121
}
113122
DataType::Int32 => {
114123
let days_array = as_int32_array(days_arg)?;
115-
compute::binary::<_, _, _, Date32Type>(
124+
compute::try_binary::<_, _, _, Date32Type>(
116125
date_array,
117126
days_array,
118-
|date, days| date + days,
127+
|date, days| {
128+
date.checked_add(days).ok_or_else(|| {
129+
ArrowError::ArithmeticOverflow("date_add".to_string())
130+
})
131+
},
119132
)?
120133
}
121134
_ => {

datafusion/spark/src/function/datetime/date_sub.rs

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use std::sync::Arc;
2121
use arrow::array::ArrayRef;
2222
use arrow::compute;
2323
use arrow::datatypes::{DataType, Date32Type};
24+
use arrow::error::ArrowError;
2425
use datafusion_common::cast::{
2526
as_date32_array, as_int16_array, as_int32_array, as_int8_array,
2627
};
@@ -90,26 +91,38 @@ fn spark_date_sub(args: &[ArrayRef]) -> Result<ArrayRef> {
9091
let result = match days_arg.data_type() {
9192
DataType::Int8 => {
9293
let days_array = as_int8_array(days_arg)?;
93-
compute::binary::<_, _, _, Date32Type>(
94+
compute::try_binary::<_, _, _, Date32Type>(
9495
date_array,
9596
days_array,
96-
|date, days| date - days as i32,
97+
|date, days| {
98+
date.checked_sub(days as i32).ok_or_else(|| {
99+
ArrowError::ArithmeticOverflow("date_sub".to_string())
100+
})
101+
},
97102
)?
98103
}
99104
DataType::Int16 => {
100105
let days_array = as_int16_array(days_arg)?;
101-
compute::binary::<_, _, _, Date32Type>(
106+
compute::try_binary::<_, _, _, Date32Type>(
102107
date_array,
103108
days_array,
104-
|date, days| date - days as i32,
109+
|date, days| {
110+
date.checked_sub(days as i32).ok_or_else(|| {
111+
ArrowError::ArithmeticOverflow("date_sub".to_string())
112+
})
113+
},
105114
)?
106115
}
107116
DataType::Int32 => {
108117
let days_array = as_int32_array(days_arg)?;
109-
compute::binary::<_, _, _, Date32Type>(
118+
compute::try_binary::<_, _, _, Date32Type>(
110119
date_array,
111120
days_array,
112-
|date, days| date - days,
121+
|date, days| {
122+
date.checked_sub(days).ok_or_else(|| {
123+
ArrowError::ArithmeticOverflow("date_sub".to_string())
124+
})
125+
},
113126
)?
114127
}
115128
_ => {

datafusion/sqllogictest/test_files/spark/datetime/date_add.slt

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,22 @@ SELECT date_sub('2016-07-30'::date, 0::int);
4545
----
4646
2016-07-30
4747

48+
query error DataFusion error: Arrow error: Arithmetic overflow: date_add
49+
SELECT date_add('2016-07-30'::date, 2147483647::int);
50+
51+
query error DataFusion error: Arrow error: Arithmetic overflow: date_sub
52+
SELECT date_sub('1969-01-01'::date, 2147483647::int);
53+
54+
query D
55+
SELECT date_add('2016-07-30'::date, 100000::int);
56+
----
57+
2290-05-15
58+
59+
query D
60+
SELECT date_sub('2016-07-30'::date, 100000::int);
61+
----
62+
1742-10-15
63+
4864
# Test with negative day values (should subtract days)
4965
query D
5066
SELECT date_add('2016-07-30'::date, -5::int);

0 commit comments

Comments
 (0)