datafusion_spark/function/bitmap/
bitmap_bit_position.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, AsArray, Int64Array};
19use arrow::datatypes::Field;
20use arrow::datatypes::{DataType, FieldRef, Int8Type, Int16Type, Int32Type, Int64Type};
21use datafusion_common::utils::take_function_args;
22use datafusion_common::{Result, internal_err};
23use datafusion_expr::{
24    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
25    Volatility,
26};
27use datafusion_functions::utils::make_scalar_function;
28use std::any::Any;
29use std::sync::Arc;
30
31/// Spark-compatible `bitmap_bit_position` expression
32/// <https://spark.apache.org/docs/latest/api/sql/index.html#bitmap_bit_position>
33#[derive(Debug, PartialEq, Eq, Hash)]
34pub struct BitmapBitPosition {
35    signature: Signature,
36}
37
38impl Default for BitmapBitPosition {
39    fn default() -> Self {
40        Self::new()
41    }
42}
43
44impl BitmapBitPosition {
45    pub fn new() -> Self {
46        Self {
47            signature: Signature::one_of(
48                vec![
49                    TypeSignature::Exact(vec![DataType::Int8]),
50                    TypeSignature::Exact(vec![DataType::Int16]),
51                    TypeSignature::Exact(vec![DataType::Int32]),
52                    TypeSignature::Exact(vec![DataType::Int64]),
53                ],
54                Volatility::Immutable,
55            ),
56        }
57    }
58}
59
60impl ScalarUDFImpl for BitmapBitPosition {
61    fn as_any(&self) -> &dyn Any {
62        self
63    }
64
65    fn name(&self) -> &str {
66        "bitmap_bit_position"
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(
78        &self,
79        args: datafusion_expr::ReturnFieldArgs,
80    ) -> Result<FieldRef> {
81        Ok(Arc::new(Field::new(
82            self.name(),
83            DataType::Int64,
84            args.arg_fields[0].is_nullable(),
85        )))
86    }
87
88    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
89        make_scalar_function(bitmap_bit_position_inner, vec![])(&args.args)
90    }
91}
92
93pub fn bitmap_bit_position_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
94    let [array] = take_function_args("bitmap_bit_position", arg)?;
95    match &array.data_type() {
96        DataType::Int8 => {
97            let result: Int64Array = array
98                .as_primitive::<Int8Type>()
99                .iter()
100                .map(|opt| opt.map(|value| bitmap_bit_position(value.into())))
101                .collect();
102            Ok(Arc::new(result))
103        }
104        DataType::Int16 => {
105            let result: Int64Array = array
106                .as_primitive::<Int16Type>()
107                .iter()
108                .map(|opt| opt.map(|value| bitmap_bit_position(value.into())))
109                .collect();
110            Ok(Arc::new(result))
111        }
112        DataType::Int32 => {
113            let result: Int64Array = array
114                .as_primitive::<Int32Type>()
115                .iter()
116                .map(|opt| opt.map(|value| bitmap_bit_position(value.into())))
117                .collect();
118            Ok(Arc::new(result))
119        }
120        DataType::Int64 => {
121            let result: Int64Array = array
122                .as_primitive::<Int64Type>()
123                .iter()
124                .map(|opt| opt.map(bitmap_bit_position))
125                .collect();
126            Ok(Arc::new(result))
127        }
128        data_type => {
129            internal_err!("bitmap_bit_position does not support {data_type}")
130        }
131    }
132}
133
134const NUM_BYTES: i64 = 4 * 1024;
135const NUM_BITS: i64 = NUM_BYTES * 8;
136
137fn bitmap_bit_position(value: i64) -> i64 {
138    if value > 0 {
139        (value - 1) % NUM_BITS
140    } else {
141        (value.wrapping_neg()) % NUM_BITS
142    }
143}
datafusion_spark/function/bitmap/bitmap_bit_position.rs

datafusion_spark/function/bitmap/
bitmap_bit_position.rs
