Skip to content

Commit 608ee58

Browse files
alambtimsaucer
andauthored
Correct return type for initcap scalar function with utf8view (#13909) (#13934)
* Set utf8view as return type when input type is the same * Verify that the returned type from call to scalar function matches the return type specified in the return_type function * Match return type to utf8view Co-authored-by: Tim Saucer <timsaucer@gmail.com>
1 parent 073a3b1 commit 608ee58

File tree

2 files changed

+12
-7
lines changed

2 files changed

+12
-7
lines changed

datafusion/functions/src/unicode/initcap.rs

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,11 @@ impl ScalarUDFImpl for InitcapFunc {
6363
}
6464

6565
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
66-
utf8_to_str_type(&arg_types[0], "initcap")
66+
if let DataType::Utf8View = arg_types[0] {
67+
Ok(DataType::Utf8View)
68+
} else {
69+
utf8_to_str_type(&arg_types[0], "initcap")
70+
}
6771
}
6872

6973
fn invoke_batch(
@@ -188,7 +192,7 @@ mod tests {
188192
use crate::unicode::initcap::InitcapFunc;
189193
use crate::utils::test::test_function;
190194
use arrow::array::{Array, StringArray, StringViewArray};
191-
use arrow::datatypes::DataType::Utf8;
195+
use arrow::datatypes::DataType::{Utf8, Utf8View};
192196
use datafusion_common::{Result, ScalarValue};
193197
use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
194198

@@ -247,7 +251,7 @@ mod tests {
247251
)))],
248252
Ok(Some("Hi Thomas")),
249253
&str,
250-
Utf8,
254+
Utf8View,
251255
StringViewArray
252256
);
253257
test_function!(
@@ -257,7 +261,7 @@ mod tests {
257261
)))],
258262
Ok(Some("Hi Thomas With M0re Than 12 Chars")),
259263
&str,
260-
Utf8,
264+
Utf8View,
261265
StringViewArray
262266
);
263267
test_function!(
@@ -270,7 +274,7 @@ mod tests {
270274
"Đẹp Đẽ Êm Ả Ñandú Árbol Олег Иванович Íslensku Þjóðarinnar Ελληνική"
271275
)),
272276
&str,
273-
Utf8,
277+
Utf8View,
274278
StringViewArray
275279
);
276280
test_function!(
@@ -280,15 +284,15 @@ mod tests {
280284
)))],
281285
Ok(Some("")),
282286
&str,
283-
Utf8,
287+
Utf8View,
284288
StringViewArray
285289
);
286290
test_function!(
287291
InitcapFunc::new(),
288292
vec![ColumnarValue::Scalar(ScalarValue::Utf8View(None))],
289293
Ok(None),
290294
&str,
291-
Utf8,
295+
Utf8View,
292296
StringViewArray
293297
);
294298

datafusion/functions/src/utils.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ pub mod test {
154154

155155
let result = result.unwrap().to_array(cardinality).expect("Failed to convert to array");
156156
let result = result.as_any().downcast_ref::<$ARRAY_TYPE>().expect("Failed to convert to type");
157+
assert_eq!(result.data_type(), &$EXPECTED_DATA_TYPE);
157158

158159
// value is correct
159160
match expected {

0 commit comments

Comments
 (0)