Skip to content

Commit 1f2f02f

Browse files
authored
Fix strpos invocation with dictionary and null (#12712)
In 1b3608d `strpos` signature was modified to indicate it supports dictionary as input argument, but the invoke method doesn't support them.
1 parent 0b2b4fb commit 1f2f02f

File tree

3 files changed

+28
-22
lines changed

3 files changed

+28
-22
lines changed

datafusion/functions/src/unicode/strpos.rs

Lines changed: 15 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ use arrow::datatypes::{ArrowNativeType, DataType, Int32Type, Int64Type};
2323

2424
use crate::string::common::StringArrayType;
2525
use crate::utils::{make_scalar_function, utf8_to_int_type};
26-
use datafusion_common::{exec_err, plan_err, Result};
26+
use datafusion_common::{exec_err, Result};
27+
use datafusion_expr::TypeSignature::Exact;
2728
use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
2829

2930
#[derive(Debug)]
@@ -40,8 +41,20 @@ impl Default for StrposFunc {
4041

4142
impl StrposFunc {
4243
pub fn new() -> Self {
44+
use DataType::*;
4345
Self {
44-
signature: Signature::user_defined(Volatility::Immutable),
46+
signature: Signature::one_of(
47+
vec![
48+
Exact(vec![Utf8, Utf8]),
49+
Exact(vec![Utf8, LargeUtf8]),
50+
Exact(vec![LargeUtf8, Utf8]),
51+
Exact(vec![LargeUtf8, LargeUtf8]),
52+
Exact(vec![Utf8View, Utf8View]),
53+
Exact(vec![Utf8View, Utf8]),
54+
Exact(vec![Utf8View, LargeUtf8]),
55+
],
56+
Volatility::Immutable,
57+
),
4558
aliases: vec![String::from("instr"), String::from("position")],
4659
}
4760
}
@@ -71,25 +84,6 @@ impl ScalarUDFImpl for StrposFunc {
7184
fn aliases(&self) -> &[String] {
7285
&self.aliases
7386
}
74-
75-
fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
76-
match arg_types {
77-
[first, second ] => {
78-
match (first, second) {
79-
(DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8, DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8) => Ok(arg_types.to_vec()),
80-
(DataType::Null, DataType::Null) => Ok(vec![DataType::Utf8, DataType::Utf8]),
81-
(DataType::Null, _) => Ok(vec![DataType::Utf8, second.to_owned()]),
82-
(_, DataType::Null) => Ok(vec![first.to_owned(), DataType::Utf8]),
83-
(DataType::Dictionary(_, value_type), DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8) => match **value_type {
84-
DataType::LargeUtf8 | DataType::Utf8View | DataType::Utf8 | DataType::Null | DataType::Binary => Ok(vec![*value_type.clone(), second.to_owned()]),
85-
_ => plan_err!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}.", **value_type),
86-
},
87-
_ => plan_err!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}.", arg_types)
88-
}
89-
},
90-
_ => plan_err!("The STRPOS/INSTR/POSITION function can only accept strings, but got {:?}", arg_types)
91-
}
92-
}
9387
}
9488

9589
fn strpos(args: &[ArrayRef]) -> Result<ArrayRef> {

datafusion/sqllogictest/test_files/functions.slt

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -553,6 +553,16 @@ SELECT strpos(arrow_cast('helloworld', 'Dictionary(Int32, Utf8)'), 'world')
553553
----
554554
6
555555

556+
query I
557+
SELECT strpos('helloworld', NULL)
558+
----
559+
NULL
560+
561+
query I
562+
SELECT strpos(arrow_cast('helloworld', 'Dictionary(Int32, Utf8)'), NULL)
563+
----
564+
NULL
565+
556566
statement ok
557567
CREATE TABLE products (
558568
product_id INT PRIMARY KEY,

datafusion/sqllogictest/test_files/scalar.slt

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1907,8 +1907,10 @@ select position('' in '')
19071907
1
19081908

19091909

1910-
query error POSITION function can only accept strings
1910+
query I
19111911
select position(1 in 1)
1912+
----
1913+
1
19121914

19131915

19141916
query I

0 commit comments

Comments
 (0)