This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 10fe3c1658 Implement native support StringView for find in set (#11970)
10fe3c1658 is described below
commit 10fe3c165862e1157224ef2fafd6917782fb255f
Author: Chojan Shang <[email protected]>
AuthorDate: Wed Aug 14 04:16:57 2024 -0700
Implement native support StringView for find in set (#11970)
* Implement native support StringView for find in set
Signed-off-by: Chojan Shang <[email protected]>
* Add more tests
Signed-off-by: Chojan Shang <[email protected]>
* Minor update
---------
Signed-off-by: Chojan Shang <[email protected]>
---
datafusion/functions/src/unicode/find_in_set.rs | 67 +++++++++++++---------
datafusion/sqllogictest/test_files/functions.slt | 19 +++++-
datafusion/sqllogictest/test_files/string_view.slt | 14 +++--
3 files changed, 69 insertions(+), 31 deletions(-)
diff --git a/datafusion/functions/src/unicode/find_in_set.rs
b/datafusion/functions/src/unicode/find_in_set.rs
index 7c864bc191..41a2b9d9e7 100644
--- a/datafusion/functions/src/unicode/find_in_set.rs
+++ b/datafusion/functions/src/unicode/find_in_set.rs
@@ -19,11 +19,11 @@ use std::any::Any;
use std::sync::Arc;
use arrow::array::{
- ArrayRef, ArrowPrimitiveType, GenericStringArray, OffsetSizeTrait,
PrimitiveArray,
+ ArrayAccessor, ArrayIter, ArrayRef, ArrowPrimitiveType, AsArray,
OffsetSizeTrait,
+ PrimitiveArray,
};
use arrow::datatypes::{ArrowNativeType, DataType, Int32Type, Int64Type};
-use datafusion_common::cast::as_generic_string_array;
use datafusion_common::{exec_err, Result};
use datafusion_expr::TypeSignature::Exact;
use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
@@ -46,7 +46,11 @@ impl FindInSetFunc {
use DataType::*;
Self {
signature: Signature::one_of(
- vec![Exact(vec![Utf8, Utf8]), Exact(vec![LargeUtf8,
LargeUtf8])],
+ vec![
+ Exact(vec![Utf8View, Utf8View]),
+ Exact(vec![Utf8, Utf8]),
+ Exact(vec![LargeUtf8, LargeUtf8]),
+ ],
Volatility::Immutable,
),
}
@@ -71,41 +75,52 @@ impl ScalarUDFImpl for FindInSetFunc {
}
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- match args[0].data_type() {
- DataType::Utf8 => {
- make_scalar_function(find_in_set::<Int32Type>, vec![])(args)
- }
- DataType::LargeUtf8 => {
- make_scalar_function(find_in_set::<Int64Type>, vec![])(args)
- }
- other => {
- exec_err!("Unsupported data type {other:?} for function
find_in_set")
- }
- }
+ make_scalar_function(find_in_set, vec![])(args)
}
}
///Returns a value in the range of 1 to N if the string str is in the string
list strlist consisting of N substrings
///A string list is a string composed of substrings separated by , characters.
-pub fn find_in_set<T: ArrowPrimitiveType>(args: &[ArrayRef]) ->
Result<ArrayRef>
-where
- T::Native: OffsetSizeTrait,
-{
+fn find_in_set(args: &[ArrayRef]) -> Result<ArrayRef> {
if args.len() != 2 {
return exec_err!(
"find_in_set was called with {} arguments. It requires 2.",
args.len()
);
}
+ match args[0].data_type() {
+ DataType::Utf8 => {
+ let string_array = args[0].as_string::<i32>();
+ let str_list_array = args[1].as_string::<i32>();
+ find_in_set_general::<Int32Type, _>(string_array, str_list_array)
+ }
+ DataType::LargeUtf8 => {
+ let string_array = args[0].as_string::<i64>();
+ let str_list_array = args[1].as_string::<i64>();
+ find_in_set_general::<Int64Type, _>(string_array, str_list_array)
+ }
+ DataType::Utf8View => {
+ let string_array = args[0].as_string_view();
+ let str_list_array = args[1].as_string_view();
+ find_in_set_general::<Int32Type, _>(string_array, str_list_array)
+ }
+ other => {
+ exec_err!("Unsupported data type {other:?} for function
find_in_set")
+ }
+ }
+}
- let str_array: &GenericStringArray<T::Native> =
- as_generic_string_array::<T::Native>(&args[0])?;
- let str_list_array: &GenericStringArray<T::Native> =
- as_generic_string_array::<T::Native>(&args[1])?;
-
- let result = str_array
- .iter()
- .zip(str_list_array.iter())
+pub fn find_in_set_general<'a, T: ArrowPrimitiveType, V: ArrayAccessor<Item =
&'a str>>(
+ string_array: V,
+ str_list_array: V,
+) -> Result<ArrayRef>
+where
+ T::Native: OffsetSizeTrait,
+{
+ let string_iter = ArrayIter::new(string_array);
+ let str_list_iter = ArrayIter::new(str_list_array);
+ let result = string_iter
+ .zip(str_list_iter)
.map(|(string, str_list)| match (string, str_list) {
(Some(string), Some(str_list)) => {
let mut res = 0;
diff --git a/datafusion/sqllogictest/test_files/functions.slt
b/datafusion/sqllogictest/test_files/functions.slt
index bea3016a21..f728942b38 100644
--- a/datafusion/sqllogictest/test_files/functions.slt
+++ b/datafusion/sqllogictest/test_files/functions.slt
@@ -1092,7 +1092,7 @@ docs.apache.com docs com
community.influxdata.com community com
arrow.apache.org arrow org
-
+# find_in_set tests
query I
SELECT find_in_set('b', 'a,b,c,d')
----
@@ -1136,6 +1136,23 @@ SELECT find_in_set(NULL, NULL)
----
NULL
+# find_in_set tests with utf8view
+query I
+SELECT find_in_set(arrow_cast('b', 'Utf8View'), 'a,b,c,d')
+----
+2
+
+
+query I
+SELECT find_in_set('a', arrow_cast('a,b,c,d,a', 'Utf8View'))
+----
+1
+
+query I
+SELECT find_in_set(arrow_cast('', 'Utf8View'), arrow_cast('a,b,c,d,a',
'Utf8View'))
+----
+0
+
# Verify that multiple calls to volatile functions like `random()` are not
combined / optimized away
query B
SELECT r FROM (SELECT r1 == r2 r, r1, r2 FROM (SELECT random()+1 r1,
random()+1 r2) WHERE r1 > 0 AND r2 > 0)
diff --git a/datafusion/sqllogictest/test_files/string_view.slt
b/datafusion/sqllogictest/test_files/string_view.slt
index 0a9b73babb..264f85ff84 100644
--- a/datafusion/sqllogictest/test_files/string_view.slt
+++ b/datafusion/sqllogictest/test_files/string_view.slt
@@ -992,18 +992,24 @@ logical_plan
02)--TableScan: test projection=[column1_utf8view]
## Ensure no casts for FIND_IN_SET
-## TODO file ticket
query TT
EXPLAIN SELECT
FIND_IN_SET(column1_utf8view, 'a,b,c,d') as c
FROM test;
----
logical_plan
-01)Projection: find_in_set(CAST(test.column1_utf8view AS Utf8),
Utf8("a,b,c,d")) AS c
+01)Projection: find_in_set(test.column1_utf8view, Utf8View("a,b,c,d")) AS c
02)--TableScan: test projection=[column1_utf8view]
-
-
+query I
+SELECT
+ FIND_IN_SET(column1_utf8view, 'a,b,c,d') as c
+FROM test;
+----
+0
+0
+0
+NULL
statement ok
drop table test;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
