This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24100-c8b4aeaad90d92071ae7b1e88637a308cf749a4c in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 6b6cc0939ef599235e0356402492b4ab2c738e55 Author: linfeng <[email protected]> AuthorDate: Sun Aug 9 10:12:10 2026 +0800 perf: preserve dictionary encoding for `btrim`, `ltrim`, and `rtrim` (#24100) ## 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 #123` indicates that this PR will close issue #123. --> - Followup to #23930 - Related to #19458 and #20935 ## 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. Please explain the problem you are trying to solve in terms of the user-visible behavior, rather than the implementation. For example, "The code in `foo.rs` doesn't handle nulls" is a symptom of the implementation. "COUNT(DISTINCT) returns wrong results when the column contains nulls" is the user-visible problem. --> Previously, coercion materialized dictionary-encoded inputs for `btrim`, `ltrim`, and `rtrim`. This lost the encoding and evaluated the function for every row instead of once per dictionary value entry. This PR extends dictionary preservation to these trim functions. ## 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. --> - Preserve dictionary encoding for `btrim`, `ltrim`, and `rtrim`. - Add tests and benchmarks. ## 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)? --> Yes, covered by slt ## 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. --> These functions now preserve dictionary encoding in their output. ### Benchmarks ``` group optimize main ----- -------------- --------- dictionary_encoding/string/cardinality_10/btrim 1.00 286.8±13.36ns ? ?/sec 137.69 39.5±4.56µs ? ?/sec dictionary_encoding/string/cardinality_10/ltrim 1.00 282.3±4.68ns ? ?/sec 151.48 42.8±8.96µs ? ?/sec dictionary_encoding/string/cardinality_10/rtrim 1.00 278.8±13.01ns ? ?/sec 119.91 33.4±3.85µs ? ?/sec dictionary_encoding/string/cardinality_100/btrim 1.00 718.4±4.53ns ? ?/sec 51.89 37.3±0.92µs ? ?/sec dictionary_encoding/string/cardinality_100/ltrim 1.00 696.4±6.33ns ? ?/sec 54.02 37.6±1.68µs ? ?/sec dictionary_encoding/string/cardinality_100/rtrim 1.00 659.2±6.91ns ? ?/sec 55.33 36.5±6.84µs ? ?/sec dictionary_encoding/string/cardinality_1000/btrim 1.00 4.5±0.03µs ? ?/sec 9.71 43.3±8.35µs ? ?/sec dictionary_encoding/string/cardinality_1000/ltrim 1.00 4.2±0.04µs ? ?/sec 9.06 37.8±4.46µs ? ?/sec dictionary_encoding/string/cardinality_1000/rtrim 1.00 3.9±0.05µs ? ?/sec 8.69 33.9±1.96µs ? ?/sec dictionary_encoding/string/cardinality_8192/btrim 1.01 37.4±1.64µs ? ?/sec 1.00 37.1±0.68µs ? ?/sec dictionary_encoding/string/cardinality_8192/ltrim 1.00 34.7±1.75µs ? ?/sec 1.01 35.1±1.55µs ? ?/sec dictionary_encoding/string/cardinality_8192/rtrim 1.00 31.9±1.18µs ? ?/sec 1.17 37.3±5.68µs ? ?/sec ``` --- .../functions/benches/dictionary_encoding.rs | 7 +- datafusion/functions/src/string/btrim.rs | 55 +++++++-------- datafusion/functions/src/string/ltrim.rs | 47 +++++++------ datafusion/functions/src/string/rtrim.rs | 47 +++++++------ datafusion/sqllogictest/test_files/functions.slt | 81 ++++++++++++++++++---- 5 files changed, 148 insertions(+), 89 deletions(-) diff --git a/datafusion/functions/benches/dictionary_encoding.rs b/datafusion/functions/benches/dictionary_encoding.rs index 05541fc10e..4ba04a4940 100644 --- a/datafusion/functions/benches/dictionary_encoding.rs +++ b/datafusion/functions/benches/dictionary_encoding.rs @@ -24,7 +24,7 @@ use arrow::datatypes::{Field, Int32Type}; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_common::config::ConfigOptions; use datafusion_expr::type_coercion::functions::fields_with_udf; -use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; const NUM_ROWS: usize = 8_192; const DICTIONARY_CARDINALITIES: [usize; 4] = [10, 100, 1_000, 8_192]; @@ -42,16 +42,19 @@ fn create_string_dictionary(cardinality: usize) -> ArrayRef { } fn benchmark_dictionary_string_udfs(c: &mut Criterion) { - let udfs: [(&str, Arc<ScalarUDF>); 6] = [ + let udfs = [ ("ascii", datafusion_functions::string::ascii()), ("bit_length", datafusion_functions::string::bit_length()), + ("btrim", datafusion_functions::string::btrim()), ( "character_length", datafusion_functions::unicode::character_length(), ), ("initcap", datafusion_functions::unicode::initcap()), + ("ltrim", datafusion_functions::string::ltrim()), ("octet_length", datafusion_functions::string::octet_length()), ("reverse", datafusion_functions::unicode::reverse()), + ("rtrim", datafusion_functions::string::rtrim()), ]; let config_options = Arc::new(ConfigOptions::default()); diff --git a/datafusion/functions/src/string/btrim.rs b/datafusion/functions/src/string/btrim.rs index 279f444d9f..82e1f7d2c7 100644 --- a/datafusion/functions/src/string/btrim.rs +++ b/datafusion/functions/src/string/btrim.rs @@ -16,30 +16,41 @@ // under the License. use crate::string::common::*; -use crate::utils::{make_scalar_function, utf8_to_str_type}; -use arrow::array::{ArrayRef, OffsetSizeTrait}; +use crate::utils::make_scalar_function; +use arrow::array::{ArrayRef, AsArray}; use arrow::datatypes::DataType; use datafusion_common::types::logical_string; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; use std::sync::Arc; /// Returns the longest string with leading and trailing characters removed. If the characters are not specified, spaces are removed. /// btrim('xyxtrimyyx', 'xyz') = 'trim' -fn btrim<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> { - let use_string_view = args[0].data_type() == &DataType::Utf8View; +fn btrim(args: &[ArrayRef]) -> Result<ArrayRef> { let args = if args.len() > 1 { let arg1 = arrow::compute::kernels::cast::cast(&args[1], args[0].data_type())?; vec![Arc::clone(&args[0]), arg1] } else { args.to_owned() }; - general_trim::<T, TrimBoth>(&args, use_string_view) + match args[0].data_type() { + DataType::Utf8 => general_trim::<i32, TrimBoth>(&args, false), + DataType::LargeUtf8 => general_trim::<i64, TrimBoth>(&args, false), + DataType::Utf8View => general_trim::<i32, TrimBoth>(&args, true), + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let trimmed = btrim(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(trimmed)) + } + other => exec_err!( + "Unsupported data type {other:?} for function btrim, expected Utf8, LargeUtf8 or Utf8View." + ), + } } #[user_doc( @@ -85,9 +96,12 @@ impl BTrimFunc { Coercion::new_exact(TypeSignatureClass::Native(logical_string())), Coercion::new_exact(TypeSignatureClass::Native(logical_string())), ]), - TypeSignature::Coercible(vec![Coercion::new_exact( - TypeSignatureClass::Native(logical_string()), - )]), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::dictionary(), + ), + ]), ], Volatility::Immutable, ), @@ -106,28 +120,11 @@ impl ScalarUDFImpl for BTrimFunc { } 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], "btrim") - } + Ok(arg_types[0].clone()) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { - match args.args[0].data_type() { - DataType::Utf8 | DataType::Utf8View => make_scalar_function( - btrim::<i32>, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - DataType::LargeUtf8 => make_scalar_function( - btrim::<i64>, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - other => exec_err!( - "Unsupported data type {other:?} for function btrim,\ - expected Utf8, LargeUtf8 or Utf8View." - ), - } + make_scalar_function(btrim, vec![Hint::Pad, Hint::AcceptsSingular])(&args.args) } fn aliases(&self) -> &[String] { diff --git a/datafusion/functions/src/string/ltrim.rs b/datafusion/functions/src/string/ltrim.rs index e49ffeb054..04e33253ed 100644 --- a/datafusion/functions/src/string/ltrim.rs +++ b/datafusion/functions/src/string/ltrim.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, OffsetSizeTrait}; +use arrow::array::{ArrayRef, AsArray}; use arrow::datatypes::DataType; use std::sync::Arc; @@ -25,22 +25,33 @@ use datafusion_common::types::logical_string; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; /// Returns the longest string with leading characters removed. If the characters are not specified, spaces are removed. /// ltrim('zzzytest', 'xyz') = 'test' -fn ltrim<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> { - let use_string_view = args[0].data_type() == &DataType::Utf8View; +fn ltrim(args: &[ArrayRef]) -> Result<ArrayRef> { let args = if args.len() > 1 { let arg1 = arrow::compute::kernels::cast::cast(&args[1], args[0].data_type())?; vec![Arc::clone(&args[0]), arg1] } else { args.to_owned() }; - general_trim::<T, TrimLeft>(&args, use_string_view) + match args[0].data_type() { + DataType::Utf8 => general_trim::<i32, TrimLeft>(&args, false), + DataType::LargeUtf8 => general_trim::<i64, TrimLeft>(&args, false), + DataType::Utf8View => general_trim::<i32, TrimLeft>(&args, true), + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let trimmed = ltrim(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(trimmed)) + } + other => exec_err!( + "Unsupported data type {other:?} for function ltrim, expected Utf8, LargeUtf8 or Utf8View." + ), + } } #[user_doc( @@ -90,9 +101,12 @@ impl LtrimFunc { Coercion::new_exact(TypeSignatureClass::Native(logical_string())), Coercion::new_exact(TypeSignatureClass::Native(logical_string())), ]), - TypeSignature::Coercible(vec![Coercion::new_exact( - TypeSignatureClass::Native(logical_string()), - )]), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::dictionary(), + ), + ]), ], Volatility::Immutable, ), @@ -114,20 +128,7 @@ impl ScalarUDFImpl for LtrimFunc { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { - match args.args[0].data_type() { - DataType::Utf8 | DataType::Utf8View => make_scalar_function( - ltrim::<i32>, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - DataType::LargeUtf8 => make_scalar_function( - ltrim::<i64>, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - other => exec_err!( - "Unsupported data type {other:?} for function ltrim,\ - expected Utf8, LargeUtf8 or Utf8View." - ), - } + make_scalar_function(ltrim, vec![Hint::Pad, Hint::AcceptsSingular])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/string/rtrim.rs b/datafusion/functions/src/string/rtrim.rs index 05ad9e8559..d1126cb241 100644 --- a/datafusion/functions/src/string/rtrim.rs +++ b/datafusion/functions/src/string/rtrim.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, OffsetSizeTrait}; +use arrow::array::{ArrayRef, AsArray}; use arrow::datatypes::DataType; use std::sync::Arc; @@ -25,22 +25,33 @@ use datafusion_common::types::logical_string; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; /// Returns the longest string with trailing characters removed. If the characters are not specified, spaces are removed. /// rtrim('testxxzx', 'xyz') = 'test' -fn rtrim<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> { - let use_string_view = args[0].data_type() == &DataType::Utf8View; +fn rtrim(args: &[ArrayRef]) -> Result<ArrayRef> { let args = if args.len() > 1 { let arg1 = arrow::compute::kernels::cast::cast(&args[1], args[0].data_type())?; vec![Arc::clone(&args[0]), arg1] } else { args.to_owned() }; - general_trim::<T, TrimRight>(&args, use_string_view) + match args[0].data_type() { + DataType::Utf8 => general_trim::<i32, TrimRight>(&args, false), + DataType::LargeUtf8 => general_trim::<i64, TrimRight>(&args, false), + DataType::Utf8View => general_trim::<i32, TrimRight>(&args, true), + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let trimmed = rtrim(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(trimmed)) + } + other => exec_err!( + "Unsupported data type {other:?} for function rtrim, expected Utf8, LargeUtf8 or Utf8View." + ), + } } #[user_doc( @@ -90,9 +101,12 @@ impl RtrimFunc { Coercion::new_exact(TypeSignatureClass::Native(logical_string())), Coercion::new_exact(TypeSignatureClass::Native(logical_string())), ]), - TypeSignature::Coercible(vec![Coercion::new_exact( - TypeSignatureClass::Native(logical_string()), - )]), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::dictionary(), + ), + ]), ], Volatility::Immutable, ), @@ -114,20 +128,7 @@ impl ScalarUDFImpl for RtrimFunc { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { - match args.args[0].data_type() { - DataType::Utf8 | DataType::Utf8View => make_scalar_function( - rtrim::<i32>, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - DataType::LargeUtf8 => make_scalar_function( - rtrim::<i64>, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - other => exec_err!( - "Unsupported data type {other:?} for function rtrim,\ - expected Utf8, LargeUtf8 or Utf8View." - ), - } + make_scalar_function(rtrim, vec![Hint::Pad, Hint::AcceptsSingular])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index 78bdeb3e15..008be05852 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -492,10 +492,11 @@ SELECT btrim(' foo ') ---- foo -query T -SELECT btrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')) +query TT +SELECT btrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')), + arrow_typeof(btrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)'))) ---- -foo +foo Dictionary(Int32, Utf8) query T SELECT initcap('foo') @@ -689,10 +690,11 @@ SELECT ltrim(' foo') ---- foo -query T -SELECT ltrim(arrow_cast(' foo', 'Dictionary(Int32, Utf8)')) +query TT +SELECT ltrim(arrow_cast(' foo', 'Dictionary(Int32, Utf8)')), + arrow_typeof(ltrim(arrow_cast(' foo', 'Dictionary(Int32, Utf8)'))) ---- -foo +foo Dictionary(Int32, Utf8) query T SELECT md5('foo') @@ -710,20 +712,75 @@ SELECT rtrim(' foo ') ---- foo -query T -SELECT rtrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')) +query TT +SELECT rtrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')), + arrow_typeof(rtrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)'))) ---- - foo + foo Dictionary(Int32, Utf8) query T SELECT trim(' foo ') ---- foo -query T -SELECT trim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')) +query TT +SELECT trim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')), + arrow_typeof(trim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)'))) ---- -foo +foo Dictionary(Int32, Utf8) + +query TTTT?T +SELECT btrim(arrow_cast(' foo ', 'Dictionary(Int32, LargeUtf8)')), + arrow_typeof(btrim(arrow_cast(' foo ', 'Dictionary(Int32, LargeUtf8)'))), + ltrim(arrow_cast( + arrow_cast(' bar', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + )), + arrow_typeof(ltrim(arrow_cast( + arrow_cast(' bar', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + ))), + rtrim(arrow_cast( + arrow_cast('baz ', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(rtrim(arrow_cast( + arrow_cast('baz ', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +foo Dictionary(Int32, LargeUtf8) bar Dictionary(Int32, Utf8View) baz Dictionary(Int32, Dictionary(UInt32, Utf8)) + +statement ok +CREATE TABLE trim_dictionary_test AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS both, + arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS leading, + arrow_cast(column3, 'Dictionary(Int32, Utf8)') AS trailing +FROM (VALUES +(' foo ', ' bar', 'baz '), +(NULL, NULL, NULL)); + +query TTTTTT +SELECT btrim(both), arrow_typeof(btrim(both)), + ltrim(leading), arrow_typeof(ltrim(leading)), + rtrim(trailing), arrow_typeof(rtrim(trailing)) +FROM trim_dictionary_test +---- +foo Dictionary(Int32, Utf8) bar Dictionary(Int32, Utf8) baz Dictionary(Int32, Utf8) +NULL Dictionary(Int32, Utf8) NULL Dictionary(Int32, Utf8) NULL Dictionary(Int32, Utf8) + +statement ok +DROP TABLE trim_dictionary_test + +query TTTTTT +SELECT btrim(arrow_cast('__foo__', 'Dictionary(Int32, Utf8)'), '_'), + arrow_typeof(btrim(arrow_cast('__foo__', 'Dictionary(Int32, Utf8)'), '_')), + ltrim(arrow_cast('__bar', 'Dictionary(Int32, Utf8)'), '_'), + arrow_typeof(ltrim(arrow_cast('__bar', 'Dictionary(Int32, Utf8)'), '_')), + rtrim(arrow_cast('baz__', 'Dictionary(Int32, Utf8)'), '_'), + arrow_typeof(rtrim(arrow_cast('baz__', 'Dictionary(Int32, Utf8)'), '_')) +---- +foo Utf8 bar Utf8 baz Utf8 # Verify that trim, ltrim, and rtrim only strip spaces by default, # not other whitespace characters (tabs, newlines, etc.) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
