datafusion_spark/function/bitwise/
bitwise_not.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::compute::kernels::bitwise;
19use arrow::datatypes::{Int16Type, Int32Type, Int64Type, Int8Type};
20use arrow::{array::*, datatypes::DataType};
21use datafusion_common::{plan_err, Result};
22use datafusion_expr::{ColumnarValue, TypeSignature, Volatility};
23use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature};
24use datafusion_functions::utils::make_scalar_function;
25use std::{any::Any, sync::Arc};
26
27#[derive(Debug, PartialEq, Eq, Hash)]
28pub struct SparkBitwiseNot {
29    signature: Signature,
30}
31
32impl Default for SparkBitwiseNot {
33    fn default() -> Self {
34        Self::new()
35    }
36}
37
38impl SparkBitwiseNot {
39    pub fn new() -> Self {
40        Self {
41            signature: Signature::one_of(
42                vec![
43                    TypeSignature::Exact(vec![DataType::Int8]),
44                    TypeSignature::Exact(vec![DataType::Int16]),
45                    TypeSignature::Exact(vec![DataType::Int32]),
46                    TypeSignature::Exact(vec![DataType::Int64]),
47                ],
48                Volatility::Immutable,
49            ),
50        }
51    }
52}
53
54impl ScalarUDFImpl for SparkBitwiseNot {
55    fn as_any(&self) -> &dyn Any {
56        self
57    }
58
59    fn name(&self) -> &str {
60        "bitwise_not"
61    }
62
63    fn signature(&self) -> &Signature {
64        &self.signature
65    }
66
67    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
68        Ok(arg_types[0].clone())
69    }
70
71    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
72        if args.args.len() != 1 {
73            return plan_err!("bitwise_not expects exactly 1 argument");
74        }
75        make_scalar_function(spark_bitwise_not, vec![])(&args.args)
76    }
77}
78
79pub fn spark_bitwise_not(args: &[ArrayRef]) -> Result<ArrayRef> {
80    let array = args[0].as_ref();
81    match array.data_type() {
82        DataType::Int8 => {
83            let result: Int8Array =
84                bitwise::bitwise_not(array.as_primitive::<Int8Type>())?;
85            Ok(Arc::new(result))
86        }
87        DataType::Int16 => {
88            let result: Int16Array =
89                bitwise::bitwise_not(array.as_primitive::<Int16Type>())?;
90            Ok(Arc::new(result))
91        }
92        DataType::Int32 => {
93            let result: Int32Array =
94                bitwise::bitwise_not(array.as_primitive::<Int32Type>())?;
95            Ok(Arc::new(result))
96        }
97        DataType::Int64 => {
98            let result: Int64Array =
99                bitwise::bitwise_not(array.as_primitive::<Int64Type>())?;
100            Ok(Arc::new(result))
101        }
102        _ => {
103            plan_err!(
104                "bitwise_not function does not support data type: {}",
105                array.data_type()
106            )
107        }
108    }
109}
datafusion_spark/function/bitwise/bitwise_not.rs

datafusion_spark/function/bitwise/
bitwise_not.rs
