datafusion_spark/function/string/
quote.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, OffsetSizeTrait, StringArray};
19use arrow::datatypes::DataType;
20use datafusion_expr::{Coercion, ColumnarValue, Signature, TypeSignatureClass};
21use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
22use datafusion_common::types::{NativeType, logical_string};
23use datafusion_common::utils::take_function_args;
24use datafusion_common::{Result, exec_err};
25use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Volatility};
26use datafusion_functions::utils::make_scalar_function;
27
28use std::sync::Arc;
29
30/// Spark-compatible `quote` expression
31/// <https://spark.apache.org/docs/latest/api/sql/index.html#quote>
32#[derive(Debug, PartialEq, Eq, Hash)]
33pub struct SparkQuote {
34    signature: Signature,
35}
36
37impl Default for SparkQuote {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl SparkQuote {
44    pub fn new() -> Self {
45        let str_coercion = Coercion::new_implicit(
46            TypeSignatureClass::Native(logical_string()),
47            vec![TypeSignatureClass::Any],
48            NativeType::String,
49        );
50        Self {
51            signature: Signature::coercible(vec![str_coercion], Volatility::Immutable),
52        }
53    }
54}
55
56impl ScalarUDFImpl for SparkQuote {
57    fn name(&self) -> &str {
58        "quote"
59    }
60
61    fn signature(&self) -> &Signature {
62        &self.signature
63    }
64
65    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
66        match &arg_types[0] {
67            DataType::LargeUtf8 => Ok(DataType::LargeUtf8),
68            _ => Ok(DataType::Utf8),
69        }
70    }
71
72    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
73        make_scalar_function(spark_quote_inner, vec![])(&args.args)
74    }
75}
76
77fn spark_quote_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
78    let [array] = take_function_args("quote", arg)?;
79    match &array.data_type() {
80        DataType::Utf8 => quote_array::<i32>(array),
81        DataType::LargeUtf8 => quote_array::<i64>(array),
82        DataType::Utf8View => quote_view(array),
83        other => {
84            exec_err!("unsupported data type {other:?} for function `quote`")
85        }
86    }
87}
88
89fn quote_array<T: OffsetSizeTrait>(array: &ArrayRef) -> Result<ArrayRef> {
90    let str_array = as_generic_string_array::<T>(array)?;
91    let result = str_array
92        .iter()
93        .map(|s| s.map(compute_quote))
94        .collect::<StringArray>();
95    Ok(Arc::new(result))
96}
97
98fn quote_view(str_view: &ArrayRef) -> Result<ArrayRef> {
99    let str_array = as_string_view_array(str_view)?;
100    let result = str_array
101        .iter()
102        .map(|opt_str| opt_str.map(compute_quote))
103        .collect::<StringArray>();
104    Ok(Arc::new(result) as ArrayRef)
105}
106
107const QUOTE_CHAR: char = '\'';
108const ESCAPE_CHAR: char = '\\';
109
110fn compute_quote(s: &str) -> String {
111    let mut quoted = String::with_capacity(s.len() + 2);
112    quoted.push(QUOTE_CHAR);
113    for c in s.chars() {
114        if c == QUOTE_CHAR {
115            quoted.push(ESCAPE_CHAR);
116        }
117        quoted.push(c);
118    }
119    quoted.push(QUOTE_CHAR);
120    quoted
121}
datafusion_spark/function/string/quote.rs

datafusion_spark/function/string/
quote.rs
