feat(query): Support bitmap_to_array function (#18927)

Author: b41sh

src/query/functions/src/scalars/bitmap.rs

@@ -192,6 +192,31 @@ pub fn register(registry: &mut FunctionRegistry) {
         }),
     );
 
+    registry.register_passthrough_nullable_1_arg::<BitmapType, ArrayType<NumberType<u64>>, _, _>(
+        "bitmap_to_array",
+        |_, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_1_arg::<BitmapType, ArrayType<NumberType<u64>>>(
+            |b, builder, ctx| {
+                if let Some(validity) = &ctx.validity {
+                    if !validity.get_bit(builder.len()) {
+                        builder.commit_row();
+                        return;
+                    }
+                }
+                match deserialize_bitmap(b) {
+                    Ok(rb) => {
+                        let ids = rb.into_iter().collect::<Vec<_>>();
+                        builder.push(ids.into());
+                    }
+                    Err(e) => {
+                        ctx.set_error(builder.len(), e.to_string());
+                        builder.commit_row();
+                    }
+                }
+            },
+        ),
+    );
+
     registry.register_passthrough_nullable_2_arg::<BitmapType, UInt64Type, BooleanType, _, _>(
         "bitmap_contains",
         |_, _, _| FunctionDomain::Full,

src/query/functions/tests/it/scalars/bitmap.rs

@@ -37,6 +37,7 @@ fn test_bitmap() {
     test_bitmap_subset_limit(file);
     test_bitmap_subset_in_range(file);
     test_bitmap_op(file);
+    test_bitmap_to_array(file);
 }
 
 fn test_build_bitmap(file: &mut impl Write) {
@@ -140,3 +141,11 @@ fn test_bitmap_op(file: &mut impl Write) {
         &[],
     );
 }
+
+fn test_bitmap_to_array(file: &mut impl Write) {
+    run_ast(file, "bitmap_to_array(build_bitmap([1,4,5]))", &[]);
+    run_ast(file, "bitmap_to_array(build_bitmap([a, b]))", &[
+        ("a", UInt16Type::from_data(vec![1u16, 2, 3])),
+        ("b", UInt16Type::from_data(vec![1u16, 5, 6])),
+    ]);
+}

src/query/functions/tests/it/scalars/testdata/bitmap.txt

@@ -209,3 +209,35 @@ output domain  : Undefined
 output         : '3'
 
 
+ast            : bitmap_to_array(build_bitmap([1,4,5]))
+raw expr       : bitmap_to_array(build_bitmap(array(1, 4, 5)))
+checked expr   : bitmap_to_array<Bitmap>(build_bitmap<Array(UInt8 NULL)>(CAST<Array(UInt8)>(array<T0=UInt8><T0, T0, T0>(1_u8, 4_u8, 5_u8) AS Array(UInt8 NULL))))
+optimized expr : [1, 4, 5]
+output type    : Array(UInt64)
+output domain  : [{1..=5}]
+output         : [1, 4, 5]
+
+
+ast            : bitmap_to_array(build_bitmap([a, b]))
+raw expr       : bitmap_to_array(build_bitmap(array(a::UInt16, b::UInt16)))
+checked expr   : bitmap_to_array<Bitmap>(build_bitmap<Array(UInt16 NULL)>(CAST<Array(UInt16)>(array<T0=UInt16><T0, T0>(a, b) AS Array(UInt16 NULL))))
+evaluation:
++--------+---------+---------+---------------+
+|        | a       | b       | Output        |
++--------+---------+---------+---------------+
+| Type   | UInt16  | UInt16  | Array(UInt64) |
+| Domain | {1..=3} | {1..=6} | Unknown       |
+| Row 0  | 1       | 1       | [1]           |
+| Row 1  | 2       | 5       | [2, 5]        |
+| Row 2  | 3       | 6       | [3, 6]        |
++--------+---------+---------+---------------+
+evaluation (internal):
++--------+------------------------------------------------------------------------+
+| Column | Data                                                                   |
++--------+------------------------------------------------------------------------+
+| a      | UInt16([1, 2, 3])                                                      |
+| b      | UInt16([1, 5, 6])                                                      |
+| Output | ArrayColumn { values: UInt64([1, 2, 5, 3, 6]), offsets: [0, 1, 3, 5] } |
++--------+------------------------------------------------------------------------+
+
+

src/query/functions/tests/it/scalars/testdata/function_list.txt

@@ -920,6 +920,8 @@ Functions overloads:
 1 bitmap_subset_in_range(Bitmap NULL, UInt64 NULL, UInt64 NULL) :: Bitmap NULL
 0 bitmap_subset_limit(Bitmap, UInt64, UInt64) :: Bitmap
 1 bitmap_subset_limit(Bitmap NULL, UInt64 NULL, UInt64 NULL) :: Bitmap NULL
+0 bitmap_to_array(Bitmap) :: Array(UInt64)
+1 bitmap_to_array(Bitmap NULL) :: Array(UInt64) NULL
 0 bitmap_xor(Bitmap, Bitmap) :: Bitmap
 1 bitmap_xor(Bitmap NULL, Bitmap NULL) :: Bitmap NULL
 0 blake3(String) :: String

tests/sqllogictests/suites/query/functions/02_0064_function_bitmap.test

@@ -8,7 +8,7 @@ SELECT to_string(to_bitmap(1024))
 ----
 1024
 
-query II
+query TII
 SELECT build_bitmap([1,4,5])::String, bitmap_count(build_bitmap([1,2,5])), bitmap_count(build_bitmap([1, 66, 69, 70]))
 ----
 1,4,5 3 4
@@ -72,12 +72,18 @@ SELECT sub_bitmap(build_bitmap([1, 2, 3, 4, 5]), 1, 3)::String;
 ----
 2,3,4
 
-query II
+query T
 SELECT bitmap_subset_limit(build_bitmap([3,5,7]), 4, 2)::String;
 ----
 5,7
 
-query TT
+query T
 SELECT bitmap_subset_in_range(build_bitmap([5,7,9]), 6, 9)::String;
 ----
 7
+
+query T
+SELECT bitmap_to_array(build_bitmap([1, 2, 3, 100, 10000]));
+----
+[1,2,3,100,10000]
+