[Enhancement](function) add to_bitmap() function with int type (#13973)

to_bitmap function only support string param only，add to_bitmap() function with int type, this can avoid convert int type to string and then convert string to int
apache · Nov 8, 2022 · e1654bc · e1654bc
1 parent 34f43ac
commit e1654bc
Show file tree

Hide file tree

Showing 6 changed files with 276 additions and 152 deletions.
diff --git a/be/src/exprs/bitmap_function.cpp b/be/src/exprs/bitmap_function.cpp
@@ -197,6 +197,36 @@ StringVal BitmapFunctions::to_bitmap_with_check(doris_udf::FunctionContext* ctx,
     return serialize(ctx, &bitmap);
 }
 
+StringVal BitmapFunctions::to_bitmap(doris_udf::FunctionContext* ctx,
+                                     const doris_udf::BigIntVal& src) {
+    BitmapValue bitmap;
+    if (LIKELY(!src.is_null && src.val >= 0)) {
+        bitmap.add(src.val);
+    }
+    return serialize(ctx, &bitmap);
+}
+
+StringVal BitmapFunctions::to_bitmap_with_check(doris_udf::FunctionContext* ctx,
+                                                const doris_udf::BigIntVal& src) {
+    BitmapValue bitmap;
+
+    if (!src.is_null) {
+        if (LIKELY(src.val >= 0)) {
+            bitmap.add(src.val);
+        } else {
+            std::stringstream ss;
+            ss << "The input: " << src.val
+               << " is not valid, to_bitmap only support bigint value from 0 to "
+                  "18446744073709551615 currently, cannot load negative values to column with"
+                  " to_bitmap MV on it.";
+            ctx->set_error(ss.str().c_str());
+            return StringVal::null();
+        }
+    }
+
+    return serialize(ctx, &bitmap);
+}
+
 StringVal BitmapFunctions::bitmap_hash(doris_udf::FunctionContext* ctx,
                                        const doris_udf::StringVal& src) {
     BitmapValue bitmap;

diff --git a/be/src/exprs/bitmap_function.h b/be/src/exprs/bitmap_function.h
@@ -69,6 +69,8 @@ class BitmapFunctions {
     static StringVal bitmap_serialize(FunctionContext* ctx, const StringVal& src);
     static StringVal to_bitmap(FunctionContext* ctx, const StringVal& src);
     static StringVal to_bitmap_with_check(FunctionContext* ctx, const StringVal& src);
+    static StringVal to_bitmap(FunctionContext* ctx, const BigIntVal& src);
+    static StringVal to_bitmap_with_check(FunctionContext* ctx, const BigIntVal& src);
     static StringVal bitmap_hash(FunctionContext* ctx, const StringVal& src);
     static StringVal bitmap_hash64(FunctionContext* ctx, const StringVal& src);
     static StringVal bitmap_or(FunctionContext* ctx, const StringVal& src, const StringVal& dst);

diff --git a/be/src/vec/functions/function_always_not_nullable.h b/be/src/vec/functions/function_always_not_nullable.h
@@ -41,56 +41,78 @@ class FunctionAlwaysNotNullable : public IFunction {
     bool use_default_implementation_for_constants() const override { return true; }
     bool use_default_implementation_for_nulls() const override { return false; }
 
-    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
-                        size_t result, size_t input_rows_count) override {
-        auto column = block.get_by_position(arguments[0]).column;
-
-        MutableColumnPtr column_result = get_return_type_impl({})->create_column();
-        column_result->resize(input_rows_count);
-
+    template <typename ColumnType, bool is_nullable>
+    Status execute_internal(const ColumnPtr& column, const DataTypePtr& data_type,
+                            MutableColumnPtr& column_result) {
         auto type_error = [&]() {
             return Status::RuntimeError("Illegal column {} of argument of function {}",
-                                        block.get_by_position(arguments[0]).column->get_name(),
-                                        get_name());
+                                        column->get_name(), get_name());
         };
-
-        if (const ColumnNullable* col_nullable =
-                    check_and_get_column<ColumnNullable>(column.get())) {
-            const ColumnString* col =
-                    check_and_get_column<ColumnString>(col_nullable->get_nested_column_ptr().get());
+        if constexpr (is_nullable) {
+            const ColumnNullable* col_nullable = check_and_get_column<ColumnNullable>(column.get());
+            const ColumnType* col =
+                    check_and_get_column<ColumnType>(col_nullable->get_nested_column_ptr().get());
             const ColumnUInt8* col_nullmap = check_and_get_column<ColumnUInt8>(
                     col_nullable->get_null_map_column_ptr().get());
 
             if (col != nullptr && col_nullmap != nullptr) {
                 if constexpr (WithReturn) {
-                    RETURN_IF_ERROR(Function::vector_nullable(col->get_chars(), col->get_offsets(),
-                                                              col_nullmap->get_data(),
-                                                              column_result));
+                    RETURN_IF_ERROR(
+                            Function::vector_nullable(col, col_nullmap->get_data(), column_result));
                 } else {
-                    Function::vector_nullable(col->get_chars(), col->get_offsets(),
-                                              col_nullmap->get_data(), column_result);
+                    Function::vector_nullable(col, col_nullmap->get_data(), column_result);
                 }
-
-                block.replace_by_position(result, std::move(column_result));
-                return Status::OK();
             } else {
                 return type_error();
             }
-        } else if (const ColumnString* col = check_and_get_column<ColumnString>(column.get())) {
+        } else {
+            const ColumnType* col = check_and_get_column<ColumnType>(column.get());
             if constexpr (WithReturn) {
-                RETURN_IF_ERROR(
-                        Function::vector(col->get_chars(), col->get_offsets(), column_result));
+                RETURN_IF_ERROR(Function::vector(col, column_result));
             } else {
-                Function::vector(col->get_chars(), col->get_offsets(), column_result);
+                Function::vector(col, column_result);
             }
-            block.replace_by_position(result, std::move(column_result));
-            return Status::OK();
+        }
+        return Status::OK();
+    }
+
+    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
+                        size_t result, size_t input_rows_count) override {
+        const ColumnPtr& column = block.get_by_position(arguments[0]).column;
+        const DataTypePtr& data_type = block.get_by_position(arguments[0]).type;
+        WhichDataType which(data_type);
+
+        MutableColumnPtr column_result = get_return_type_impl({})->create_column();
+        column_result->resize(input_rows_count);
+
+        auto type_error = [&]() {
+            return Status::RuntimeError("Illegal column {} of argument of function {}",
+                                        block.get_by_position(arguments[0]).column->get_name(),
+                                        get_name());
+        };
+        Status status = Status::OK();
+        if (which.is_nullable()) {
+            const DataTypePtr& nested_data_type =
+                    static_cast<const DataTypeNullable*>(data_type.get())->get_nested_type();
+            WhichDataType nested_which(nested_data_type);
+            if (nested_which.is_string_or_fixed_string()) {
+                status = execute_internal<ColumnString, true>(column, data_type, column_result);
+            } else if (nested_which.is_int64()) {
+                status = execute_internal<ColumnInt64, true>(column, data_type, column_result);
+            } else {
+                return type_error();
+            }
+        } else if (which.is_string_or_fixed_string()) {
+            status = execute_internal<ColumnString, false>(column, data_type, column_result);
+        } else if (which.is_int64()) {
+            status = execute_internal<ColumnInt64, false>(column, data_type, column_result);
         } else {
             return type_error();
         }
-
-        block.replace_by_position(result, std::move(column_result));
-        return Status::OK();
+        if (status.ok()) {
+            block.replace_by_position(result, std::move(column_result));
+        }
+        return status;
     }
 };