[SR GeekCamp] Support function array_to_bitmap in StarRocks (StarRock…

…s#5602)
minchowang · May 20, 2022 · d2e666c · d2e666c
1 parent 2583b79
commit d2e666c
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 1 deletion.
diff --git a/be/src/exprs/vectorized/bitmap_functions.cpp b/be/src/exprs/vectorized/bitmap_functions.cpp
@@ -6,9 +6,11 @@
 #include "column/column_builder.h"
 #include "column/column_helper.h"
 #include "column/column_viewer.h"
+#include "column/nullable_column.h"
 #include "exprs/base64.h"
 #include "exprs/vectorized/binary_function.h"
 #include "exprs/vectorized/unary_function.h"
+#include "gutil/casts.h"
 #include "gutil/strings/split.h"
 #include "gutil/strings/substitute.h"
 #include "util/string_parser.hpp"
@@ -333,6 +335,51 @@ ColumnPtr BitmapFunctions::bitmap_to_array(FunctionContext* context, const starr
     }
 }
 
+ColumnPtr BitmapFunctions::array_to_bitmap(FunctionContext* context, const starrocks::vectorized::Columns& columns) {
+    size_t size = columns[0]->size();
+    ColumnBuilder<TYPE_OBJECT> builder(size);
+    const constexpr PrimitiveType TYPE = TYPE_BIGINT;
+
+    Column* data_column = ColumnHelper::get_data_column(columns[0].get());
+    NullData::pointer null_data = columns[0]->is_nullable()
+                                          ? down_cast<NullableColumn*>(columns[0].get())->null_column_data().data()
+                                          : nullptr;
+    ArrayColumn* array_column = down_cast<ArrayColumn*>(data_column);
+
+    RunTimeColumnType<TYPE>::Container& element_container =
+            array_column->elements_column()->is_nullable()
+                    ? down_cast<RunTimeColumnType<TYPE>*>(
+                              down_cast<NullableColumn*>(array_column->elements_column().get())->data_column().get())
+                              ->get_data()
+                    : down_cast<RunTimeColumnType<TYPE>*>(array_column->elements_column().get())->get_data();
+    const auto& offsets = array_column->offsets_column()->get_data();
+
+    NullColumn::Container::pointer element_null_data =
+            array_column->elements_column()->is_nullable()
+                    ? down_cast<NullableColumn*>(array_column->elements_column().get())->null_column_data().data()
+                    : nullptr;
+
+    for (int row = 0; row < size; ++row) {
+        uint32_t offset = offsets[row];
+        uint32_t length = offsets[row + 1] - offsets[row];
+        if (null_data && null_data[row]) {
+            builder.append_null();
+            continue;
+        }
+        // build bitmap
+        BitmapValue bitmap;
+        for (int j = offset; j < offset + length; j++) {
+            if (element_null_data && element_null_data[j]) {
+                continue;
+            }
+            bitmap.add(element_container[j]);
+        }
+        // append bitmap
+        builder.append(std::move(bitmap));
+    }
+    return builder.build(ColumnHelper::is_all_const(columns));
+}
+
 ColumnPtr BitmapFunctions::bitmap_max(FunctionContext* context, const starrocks::vectorized::Columns& columns) {
     ColumnViewer<TYPE_OBJECT> viewer(columns[0]);
 

diff --git a/be/src/exprs/vectorized/bitmap_functions.h b/be/src/exprs/vectorized/bitmap_functions.h
@@ -103,10 +103,17 @@ class BitmapFunctions {
     /**
      * @param: 
      * @paramType columns: [TYPE_OBJECT]
-     * @return ARRAY_VARCHAR
+     * @return ARRAY_BIGINT
      */
     DEFINE_VECTORIZED_FN(bitmap_to_array);
 
+    /**
+     * @param:
+     * @paramType columns: [TYPE_OBJECT]
+     * @return TYPE_OBJECT
+     */
+    DEFINE_VECTORIZED_FN(array_to_bitmap);
+
     /**
      * @param:
      * @paramType columns: [TYPE_OBJECT]

diff --git a/gensrc/script/vectorized/vectorized_functions.py b/gensrc/script/vectorized/vectorized_functions.py
@@ -473,6 +473,7 @@
     [90600, 'bitmap_max', 'BIGINT', ['BITMAP'], 'BitmapFunctions::bitmap_max', False],
     [90700, 'bitmap_min', 'BIGINT', ['BITMAP'], 'BitmapFunctions::bitmap_min', False],
     [90800, 'base64_to_bitmap', 'BITMAP', ['VARCHAR'], 'BitmapFunctions::base64_to_bitmap', False],
+    [90900, 'array_to_bitmap', 'ARRAY_BIGINT', ['BITMAP'], 'BitmapFunctions::array_to_bitmap', False],
 
     # hash function
     [100010, 'murmur_hash3_32', 'INT', ['VARCHAR', '...'], 'HashFunctions::murmur_hash3_32'],