feat: Implement Spark bitmap_bit_position function

datafusion/spark/src/function/bitmap/bitmap_bit_position.rs

@@ -0,0 +1,141 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{ArrayRef, AsArray, Int64Array};
+use arrow::datatypes::Field;
+use arrow::datatypes::{DataType, FieldRef, Int8Type, Int16Type, Int32Type, Int64Type};
+use datafusion::logical_expr::{ColumnarValue, Signature, TypeSignature, Volatility};
+use datafusion_common::utils::take_function_args;
+use datafusion_common::{Result, internal_err};
+use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_functions::utils::make_scalar_function;
+use std::any::Any;
+use std::sync::Arc;
+
+/// Spark-compatible `bitmap_bit_position` expression
+/// <https://spark.apache.org/docs/latest/api/sql/index.html#bitmap_bit_position>
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct BitmapBitPosition {
+    signature: Signature,
+}
+
+impl Default for BitmapBitPosition {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl BitmapBitPosition {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::Exact(vec![DataType::Int8]),
+                    TypeSignature::Exact(vec![DataType::Int16]),
+                    TypeSignature::Exact(vec![DataType::Int32]),
+                    TypeSignature::Exact(vec![DataType::Int64]),
+                ],
+                Volatility::Immutable,
+            ),
+        }
+    }
+}
+
+impl ScalarUDFImpl for BitmapBitPosition {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "bitmap_bit_position"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        internal_err!("return_field_from_args should be used instead")
+    }
+
+    fn return_field_from_args(
+        &self,
+        args: datafusion_expr::ReturnFieldArgs,
+    ) -> Result<FieldRef> {
+        Ok(Arc::new(Field::new(
+            self.name(),
+            DataType::Int64,
+            args.arg_fields[0].is_nullable(),
+        )))
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        make_scalar_function(bitmap_bit_position_inner, vec![])(&args.args)
+    }
+}
+
+pub fn bitmap_bit_position_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
+    let [array] = take_function_args("bitmap_bit_position", arg)?;
+    match &array.data_type() {
+        DataType::Int8 => {
+            let result: Int64Array = array
+                .as_primitive::<Int8Type>()
+                .iter()
+                .map(|opt| opt.map(|value| bitmap_bit_position(value.into())))
+                .collect();
+            Ok(Arc::new(result))
+        }
+        DataType::Int16 => {
+            let result: Int64Array = array
+                .as_primitive::<Int16Type>()
+                .iter()
+                .map(|opt| opt.map(|value| bitmap_bit_position(value.into())))
+                .collect();
+            Ok(Arc::new(result))
+        }
+        DataType::Int32 => {
+            let result: Int64Array = array
+                .as_primitive::<Int32Type>()
+                .iter()
+                .map(|opt| opt.map(|value| bitmap_bit_position(value.into())))
+                .collect();
+            Ok(Arc::new(result))
+        }
+        DataType::Int64 => {
+            let result: Int64Array = array
+                .as_primitive::<Int64Type>()
+                .iter()
+                .map(|opt| opt.map(bitmap_bit_position))
+                .collect();
+            Ok(Arc::new(result))
+        }
+        data_type => {
+            internal_err!("bitmap_bit_position does not support {data_type}")
+        }
+    }
+}
+
+const NUM_BYTES: i64 = 4 * 1024;
+const NUM_BITS: i64 = NUM_BYTES * 8;
+
+fn bitmap_bit_position(value: i64) -> i64 {
+    if value > 0 {
+        (value - 1) % NUM_BITS
+    } else {
+        (value.wrapping_neg()) % NUM_BITS
+    }
+}

datafusion/spark/src/function/bitmap/mod.rs

@@ -15,13 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
+pub mod bitmap_bit_position;
 pub mod bitmap_count;
 
 use datafusion_expr::ScalarUDF;
 use datafusion_functions::make_udf_function;
 use std::sync::Arc;
 
 make_udf_function!(bitmap_count::BitmapCount, bitmap_count);
+make_udf_function!(bitmap_bit_position::BitmapBitPosition, bitmap_bit_position);
 
 pub mod expr_fn {
     use datafusion_functions::export_functions;
@@ -31,8 +33,13 @@ pub mod expr_fn {
         "Returns the number of set bits in the input bitmap.",
         arg
     ));
+    export_functions!((
+        bitmap_bit_position,
+        "Returns the bit position for the given input child expression.",
+        arg
+    ));
 }
 
 pub fn functions() -> Vec<Arc<ScalarUDF>> {
-    vec![bitmap_count()]
+    vec![bitmap_count(), bitmap_bit_position()]
 }

datafusion/sqllogictest/test_files/spark/bitmap/bitmap_bit_position.slt

@@ -0,0 +1,112 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+
+#   http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+query I
+SELECT bitmap_bit_position(arrow_cast(1, 'Int8'));
+----
+0
+
+query I
+SELECT bitmap_bit_position(arrow_cast(3, 'Int8'));
+----
+2
+
+query I
+SELECT bitmap_bit_position(arrow_cast(7, 'Int8'));
+----
+6
+
+query I
+SELECT bitmap_bit_position(arrow_cast(15, 'Int8'));
+----
+14
+
+query I
+SELECT bitmap_bit_position(arrow_cast(-1, 'Int8'));
+----
+1
+
+query I
+SELECT bitmap_bit_position(arrow_cast(256, 'Int16'));
+----
+255
+
+query I
+SELECT bitmap_bit_position(arrow_cast(1024, 'Int16'));
+----
+1023
+
+query I
+SELECT bitmap_bit_position(arrow_cast(-32768, 'Int16'));
+----
+0
+
+query I
+SELECT bitmap_bit_position(arrow_cast(16384, 'Int16'));
+----
+16383
+
+query I
+SELECT bitmap_bit_position(arrow_cast(-1, 'Int16'));
+----
+1
+
+query I
+SELECT bitmap_bit_position(arrow_cast(65536, 'Int32'));
+----
+32767
+
+query I
+SELECT bitmap_bit_position(arrow_cast(1048576, 'Int32'));
+----
+32767
+
+query I
+SELECT bitmap_bit_position(arrow_cast(-2147483648, 'Int32'));
+----
+0
+
+query I
+SELECT bitmap_bit_position(arrow_cast(1073741824, 'Int32'));
+----
+32767
+
+query I
+SELECT bitmap_bit_position(arrow_cast(-1, 'Int32'));
+----
+1
+
+query I
+SELECT bitmap_bit_position(arrow_cast(4294967296, 'Int64'));
+----
+32767
+
+query I
+SELECT bitmap_bit_position(arrow_cast(-1, 'Int64'));
+----
+1
+
+query I
+SELECT bitmap_bit_position(arrow_cast(-9223372036854775808, 'Int64'));
+----
+0
+
+query I
+SELECT bitmap_bit_position(arrow_cast(9223372036854775807, 'Int64'));
+----
+32766