Skip to content

Always use StringViewArray as output of substr #14498

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

Merged
merged 5 commits into from
Feb 8, 2025
Merged

Always use StringViewArray as output of substr #14498

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5d47437
change substr return type to utf8view (bin/sqllogicaltests.rs to fix)
Kev1n8 Jan 25, 2025
08210dc
fix sqllogictest/string_view
Kev1n8 Feb 5, 2025
bca6153
fix clippy
Kev1n8 Feb 5, 2025
9a2a4b8
fix tcph benchmark result
Kev1n8 Feb 5, 2025
161454e
Merge branch 'main' into substr-always-output-utf8view
Kev1n8 Feb 7, 2025
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
118 changes: 57 additions & 61 deletions datafusion/functions/src/unicode/substr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,10 @@ use std::any::Any;
use std::sync::Arc;

use crate::strings::make_and_append_view;
use crate::utils::{make_scalar_function, utf8_to_str_type};
use crate::utils::make_scalar_function;
use arrow::array::{
Array, ArrayIter, ArrayRef, AsArray, GenericStringBuilder, Int64Array,
NullBufferBuilder, OffsetSizeTrait, StringArrayType, StringViewArray,
Array, ArrayIter, ArrayRef, AsArray, Int64Array, NullBufferBuilder, StringArrayType,
StringViewArray, StringViewBuilder,
};
use arrow::buffer::ScalarBuffer;
use arrow::datatypes::DataType;
Expand Down Expand Up @@ -90,12 +90,9 @@ impl ScalarUDFImpl for SubstrFunc {
&self.signature
}

fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
if arg_types[0] == DataType::Utf8View {
Ok(DataType::Utf8View)
} else {
utf8_to_str_type(&arg_types[0], "substr")
}
// `SubstrFunc` always generates `Utf8View` output for its efficiency.
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Utf8View)
}

