datafusion_spark/function/string/
make_valid_utf8.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, LargeStringArray, StringArray};
19use arrow::datatypes::{DataType, Field, FieldRef};
20use datafusion_common::cast::{
21    as_binary_array, as_binary_view_array, as_large_binary_array,
22};
23use datafusion_common::utils::take_function_args;
24use datafusion_common::{Result, internal_err};
25use datafusion_expr::{
26    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature,
27    Volatility,
28};
29use datafusion_functions::utils::make_scalar_function;
30use std::sync::Arc;
31
32/// Spark-compatible `make_valid_utf8` expression
33/// <https://spark.apache.org/docs/latest/api/sql/index.html#make_valid_utf8>
34#[derive(Debug, PartialEq, Eq, Hash)]
35pub struct SparkMakeValidUtf8 {
36    signature: Signature,
37}
38
39impl Default for SparkMakeValidUtf8 {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl SparkMakeValidUtf8 {
46    pub fn new() -> Self {
47        Self {
48            signature: Signature::uniform(
49                1,
50                vec![
51                    DataType::Utf8,
52                    DataType::LargeUtf8,
53                    DataType::Utf8View,
54                    DataType::Binary,
55                    DataType::BinaryView,
56                    DataType::LargeBinary,
57                ],
58                Volatility::Immutable,
59            ),
60        }
61    }
62}
63
64impl ScalarUDFImpl for SparkMakeValidUtf8 {
65    fn name(&self) -> &str {
66        "make_valid_utf8"
67    }
68
69    fn signature(&self) -> &Signature {
70        &self.signature
71    }
72
73    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
74        internal_err!("return_field_from_args should be used instead")
75    }
76
77    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
78        let [make_valid_utf8] = take_function_args(self.name(), args.arg_fields)?;
79        let return_type = match make_valid_utf8.data_type() {
80            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
81                Ok(make_valid_utf8.data_type().clone())
82            }
83            DataType::Binary | DataType::BinaryView => Ok(DataType::Utf8),
84            DataType::LargeBinary => Ok(DataType::LargeUtf8),
85            data_type => internal_err!("make_valid_utf8 does not support: {data_type}"),
86        }?;
87        Ok(Arc::new(Field::new(
88            self.name(),
89            return_type,
90            make_valid_utf8.is_nullable(),
91        )))
92    }
93
94    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
95        make_scalar_function(spark_make_valid_utf8_inner, vec![])(&args.args)
96    }
97}
98
99fn spark_make_valid_utf8_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
100    let array = &args[0];
101    match &array.data_type() {
102        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(array.to_owned()),
103        DataType::Binary => Ok(Arc::new(
104            as_binary_array(&array)?
105                .iter()
106                .map(|x| x.map(String::from_utf8_lossy))
107                .collect::<StringArray>(),
108        )),
109        DataType::BinaryView => Ok(Arc::new(
110            as_binary_view_array(&array)?
111                .iter()
112                .map(|x| x.map(String::from_utf8_lossy))
113                .collect::<StringArray>(),
114        )),
115        DataType::LargeBinary => Ok(Arc::new(
116            as_large_binary_array(&array)?
117                .iter()
118                .map(|x| x.map(String::from_utf8_lossy))
119                .collect::<LargeStringArray>(),
120        )),
121        data_type => {
122            internal_err!("make_valid_utf8 does not support: {data_type}")
123        }
124    }
125}
datafusion_spark/function/string/make_valid_utf8.rs

datafusion_spark/function/string/
make_valid_utf8.rs
