From f98b6e9ec45b72540b6d09c4a875ebe60561c367 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Thu, 23 Oct 2025 22:08:45 +0400 Subject: [PATCH 1/7] Fix bitmap_count --- .../spark/src/function/bitmap/bitmap_count.rs | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/datafusion/spark/src/function/bitmap/bitmap_count.rs b/datafusion/spark/src/function/bitmap/bitmap_count.rs index 15bd33229a3d5..1df5db3409f76 100644 --- a/datafusion/spark/src/function/bitmap/bitmap_count.rs +++ b/datafusion/spark/src/function/bitmap/bitmap_count.rs @@ -19,13 +19,14 @@ use std::any::Any; use std::sync::Arc; use arrow::array::{ - Array, ArrayRef, BinaryArray, BinaryViewArray, FixedSizeBinaryArray, Int64Array, - LargeBinaryArray, + as_dictionary_array, Array, ArrayRef, BinaryArray, BinaryViewArray, + FixedSizeBinaryArray, Int64Array, LargeBinaryArray, }; -use arrow::datatypes::DataType; use arrow::datatypes::DataType::{ - Binary, BinaryView, FixedSizeBinary, Int64, LargeBinary, + Binary, BinaryView, Dictionary, FixedSizeBinary, Int64, LargeBinary, }; +use arrow::datatypes::{DataType, Int32Type}; +use datafusion_common::cast::as_binary_array; use datafusion_common::utils::take_function_args; use datafusion_common::{internal_err, Result}; use datafusion_expr::{ @@ -100,6 +101,22 @@ pub fn bitmap_count_inner(arg: &[ArrayRef]) -> Result { FixedSizeBinary(_size) => { downcast_and_count_ones!(input_array, FixedSizeBinaryArray) } + Dictionary(k, v) if k.as_ref() == &DataType::Int32 && v.as_ref() == &Binary => { + let dict_array = as_dictionary_array::(input_array); + let binary_array = as_binary_array(dict_array.values())?; + + let result: Int64Array = dict_array + .keys() + .iter() + .map(|key| { + key.and_then(|k| { + binary_count_ones(Some(binary_array.value(k as usize))) + }) + }) + .collect(); + + Ok(result) + } data_type => { internal_err!("bitmap_count does not support {data_type}") } @@ -113,6 +130,7 @@ mod tests { use crate::function::bitmap::bitmap_count::BitmapCount; use crate::function::utils::test::test_scalar_function; use arrow::array::{Array, Int64Array}; + use arrow::datatypes::DataType; use arrow::datatypes::DataType::Int64; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; From 621acb89cac2c84ac8ba1af05cce8512bfd475d6 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Fri, 24 Oct 2025 20:18:28 +0400 Subject: [PATCH 2/7] add tests --- .../spark/src/function/bitmap/bitmap_count.rs | 34 +++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/datafusion/spark/src/function/bitmap/bitmap_count.rs b/datafusion/spark/src/function/bitmap/bitmap_count.rs index 1df5db3409f76..2858601a4d8f4 100644 --- a/datafusion/spark/src/function/bitmap/bitmap_count.rs +++ b/datafusion/spark/src/function/bitmap/bitmap_count.rs @@ -130,10 +130,13 @@ mod tests { use crate::function::bitmap::bitmap_count::BitmapCount; use crate::function::utils::test::test_scalar_function; use arrow::array::{Array, Int64Array}; - use arrow::datatypes::DataType; use arrow::datatypes::DataType::Int64; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::config::ConfigOptions; use datafusion_common::{Result, ScalarValue}; - use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; + use datafusion_expr::ColumnarValue::Scalar; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; + use std::sync::Arc; macro_rules! test_bitmap_count_binary_invoke { ($INPUT:expr, $EXPECTED:expr) => { @@ -189,4 +192,31 @@ mod tests { ); Ok(()) } + + #[test] + fn test_dictionary_encoded_bitmap_count_invoke() -> Result<()> { + let dict = Scalar(ScalarValue::Dictionary( + Box::new(DataType::Int32), + Box::new(ScalarValue::Binary(Some(vec![0xFFu8, 0xFFu8]))), + )); + + let arg_fields = vec![Field::new( + "a", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)), + true, + ) + .into()]; + let args = ScalarFunctionArgs { + args: vec![dict.clone()], + arg_fields, + number_rows: 1, + return_field: Field::new("f", Int64, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }; + let udf = BitmapCount::new(); + let actual = udf.invoke_with_args(args)?; + let expect = Scalar(ScalarValue::Int64(Some(16))); + assert_eq!(*actual.into_array(1)?, *expect.into_array(1)?); + Ok(()) + } } From 074aa10eb4a8ada8237093cb2e810f838329775d Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 25 Oct 2025 21:26:07 +0400 Subject: [PATCH 3/7] fix MR issues --- .../spark/src/function/bitmap/bitmap_count.rs | 45 ++++++++++--------- .../test_files/spark/bitmap/bitmap_count.slt | 8 ++++ 2 files changed, 33 insertions(+), 20 deletions(-) diff --git a/datafusion/spark/src/function/bitmap/bitmap_count.rs b/datafusion/spark/src/function/bitmap/bitmap_count.rs index 2858601a4d8f4..56a9c5edb812c 100644 --- a/datafusion/spark/src/function/bitmap/bitmap_count.rs +++ b/datafusion/spark/src/function/bitmap/bitmap_count.rs @@ -23,10 +23,9 @@ use arrow::array::{ FixedSizeBinaryArray, Int64Array, LargeBinaryArray, }; use arrow::datatypes::DataType::{ - Binary, BinaryView, Dictionary, FixedSizeBinary, Int64, LargeBinary, + Binary, BinaryView, Dictionary, FixedSizeBinary, LargeBinary, }; -use arrow::datatypes::{DataType, Int32Type}; -use datafusion_common::cast::as_binary_array; +use arrow::datatypes::{DataType, Int16Type, Int32Type, Int64Type, Int8Type}; use datafusion_common::utils::take_function_args; use datafusion_common::{internal_err, Result}; use datafusion_expr::{ @@ -72,7 +71,7 @@ impl ScalarUDFImpl for BitmapCount { } fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(Int64) + Ok(DataType::Int64) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -91,6 +90,17 @@ macro_rules! downcast_and_count_ones { }}; } +macro_rules! downcast_dict_and_count_ones { + ($input_dict:expr, $key_array_type:ident) => {{ + let dict_array = as_dictionary_array::<$key_array_type>($input_dict); + let array = dict_array.downcast_dict::().unwrap(); + Ok(array + .into_iter() + .map(binary_count_ones) + .collect::()) + }}; +} + pub fn bitmap_count_inner(arg: &[ArrayRef]) -> Result { let [input_array] = take_function_args("bitmap_count", arg)?; @@ -101,22 +111,17 @@ pub fn bitmap_count_inner(arg: &[ArrayRef]) -> Result { FixedSizeBinary(_size) => { downcast_and_count_ones!(input_array, FixedSizeBinaryArray) } - Dictionary(k, v) if k.as_ref() == &DataType::Int32 && v.as_ref() == &Binary => { - let dict_array = as_dictionary_array::(input_array); - let binary_array = as_binary_array(dict_array.values())?; - - let result: Int64Array = dict_array - .keys() - .iter() - .map(|key| { - key.and_then(|k| { - binary_count_ones(Some(binary_array.value(k as usize))) - }) - }) - .collect(); - - Ok(result) - } + Dictionary(k, v) if v.as_ref() == &Binary => match k.as_ref() { + DataType::Int8 => downcast_dict_and_count_ones!(input_array, Int8Type), + DataType::Int16 => downcast_dict_and_count_ones!(input_array, Int16Type), + DataType::Int32 => downcast_dict_and_count_ones!(input_array, Int32Type), + DataType::Int64 => downcast_dict_and_count_ones!(input_array, Int64Type), + data_type => { + internal_err!( + "bitmap_count does not support Dictionary({data_type}, Binary)" + ) + } + }, data_type => { internal_err!("bitmap_count does not support {data_type}") } diff --git a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt index 2789efef7bf36..4c66b942195c5 100644 --- a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt +++ b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt @@ -59,3 +59,11 @@ SELECT bitmap_count(arrow_cast(a, 'FixedSizeBinary(2)')) FROM (VALUES (X'1010'), 5 16 NULL + +query I +SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int32, Binary)')) FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); +---- +2 +5 +16 +NULL From ff644ab0f2efeb96c9462f16d5ca51db9c2b226c Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 25 Oct 2025 21:26:53 +0400 Subject: [PATCH 4/7] fix warning --- datafusion/spark/src/function/bitmap/bitmap_count.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/spark/src/function/bitmap/bitmap_count.rs b/datafusion/spark/src/function/bitmap/bitmap_count.rs index 56a9c5edb812c..f882dbbab4221 100644 --- a/datafusion/spark/src/function/bitmap/bitmap_count.rs +++ b/datafusion/spark/src/function/bitmap/bitmap_count.rs @@ -134,7 +134,7 @@ pub fn bitmap_count_inner(arg: &[ArrayRef]) -> Result { mod tests { use crate::function::bitmap::bitmap_count::BitmapCount; use crate::function::utils::test::test_scalar_function; - use arrow::array::{Array, Int64Array}; + use arrow::array::Array; use arrow::datatypes::DataType::Int64; use arrow::datatypes::{DataType, Field}; use datafusion_common::config::ConfigOptions; From 4644c3ac33742f9e8648da5586d2ee0c14d082f6 Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 25 Oct 2025 21:34:41 +0400 Subject: [PATCH 5/7] fix compile --- datafusion/spark/src/function/bitmap/bitmap_count.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/spark/src/function/bitmap/bitmap_count.rs b/datafusion/spark/src/function/bitmap/bitmap_count.rs index f882dbbab4221..56a9c5edb812c 100644 --- a/datafusion/spark/src/function/bitmap/bitmap_count.rs +++ b/datafusion/spark/src/function/bitmap/bitmap_count.rs @@ -134,7 +134,7 @@ pub fn bitmap_count_inner(arg: &[ArrayRef]) -> Result { mod tests { use crate::function::bitmap::bitmap_count::BitmapCount; use crate::function::utils::test::test_scalar_function; - use arrow::array::Array; + use arrow::array::{Array, Int64Array}; use arrow::datatypes::DataType::Int64; use arrow::datatypes::{DataType, Field}; use datafusion_common::config::ConfigOptions; From 3678b9128070898f84a4501d6ffbce27a01f555d Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 25 Oct 2025 22:07:22 +0400 Subject: [PATCH 6/7] test slt --- .../sqllogictest/test_files/spark/bitmap/bitmap_count.slt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt index 4c66b942195c5..1f6e9ad4d8afe 100644 --- a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt +++ b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt @@ -64,6 +64,6 @@ query I SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int32, Binary)')) FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); ---- 2 -5 +6 16 NULL From 3af9bd9d219c28f9ab51faba8effc422d912e77a Mon Sep 17 00:00:00 2001 From: Kazantsev Maksim Date: Sat, 25 Oct 2025 22:26:18 +0400 Subject: [PATCH 7/7] test slt --- .../test_files/spark/bitmap/bitmap_count.slt | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt index 1f6e9ad4d8afe..39dca512226b2 100644 --- a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt +++ b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt @@ -64,6 +64,30 @@ query I SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int32, Binary)')) FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); ---- 2 -6 +5 +16 +NULL + +query I +SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int8, Binary)')) FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); +---- +2 +5 +16 +NULL + +query I +SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int16, Binary)')) FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); +---- +2 +5 +16 +NULL + +query I +SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int64, Binary)')) FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); +---- +2 +5 16 NULL