diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 58c5ccd2d842a..d0ce0d0be3b15 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -398,3 +398,8 @@ required-features = ["math_expressions"] harness = false name = "round" required-features = ["math_expressions"] + +[[bench]] +harness = false +name = "dictionary_encoding" +required-features = ["string_expressions"] diff --git a/datafusion/functions/benches/dictionary_encoding.rs b/datafusion/functions/benches/dictionary_encoding.rs new file mode 100644 index 0000000000000..3afc1d5eb4c19 --- /dev/null +++ b/datafusion/functions/benches/dictionary_encoding.rs @@ -0,0 +1,95 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, DictionaryArray}; +use arrow::compute::cast; +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}; + +const NUM_ROWS: usize = 8_192; +const DICTIONARY_CARDINALITIES: [usize; 4] = [10, 100, 1_000, 8_192]; + +fn create_string_dictionary(cardinality: usize) -> ArrayRef { + let values = (0..NUM_ROWS) + .map(|index| Some(format!("value_{:04}", index % cardinality))) + .collect::>(); + Arc::new( + values + .iter() + .map(|value| value.as_deref()) + .collect::>(), + ) +} + +fn benchmark_dictionary_string_udfs(c: &mut Criterion) { + let udfs: [(&str, Arc); 3] = [ + ("ascii", datafusion_functions::string::ascii()), + ("bit_length", datafusion_functions::string::bit_length()), + ("octet_length", datafusion_functions::string::octet_length()), + ]; + let config_options = Arc::new(ConfigOptions::default()); + + for cardinality in DICTIONARY_CARDINALITIES { + let dictionary = create_string_dictionary(cardinality); + let mut group = c.benchmark_group(format!( + "dictionary_encoding/string/cardinality_{cardinality}" + )); + for (name, udf) in &udfs { + let input_field = + Field::new("a", dictionary.data_type().clone(), false).into(); + let coerced_field = fields_with_udf(&[input_field], udf.as_ref()) + .unwrap() + .into_iter() + .next() + .unwrap(); + let coerced_type = coerced_field.data_type(); + let return_type = + udf.return_type(std::slice::from_ref(coerced_type)).unwrap(); + let return_field = Field::new("f", return_type, false).into(); + let input = if dictionary.data_type() == coerced_type { + Arc::clone(&dictionary) + } else { + cast(dictionary.as_ref(), coerced_type).unwrap() + }; + + group.bench_function(*name, |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::clone(&input))], + arg_fields: vec![Arc::clone(&coerced_field)], + number_rows: NUM_ROWS, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + } + group.finish(); + } +} + +criterion_group!(benches, benchmark_dictionary_string_udfs); +criterion_main!(benches); diff --git a/datafusion/functions/src/string/ascii.rs b/datafusion/functions/src/string/ascii.rs index 4447d1f174660..db539a4d11719 100644 --- a/datafusion/functions/src/string/ascii.rs +++ b/datafusion/functions/src/string/ascii.rs @@ -15,13 +15,16 @@ // specific language governing permissions and limitations // under the License. +use crate::utils::transform_leaf_type_preserving_encoding; use arrow::array::{ArrayRef, AsArray, Int32Array, StringArrayType}; use arrow::datatypes::DataType; use arrow::error::ArrowError; use datafusion_common::types::logical_string; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue, internal_err}; -use datafusion_expr::{ColumnarValue, Documentation, TypeSignatureClass}; +use datafusion_expr::{ + ColumnarValue, Documentation, EncodingPreservation, TypeSignatureClass, +}; use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; use datafusion_expr_common::signature::Coercion; use datafusion_macros::user_doc; @@ -63,9 +66,10 @@ impl AsciiFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -81,8 +85,8 @@ impl ScalarUDFImpl for AsciiFunc { &self.signature } - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Int32) + fn return_type(&self, arg_types: &[DataType]) -> Result { + transform_leaf_type_preserving_encoding(&arg_types[0], &|_| Ok(DataType::Int32)) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -90,24 +94,7 @@ impl ScalarUDFImpl for AsciiFunc { match arg { ColumnarValue::Scalar(scalar) => { - if scalar.is_null() { - return Ok(ColumnarValue::Scalar(ScalarValue::Int32(None))); - } - - match scalar { - ScalarValue::Utf8(Some(s)) - | ScalarValue::LargeUtf8(Some(s)) - | ScalarValue::Utf8View(Some(s)) => { - let result = first_char_code(&s); - Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(result)))) - } - _ => { - internal_err!( - "Unexpected data type {:?} for function ascii", - scalar.data_type() - ) - } - } + Ok(ColumnarValue::Scalar(ascii_scalar(&scalar)?)) } ColumnarValue::Array(array) => Ok(ColumnarValue::Array(ascii(&[array])?)), } @@ -118,6 +105,24 @@ impl ScalarUDFImpl for AsciiFunc { } } +fn ascii_scalar(scalar: &ScalarValue) -> Result { + match scalar { + ScalarValue::Utf8(value) + | ScalarValue::LargeUtf8(value) + | ScalarValue::Utf8View(value) => { + Ok(ScalarValue::Int32(value.as_deref().map(first_char_code))) + } + ScalarValue::Dictionary(key_type, value) => Ok(ScalarValue::Dictionary( + key_type.clone(), + Box::new(ascii_scalar(value)?), + )), + _ => internal_err!( + "Unexpected data type {:?} for function ascii", + scalar.data_type() + ), + } +} + /// Returns the Unicode scalar value of the first character of `s`, or 0 when /// `s` is empty. Reads the leading byte first so the common all-ASCII case /// avoids constructing a `char` iterator and decoding a multi-byte sequence. @@ -184,6 +189,11 @@ pub fn ascii(args: &[ArrayRef]) -> Result { let string_array = args[0].as_string_view(); Ok(calculate_ascii(&string_array)?) } + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let converted = ascii(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(converted)) + } _ => internal_err!("Unsupported data type"), } } diff --git a/datafusion/functions/src/string/bit_length.rs b/datafusion/functions/src/string/bit_length.rs index 76d8bb73bba87..4af22f5db5b5f 100644 --- a/datafusion/functions/src/string/bit_length.rs +++ b/datafusion/functions/src/string/bit_length.rs @@ -18,13 +18,13 @@ use arrow::compute::kernels::length::bit_length; use arrow::datatypes::DataType; -use crate::utils::utf8_to_int_type; +use crate::utils::{transform_leaf_type_preserving_encoding, utf8_to_int_type}; use datafusion_common::types::logical_string; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -59,9 +59,10 @@ impl BitLengthFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -78,7 +79,9 @@ impl ScalarUDFImpl for BitLengthFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - utf8_to_int_type(&arg_types[0], "bit_length") + transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| { + utf8_to_int_type(data_type, "bit_length") + }) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -86,18 +89,7 @@ impl ScalarUDFImpl for BitLengthFunc { match array { ColumnarValue::Array(v) => Ok(ColumnarValue::Array(bit_length(v.as_ref())?)), - ColumnarValue::Scalar(v) => match v { - ScalarValue::Utf8(v) => Ok(ColumnarValue::Scalar(ScalarValue::Int32( - v.as_ref().map(|x| (x.len() * 8) as i32), - ))), - ScalarValue::LargeUtf8(v) => Ok(ColumnarValue::Scalar( - ScalarValue::Int64(v.as_ref().map(|x| (x.len() * 8) as i64)), - )), - ScalarValue::Utf8View(v) => Ok(ColumnarValue::Scalar( - ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32)), - )), - _ => unreachable!("bit length"), - }, + ColumnarValue::Scalar(v) => Ok(ColumnarValue::Scalar(bit_length_scalar(v))), } } @@ -105,3 +97,21 @@ impl ScalarUDFImpl for BitLengthFunc { self.doc() } } + +fn bit_length_scalar(value: &ScalarValue) -> ScalarValue { + match value { + ScalarValue::Utf8(v) => { + ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32)) + } + ScalarValue::LargeUtf8(v) => { + ScalarValue::Int64(v.as_ref().map(|x| (x.len() * 8) as i64)) + } + ScalarValue::Utf8View(v) => { + ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32)) + } + ScalarValue::Dictionary(key_type, value) => { + ScalarValue::Dictionary(key_type.clone(), Box::new(bit_length_scalar(value))) + } + _ => unreachable!("bit length"), + } +} diff --git a/datafusion/functions/src/string/octet_length.rs b/datafusion/functions/src/string/octet_length.rs index ecffb2a6de7af..02df262ee27aa 100644 --- a/datafusion/functions/src/string/octet_length.rs +++ b/datafusion/functions/src/string/octet_length.rs @@ -18,13 +18,13 @@ use arrow::compute::kernels::length::length; use arrow::datatypes::DataType; -use crate::utils::utf8_to_int_type; +use crate::utils::{transform_leaf_type_preserving_encoding, utf8_to_int_type}; use datafusion_common::types::logical_string; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -59,9 +59,10 @@ impl OctetLengthFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -78,7 +79,9 @@ impl ScalarUDFImpl for OctetLengthFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - utf8_to_int_type(&arg_types[0], "octet_length") + transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| { + utf8_to_int_type(data_type, "octet_length") + }) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -86,18 +89,7 @@ impl ScalarUDFImpl for OctetLengthFunc { match array { ColumnarValue::Array(v) => Ok(ColumnarValue::Array(length(v.as_ref())?)), - ColumnarValue::Scalar(v) => match v { - ScalarValue::Utf8(v) => Ok(ColumnarValue::Scalar(ScalarValue::Int32( - v.as_ref().map(|x| x.len() as i32), - ))), - ScalarValue::LargeUtf8(v) => Ok(ColumnarValue::Scalar( - ScalarValue::Int64(v.as_ref().map(|x| x.len() as i64)), - )), - ScalarValue::Utf8View(v) => Ok(ColumnarValue::Scalar( - ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)), - )), - _ => unreachable!("OctetLengthFunc"), - }, + ColumnarValue::Scalar(v) => Ok(ColumnarValue::Scalar(octet_length_scalar(v))), } } @@ -106,6 +98,23 @@ impl ScalarUDFImpl for OctetLengthFunc { } } +fn octet_length_scalar(value: &ScalarValue) -> ScalarValue { + match value { + ScalarValue::Utf8(v) => ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)), + ScalarValue::LargeUtf8(v) => { + ScalarValue::Int64(v.as_ref().map(|x| x.len() as i64)) + } + ScalarValue::Utf8View(v) => { + ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)) + } + ScalarValue::Dictionary(key_type, value) => ScalarValue::Dictionary( + key_type.clone(), + Box::new(octet_length_scalar(value)), + ), + _ => unreachable!("OctetLengthFunc"), + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/datafusion/functions/src/utils.rs b/datafusion/functions/src/utils.rs index f42ecc789babd..b93bdb0b0d3bb 100644 --- a/datafusion/functions/src/utils.rs +++ b/datafusion/functions/src/utils.rs @@ -74,6 +74,28 @@ get_optimal_return_type!(utf8_to_str_type, DataType::LargeUtf8, DataType::Utf8); // `utf8_to_int_type`: returns either a Int32 or Int64 based on the input type size. get_optimal_return_type!(utf8_to_int_type, DataType::Int64, DataType::Int32); +/// Transforms the leaf type while preserving supported encoding containers. +/// +/// Keep encoded type handling centralized here so additional encodings can be +/// supported without changing each function's return type implementation. +pub(crate) fn transform_leaf_type_preserving_encoding( + arg_type: &DataType, + transform: &F, +) -> Result +where + F: Fn(&DataType) -> Result, +{ + match arg_type { + DataType::Dictionary(key_type, value_type) => Ok(DataType::Dictionary( + key_type.clone(), + Box::new(transform_leaf_type_preserving_encoding( + value_type, transform, + )?), + )), + _ => transform(arg_type), + } +} + /// Creates a scalar function implementation for the given function. /// * `inner` - the function to be executed /// * `hints` - hints to be used when expanding scalars to arrays diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index 98edfa189d3e3..78045936a1893 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -507,6 +507,28 @@ SELECT initcap(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- Foo +query ? +SELECT ascii(arrow_cast('é', 'Dictionary(Int32, Utf8)')) +---- +233 + +query T +SELECT arrow_typeof(ascii(arrow_cast('é', 'Dictionary(Int32, Utf8)'))) +---- +Dictionary(Int32, Int32) + +query ?T +SELECT ascii(arrow_cast( + arrow_cast('💯', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(ascii(arrow_cast( + arrow_cast('💯', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +128175 Dictionary(Int32, Dictionary(UInt32, Int32)) + query I SELECT instr('foobarbar', 'bar') ---- @@ -639,11 +661,28 @@ SELECT bit_length('foo') ---- 24 -query I +query ? SELECT bit_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- 24 +query T +SELECT arrow_typeof(bit_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)'))) +---- +Dictionary(Int32, Int32) + +query ?T +SELECT bit_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(bit_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +16 Dictionary(Int32, Dictionary(UInt32, Int32)) + query I SELECT character_length('foo') ---- @@ -659,11 +698,77 @@ SELECT octet_length('foo') ---- 3 -query I +query ? SELECT octet_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- 3 +query T +SELECT arrow_typeof(octet_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)'))) +---- +Dictionary(Int32, Int32) + +query ?T +SELECT octet_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(octet_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +2 Dictionary(Int32, Dictionary(UInt32, Int32)) + +statement ok +CREATE TABLE string_length_dictionary_test AS +SELECT column1 AS id, + arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS dict_col, + arrow_cast( + arrow_cast(column2, 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ) AS nested_dict_col +FROM (VALUES +(1, 'foo'), +(2, 'é'), +(3, NULL)); + +query ??TT +SELECT bit_length(dict_col), bit_length(nested_dict_col), + arrow_typeof(bit_length(dict_col)), + arrow_typeof(bit_length(nested_dict_col)) +FROM string_length_dictionary_test +ORDER BY id +---- +24 24 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +16 16 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) + +query ??TT +SELECT octet_length(dict_col), octet_length(nested_dict_col), + arrow_typeof(octet_length(dict_col)), + arrow_typeof(octet_length(nested_dict_col)) +FROM string_length_dictionary_test +ORDER BY id +---- +3 3 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +2 2 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) + +query ??TT +SELECT ascii(dict_col), ascii(nested_dict_col), + arrow_typeof(ascii(dict_col)), + arrow_typeof(ascii(nested_dict_col)) +FROM string_length_dictionary_test +ORDER BY id +---- +102 102 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +233 233 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) + +statement ok +DROP TABLE string_length_dictionary_test + query I SELECT strpos('helloworld', 'world') ---- diff --git a/datafusion/sqllogictest/test_files/string/string_literal.slt b/datafusion/sqllogictest/test_files/string/string_literal.slt index 81aaf48629998..c175f52a35f99 100644 --- a/datafusion/sqllogictest/test_files/string/string_literal.slt +++ b/datafusion/sqllogictest/test_files/string/string_literal.slt @@ -1879,7 +1879,7 @@ SELECT ---- 48 176 32 40 -query IIII +query ???? SELECT bit_length(arrow_cast('Andrew', 'Dictionary(Int32, Utf8)')), bit_length(arrow_cast('datafusion数据融合', 'Dictionary(Int32, Utf8)')), diff --git a/datafusion/sqllogictest/test_files/string/string_query.slt.part b/datafusion/sqllogictest/test_files/string/string_query.slt.part index 9fcacbaa54921..9231ec7b9c976 100644 --- a/datafusion/sqllogictest/test_files/string/string_query.slt.part +++ b/datafusion/sqllogictest/test_files/string/string_query.slt.part @@ -645,10 +645,10 @@ drop table test_lowercase; query IIII SELECT - ASCII(ascii_1) as c1, - ASCII(ascii_2) as c2, - ASCII(unicode_1) as c3, - ASCII(unicode_2) as c4 + arrow_cast(ASCII(ascii_1), 'Int32') as c1, + arrow_cast(ASCII(ascii_2), 'Int32') as c2, + arrow_cast(ASCII(unicode_1), 'Int32') as c3, + arrow_cast(ASCII(unicode_2), 'Int32') as c4 FROM test_basic_operator; ---- 65 88 100 128293 @@ -1275,7 +1275,12 @@ NULL NULL # -------------------------------------- query IIII -select bit_length(ascii_1), bit_length(ascii_2), bit_length(unicode_1), bit_length(unicode_2) from test_basic_operator; +select + arrow_cast(bit_length(ascii_1), 'Int64'), + arrow_cast(bit_length(ascii_2), 'Int64'), + arrow_cast(bit_length(unicode_1), 'Int64'), + arrow_cast(bit_length(unicode_2), 'Int64') +from test_basic_operator; ---- 48 8 144 32 72 72 176 176