datafusion_functions/string/
contains.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow::array::{Array, ArrayRef, Scalar};
19use arrow::compute::contains as arrow_contains;
20use arrow::datatypes::DataType;
21use arrow::datatypes::DataType::{Boolean, LargeUtf8, Utf8, Utf8View};
22use datafusion_common::types::logical_string;
23use datafusion_common::{Result, exec_err};
24use datafusion_expr::binary::{binary_to_string_coercion, string_coercion};
25use datafusion_expr::{
26    Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
27    TypeSignatureClass, Volatility,
28};
29use datafusion_macros::user_doc;
30use std::any::Any;
31use std::sync::Arc;
32
33#[user_doc(
34    doc_section(label = "String Functions"),
35    description = "Return true if search_str is found within string (case-sensitive).",
36    syntax_example = "contains(str, search_str)",
37    sql_example = r#"```sql
38> select contains('the quick brown fox', 'row');
39+---------------------------------------------------+
40| contains(Utf8("the quick brown fox"),Utf8("row")) |
41+---------------------------------------------------+
42| true                                              |
43+---------------------------------------------------+
44```"#,
45    standard_argument(name = "str", prefix = "String"),
46    argument(name = "search_str", description = "The string to search for in str.")
47)]
48#[derive(Debug, PartialEq, Eq, Hash)]
49pub struct ContainsFunc {
50    signature: Signature,
51}
52
53impl Default for ContainsFunc {
54    fn default() -> Self {
55        ContainsFunc::new()
56    }
57}
58
59impl ContainsFunc {
60    pub fn new() -> Self {
61        Self {
62            signature: Signature::coercible(
63                vec![
64                    Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
65                    Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
66                ],
67                Volatility::Immutable,
68            ),
69        }
70    }
71}
72
73impl ScalarUDFImpl for ContainsFunc {
74    fn as_any(&self) -> &dyn Any {
75        self
76    }
77
78    fn name(&self) -> &str {
79        "contains"
80    }
81
82    fn signature(&self) -> &Signature {
83        &self.signature
84    }
85
86    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
87        Ok(Boolean)
88    }
89
90    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
91        contains(args.args.as_slice())
92    }
93
94    fn documentation(&self) -> Option<&Documentation> {
95        self.doc()
96    }
97}
98
99fn to_array(value: &ColumnarValue) -> Result<(ArrayRef, bool)> {
100    match value {
101        ColumnarValue::Array(array) => Ok((Arc::clone(array), false)),
102        ColumnarValue::Scalar(scalar) => Ok((scalar.to_array()?, true)),
103    }
104}
105
106/// Helper to call arrow_contains with proper Datum handling.
107/// When an argument is marked as scalar, we wrap it in `Scalar` to tell arrow's
108/// kernel to use the optimized single-value code path instead of iterating.
109fn call_arrow_contains(
110    haystack: &ArrayRef,
111    haystack_is_scalar: bool,
112    needle: &ArrayRef,
113    needle_is_scalar: bool,
114) -> Result<ColumnarValue> {
115    // Arrow's Datum trait is implemented for ArrayRef, Arc<dyn Array>, and Scalar<T>
116    // We pass ArrayRef directly when not scalar, or wrap in Scalar when it is
117    let result = match (haystack_is_scalar, needle_is_scalar) {
118        (false, false) => arrow_contains(haystack, needle)?,
119        (false, true) => arrow_contains(haystack, &Scalar::new(Arc::clone(needle)))?,
120        (true, false) => arrow_contains(&Scalar::new(Arc::clone(haystack)), needle)?,
121        (true, true) => arrow_contains(
122            &Scalar::new(Arc::clone(haystack)),
123            &Scalar::new(Arc::clone(needle)),
124        )?,
125    };
126
127    // If both inputs were scalar, return a scalar result
128    if haystack_is_scalar && needle_is_scalar {
129        let scalar = datafusion_common::ScalarValue::try_from_array(&result, 0)?;
130        Ok(ColumnarValue::Scalar(scalar))
131    } else {
132        Ok(ColumnarValue::Array(Arc::new(result)))
133    }
134}
135
136/// use `arrow::compute::contains` to do the calculation for contains
137fn contains(args: &[ColumnarValue]) -> Result<ColumnarValue> {
138    let (haystack, haystack_is_scalar) = to_array(&args[0])?;
139    let (needle, needle_is_scalar) = to_array(&args[1])?;
140
141    if let Some(coercion_data_type) =
142        string_coercion(haystack.data_type(), needle.data_type()).or_else(|| {
143            binary_to_string_coercion(haystack.data_type(), needle.data_type())
144        })
145    {
146        let haystack = if haystack.data_type() == &coercion_data_type {
147            haystack
148        } else {
149            arrow::compute::kernels::cast::cast(&haystack, &coercion_data_type)?
150        };
151        let needle = if needle.data_type() == &coercion_data_type {
152            needle
153        } else {
154            arrow::compute::kernels::cast::cast(&needle, &coercion_data_type)?
155        };
156
157        match coercion_data_type {
158            Utf8View | Utf8 | LargeUtf8 => call_arrow_contains(
159                &haystack,
160                haystack_is_scalar,
161                &needle,
162                needle_is_scalar,
163            ),
164            other => {
165                exec_err!("Unsupported data type {other:?} for function `contains`.")
166            }
167        }
168    } else {
169        exec_err!(
170            "Unsupported data type {}, {:?} for function `contains`.",
171            args[0].data_type(),
172            args[1].data_type()
173        )
174    }
175}
176
177#[cfg(test)]
178mod test {
179    use super::ContainsFunc;
180    use crate::expr_fn::contains;
181    use arrow::array::{BooleanArray, StringArray};
182    use arrow::datatypes::{DataType, Field};
183    use datafusion_common::ScalarValue;
184    use datafusion_common::config::ConfigOptions;
185    use datafusion_expr::{ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDFImpl};
186    use std::sync::Arc;
187
188    #[test]
189    fn test_contains_udf() {
190        let udf = ContainsFunc::new();
191        let array = ColumnarValue::Array(Arc::new(StringArray::from(vec![
192            Some("xxx?()"),
193            Some("yyy?()"),
194        ])));
195        let scalar = ColumnarValue::Scalar(ScalarValue::Utf8(Some("x?(".to_string())));
196        let arg_fields = vec![
197            Field::new("a", DataType::Utf8, true).into(),
198            Field::new("a", DataType::Utf8, true).into(),
199        ];
200
201        let args = ScalarFunctionArgs {
202            args: vec![array, scalar],
203            arg_fields,
204            number_rows: 2,
205            return_field: Field::new("f", DataType::Boolean, true).into(),
206            config_options: Arc::new(ConfigOptions::default()),
207        };
208
209        let actual = udf.invoke_with_args(args).unwrap();
210        let expect = ColumnarValue::Array(Arc::new(BooleanArray::from(vec![
211            Some(true),
212            Some(false),
213        ])));
214        assert_eq!(
215            *actual.into_array(2).unwrap(),
216            *expect.into_array(2).unwrap()
217        );
218    }
219
220    #[test]
221    fn test_contains_api() {
222        let expr = contains(
223            Expr::Literal(
224                ScalarValue::Utf8(Some("the quick brown fox".to_string())),
225                None,
226            ),
227            Expr::Literal(ScalarValue::Utf8(Some("row".to_string())), None),
228        );
229        assert_eq!(
230            expr.to_string(),
231            "contains(Utf8(\"the quick brown fox\"), Utf8(\"row\"))"
232        );
233    }
234}
datafusion_functions/string/contains.rs

datafusion_functions/string/
contains.rs
