datafusion_spark/function/bitmap/
bitmap_count.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::sync::Arc;
19
20use arrow::array::{
21    Array, ArrayRef, BinaryArray, BinaryViewArray, FixedSizeBinaryArray, Int64Array,
22    LargeBinaryArray, as_dictionary_array,
23};
24use arrow::datatypes::DataType::{
25    Binary, BinaryView, Dictionary, FixedSizeBinary, LargeBinary,
26};
27use arrow::datatypes::{DataType, FieldRef, Int8Type, Int16Type, Int32Type, Int64Type};
28use datafusion_common::utils::take_function_args;
29use datafusion_common::{Result, internal_err};
30use datafusion_expr::{
31    Coercion, ColumnarValue, EncodingPreservation, ScalarFunctionArgs, ScalarUDFImpl,
32    Signature, TypeSignatureClass, Volatility,
33};
34use datafusion_functions::downcast_arg;
35use datafusion_functions::utils::make_scalar_function;
36
37#[derive(Debug, PartialEq, Eq, Hash)]
38pub struct BitmapCount {
39    signature: Signature,
40}
41
42impl Default for BitmapCount {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl BitmapCount {
49    pub fn new() -> Self {
50        Self {
51            signature: Signature::coercible(
52                vec![
53                    Coercion::new_exact(TypeSignatureClass::Binary)
54                        .with_encoding_preservation(EncodingPreservation::dictionary()),
55                ],
56                Volatility::Immutable,
57            ),
58        }
59    }
60}
61
62impl ScalarUDFImpl for BitmapCount {
63    fn name(&self) -> &str {
64        "bitmap_count"
65    }
66
67    fn signature(&self) -> &Signature {
68        &self.signature
69    }
70
71    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
72        internal_err!("return_field_from_args should be used instead")
73    }
74
75    fn return_field_from_args(
76        &self,
77        args: datafusion_expr::ReturnFieldArgs,
78    ) -> Result<FieldRef> {
79        use arrow::datatypes::Field;
80        // bitmap_count returns Int64 with the same nullability as the input
81        Ok(Arc::new(Field::new(
82            args.arg_fields[0].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_count_inner, vec![])(&args.args)
90    }
91}
92
93fn binary_count_ones(opt: Option<&[u8]>) -> Option<i64> {
94    opt.map(|value| value.iter().map(|b| b.count_ones() as i64).sum())
95}
96
97macro_rules! downcast_and_count_ones {
98    ($input_array:expr, $array_type:ident) => {{
99        let arr = downcast_arg!($input_array, $array_type);
100        Ok(arr.iter().map(binary_count_ones).collect::<Int64Array>())
101    }};
102}
103
104macro_rules! downcast_dict_and_count_ones {
105    ($input_dict:expr, $key_array_type:ident) => {{
106        let dict_array = as_dictionary_array::<$key_array_type>($input_dict);
107        let array = dict_array.downcast_dict::<BinaryArray>().unwrap();
108        Ok(array
109            .into_iter()
110            .map(binary_count_ones)
111            .collect::<Int64Array>())
112    }};
113}
114
115pub fn bitmap_count_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
116    let [input_array] = take_function_args("bitmap_count", arg)?;
117
118    let res: Result<Int64Array> = match &input_array.data_type() {
119        Binary => downcast_and_count_ones!(input_array, BinaryArray),
120        BinaryView => downcast_and_count_ones!(input_array, BinaryViewArray),
121        LargeBinary => downcast_and_count_ones!(input_array, LargeBinaryArray),
122        FixedSizeBinary(_size) => {
123            downcast_and_count_ones!(input_array, FixedSizeBinaryArray)
124        }
125        Dictionary(k, v) if v.as_ref() == &Binary => match k.as_ref() {
126            DataType::Int8 => downcast_dict_and_count_ones!(input_array, Int8Type),
127            DataType::Int16 => downcast_dict_and_count_ones!(input_array, Int16Type),
128            DataType::Int32 => downcast_dict_and_count_ones!(input_array, Int32Type),
129            DataType::Int64 => downcast_dict_and_count_ones!(input_array, Int64Type),
130            data_type => {
131                internal_err!(
132                    "bitmap_count does not support Dictionary({data_type}, Binary)"
133                )
134            }
135        },
136        data_type => {
137            internal_err!("bitmap_count does not support {data_type}")
138        }
139    };
140
141    Ok(Arc::new(res?))
142}
143
144#[cfg(test)]
145mod tests {
146    use crate::function::bitmap::bitmap_count::BitmapCount;
147    use crate::function::utils::test::test_scalar_function;
148    use arrow::array::{Array, Int64Array};
149    use arrow::datatypes::DataType::Int64;
150    use arrow::datatypes::{DataType, Field};
151    use datafusion_common::config::ConfigOptions;
152    use datafusion_common::{Result, ScalarValue};
153    use datafusion_expr::ColumnarValue::Scalar;
154    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
155    use std::sync::Arc;
156
157    macro_rules! test_bitmap_count_binary_invoke {
158        ($INPUT:expr, $EXPECTED:expr) => {
159            test_scalar_function!(
160                BitmapCount::new(),
161                vec![ColumnarValue::Scalar(ScalarValue::Binary($INPUT))],
162                $EXPECTED,
163                i64,
164                Int64,
165                Int64Array
166            );
167
168            test_scalar_function!(
169                BitmapCount::new(),
170                vec![ColumnarValue::Scalar(ScalarValue::LargeBinary($INPUT))],
171                $EXPECTED,
172                i64,
173                Int64,
174                Int64Array
175            );
176
177            test_scalar_function!(
178                BitmapCount::new(),
179                vec![ColumnarValue::Scalar(ScalarValue::BinaryView($INPUT))],
180                $EXPECTED,
181                i64,
182                Int64,
183                Int64Array
184            );
185
186            test_scalar_function!(
187                BitmapCount::new(),
188                vec![ColumnarValue::Scalar(ScalarValue::FixedSizeBinary(
189                    $INPUT.map(|a| a.len()).unwrap_or(0) as i32,
190                    $INPUT
191                ))],
192                $EXPECTED,
193                i64,
194                Int64,
195                Int64Array
196            );
197        };
198    }
199
200    #[test]
201    fn test_bitmap_count_invoke() -> Result<()> {
202        test_bitmap_count_binary_invoke!(None::<Vec<u8>>, Ok(None));
203        test_bitmap_count_binary_invoke!(Some(vec![0x0Au8]), Ok(Some(2)));
204        test_bitmap_count_binary_invoke!(Some(vec![0xFFu8, 0xFFu8]), Ok(Some(16)));
205        test_bitmap_count_binary_invoke!(
206            Some(vec![0x0Au8, 0xB0u8, 0xCDu8]),
207            Ok(Some(10))
208        );
209        Ok(())
210    }
211
212    #[test]
213    fn test_dictionary_encoded_bitmap_count_invoke() -> Result<()> {
214        let dict = Scalar(ScalarValue::Dictionary(
215            Box::new(DataType::Int32),
216            Box::new(ScalarValue::Binary(Some(vec![0xFFu8, 0xFFu8]))),
217        ));
218
219        let arg_fields = vec![
220            Field::new(
221                "a",
222                DataType::Dictionary(
223                    Box::new(DataType::Int32),
224                    Box::new(DataType::Binary),
225                ),
226                true,
227            )
228            .into(),
229        ];
230        let args = ScalarFunctionArgs {
231            args: vec![dict.clone()],
232            arg_fields,
233            number_rows: 1,
234            return_field: Field::new("f", Int64, true).into(),
235            config_options: Arc::new(ConfigOptions::default()),
236        };
237        let udf = BitmapCount::new();
238        let actual = udf.invoke_with_args(args)?;
239        let expect = Scalar(ScalarValue::Int64(Some(16)));
240        assert_eq!(*actual.into_array(1)?, *expect.into_array(1)?);
241        Ok(())
242    }
243
244    #[test]
245    fn test_bitmap_count_nullability() -> Result<()> {
246        use datafusion_expr::ReturnFieldArgs;
247
248        let bitmap_count = BitmapCount::new();
249
250        // Test with non-nullable binary field
251        let non_nullable_field = Arc::new(Field::new("bin", DataType::Binary, false));
252
253        let result = bitmap_count.return_field_from_args(ReturnFieldArgs {
254            arg_fields: &[Arc::clone(&non_nullable_field)],
255            scalar_arguments: &[None],
256        })?;
257
258        // The result should not be nullable (same as input)
259        assert!(!result.is_nullable());
260        assert_eq!(result.data_type(), &Int64);
261
262        // Test with nullable binary field
263        let nullable_field = Arc::new(Field::new("bin", DataType::Binary, true));
264
265        let result = bitmap_count.return_field_from_args(ReturnFieldArgs {
266            arg_fields: &[Arc::clone(&nullable_field)],
267            scalar_arguments: &[None],
268        })?;
269
270        // The result should be nullable (same as input)
271        assert!(result.is_nullable());
272        assert_eq!(result.data_type(), &Int64);
273
274        Ok(())
275    }
276}
datafusion_spark/function/bitmap/bitmap_count.rs

datafusion_spark/function/bitmap/
bitmap_count.rs
