datafusion_spark/function/url/
url_encode.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 std::sync::Arc;
19
20use arrow::array::{
21    Array, ArrayRef, LargeStringBuilder, StringBuilder, StringViewBuilder,
22};
23use arrow::datatypes::DataType;
24use datafusion_common::cast::{
25    as_large_string_array, as_string_array, as_string_view_array,
26};
27use datafusion_common::{Result, exec_err, plan_err};
28use datafusion_expr::{
29    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
30};
31use datafusion_functions::utils::make_scalar_function;
32use url::form_urlencoded::byte_serialize;
33
34#[derive(Debug, PartialEq, Eq, Hash)]
35pub struct UrlEncode {
36    signature: Signature,
37}
38
39impl Default for UrlEncode {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl UrlEncode {
46    pub fn new() -> Self {
47        Self {
48            signature: Signature::string(1, Volatility::Immutable),
49        }
50    }
51}
52
53impl ScalarUDFImpl for UrlEncode {
54    fn name(&self) -> &str {
55        "url_encode"
56    }
57
58    fn signature(&self) -> &Signature {
59        &self.signature
60    }
61
62    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
63        if arg_types.len() != 1 {
64            return plan_err!(
65                "{} expects 1 argument, but got {}",
66                self.name(),
67                arg_types.len()
68            );
69        }
70        // As the type signature is already checked, we can safely return the type of the first argument
71        Ok(arg_types[0].clone())
72    }
73
74    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
75        let ScalarFunctionArgs { args, .. } = args;
76        make_scalar_function(spark_url_encode, vec![])(&args)
77    }
78}
79
80/// Core implementation of URL encoding function.
81///
82/// # Arguments
83///
84/// * `args` - A slice containing exactly one ArrayRef with the strings to encode
85///
86/// # Returns
87///
88/// * `Ok(ArrayRef)` - A new array of the same type containing encoded strings
89/// * `Err(DataFusionError)` - If invalid arguments are provided
90///
91fn spark_url_encode(args: &[ArrayRef]) -> Result<ArrayRef> {
92    if args.len() != 1 {
93        return exec_err!("`url_encode` expects 1 argument");
94    }
95
96    // The percent-encoded form of each value is assembled in a single scratch buffer
97    // reused across rows, rather than allocating a `String` per row.
98    macro_rules! encode_all {
99        ($array:expr, $builder:expr) => {{
100            let array = $array;
101            let mut builder = $builder;
102            let mut encoded = String::new();
103            for value in array.iter() {
104                match value {
105                    Some(value) => {
106                        encoded.clear();
107                        encoded.extend(byte_serialize(value.as_bytes()));
108                        builder.append_value(&encoded);
109                    }
110                    None => builder.append_null(),
111                }
112            }
113            Ok(Arc::new(builder.finish()) as ArrayRef)
114        }};
115    }
116
117    match &args[0].data_type() {
118        DataType::Utf8 => {
119            let array = as_string_array(&args[0])?;
120            let builder =
121                StringBuilder::with_capacity(array.len(), array.value_data().len());
122            encode_all!(array, builder)
123        }
124        DataType::LargeUtf8 => {
125            let array = as_large_string_array(&args[0])?;
126            let builder =
127                LargeStringBuilder::with_capacity(array.len(), array.value_data().len());
128            encode_all!(array, builder)
129        }
130        DataType::Utf8View => {
131            let array = as_string_view_array(&args[0])?;
132            let builder = StringViewBuilder::with_capacity(array.len());
133            encode_all!(array, builder)
134        }
135        other => exec_err!("`url_encode`: Expr must be STRING, got {other:?}"),
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use arrow::array::{LargeStringArray, StringArray, StringViewArray};
143
144    const INPUT: [Option<&str>; 5] = [
145        Some("https://spark.apache.org"),
146        Some("inva lid://user:pass@host/file\\;param?query\\;p2"),
147        Some("你好"),
148        Some(""),
149        None,
150    ];
151
152    const EXPECTED: [Option<&str>; 5] = [
153        Some("https%3A%2F%2Fspark.apache.org"),
154        Some("inva+lid%3A%2F%2Fuser%3Apass%40host%2Ffile%5C%3Bparam%3Fquery%5C%3Bp2"),
155        Some("%E4%BD%A0%E5%A5%BD"),
156        Some(""),
157        None,
158    ];
159
160    #[test]
161    fn test_encode_utf8() -> Result<()> {
162        let input = Arc::new(StringArray::from(INPUT.to_vec())) as ArrayRef;
163        let result = spark_url_encode(&[input])?;
164        let result = as_string_array(&result)?;
165        assert_eq!(&StringArray::from(EXPECTED.to_vec()), result);
166        Ok(())
167    }
168
169    #[test]
170    fn test_encode_large_utf8() -> Result<()> {
171        let input = Arc::new(LargeStringArray::from(INPUT.to_vec())) as ArrayRef;
172        let result = spark_url_encode(&[input])?;
173        let result = as_large_string_array(&result)?;
174        assert_eq!(&LargeStringArray::from(EXPECTED.to_vec()), result);
175        Ok(())
176    }
177
178    #[test]
179    fn test_encode_utf8_view() -> Result<()> {
180        let input = Arc::new(StringViewArray::from(INPUT.to_vec())) as ArrayRef;
181        let result = spark_url_encode(&[input])?;
182        let result = as_string_view_array(&result)?;
183        assert_eq!(&StringViewArray::from(EXPECTED.to_vec()), result);
184        Ok(())
185    }
186}
datafusion_spark/function/url/url_encode.rs

datafusion_spark/function/url/
url_encode.rs
