datafusion_spark/function/url/
parse_url.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, AsArray, LargeStringArray, StringArray, StringArrayType,
22    StringViewArray, new_null_array,
23};
24use arrow::datatypes::DataType;
25use datafusion_common::cast::{
26    as_large_string_array, as_string_array, as_string_view_array,
27};
28use datafusion_common::{Result, exec_datafusion_err, exec_err};
29use datafusion_expr::{
30    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
31    Volatility,
32};
33use datafusion_functions::utils::make_scalar_function;
34use url::{ParseError, Url};
35
36#[derive(Debug, PartialEq, Eq, Hash)]
37pub struct ParseUrl {
38    signature: Signature,
39}
40
41impl Default for ParseUrl {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl ParseUrl {
48    pub fn new() -> Self {
49        Self {
50            signature: Signature::one_of(
51                vec![TypeSignature::String(2), TypeSignature::String(3)],
52                Volatility::Immutable,
53            ),
54        }
55    }
56    /// Parses a URL and extracts the specified component.
57    ///
58    /// This function takes a URL string and extracts different parts of it based on the
59    /// `part` parameter. For query parameters, an optional `key` can be specified to
60    /// extract a specific query parameter value.
61    ///
62    /// # Arguments
63    ///
64    /// * `value` - The URL string to parse
65    /// * `part` - The component of the URL to extract. Valid values are:
66    ///   - `"HOST"` - The hostname (e.g., "example.com")
67    ///   - `"PATH"` - The path portion (e.g., "/path/to/resource")
68    ///   - `"QUERY"` - The query string or a specific query parameter
69    ///   - `"REF"` - The fragment/anchor (the part after #)
70    ///   - `"PROTOCOL"` - The URL scheme (e.g., "https", "http")
71    ///   - `"FILE"` - The path with query string (e.g., "/path?query=value")
72    ///   - `"AUTHORITY"` - The authority component (host:port)
73    ///   - `"USERINFO"` - The user information (username:password)
74    /// * `key` - Optional parameter used only with `"QUERY"`. When provided, extracts
75    ///   the value of the specific query parameter with this key name.
76    ///
77    /// # Returns
78    ///
79    /// * `Ok(Some(String))` - The extracted URL component as a string
80    /// * `Ok(None)` - If the requested component doesn't exist
81    /// * `Err(DataFusionError)` - If the URL is malformed and cannot be parsed
82    fn parse(value: &str, part: &str, key: Option<&str>) -> Result<Option<String>> {
83        let url: std::result::Result<Url, ParseError> = Url::parse(value);
84        if let Err(ParseError::RelativeUrlWithoutBase) = url {
85            return if !value.contains("://") {
86                // Schemeless URLs are treated as relative URIs (like java.net.URI).
87                // Manually parse path, query, and fragment components.
88                let (without_fragment, fragment) = match value.split_once('#') {
89                    Some((before, frag)) => (before, Some(frag)),
90                    None => (value, None),
91                };
92                let (path, query) = match without_fragment.split_once('?') {
93                    Some((p, q)) => (p, Some(q)),
94                    None => (without_fragment, None),
95                };
96                Ok(match part {
97                    "PATH" => Some(path.to_string()),
98                    "QUERY" => match key {
99                        None => query.map(String::from),
100                        Some(key) => Self::query_value(query, key).map(String::from),
101                    },
102                    "REF" => fragment.map(String::from),
103                    "FILE" => {
104                        // FILE = path + query (without fragment)
105                        Some(without_fragment.to_string())
106                    }
107                    // HOST, PROTOCOL, AUTHORITY, USERINFO → NULL
108                    _ => None,
109                })
110            } else {
111                Err(exec_datafusion_err!(
112                    "The url is invalid: {value}. Use `try_parse_url` to tolerate invalid URL and return NULL instead. SQLSTATE: 22P02"
113                ))
114            };
115        };
116        url.map_err(|e| exec_datafusion_err!("{e:?}"))
117            .map(|url| match part {
118                "HOST" => url.host_str().map(String::from),
119                "PATH" => {
120                    let path = Self::path(value, &url);
121                    Some(path.to_string())
122                }
123                "QUERY" => match key {
124                    None => url.query().map(String::from),
125                    Some(key) => Self::query_value(url.query(), key).map(String::from),
126                },
127                "REF" => url.fragment().map(String::from),
128                "PROTOCOL" => Some(url.scheme().to_string()),
129                "FILE" => {
130                    let path = Self::path(value, &url);
131                    match url.query() {
132                        Some(query) => Some(format!("{path}?{query}")),
133                        None => Some(path.to_string()),
134                    }
135                }
136                "AUTHORITY" => Some(url.authority().to_string()),
137                "USERINFO" => {
138                    let username = url.username();
139                    if username.is_empty() {
140                        return None;
141                    }
142                    match url.password() {
143                        Some(password) => Some(format!("{username}:{password}")),
144                        None => Some(username.to_string()),
145                    }
146                }
147                _ => None,
148            })
149    }
150
151    fn path<'a>(value: &str, url: &'a Url) -> &'a str {
152        let path = url.path();
153        if path == "/" && Self::absolute_url_has_empty_path(value) {
154            ""
155        } else {
156            path
157        }
158    }
159
160    fn absolute_url_has_empty_path(value: &str) -> bool {
161        let Some(authority_start) = value.find("://").map(|index| index + 3) else {
162            return false;
163        };
164        let after_authority = &value[authority_start..];
165        match after_authority.find(['/', '?', '#']) {
166            None => true,
167            Some(index) => matches!(after_authority.as_bytes()[index], b'?' | b'#'),
168        }
169    }
170
171    fn query_value<'a>(query: Option<&'a str>, key: &str) -> Option<&'a str> {
172        query.and_then(|query| {
173            query
174                .split('&')
175                .filter_map(|pair| pair.split_once('='))
176                .find(|(query_key, _)| *query_key == key)
177                .map(|(_, value)| value)
178        })
179    }
180}
181
182impl ScalarUDFImpl for ParseUrl {
183    fn name(&self) -> &str {
184        "parse_url"
185    }
186
187    fn signature(&self) -> &Signature {
188        &self.signature
189    }
190
191    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
192        Ok(arg_types[0].clone())
193    }
194
195    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
196        let ScalarFunctionArgs { args, .. } = args;
197        make_scalar_function(spark_parse_url, vec![])(&args)
198    }
199}
200
201/// Core implementation of URL parsing function.
202///
203/// # Arguments
204///
205/// * `args` - A slice of ArrayRef containing the input arrays:
206///   - `args[0]` - URL array: The URLs to parse
207///   - `args[1]` - Part array: The URL components to extract (HOST, PATH, QUERY, etc.)
208///   - `args[2]` - Key array (optional): For QUERY part, the specific parameter names to extract
209///
210/// # Return Value
211///
212/// Returns `Result<ArrayRef>` containing:
213/// - A string array with extracted URL components
214/// - `None` values where extraction failed or component doesn't exist
215/// - The output array type (StringArray or LargeStringArray) is determined by input types
216fn spark_parse_url(args: &[ArrayRef]) -> Result<ArrayRef> {
217    spark_handled_parse_url(args, |x| x)
218}
219
220pub fn spark_handled_parse_url(
221    args: &[ArrayRef],
222    handler_err: impl Fn(Result<Option<String>>) -> Result<Option<String>>,
223) -> Result<ArrayRef> {
224    if args.len() < 2 || args.len() > 3 {
225        return exec_err!(
226            "{} expects 2 or 3 arguments, but got {}",
227            "`parse_url`",
228            args.len()
229        );
230    }
231    // Required arguments
232    let url = &args[0];
233    let part = &args[1];
234
235    if args.len() == 3 {
236        // In this case, the 'key' argument is passed
237        let key = &args[2];
238
239        match (url.data_type(), part.data_type(), key.data_type()) {
240            (DataType::Utf8, DataType::Utf8, DataType::Utf8) => {
241                process_parse_url::<_, _, _, StringArray>(
242                    as_string_array(url)?,
243                    as_string_array(part)?,
244                    as_string_array(key)?,
245                    handler_err,
246                    true,
247                )
248            }
249            (DataType::Utf8View, DataType::Utf8View, DataType::Utf8View) => {
250                process_parse_url::<_, _, _, StringViewArray>(
251                    as_string_view_array(url)?,
252                    as_string_view_array(part)?,
253                    as_string_view_array(key)?,
254                    handler_err,
255                    true,
256                )
257            }
258            (DataType::LargeUtf8, DataType::LargeUtf8, DataType::LargeUtf8) => {
259                process_parse_url::<_, _, _, LargeStringArray>(
260                    as_large_string_array(url)?,
261                    as_large_string_array(part)?,
262                    as_large_string_array(key)?,
263                    handler_err,
264                    true,
265                )
266            }
267            _ => exec_err!(
268                "`parse_url` expects STRING arguments, got ({}, {}, {})",
269                url.data_type(),
270                part.data_type(),
271                key.data_type()
272            ),
273        }
274    } else {
275        // The 'key' argument is omitted, assume all values are null.
276        // `new_null_array` allocates the null array outright, rather than
277        // appending one null per row through a builder.
278        let key_array = new_null_array(&DataType::Utf8, args[0].len());
279        let key = key_array.as_string::<i32>();
280
281        match (url.data_type(), part.data_type()) {
282            (DataType::Utf8, DataType::Utf8) => {
283                process_parse_url::<_, _, _, StringArray>(
284                    as_string_array(url)?,
285                    as_string_array(part)?,
286                    key,
287                    handler_err,
288                    false,
289                )
290            }
291            (DataType::Utf8View, DataType::Utf8View) => {
292                process_parse_url::<_, _, _, StringViewArray>(
293                    as_string_view_array(url)?,
294                    as_string_view_array(part)?,
295                    key,
296                    handler_err,
297                    false,
298                )
299            }
300            (DataType::LargeUtf8, DataType::LargeUtf8) => {
301                process_parse_url::<_, _, _, LargeStringArray>(
302                    as_large_string_array(url)?,
303                    as_large_string_array(part)?,
304                    key,
305                    handler_err,
306                    false,
307                )
308            }
309            _ => exec_err!(
310                "`parse_url` expects STRING arguments, got ({}, {})",
311                url.data_type(),
312                part.data_type()
313            ),
314        }
315    }
316}
317
318fn process_parse_url<'a, A, B, C, T>(
319    url_array: &'a A,
320    part_array: &'a B,
321    key_array: &'a C,
322    handle: impl Fn(Result<Option<String>>) -> Result<Option<String>>,
323    has_key_arg: bool,
324) -> Result<ArrayRef>
325where
326    &'a A: StringArrayType<'a>,
327    &'a B: StringArrayType<'a>,
328    &'a C: StringArrayType<'a>,
329    T: Array + FromIterator<Option<String>> + 'static,
330{
331    url_array
332        .iter()
333        .zip(part_array.iter())
334        .zip(key_array.iter())
335        .map(|((url, part), key)| {
336            // Spark returns NULL when the third argument is explicitly NULL
337            if has_key_arg && key.is_none() {
338                return Ok(None);
339            }
340            if let (Some(url), Some(part)) = (url, part) {
341                handle(ParseUrl::parse(url, part, key))
342            } else {
343                Ok(None)
344            }
345        })
346        .collect::<Result<T>>()
347        .map(|array| Arc::new(array) as ArrayRef)
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353    use arrow::array::Int32Array;
354    use std::array::from_ref;
355
356    fn sa(vals: &[Option<&str>]) -> ArrayRef {
357        Arc::new(StringArray::from(vals.to_vec())) as ArrayRef
358    }
359
360    #[test]
361    fn test_parse_host() -> Result<()> {
362        let got = ParseUrl::parse("https://example.com/a?x=1", "HOST", None)?;
363        assert_eq!(got, Some("example.com".to_string()));
364        Ok(())
365    }
366
367    #[test]
368    fn test_parse_query_no_key_vs_with_key() -> Result<()> {
369        let got_all = ParseUrl::parse("https://ex.com/p?a=1&b=2", "QUERY", None)?;
370        assert_eq!(got_all, Some("a=1&b=2".to_string()));
371
372        let got_a = ParseUrl::parse("https://ex.com/p?a=1&b=2", "QUERY", Some("a"))?;
373        assert_eq!(got_a, Some("1".to_string()));
374
375        let got_c = ParseUrl::parse("https://ex.com/p?a=1&b=2", "QUERY", Some("c"))?;
376        assert_eq!(got_c, None);
377        Ok(())
378    }
379
380    #[test]
381    fn test_parse_ref_protocol_userinfo_file_authority() -> Result<()> {
382        let url = "ftp://user:pwd@ftp.example.com:21/files?x=1#frag";
383        assert_eq!(ParseUrl::parse(url, "REF", None)?, Some("frag".to_string()));
384        assert_eq!(
385            ParseUrl::parse(url, "PROTOCOL", None)?,
386            Some("ftp".to_string())
387        );
388        assert_eq!(
389            ParseUrl::parse(url, "USERINFO", None)?,
390            Some("user:pwd".to_string())
391        );
392        assert_eq!(
393            ParseUrl::parse(url, "FILE", None)?,
394            Some("/files?x=1".to_string())
395        );
396        assert_eq!(
397            ParseUrl::parse(url, "AUTHORITY", None)?,
398            Some("user:pwd@ftp.example.com".to_string())
399        );
400        Ok(())
401    }
402
403    #[test]
404    fn test_parse_path_empty_vs_root() -> Result<()> {
405        assert_eq!(
406            ParseUrl::parse("https://example.com", "PATH", None)?,
407            Some("".to_string())
408        );
409        assert_eq!(
410            ParseUrl::parse("https://example.com/", "PATH", None)?,
411            Some("/".to_string())
412        );
413        assert_eq!(
414            ParseUrl::parse("https://ex.com/dir%20/pa%20th.HTML", "PATH", None)?,
415            Some("/dir%20/pa%20th.HTML".to_string())
416        );
417        Ok(())
418    }
419
420    #[test]
421    fn test_parse_query_key_is_raw() -> Result<()> {
422        let url = "https://use%20r:pas%20s@example.com/dir%20/pa%20th.HTML?query=x%20y&q2=2#Ref%20two";
423        assert_eq!(
424            ParseUrl::parse(url, "QUERY", None)?,
425            Some("query=x%20y&q2=2".to_string())
426        );
427        assert_eq!(
428            ParseUrl::parse(url, "QUERY", Some("query"))?,
429            Some("x%20y".to_string())
430        );
431        assert_eq!(
432            ParseUrl::parse("http://ex.com?key=", "QUERY", Some("key"))?,
433            Some("".to_string())
434        );
435        assert_eq!(
436            ParseUrl::parse("http://ex.com?keyonly", "QUERY", Some("keyonly"))?,
437            None
438        );
439        assert_eq!(
440            ParseUrl::parse("http://ex.com?a=1&a=2", "QUERY", Some("a"))?,
441            Some("1".to_string())
442        );
443        assert_eq!(
444            ParseUrl::parse("http://ex.com?a%20b=1", "QUERY", Some("a b"))?,
445            None
446        );
447        Ok(())
448    }
449
450    #[test]
451    fn test_parse_empty_path_file() -> Result<()> {
452        assert_eq!(ParseUrl::parse("", "PATH", None)?, Some("".to_string()));
453        assert_eq!(
454            ParseUrl::parse("http://example.com", "FILE", None)?,
455            Some("".to_string())
456        );
457        assert_eq!(
458            ParseUrl::parse("http://example.com?foo=bar", "FILE", None)?,
459            Some("?foo=bar".to_string())
460        );
461        assert_eq!(
462            ParseUrl::parse("http://example.com#fragment", "FILE", None)?,
463            Some("".to_string())
464        );
465        assert_eq!(
466            ParseUrl::parse("http://example.com/?foo=bar", "FILE", None)?,
467            Some("/?foo=bar".to_string())
468        );
469        assert_eq!(
470            ParseUrl::parse("http://ex.com/?", "FILE", None)?,
471            Some("/?".to_string())
472        );
473        assert_eq!(
474            ParseUrl::parse("http://ex.com?", "FILE", None)?,
475            Some("?".to_string())
476        );
477        Ok(())
478    }
479
480    #[test]
481    fn test_parse_schemeless_url() -> Result<()> {
482        // Spark's java.net.URI treats schemeless strings as relative URIs.
483        // Simple schemeless string: no query, no fragment.
484        assert_eq!(
485            ParseUrl::parse("notaurl", "PATH", None)?,
486            Some("notaurl".to_string())
487        );
488        assert_eq!(
489            ParseUrl::parse("notaurl", "FILE", None)?,
490            Some("notaurl".to_string())
491        );
492        assert_eq!(ParseUrl::parse("notaurl", "HOST", None)?, None);
493        assert_eq!(ParseUrl::parse("notaurl", "PROTOCOL", None)?, None);
494        assert_eq!(ParseUrl::parse("notaurl", "QUERY", None)?, None);
495        assert_eq!(ParseUrl::parse("notaurl", "REF", None)?, None);
496        assert_eq!(ParseUrl::parse("notaurl", "AUTHORITY", None)?, None);
497        assert_eq!(ParseUrl::parse("notaurl", "USERINFO", None)?, None);
498
499        // Schemeless URL with query string
500        assert_eq!(
501            ParseUrl::parse("notaurl?key=value", "PATH", None)?,
502            Some("notaurl".to_string())
503        );
504        assert_eq!(
505            ParseUrl::parse("notaurl?key=value", "FILE", None)?,
506            Some("notaurl?key=value".to_string())
507        );
508        assert_eq!(
509            ParseUrl::parse("notaurl?key=value", "QUERY", None)?,
510            Some("key=value".to_string())
511        );
512        assert_eq!(
513            ParseUrl::parse("notaurl?key=value", "QUERY", Some("key"))?,
514            Some("value".to_string())
515        );
516        assert_eq!(
517            ParseUrl::parse("notaurl?key=value", "QUERY", Some("missing"))?,
518            None
519        );
520        assert_eq!(ParseUrl::parse("notaurl?key=value", "HOST", None)?, None);
521        assert_eq!(
522            ParseUrl::parse("notaurl?key=value", "PROTOCOL", None)?,
523            None
524        );
525
526        // Schemeless URL with fragment
527        assert_eq!(
528            ParseUrl::parse("notaurl#reference", "REF", None)?,
529            Some("reference".to_string())
530        );
531        assert_eq!(
532            ParseUrl::parse("notaurl#reference", "PATH", None)?,
533            Some("notaurl".to_string())
534        );
535        assert_eq!(
536            ParseUrl::parse("notaurl#reference", "FILE", None)?,
537            Some("notaurl".to_string())
538        );
539
540        // Schemeless URL with both query and fragment
541        assert_eq!(
542            ParseUrl::parse("notaurl?a=1&b=2#frag", "PATH", None)?,
543            Some("notaurl".to_string())
544        );
545        assert_eq!(
546            ParseUrl::parse("notaurl?a=1&b=2#frag", "QUERY", None)?,
547            Some("a=1&b=2".to_string())
548        );
549        assert_eq!(
550            ParseUrl::parse("notaurl?a=1&b=2#frag", "QUERY", Some("b"))?,
551            Some("2".to_string())
552        );
553        assert_eq!(
554            ParseUrl::parse("notaurl?a=1&b=2#frag", "REF", None)?,
555            Some("frag".to_string())
556        );
557        assert_eq!(
558            ParseUrl::parse("notaurl?a=1&b=2#frag", "FILE", None)?,
559            Some("notaurl?a=1&b=2".to_string())
560        );
561        Ok(())
562    }
563
564    #[test]
565    fn test_spark_utf8_two_args() -> Result<()> {
566        let urls = sa(&[Some("https://example.com/a?x=1"), Some("https://ex.com/")]);
567        let parts = sa(&[Some("HOST"), Some("PATH")]);
568
569        let out = spark_handled_parse_url(&[urls, parts], |x| x)?;
570        let out_sa = out.as_any().downcast_ref::<StringArray>().unwrap();
571
572        assert_eq!(out_sa.len(), 2);
573        assert_eq!(out_sa.value(0), "example.com");
574        assert_eq!(out_sa.value(1), "/");
575        Ok(())
576    }
577
578    #[test]
579    fn test_spark_utf8_three_args_query_key() -> Result<()> {
580        let urls = sa(&[
581            Some("https://example.com/a?x=1&y=2"),
582            Some("https://ex.com/?a=1"),
583        ]);
584        let parts = sa(&[Some("QUERY"), Some("QUERY")]);
585        let keys = sa(&[Some("y"), Some("b")]);
586
587        let out = spark_handled_parse_url(&[urls, parts, keys], |x| x)?;
588        let out_sa = out.as_any().downcast_ref::<StringArray>().unwrap();
589
590        assert_eq!(out_sa.len(), 2);
591        assert_eq!(out_sa.value(0), "2");
592        assert!(out_sa.is_null(1));
593        Ok(())
594    }
595
596    #[test]
597    fn test_spark_userinfo_and_nulls() -> Result<()> {
598        let urls = sa(&[
599            Some("ftp://user:pwd@ftp.example.com:21/files"),
600            Some("https://example.com"),
601            None,
602        ]);
603        let parts = sa(&[Some("USERINFO"), Some("USERINFO"), Some("USERINFO")]);
604
605        let out = spark_handled_parse_url(&[urls, parts], |x| x)?;
606        let out_sa = out.as_any().downcast_ref::<StringArray>().unwrap();
607
608        assert_eq!(out_sa.len(), 3);
609        assert_eq!(out_sa.value(0), "user:pwd");
610        assert!(out_sa.is_null(1));
611        assert!(out_sa.is_null(2));
612        Ok(())
613    }
614
615    #[test]
616    fn test_invalid_arg_count() {
617        let urls = sa(&[Some("https://example.com")]);
618        let err = spark_handled_parse_url(from_ref(&urls), |x| x).unwrap_err();
619        assert!(format!("{err}").contains("expects 2 or 3 arguments"));
620
621        let parts = sa(&[Some("HOST")]);
622        let keys = sa(&[Some("x")]);
623        let err =
624            spark_handled_parse_url(&[urls, parts, keys, sa(&[Some("extra")])], |x| x)
625                .unwrap_err();
626        assert!(format!("{err}").contains("expects 2 or 3 arguments"));
627    }
628
629    #[test]
630    fn test_non_string_types_error() {
631        let urls = sa(&[Some("https://example.com")]);
632        let bad_part = Arc::new(Int32Array::from(vec![1])) as ArrayRef;
633
634        let err = spark_handled_parse_url(&[urls, bad_part], |x| x).unwrap_err();
635        let msg = format!("{err}");
636        assert!(msg.contains("expects STRING arguments"));
637    }
638}
datafusion_spark/function/url/parse_url.rs

datafusion_spark/function/url/
parse_url.rs
