datafusion_comet_spark_expr/string_funcs/
chr.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::{any::Any, sync::Arc};
19
20use arrow::{
21    array::{ArrayRef, StringArray},
22    datatypes::{
23        DataType,
24        DataType::{Int64, Utf8},
25    },
26};
27
28use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
29use datafusion_common::{cast::as_int64_array, exec_err, Result, ScalarValue};
30
31fn chr(args: &[ArrayRef]) -> Result<ArrayRef> {
32    let integer_array = as_int64_array(&args[0])?;
33
34    // first map is the iterator, second is for the `Option<_>`
35    let result = integer_array
36        .iter()
37        .map(|integer: Option<i64>| {
38            integer
39                .map(|integer| {
40                    if integer < 0 {
41                        return Ok("".to_string()); // Return empty string for negative integers
42                    }
43                    match core::char::from_u32((integer % 256) as u32) {
44                        Some(ch) => Ok(ch.to_string()),
45                        None => {
46                            exec_err!("requested character not compatible for encoding.")
47                        }
48                    }
49                })
50                .transpose()
51        })
52        .collect::<Result<StringArray>>()?;
53
54    Ok(Arc::new(result) as ArrayRef)
55}
56
57/// Spark-compatible `chr` expression
58#[derive(Debug)]
59pub struct SparkChrFunc {
60    signature: Signature,
61}
62
63impl Default for SparkChrFunc {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69impl SparkChrFunc {
70    pub fn new() -> Self {
71        Self {
72            signature: Signature::uniform(1, vec![Int64], Volatility::Immutable),
73        }
74    }
75}
76
77impl ScalarUDFImpl for SparkChrFunc {
78    fn as_any(&self) -> &dyn Any {
79        self
80    }
81
82    fn name(&self) -> &str {
83        "chr"
84    }
85
86    fn signature(&self) -> &Signature {
87        &self.signature
88    }
89
90    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
91        Ok(Utf8)
92    }
93
94    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
95        spark_chr(args)
96    }
97}
98
99/// Returns the ASCII character having the binary equivalent to the input expression.
100/// E.g., chr(65) = 'A'.
101/// Compatible with Apache Spark's Chr function
102fn spark_chr(args: &[ColumnarValue]) -> Result<ColumnarValue> {
103    let array = args[0].clone();
104    match array {
105        ColumnarValue::Array(array) => {
106            let array = chr(&[array])?;
107            Ok(ColumnarValue::Array(array))
108        }
109        ColumnarValue::Scalar(ScalarValue::Int64(Some(value))) => {
110            if value < 0 {
111                Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
112                    "".to_string(),
113                ))))
114            } else {
115                match core::char::from_u32((value % 256) as u32) {
116                    Some(ch) => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
117                        ch.to_string(),
118                    )))),
119                    None => exec_err!("requested character was incompatible for encoding."),
120                }
121            }
122        }
123        _ => exec_err!("The argument must be an Int64 array or scalar."),
124    }
125}
datafusion_comet_spark_expr/string_funcs/chr.rs

datafusion_comet_spark_expr/string_funcs/
chr.rs
