Skip to main content

datafusion_spark/function/url/
try_url_decode.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::ArrayRef;
19use arrow::datatypes::DataType;
20
21use datafusion_common::Result;
22use datafusion_expr::{
23    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
24};
25use datafusion_functions::utils::make_scalar_function;
26
27use crate::function::url::url_decode::{
28    OnDecodeError, UrlDecode, spark_handled_url_decode,
29};
30
31#[derive(Debug, PartialEq, Eq, Hash)]
32pub struct TryUrlDecode {
33    signature: Signature,
34    url_decoder: UrlDecode,
35}
36
37impl Default for TryUrlDecode {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl TryUrlDecode {
44    pub fn new() -> Self {
45        Self {
46            signature: Signature::string(1, Volatility::Immutable),
47            url_decoder: UrlDecode::new(),
48        }
49    }
50}
51
52impl ScalarUDFImpl for TryUrlDecode {
53    fn name(&self) -> &str {
54        "try_url_decode"
55    }
56
57    fn signature(&self) -> &Signature {
58        &self.signature
59    }
60
61    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
62        self.url_decoder.return_type(arg_types)
63    }
64
65    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
66        let ScalarFunctionArgs { args, .. } = args;
67        make_scalar_function(spark_try_url_decode, vec![])(&args)
68    }
69}
70
71fn spark_try_url_decode(args: &[ArrayRef]) -> Result<ArrayRef> {
72    spark_handled_url_decode(args, OnDecodeError::Null)
73}
74
75#[cfg(test)]
76mod tests {
77    use std::sync::Arc;
78
79    use arrow::array::StringArray;
80    use datafusion_common::cast::as_string_array;
81
82    use super::*;
83
84    #[test]
85    fn test_try_decode_error_handled() -> Result<()> {
86        let input = Arc::new(StringArray::from(vec![
87            Some("http%3A%2F%2spark.apache.org"), // '%2s' is not a valid percent encoded character
88            // Valid cases
89            Some("https%3A%2F%2Fspark.apache.org"),
90            None,
91        ]));
92
93        let expected =
94            StringArray::from(vec![None, Some("https://spark.apache.org"), None]);
95
96        let result = spark_try_url_decode(&[input as ArrayRef])?;
97        let result = as_string_array(&result)?;
98
99        assert_eq!(&expected, result);
100        Ok(())
101    }
102}