fn invoke_batch(
Expand Down Expand Up @@ -189,11 +186,11 @@ pub fn substr(args: &[ArrayRef]) -> Result<ArrayRef> {
match args[0].data_type() {
DataType::Utf8 => {
let string_array = args[0].as_string::<i32>();
string_substr::<_, i32>(string_array, &args[1..])
string_substr::<_>(string_array, &args[1..])
}
DataType::LargeUtf8 => {
let string_array = args[0].as_string::<i64>();
string_substr::<_, i64>(string_array, &args[1..])
string_substr::<_>(string_array, &args[1..])
}
DataType::Utf8View => {
let string_array = args[0].as_string_view();
Expand Down Expand Up @@ -429,10 +426,9 @@ fn string_view_substr(
}
}

fn string_substr<'a, V, T>(string_array: V, args: &[ArrayRef]) -> Result<ArrayRef>
fn string_substr<'a, V>(string_array: V, args: &[ArrayRef]) -> Result<ArrayRef>
where
V: StringArrayType<'a>,
T: OffsetSizeTrait,
{
let start_array = as_int64_array(&args[0])?;
let count_array_opt = if args.len() == 2 {
Expand All @@ -447,7 +443,7 @@ where
match args.len() {
1 => {
let iter = ArrayIter::new(string_array);
let mut result_builder = GenericStringBuilder::<T>::new();
let mut result_builder = StringViewBuilder::new();
for (string, start) in iter.zip(start_array.iter()) {
match (string, start) {
(Some(string), Some(start)) => {
Expand All @@ -470,7 +466,7 @@ where
2 => {
let iter = ArrayIter::new(string_array);
let count_array = count_array_opt.unwrap();
let mut result_builder = GenericStringBuilder::<T>::new();
let mut result_builder = StringViewBuilder::new();

for ((string, start), count) in
iter.zip(start_array.iter()).zip(count_array.iter())
Expand Down Expand Up @@ -512,8 +508,8 @@ where

#[cfg(test)]
mod tests {
use arrow::array::{Array, StringArray, StringViewArray};
use arrow::datatypes::DataType::{Utf8, Utf8View};
use arrow::array::{Array, StringViewArray};
use arrow::datatypes::DataType::Utf8View;

use datafusion_common::{exec_err, Result, ScalarValue};
use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
Expand Down Expand Up @@ -623,8 +619,8 @@ mod tests {
],
Ok(Some("alphabet")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -634,8 +630,8 @@ mod tests {
],
Ok(Some("ésoj")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -645,8 +641,8 @@ mod tests {
],
Ok(Some("joséésoj")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -656,8 +652,8 @@ mod tests {
],
Ok(Some("alphabet")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -667,8 +663,8 @@ mod tests {
],
Ok(Some("lphabet")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -678,8 +674,8 @@ mod tests {
],
Ok(Some("phabet")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -689,8 +685,8 @@ mod tests {
],
Ok(Some("alphabet")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -700,8 +696,8 @@ mod tests {
],
Ok(Some("")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -711,8 +707,8 @@ mod tests {
],
Ok(None),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -723,8 +719,8 @@ mod tests {
],
Ok(Some("ph")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -735,8 +731,8 @@ mod tests {
],
Ok(Some("phabet")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -747,8 +743,8 @@ mod tests {
],
Ok(Some("alph")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
// starting from 5 (10 + -5)
test_function!(
Expand All @@ -760,8 +756,8 @@ mod tests {
],
Ok(Some("alph")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
// starting from -1 (4 + -5)
test_function!(
Expand All @@ -773,8 +769,8 @@ mod tests {
],
Ok(Some("")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
// starting from 0 (5 + -5)
test_function!(
Expand All @@ -786,8 +782,8 @@ mod tests {
],
Ok(Some("")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -798,8 +794,8 @@ mod tests {
],
Ok(None),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -810,8 +806,8 @@ mod tests {
],
Ok(None),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -822,8 +818,8 @@ mod tests {
],
exec_err!("negative substring length not allowed: substr(<str>, 1, -1)"),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -834,8 +830,8 @@ mod tests {
],
Ok(Some("és")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
#[cfg(not(feature = "unicode_expressions"))]
test_function!(
Expand All @@ -848,8 +844,8 @@ mod tests {
"function substr requires compilation with feature flag: unicode_expressions."
),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -859,8 +855,8 @@ mod tests {
],
Ok(Some("abc")),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);
test_function!(
SubstrFunc::new(),
Expand All @@ -871,8 +867,8 @@ mod tests {
],
exec_err!("negative overflow when calculating skip value"),
&str,
Utf8,
StringArray
Utf8View,
StringViewArray
);

Ok(())
Expand Down
2 changes: 1 addition & 1 deletion datafusion/sqllogictest/test_files/string/string_view.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,7 @@ EXPLAIN SELECT
FROM test;
----
logical_plan
01)Projection: starts_with(test.column1_utf8, substr(test.column1_utf8, Int64(1), Int64(2))) AS c1, starts_with(test.column1_large_utf8, substr(test.column1_large_utf8, Int64(1), Int64(2))) AS c2, starts_with(test.column1_utf8view, substr(test.column1_utf8view, Int64(1), Int64(2))) AS c3
01)Projection: starts_with(CAST(test.column1_utf8 AS Utf8View), substr(test.column1_utf8, Int64(1), Int64(2))) AS c1, starts_with(CAST(test.column1_large_utf8 AS Utf8View), substr(test.column1_large_utf8, Int64(1), Int64(2))) AS c2, starts_with(test.column1_utf8view, substr(test.column1_utf8view, Int64(1), Int64(2))) AS c3
02)--TableScan: test projection=[column1_utf8, column1_large_utf8, column1_utf8view]

query BBB
Expand Down
12 changes: 6 additions & 6 deletions datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,15 +64,15 @@ logical_plan
06)----------Inner Join: Filter: CAST(customer.c_acctbal AS Decimal128(19, 6)) > __scalar_sq_2.avg(customer.c_acctbal)
07)------------Projection: customer.c_phone, customer.c_acctbal
08)--------------LeftAnti Join: customer.c_custkey = __correlated_sq_1.o_custkey
09)----------------Filter: substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8("13"), Utf8("31"), Utf8("23"), Utf8("29"), Utf8("30"), Utf8("18"), Utf8("17")])
10)------------------TableScan: customer projection=[c_custkey, c_phone, c_acctbal], partial_filters=[substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8("13"), Utf8("31"), Utf8("23"), Utf8("29"), Utf8("30"), Utf8("18"), Utf8("17")])]
09)----------------Filter: substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")])
10)------------------TableScan: customer projection=[c_custkey, c_phone, c_acctbal], partial_filters=[substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")])]
11)----------------SubqueryAlias: __correlated_sq_1
12)------------------TableScan: orders projection=[o_custkey]
13)------------SubqueryAlias: __scalar_sq_2
14)--------------Aggregate: groupBy=[[]], aggr=[[avg(customer.c_acctbal)]]
15)----------------Projection: customer.c_acctbal
16)------------------Filter: customer.c_acctbal > Decimal128(Some(0),15,2) AND substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8("13"), Utf8("31"), Utf8("23"), Utf8("29"), Utf8("30"), Utf8("18"), Utf8("17")])
17)--------------------TableScan: customer projection=[c_phone, c_acctbal], partial_filters=[customer.c_acctbal > Decimal128(Some(0),15,2), substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8("13"), Utf8("31"), Utf8("23"), Utf8("29"), Utf8("30"), Utf8("18"), Utf8("17")])]
16)------------------Filter: customer.c_acctbal > Decimal128(Some(0),15,2) AND substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")])
17)--------------------TableScan: customer projection=[c_phone, c_acctbal], partial_filters=[customer.c_acctbal > Decimal128(Some(0),15,2), substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")])]
physical_plan
01)SortPreservingMergeExec: [cntrycode@0 ASC NULLS LAST]
02)--SortExec: expr=[cntrycode@0 ASC NULLS LAST], preserve_partitioning=[true]
Expand All @@ -90,7 +90,7 @@ physical_plan
14)--------------------------CoalesceBatchesExec: target_batch_size=8192
15)----------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4
16)------------------------------CoalesceBatchesExec: target_batch_size=8192
17)--------------------------------FilterExec: Use substr(c_phone@1, 1, 2) IN (SET) ([Literal { value: Utf8("13") }, Literal { value: Utf8("31") }, Literal { value: Utf8("23") }, Literal { value: Utf8("29") }, Literal { value: Utf8("30") }, Literal { value: Utf8("18") }, Literal { value: Utf8("17") }])
17)--------------------------------FilterExec: substr(c_phone@1, 1, 2) IN ([Literal { value: Utf8View("13") }, Literal { value: Utf8View("31") }, Literal { value: Utf8View("23") }, Literal { value: Utf8View("29") }, Literal { value: Utf8View("30") }, Literal { value: Utf8View("18") }, Literal { value: Utf8View("17") }])
18)----------------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
19)------------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_phone, c_acctbal], file_type=csv, has_header=false
20)--------------------------CoalesceBatchesExec: target_batch_size=8192
Expand All @@ -100,6 +100,6 @@ physical_plan
24)----------------------CoalescePartitionsExec
25)------------------------AggregateExec: mode=Partial, gby=[], aggr=[avg(customer.c_acctbal)]
26)--------------------------CoalesceBatchesExec: target_batch_size=8192
27)----------------------------FilterExec: c_acctbal@1 > Some(0),15,2 AND Use substr(c_phone@0, 1, 2) IN (SET) ([Literal { value: Utf8("13") }, Literal { value: Utf8("31") }, Literal { value: Utf8("23") }, Literal { value: Utf8("29") }, Literal { value: Utf8("30") }, Literal { value: Utf8("18") }, Literal { value: Utf8("17") }]), projection=[c_acctbal@1]
27)----------------------------FilterExec: c_acctbal@1 > Some(0),15,2 AND substr(c_phone@0, 1, 2) IN ([Literal { value: Utf8View("13") }, Literal { value: Utf8View("31") }, Literal { value: Utf8View("23") }, Literal { value: Utf8View("29") }, Literal { value: Utf8View("30") }, Literal { value: Utf8View("18") }, Literal { value: Utf8View("17") }]), projection=[c_acctbal@1]
28)------------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
29)--------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_phone, c_acctbal], file_type=csv, has_header=false