Skip to content

Commit 04ef3c7

Browse files
kazantsev-maksimKazantsev Maksim
kazantsev-maksim
and
Kazantsev Maksim
authored
Spark quote function implementation (#22642)
## Which issue does this PR close? - N/A ## Rationale for this change Add new spark function: https://spark.apache.org/docs/latest/api/sql/index.html#quote ## What changes are included in this PR? - Implementation - SLT tests ## Are these changes tested? Yes, tests added as part of this PR. ## Are there any user-facing changes? No, these are new function. --------- Co-authored-by: Kazantsev Maksim <mn.kazantsev@gmail.com>
1 parent 6fdef65 commit 04ef3c7

3 files changed

Lines changed: 290 additions & 0 deletions

File tree

‎datafusion/spark/src/function/string/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ pub mod length;
2727
pub mod like;
2828
pub mod luhn_check;
2929
pub mod make_valid_utf8;
30+
pub mod quote;
3031
pub mod soundex;
3132
pub mod space;
3233
pub mod substring;
@@ -51,6 +52,7 @@ make_udf_function!(base64::SparkUnBase64, unbase64);
5152
make_udf_function!(soundex::SparkSoundex, soundex);
5253
make_udf_function!(make_valid_utf8::SparkMakeValidUtf8, make_valid_utf8);
5354
make_udf_function!(is_valid_utf8::SparkIsValidUtf8, is_valid_utf8);
55+
make_udf_function!(quote::SparkQuote, quote);
5456

5557
pub mod expr_fn {
5658
use datafusion_functions::export_functions;
@@ -127,6 +129,11 @@ pub mod expr_fn {
127129
"Returns the original string if str is a valid UTF-8 string, otherwise returns a new string whose invalid UTF8 byte sequences are replaced using the UNICODE replacement character U+FFFD.",
128130
str
129131
));
132+
export_functions!((
133+
quote,
134+
"Returns str enclosed by single quotes and each instance of single quote in it is preceded by a backslash",
135+
str
136+
));
130137
}
131138

132139
pub fn functions() -> Vec<Arc<ScalarUDF>> {
@@ -147,5 +154,6 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
147154
soundex(),
148155
make_valid_utf8(),
149156
is_valid_utf8(),
157+
quote(),
150158
]
151159
}

‎datafusion/spark/src/function/string/quote.rs‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
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+
18+
use arrow::array::{ArrayRef, OffsetSizeTrait, StringArray};
19+
use arrow::datatypes::DataType;
20+
use datafusion::logical_expr::{Coercion, ColumnarValue, Signature, TypeSignatureClass};
21+
use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
22+
use datafusion_common::types::{NativeType, logical_string};
23+
use datafusion_common::utils::take_function_args;
24+
use datafusion_common::{Result, exec_err};
25+
use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Volatility};
26+
use datafusion_functions::utils::make_scalar_function;
27+
28+
use 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)]
33+
pub struct SparkQuote {
34+
signature: Signature,
35+
}
36+
37+
impl Default for SparkQuote {
38+
fn default() -> Self {
39+
Self::new()
40+
}
41+
}
42+
43+
impl 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+
56+
impl 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+
77+
fn 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+
89+
fn 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+
98+
fn 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+
107+
const QUOTE_CHAR: char = '\'';
108+
const ESCAPE_CHAR: char = '\\';
109+
110+
fn 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/sqllogictest/test_files/spark/string/quote.slt‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
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+
18+
query T
19+
SELECT quote(arrow_cast(127, 'Int8'));
20+
----
21+
'127'
22+
23+
query T
24+
SELECT quote(arrow_cast(-128, 'Int8'));
25+
----
26+
'-128'
27+
28+
query T
29+
SELECT quote(arrow_cast(32767, 'Int16'));
30+
----
31+
'32767'
32+
33+
query T
34+
SELECT quote(arrow_cast(-32768, 'Int16'));
35+
----
36+
'-32768'
37+
38+
query T
39+
SELECT quote(arrow_cast(2147483647, 'Int32'));
40+
----
41+
'2147483647'
42+
43+
query T
44+
SELECT quote(arrow_cast(-2147483648, 'Int32'));
45+
----
46+
'-2147483648'
47+
48+
query T
49+
SELECT quote(arrow_cast(9223372036854775807, 'Int64'));
50+
----
51+
'9223372036854775807'
52+
53+
query T
54+
SELECT quote(arrow_cast(-9223372036854775808, 'Int64'));
55+
----
56+
'-9223372036854775808'
57+
58+
query T
59+
SELECT quote(arrow_cast(3.14, 'Float32'));
60+
----
61+
'3.14'
62+
63+
query T
64+
SELECT quote(arrow_cast(2.718281828459045, 'Float64'));
65+
----
66+
'2.718281828459045'
67+
68+
query T
69+
SELECT quote(arrow_cast(0, 'UInt8'));
70+
----
71+
'0'
72+
73+
query T
74+
SELECT quote(arrow_cast(255, 'UInt8'));
75+
----
76+
'255'
77+
78+
query T
79+
SELECT quote(arrow_cast(65535, 'UInt16'));
80+
----
81+
'65535'
82+
83+
query T
84+
SELECT quote(arrow_cast(4294967295, 'UInt32'));
85+
----
86+
'4294967295'
87+
88+
query T
89+
SELECT quote(arrow_cast(18446744073709551615, 'UInt64'));
90+
----
91+
'18446744073709551615'
92+
93+
query T
94+
SELECT quote('special chars: !@#$%^&*()');
95+
----
96+
'special chars: !@#$%^&*()'
97+
98+
query T
99+
SELECT quote('tab\tseparated');
100+
----
101+
'tab\tseparated'
102+
103+
query T
104+
SELECT quote('carriage\rreturn');
105+
----
106+
'carriage\rreturn'
107+
108+
query T
109+
SELECT quote('backslash\\test');
110+
----
111+
'backslash\\test'
112+
113+
query T
114+
SELECT quote('quote\"inside\"');
115+
----
116+
'quote\"inside\"'
117+
118+
query T
119+
SELECT quote('mixed\nescape\tchars\r\n');
120+
----
121+
'mixed\nescape\tchars\r\n'
122+
123+
query T
124+
SELECT quote('unicode: 你好, 世界');
125+
----
126+
'unicode: 你好, 世界'
127+
128+
query T
129+
SELECT quote('emoji: 😀🎉❤️🚀');
130+
----
131+
'emoji: 😀🎉❤️🚀'
132+
133+
query T
134+
SELECT quote(arrow_cast('2024-01-15', 'Date32'));
135+
----
136+
'2024-01-15'
137+
138+
query T
139+
SELECT quote(arrow_cast('2024-01-15T12:30:45', 'Timestamp(µs)'));
140+
----
141+
'2024-01-15T12:30:45'
142+
143+
query T
144+
SELECT quote('special\n\t\r');
145+
----
146+
'special\n\t\r'
147+
148+
query T
149+
SELECT quote('a''b');
150+
----
151+
'a\'b'
152+
153+
query T
154+
SELECT quote('it''s a ''test''');
155+
----
156+
'it\'s a \'test\''
157+
158+
query T
159+
SELECT quote('''');
160+
----
161+
'\''

0 commit comments

Comments
 (0)