datafusion_functions/string/
bit_length.rs1use arrow::compute::kernels::length::bit_length;
19use arrow::datatypes::DataType;
20use std::any::Any;
21
22use crate::utils::utf8_to_int_type;
23use datafusion_common::{utils::take_function_args, Result, ScalarValue};
24use datafusion_expr::{ColumnarValue, Documentation, Volatility};
25use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
26use datafusion_macros::user_doc;
27
28#[user_doc(
29 doc_section(label = "String Functions"),
30 description = "Returns the bit length of a string.",
31 syntax_example = "bit_length(str)",
32 sql_example = r#"```sql
33> select bit_length('datafusion');
34+--------------------------------+
35| bit_length(Utf8("datafusion")) |
36+--------------------------------+
37| 80 |
38+--------------------------------+
39```"#,
40 standard_argument(name = "str", prefix = "String"),
41 related_udf(name = "length"),
42 related_udf(name = "octet_length")
43)]
44#[derive(Debug)]
45pub struct BitLengthFunc {
46 signature: Signature,
47}
48
49impl Default for BitLengthFunc {
50 fn default() -> Self {
51 Self::new()
52 }
53}
54
55impl BitLengthFunc {
56 pub fn new() -> Self {
57 Self {
58 signature: Signature::string(1, Volatility::Immutable),
59 }
60 }
61}
62
63impl ScalarUDFImpl for BitLengthFunc {
64 fn as_any(&self) -> &dyn Any {
65 self
66 }
67
68 fn name(&self) -> &str {
69 "bit_length"
70 }
71
72 fn signature(&self) -> &Signature {
73 &self.signature
74 }
75
76 fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
77 utf8_to_int_type(&arg_types[0], "bit_length")
78 }
79
80 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
81 let [array] = take_function_args(self.name(), &args.args)?;
82
83 match array {
84 ColumnarValue::Array(v) => Ok(ColumnarValue::Array(bit_length(v.as_ref())?)),
85 ColumnarValue::Scalar(v) => match v {
86 ScalarValue::Utf8(v) => Ok(ColumnarValue::Scalar(ScalarValue::Int32(
87 v.as_ref().map(|x| (x.len() * 8) as i32),
88 ))),
89 ScalarValue::LargeUtf8(v) => Ok(ColumnarValue::Scalar(
90 ScalarValue::Int64(v.as_ref().map(|x| (x.len() * 8) as i64)),
91 )),
92 ScalarValue::Utf8View(v) => Ok(ColumnarValue::Scalar(
93 ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32)),
94 )),
95 _ => unreachable!("bit length"),
96 },
97 }
98 }
99
100 fn documentation(&self) -> Option<&Documentation> {
101 self.doc()
102 }
103}